From e434b7f1a9a1f2a8961eb9f24a7de4da3f895fdb Mon Sep 17 00:00:00 2001
From: Alexandre Beaulieu <alex@segfault.me>
Date: Fri, 21 Aug 2020 10:19:52 -0400
Subject: [PATCH] refactor(#94): Changed timing of StartTLS.

---
 pyrdp/mitm/RDPMITM.py  | 52 ++++++++++++++++++++++++++++++++++--------
 pyrdp/mitm/X224MITM.py |  9 +++++---
 pyrdp/mitm/config.py   |  7 ++++++
 3 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/pyrdp/mitm/RDPMITM.py b/pyrdp/mitm/RDPMITM.py
index 3a11fdf1f..84f2cf125 100644
--- a/pyrdp/mitm/RDPMITM.py
+++ b/pyrdp/mitm/RDPMITM.py
@@ -1,11 +1,12 @@
 #
 # This file is part of the PyRDP project.
-# Copyright (C) 2019 GoSecure Inc.
+# Copyright (C) 2019-2020 GoSecure Inc.
 # Licensed under the GPLv3 or later.
 #
 
 import asyncio
 import datetime
+import typing
 
 from twisted.internet import reactor
 from twisted.internet.protocol import Protocol
@@ -137,9 +138,11 @@ def __init__(self, mainLogger: SessionLogger, crawlerLogger: SessionLogger, conf
 
         self.player.player.addObserver(LayerLogger(self.attackerLog))
 
-        self.config.outDir.mkdir(parents=True, exist_ok=True)
-        self.config.replayDir.mkdir(exist_ok=True)
-        self.config.fileDir.mkdir(exist_ok=True)
+        self.ensureOutDir()
+        # if config.certificateFileName == "auto":
+        #     self.certCache: CertificateCache = CertificateCache(self.config.certDir)
+        # else:
+        #     self.certCache = None
 
         self.state.securitySettings.addObserver(RC4LoggingObserver(self.rc4Log))
 
@@ -213,17 +216,40 @@ async def connectToServer(self):
             except asyncio.TimeoutError:
                 self.log.error("Failed to connect to recording host: timeout expired")
 
-    def startTLS(self):
+    def doClientTls(self):
+        cert = self.server.tcp.transport.getPeerCertificate()
+        if not cert:
+            # Wait for server certificate
+            reactor.callLater(1, self.doClientTls)
+
+        # Clone certificate if necessary.
+        if self.certs:
+            contextForClient = ServerTLSContext(self.config.privateKeyFileName, self.config.certificateFileName)
+        else:
+            # No automated certificate cloning. Use the specified certificate.
+            contextForClient = ServerTLSContext(self.config.privateKeyFileName, self.config.certificateFileName)
+
+        # Establish TLS tunnel with the client
+        self.onTlsReady()
+        self.client.tcp.startTLS(contextForClient)
+        self.onTlsReady = None
+
+        # Add unknown packet handlers.
+        self.client.segmentation.addObserver(PacketForwarder(self.server.tcp))
+        self.server.segmentation.addObserver(PacketForwarder(self.client.tcp))
+
+    def startTLS(self, onTlsReady: typing.Callable[[], None]):
         """
         Execute a startTLS on both the client and server side.
         """
-        contextForClient = ServerTLSContext(self.config.privateKeyFileName, self.config.certificateFileName)
-        contextForServer = ClientTLSContext()
+        self.onTlsReady = onTlsReady
 
-        self.client.tcp.startTLS(contextForClient)
+        # Establish TLS tunnel with target server...
+        contextForServer = ClientTLSContext()
         self.server.tcp.startTLS(contextForServer)
-        self.client.segmentation.addObserver(PacketForwarder(self.server.tcp))
-        self.server.segmentation.addObserver(PacketForwarder(self.client.tcp))
+
+        # Establish TLS tunnel with client.
+        reactor.callLater(1, self.doClientTls)
 
     def buildChannel(self, client: MCSServerChannel, server: MCSClientChannel):
         """
@@ -420,3 +446,9 @@ def enableForwarding():
             enableForwarding
         ])
         sequencer.run()
+
+    def ensureOutDir(self):
+        self.config.outDir.mkdir(parents=True, exist_ok=True)
+        self.config.replayDir.mkdir(exist_ok=True)
+        self.config.fileDir.mkdir(exist_ok=True)
+        self.config.certDir.mkdir(exist_ok=True)
diff --git a/pyrdp/mitm/X224MITM.py b/pyrdp/mitm/X224MITM.py
index 9fd5f6131..a30f86f8e 100644
--- a/pyrdp/mitm/X224MITM.py
+++ b/pyrdp/mitm/X224MITM.py
@@ -17,7 +17,7 @@
 
 
 class X224MITM:
-    def __init__(self, client: X224Layer, server: X224Layer, log: LoggerAdapter, state: RDPMITMState, connector: typing.Coroutine, startTLSCallback: typing.Callable[[], None]):
+    def __init__(self, client: X224Layer, server: X224Layer, log: LoggerAdapter, state: RDPMITMState, connector: typing.Coroutine, startTLSCallback: typing.Callable[[typing.Callable[[], None]], None]):
         """
 
         :param client: X224 layer for the client side
@@ -112,13 +112,16 @@ def onConnectionConfirm(self, pdu: X224ConnectionConfirmPDU):
             payload = pdu.payload
         else:
             payload = parser.write(NegotiationResponsePDU(NegotiationType.TYPE_RDP_NEG_RSP, 0x00, response.selectedProtocols))
-        self.client.sendConnectionConfirm(payload, source=0x1234)
 
         # FIXME: This should be done based on what authentication method the server selected, not on what
         #        the client supports.
         if self.originalRequest.tlsSupported:
-            self.startTLSCallback()
+            # If a TLS tunnel is requested, then we establish the server-side tunnel before
+            # replying to the client, so that we can clone the certificate if needed.
+            self.startTLSCallback(lambda: self.client.sendConnectionConfirm(payload, source=0x1234))
             self.state.useTLS = True
+        else:
+            self.client.sendConnectionConfirm(payload, source=0x1234)
 
     def onClientDisconnectRequest(self, pdu: X224DisconnectRequestPDU):
         self.server.sendPDU(pdu)
diff --git a/pyrdp/mitm/config.py b/pyrdp/mitm/config.py
index f365ad848..a88d28fd2 100644
--- a/pyrdp/mitm/config.py
+++ b/pyrdp/mitm/config.py
@@ -96,6 +96,13 @@ def fileDir(self) -> Path:
         """
         return self.outDir / "files"
 
+    @property
+    def certDir(self) -> Path:
+        """
+        Get the directory for dynamically generated certificates.
+        """
+        return self.outDir / "certs"
+
 
 """
 The default MITM configuration.