Merge pull request #78 from speidy/mitm

rdpy-rdpmitm: pep8 fixes, argument parsing improvements

bin/rdpy-rdpmitm.py

@@ -29,7 +29,10 @@ Client RDP -> | ProxyServer | ProxyClient | -> Server RDP
                    -----------------
 """
 
-import sys, os, getopt, time
+import sys
+import os
+import argparse
+import time
 
 from rdpy.core import log, error, rss
 from rdpy.protocol.rdp import rdp
@@ -37,10 +40,12 @@ from twisted.internet import reactor
 
 log._LOG_LEVEL = log.Level.INFO
 
+
 class ProxyServer(rdp.RDPServerObserver):
     """
     @summary: Server side of proxy
     """
+
     def __init__(self, controller, target, clientSecurityLevel, rssRecorder):
         """
         @param controller: {RDPServerController}
@@ -71,7 +76,8 @@ class ProxyServer(rdp.RDPServerObserver):
         if self._client is None:
             # try a connection
             domain, username, password = self._controller.getCredentials()
-            self._rss.credentials(username, password, domain, self._controller.getHostname())
+            self._rss.credentials(username, password,
+                                  domain, self._controller.getHostname())
 
             width, height = self._controller.getScreen()
             self._rss.screen(width, height, self._controller.getColorDepth())
@@ -102,7 +108,8 @@ class ProxyServer(rdp.RDPServerObserver):
         """
         if self._client is None:
             return
-        self._client._controller.sendKeyEventScancode(code, isPressed, isExtended)
+        self._client._controller.sendKeyEventScancode(
+            code, isPressed, isExtended)
         self._rss.keyScancode(code, isPressed)
 
     def onKeyEventUnicode(self, code, isPressed):
@@ -130,10 +137,12 @@ class ProxyServer(rdp.RDPServerObserver):
             return
         self._client._controller.sendPointerEvent(x, y, button, isPressed)
 
+
 class ProxyServerFactory(rdp.ServerFactory):
     """
     @summary: Factory on listening events
     """
+
     def __init__(self, target, ouputDir, privateKeyFilePath, certificateFilePath, clientSecurity):
         """
         @param target: {tuple(ip, prt)}
@@ -141,7 +150,8 @@ class ProxyServerFactory(rdp.ServerFactory):
         @param certificateFilePath: {str} file contain server certificate (if none -> back to standard RDP security)
         @param clientSecurity: {str(ssl|rdp)} security layer use in client connection side
         """
-        rdp.ServerFactory.__init__(self, 16, privateKeyFilePath, certificateFilePath)
+        rdp.ServerFactory.__init__(
+            self, 16, privateKeyFilePath, certificateFilePath)
         self._target = target
         self._ouputDir = ouputDir
         self._clientSecurity = clientSecurity
@@ -157,10 +167,12 @@ class ProxyServerFactory(rdp.ServerFactory):
         self._uniqueId += 1
         return ProxyServer(controller, self._target, self._clientSecurity, rss.createRecorder(os.path.join(self._ouputDir, "%s_%s_%s.rss" % (time.strftime('%Y%m%d%H%M%S'), addr.host, self._uniqueId))))
 
+
 class ProxyClient(rdp.RDPClientObserver):
     """
     @summary: Client side of proxy
     """
+
     def __init__(self, controller, server):
         """
         @param controller: {rdp.RDPClientController}
@@ -177,7 +189,8 @@ class ProxyClient(rdp.RDPClientObserver):
         """
         self._server.setClient(self)
         # maybe color depth change
-        self._server._controller.setColorDepth(self._controller.getColorDepth())
+        self._server._controller.setColorDepth(
+            self._controller.getColorDepth())
 
     def onSessionReady(self):
         """
@@ -209,13 +222,17 @@ class ProxyClient(rdp.RDPClientObserver):
         @param data: {str} bitmap data
         @see: rdp.RDPClientObserver.onUpdate
         """
-        self._server._rss.update(destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, rss.UpdateFormat.BMP if isCompress else rss.UpdateFormat.RAW, data)
+        self._server._rss.update(destLeft, destTop, destRight, destBottom, width, height,
-        self._server._controller.sendUpdate(destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, isCompress, data)
+                                 bitsPerPixel, rss.UpdateFormat.BMP if isCompress else rss.UpdateFormat.RAW, data)
+        self._server._controller.sendUpdate(
+            destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, isCompress, data)
+
 
 class ProxyClientFactory(rdp.ClientFactory):
     """
     @summary: Factory for proxy client
     """
+
     def __init__(self, server, width, height, domain, username, password, security):
         """
         @param server: {ProxyServer}
@@ -252,59 +269,55 @@ class ProxyClientFactory(rdp.ClientFactory):
         controller.setPerformanceSession()
         return ProxyClient(controller, self._server)
 
-def help():
-    """
-    @summary: Print help in console
-    """
-    print """
-    Usage:  rdpy-rdpmitm.py -o output_directory target
-            [-l listen_port default 3389] 
-            [-k private_key_file_path (mandatory for SSL)] 
-            [-c certificate_file_path (mandatory for SSL)] 
-            [-o output directory for recoded files] 
-            [-r RDP standard security (XP or server 2003 client or older)] 
-            [-n For NLA Client authentication (need to provide credentials)] 
-    """
 
 def parseIpPort(interface, defaultPort="3389"):
     if ':' in interface:
-        return interface.split(':')
+        s = interface.split(':')
+        return s[0], int(s[1])
     else:
-        return interface, defaultPort
+        return interface, int(defaultPort)
+
+
+def isDirectory(outputDirectory):
+    if outputDirectory is None or not os.path.dirname(outputDirectory):
+        log.error("{} is an invalid output directory or directory doesn't exist".format(
+                  outputDirectory))
+    return outputDirectory
+
+
+def mapSecurityLayer(layer):
+    return {
+        "rdp": rdp.SecurityLevel.RDP_LEVEL_RDP,
+        "tls": rdp.SecurityLevel.RDP_LEVEL_SSL,
+        "nla": rdp.SecurityLevel.RDP_LEVEL_NLA
+    }[layer]
+
 
 if __name__ == '__main__':
-    listen = "3389"
+    p = argparse.ArgumentParser(
-    privateKeyFilePath = None
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    certificateFilePath = None
-    ouputDirectory = None
-    #for anonymous authentication
-    clientSecurity = rdp.SecurityLevel.RDP_LEVEL_SSL
 
-    try:
+    p.add_argument('-l', '--listen', type=parseIpPort, default="0.0.0.0:3389",
-        opts, args = getopt.getopt(sys.argv[1:], "hl:k:c:o:rn")
+                   help="<addr>[:<port>] to bind the server")
-    except getopt.GetoptError:
+    p.add_argument('-t', '--target', type=parseIpPort, required=True,
-        help()
+                   help="<addr>[:<port>] of the target you want to connect to via proxy")
-    for opt, arg in opts:
+    p.add_argument('-o', '--output', type=isDirectory,
-        if opt == "-h":
+                   help="output directory", required=True)
-            help()
+    p.add_argument('-s', '--sec', choices=["rdp", "tls", "nla"],
-            sys.exit()
+                   default="rdp", help="set protocol security layer")
-        elif opt == "-l":
+    ssl = p.add_argument_group()
-            listen = arg
+    ssl.add_argument('-c', '--certificate', help="certificate for TLS connections")
-        elif opt == "-k":
+    ssl.add_argument('-k', '--key', help="private key of the given certificate for TLS connections")
-            privateKeyFilePath = arg
-        elif opt == "-c":
-            certificateFilePath = arg
-        elif opt == "-o":
-            ouputDirectory = arg
-        elif opt == "-r":
-            clientSecurity = rdp.SecurityLevel.RDP_LEVEL_RDP
-        elif opt == "-n":
-            clientSecurity = rdp.SecurityLevel.RDP_LEVEL_NLA
 
-    if ouputDirectory is None or not os.path.dirname(ouputDirectory):
+    args = p.parse_args()
-        log.error("%s is an invalid output directory"%ouputDirectory)
+
-        help()
+    if args.certificate and args.key and not args.sec == "nla":
-        sys.exit()
+        args.sec = "tls"
+
+    log.info("running server on {addr}, using {sec} security layer, proxying to {target}".format(
+             addr=args.listen, sec=args.sec.upper(), target=args.target))
+    reactor.listenTCP(args.listen[1], ProxyServerFactory(
+        args.target, args.output, args.key, args.certificate, mapSecurityLayer(args.sec)),
+        interface=args.listen[0])
 
-    reactor.listenTCP(int(listen), ProxyServerFactory(parseIpPort(args[0]), ouputDirectory, privateKeyFilePath, certificateFilePath, clientSecurity))
     reactor.run()