From c4d071a2e00e6274173f3a84d1104aded1251067 Mon Sep 17 00:00:00 2001
From: citronneur <citronneur@gmail.com>
Date: Tue, 2 Dec 2014 21:32:52 +0100
Subject: [PATCH] Bug fix in screenshot script

---
 bin/rdpy-rdpscreenshot.py | 37 ++++++++++++++++++-------------------
 bin/rdpy-vncscreenshot.py | 21 ++++++++++++---------
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/bin/rdpy-rdpscreenshot.py b/bin/rdpy-rdpscreenshot.py
index cd12f97..d885df7 100755
--- a/bin/rdpy-rdpscreenshot.py
+++ b/bin/rdpy-rdpscreenshot.py
@@ -99,19 +99,20 @@ class RDPScreenShotFactory(rdp.ClientFactory):
                 controller.setScreen(width, height);
                 self._buffer = QtGui.QImage(width, height, QtGui.QImage.Format_RGB32)
                 self._path = path
-                self._hasUpdated = True
-                self._brandWidthTask = task.LoopingCall(self.checkUpdate)
-                self._brandWidthTask.start(timeout) # call every second
+                self._timeout = timeout
+                self._startTimeout = False
                 
             def onUpdate(self, destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, isCompress, data):
                 """
                 @summary: callback use when bitmap is received 
                 """
-                self._hasUpdated = True
                 image = RDPBitmapToQtImage(destLeft, width, height, bitsPerPixel, isCompress, data);
                 with QtGui.QPainter(self._buffer) as qp:
                 #draw image
                     qp.drawImage(destLeft, destTop, image, 0, 0, destRight - destLeft + 1, destBottom - destTop + 1)
+                if not self._startTimeout:
+                    self._startTimeout = False
+                    reactor.callLater(self._timeout, self.checkUpdate)
                    
             def onReady(self):
                 """
@@ -127,12 +128,7 @@ class RDPScreenShotFactory(rdp.ClientFactory):
                 self._buffer.save(self._path)
                 
             def checkUpdate(self):
-                if not self._hasUpdated:
-                    log.info("close connection on timeout without updating orders")
-                    self._controller.close();
-                    self._brandWidthTask.stop()
-                    return
-                self._hasUpdated = False
+                self._controller.close();
         
         return ScreenShotObserver(controller, self._width, self._height, self._path, self._timeout)
         
@@ -147,8 +143,8 @@ if __name__ == '__main__':
     #default script argument
     width = 1024
     height = 800
-    path = "/tmp/rdpy-rdpscreenshot.jpg"
-    timeout = 2.0
+    path = "/tmp/"
+    timeout = 5.0
     
     try:
         opts, args = getopt.getopt(sys.argv[1:], "hw:l:o:t:")
@@ -166,12 +162,7 @@ if __name__ == '__main__':
             path = arg
         elif opt == "-t":
             timeout = float(arg)
-            
-    if ':' in args[0]:
-        ip, port = args[0].split(':')
-    else:
-        ip, port = args[0], "3389"
-        
+    
     #create application
     app = QtGui.QApplication(sys.argv)
     
@@ -180,6 +171,14 @@ if __name__ == '__main__':
     qt4reactor.install()
     
     from twisted.internet import reactor
-    reactor.connectTCP(ip, int(port), RDPScreenShotFactory(width, height, path, timeout))
+        
+    for arg in args:      
+        if ':' in arg:
+            ip, port = arg.split(':')
+        else:
+            ip, port = arg, "3389"
+            
+        reactor.connectTCP(ip, int(port), RDPScreenShotFactory(width, height, path + "%s.jpg"%ip, timeout))
+        
     reactor.runReturn()
     app.exec_()
\ No newline at end of file
diff --git a/bin/rdpy-vncscreenshot.py b/bin/rdpy-vncscreenshot.py
index d2bb21d..e9040a4 100755
--- a/bin/rdpy-vncscreenshot.py
+++ b/bin/rdpy-vncscreenshot.py
@@ -138,7 +138,7 @@ def help():
         
 if __name__ == '__main__':
     #default script argument
-    path = "/tmp/rdpy-vncscreenshot.jpg"
+    path = "/tmp/"
     password = ""
     
     try:
@@ -154,20 +154,23 @@ if __name__ == '__main__':
         elif opt == "-p":
             password = arg
         
-            
-    if ':' in args[0]:
-        ip, port = args[0].split(':')
-    else:
-        ip, port = args[0], "5900"
-        
     #create application
     app = QtGui.QApplication(sys.argv)
     
     #add qt4 reactor
     import qt4reactor
     qt4reactor.install()
-    
     from twisted.internet import reactor
-    reactor.connectTCP(ip, int(port), RFBScreenShotFactory(password, path))
+
+    
+    for arg in args:      
+        if ':' in arg:
+            ip, port = arg.split(':')
+        else:
+            ip, port = arg, "5900"
+        
+        reactor.connectTCP(ip, int(port), RFBScreenShotFactory(password, path + "%s.jpg"%ip))
+        
+    
     reactor.runReturn()
     app.exec_()
\ No newline at end of file