From d32bb1e22be00e05835b17bb3ce428eb745f3453 Mon Sep 17 00:00:00 2001
From: citronneur <citronneur@gmail.com>
Date: Sat, 12 Jul 2014 15:46:59 +0200
Subject: [PATCH] bug fix in tpdu layer for server side

---
 rdpy/protocol/rdp/rdp.py  | 34 ++++++++++++++++++++++++++++++----
 rdpy/protocol/rdp/tpdu.py | 26 +++++++++++++++++++++++---
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/rdpy/protocol/rdp/rdp.py b/rdpy/protocol/rdp/rdp.py
index a876b7a..af098b7 100644
--- a/rdpy/protocol/rdp/rdp.py
+++ b/rdpy/protocol/rdp/rdp.py
@@ -36,12 +36,12 @@ class RDPClientController(pdu.PDUClientListener):
         """
         #list of observer
         self._clientObserver = []
-        #transport layer
+        #PDU layer
         self._pduLayer = pdu.PDU(self)
         #multi channel service
         self._mcsLayer = mcs.MCS(LayerMode.CLIENT, self._pduLayer)
         #transport pdu layer
-        self._tpduLayer = tpdu.TPDU(LayerMode.CLIENT, self._mcsLayer)
+        self._tpduLayer = tpdu.createClient(self._mcsLayer)
         #transport packet (protocol layer)
         self._tpktLayer = tpkt.TPKT(self._tpduLayer, self._pduLayer)
         
@@ -186,7 +186,30 @@ class RDPClientController(pdu.PDUClientListener):
             
         except InvalidValue:
             print "try send bad key event"
-                
+
+class RDPServerController(pdu.PDUServerListener):
+    """
+    Controller use in server side mode
+    """               
+    def __init__(self, privateKeyFileName, certificateFileName):
+        """
+        @param privateKeyFileName: file contain server private key
+        @param certficiateFileName: file that contain public key
+        """
+        self._pduLayer = pdu.PDU(self)
+        #multi channel service
+        self._mcsLayer = mcs.MCS(LayerMode.SERVER, self._pduLayer)
+        #transport pdu layer
+        self._tpduLayer = tpdu.createServer(self._mcsLayer, privateKeyFileName, certificateFileName)
+        #transport packet (protocol layer)
+        self._tpktLayer = tpkt.TPKT(self._tpduLayer, self._pduLayer)
+        
+    def getProtocol(self):
+        """
+        @return: the twisted protocol layer
+        in RDP case is TPKT layer
+        """
+        return self._tpktLayer;
 
 class ClientFactory(protocol.Factory):
     """
@@ -225,12 +248,15 @@ class ServerFactory(protocol.Factory):
         Function call from twisted and build rdp protocol stack
         @param addr: destination address
         """
+        controller = RDPServerController(self._privateKeyFileName, self._certificateFileName)
+        return controller.getProtocol()
         #pduLayer = pdu.PDU(pdu.PDUServerListener())
         #return tpkt.TPKT(tpdu.createServer(mcs.createServer(pduLayer), self._privateKeyFileName, self._certificateFileName));
     
-    def buildObserver(self):
+    def buildObserver(self, controller):
         """
         Build observer use for connection
+        @param controller: RDP stack controller
         """
         raise CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "buildObserver", "ServerFactory")) 
         
diff --git a/rdpy/protocol/rdp/tpdu.py b/rdpy/protocol/rdp/tpdu.py
index f217628..b44febc 100644
--- a/rdpy/protocol/rdp/tpdu.py
+++ b/rdpy/protocol/rdp/tpdu.py
@@ -177,7 +177,7 @@ class TPDU(LayerAutomata, StreamSender):
         
         self.setNextState(self.recvData)
         #connection is done send to presentation
-        LayerAutomata.connect(self)
+        self._presentation.connect(self)
         
     def recvConnectionRequest(self, data):
         """
@@ -231,17 +231,19 @@ class TPDU(LayerAutomata, StreamSender):
     def sendConnectionConfirm(self):
         """
         Write connection confirm message
+        Start TLS connection
         Next state is recvData
         @see : http://msdn.microsoft.com/en-us/library/cc240501.aspx
         """
         message = TPDUConnectMessage(MessageType.X224_TPDU_CONNECTION_CONFIRM)
-        message.protocolNeg.code.value = NegociationType.TYPE_RDP_NEG_REQ
+        message.protocolNeg.code.value = NegociationType.TYPE_RDP_NEG_RSP
         message.protocolNeg.selectedProtocol.value = self._selectedProtocol
         self._transport.send(message)
         #_transport is TPKT and transport is TCP layer of twisted
         self._transport.transport.startTLS(ServerTLSContext(self._serverPrivateKeyFileName, self._serverCertificateFileName))
         #connection is done send to presentation
-        LayerAutomata.connect(self)
+        self.setNextState(self.recvData)
+        self._presentation.connect()
         
     def send(self, message):
         """
@@ -250,6 +252,24 @@ class TPDU(LayerAutomata, StreamSender):
         @param message: network.Type message
         """
         self._transport.send((TPDUDataHeader(), message))
+        
+def createClient(mcsLayer):
+    """
+    Factory for client TPDU automata
+    @param mcsLayer: presentation layer of TPDU
+    """
+    return TPDU(LayerMode.CLIENT, mcsLayer)
+
+def createServer(mcsLayer, privateKeyFileName, certificateFileName):
+    """
+    Factory for server TPDU automata
+    @param mcsLayer: presentation layer of TPDU
+    @param privateKeyFileName: file contain server private key
+    @param certficiateFileName: file that contain public key
+    """
+    tpduLayer = TPDU(LayerMode.SERVER, mcsLayer)
+    tpduLayer.initTLSServerInfos(privateKeyFileName, certificateFileName)
+    return tpduLayer
 
 #open ssl needed
 from twisted.internet import ssl