Merge pull request #136 from GoSecure/scanner_support

Added logger to TwisterTCPLayer for error handling and made PyRDP more "scanner friendly"
GoSecure · Aug 1, 2019 · 7cb6ac5 · 7cb6ac5
2 parents e28a31b + 7f4d2f2
commit 7cb6ac5
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/pyrdp/layer/tcp.py b/pyrdp/layer/tcp.py
@@ -42,6 +42,7 @@ class TwistedTCPLayer(IntermediateLayer, Protocol):
     """
 
     def __init__(self):
+        self.log = logging.getLogger(LOGGER_NAMES.PYRDP)
         super().__init__(TCPParser())
         self.connectedEvent = asyncio.Event()
         self.logSSLRequired = False
@@ -91,9 +92,8 @@ def dataReceived(self, data: bytes):
         except KeyboardInterrupt:
             raise
         except Exception as e:
-            log = logging.getLogger(LOGGER_NAMES.PYRDP)
-            log.exception(e)
-            log.error("Exception occurred when receiving: %(data)s" , {"data": hexlify(data).decode()})
+            self.log.exception(e)
+            self.log.error("Exception occurred when receiving: %(data)s" , {"data": hexlify(data).decode()})
             raise
 
     def sendBytes(self, data: bytes):

diff --git a/pyrdp/mitm/TCPMITM.py b/pyrdp/mitm/TCPMITM.py
@@ -39,6 +39,10 @@ def __init__(self, client: TwistedTCPLayer, server: TwistedTCPLayer, attacker: T
         self.recorder = recorder
         self.serverConnector = serverConnector
 
+        # Allows a lower layer to raise error tagged with the correct sessionID
+        self.client.log = log
+        self.server.log = log
+
         self.clientObserver = self.client.createObserver(
             onConnection = self.onClientConnection,
             onDisconnection = self.onClientDisconnection,

diff --git a/pyrdp/parser/rdp/connection.py b/pyrdp/parser/rdp/connection.py
@@ -110,7 +110,10 @@ def parseClientCoreData(self, stream: BytesIO) -> ClientCoreData:
             core.postBeta2ColorDepth = Uint16LE.unpack(stream)
             core.clientProductId = Uint16LE.unpack(stream)
             core.serialNumber = Uint32LE.unpack(stream)
-            core.highColorDepth = HighColorDepth(Uint16LE.unpack(stream))
+
+            # Should match HighColorDepth enum most of the time, but in order to support scanners and we script, we have to loosely accept this one
+            # Anyway, the server will reject it and enforce another one
+            core.highColorDepth = Uint16LE.unpack(stream)
             core.supportedColorDepths = Uint16LE.unpack(stream)
             core.earlyCapabilityFlags = Uint16LE.unpack(stream)
             core.clientDigProductId = decodeUTF16LE(stream.read(64))