From 44d99eaca42d65df07c889d026e376f404960d35 Mon Sep 17 00:00:00 2001
From: speyrefitte <citronneur@gmail.com>
Date: Tue, 2 Dec 2014 17:00:42 +0100
Subject: [PATCH] Eyewitness integration

---
 bin/rdpy-rdpscreenshot.py | 17 ++++++++++++-----
 bin/rdpy-vncscreenshot.py | 14 ++++++++++----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/bin/rdpy-rdpscreenshot.py b/bin/rdpy-rdpscreenshot.py
index 01791ff..cd12f97 100755
--- a/bin/rdpy-rdpscreenshot.py
+++ b/bin/rdpy-rdpscreenshot.py
@@ -38,6 +38,7 @@ class RDPScreenShotFactory(rdp.ClientFactory):
     """
     @summary: Factory for screenshot exemple
     """
+    __INSTANCE__ = 0
     def __init__(self, width, height, path, timeout):
         """
         @param width: width of screen
@@ -45,6 +46,7 @@ class RDPScreenShotFactory(rdp.ClientFactory):
         @param path: path of output screenshot
         @param timeout: close connection after timeout s without any updating
         """
+        RDPScreenShotFactory.__INSTANCE__ += 1
         self._width = width
         self._height = height
         self._path = path
@@ -57,8 +59,10 @@ class RDPScreenShotFactory(rdp.ClientFactory):
         @param reason: str use to advertise reason of lost connection
         """
         log.info("connection lost : %s"%reason)
-        reactor.stop()
-        app.exit()
+        RDPScreenShotFactory.__INSTANCE__ -= 1
+        if(RDPScreenShotFactory.__INSTANCE__ == 0):
+            reactor.stop()
+            app.exit()
         
     def clientConnectionFailed(self, connector, reason):
         """
@@ -66,9 +70,11 @@ class RDPScreenShotFactory(rdp.ClientFactory):
         @param connector: twisted connector use for rdp connection (use reconnect to restart connection)
         @param reason: str use to advertise reason of lost connection
         """
-        log.info("connection failes : %s"%reason)
-        reactor.stop()
-        app.exit()
+        log.info("connection failed : %s"%reason)
+        RDPScreenShotFactory.__INSTANCE__ -= 1
+        if(RDPScreenShotFactory.__INSTANCE__ == 0):
+            reactor.stop()
+            app.exit()
         
         
     def buildObserver(self, controller, addr):
@@ -124,6 +130,7 @@ class RDPScreenShotFactory(rdp.ClientFactory):
                 if not self._hasUpdated:
                     log.info("close connection on timeout without updating orders")
                     self._controller.close();
+                    self._brandWidthTask.stop()
                     return
                 self._hasUpdated = False
         
diff --git a/bin/rdpy-vncscreenshot.py b/bin/rdpy-vncscreenshot.py
index 89486f2..d2bb21d 100755
--- a/bin/rdpy-vncscreenshot.py
+++ b/bin/rdpy-vncscreenshot.py
@@ -37,11 +37,13 @@ class RFBScreenShotFactory(rfb.ClientFactory):
     """
     @summary: Factory for screenshot exemple
     """
+    __INSTANCE__ = 0
     def __init__(self, password, path):
         """
         @param password: password for VNC authentication
         @param path: path of output screenshot
         """
+        RFBScreenShotFactory.__INSTANCE__ += 1
         self._path = path
         self._password = password
         
@@ -52,8 +54,10 @@ class RFBScreenShotFactory(rfb.ClientFactory):
         @param reason: str use to advertise reason of lost connection
         """
         log.info("connection lost : %s"%reason)
-        reactor.stop()
-        app.exit()
+        RFBScreenShotFactory.__INSTANCE__ -= 1
+        if(RFBScreenShotFactory.__INSTANCE__ == 0):
+            reactor.stop()
+            app.exit()
         
     def clientConnectionFailed(self, connector, reason):
         """
@@ -62,8 +66,10 @@ class RFBScreenShotFactory(rfb.ClientFactory):
         @param reason: str use to advertise reason of lost connection
         """
         log.info("connection failed : %s"%reason)
-        reactor.stop()
-        app.exit()
+        RFBScreenShotFactory.__INSTANCE__ -= 1
+        if(RFBScreenShotFactory.__INSTANCE__ == 0):
+            reactor.stop()
+            app.exit()
         
         
     def buildObserver(self, controller, addr):