add test + ssl server connection

rdpy/protocol/rdp/tpdu.py

@@ -56,12 +56,12 @@ class TPDUConnectMessage(CompositeType):
     '''
     header of TPDU connection messages 
     '''
-    def __init__(self):
+    def __init__(self, code):
         CompositeType.__init__(self)
         self.len = UInt8(lambda:sizeof(self) - 1)
-        self.code = UInt8()
+        self.code = UInt8(code.value, constant = True)
         self.padding = (UInt16Be(), UInt16Be(), UInt8())
-        #read if there is enought data
+        #read if there is enough data
         self.protocolNeg = Negotiation(optional = True)
         
 class TPDUDataHeader(CompositeType):
@@ -71,7 +71,7 @@ class TPDUDataHeader(CompositeType):
     def __init__(self):
         CompositeType.__init__(self)
         self.header = UInt8(2, constant = True)
-        self.messageType = MessageType.X224_TPDU_DATA
+        self.messageType = UInt8(MessageType.X224_TPDU_DATA.value, constant = True)
         self.separator = UInt8(0x80, constant = True)
     
 class Negotiation(CompositeType):
@@ -127,10 +127,9 @@ class TPDU(LayerAutomata):
         @see: response -> http://msdn.microsoft.com/en-us/library/cc240506.aspx
         @see: failure ->http://msdn.microsoft.com/en-us/library/cc240507.aspx
         '''
-        message = TPDUConnectMessage()
+        message = TPDUConnectMessage(MessageType.X224_TPDU_CONNECTION_CONFIRM)
         data.readType(message)
-        if message.code != MessageType.X224_TPDU_CONNECTION_CONFIRM:
-            raise InvalidExpectedDataException("invalid TPDU header code X224_TPDU_CONNECTION_CONFIRM != %d"%message.code)
+        
         #check presence of negotiation response
         if not message.protocolNeg._is_readed:
             raise InvalidExpectedDataException("server must support negotiation protocol to use SSL")
@@ -157,10 +156,8 @@ class TPDU(LayerAutomata):
         @param data: stream
         @see : http://msdn.microsoft.com/en-us/library/cc240470.aspx
         '''
-        message = TPDUConnectMessage()
+        message = TPDUConnectMessage(MessageType.X224_TPDU_CONNECTION_REQUEST)
         data.readType(message)
-        if message.code != MessageType.X224_TPDU_CONNECTION_REQUEST:
-            raise InvalidExpectedDataException("Expect connection packet")
         
         if not message.protocolNeg._is_readed or message.protocolNeg.failureCode._is_readed:
             raise InvalidExpectedDataException("Too older rdp client")
@@ -187,12 +184,7 @@ class TPDU(LayerAutomata):
         '''
         header = TPDUDataHeader()
         data.readType(header)
-        if header.messageType == MessageType.X224_TPDU_DATA:
-            LayerAutomata.recv(self, data)
-        elif header.messageType == MessageType.X224_TPDU_ERROR:
-            raise Exception("receive error from tpdu layer")
-        else:
-            raise InvalidExpectedDataException("unknow tpdu code %s"%header.messageType)
+        LayerAutomata.recv(self, data)
         
     def sendConnectionRequest(self):
         '''
@@ -200,8 +192,7 @@ class TPDU(LayerAutomata):
         next state is recvConnectionConfirm
         @see: http://msdn.microsoft.com/en-us/library/cc240500.aspx
         '''
-        message = TPDUConnectMessage()
-        message.code = MessageType.X224_TPDU_CONNECTION_REQUEST
+        message = TPDUConnectMessage(MessageType.X224_TPDU_CONNECTION_REQUEST)
         message.protocolNeg.code = NegociationType.TYPE_RDP_NEG_REQ
         message.protocolNeg.selectedProtocol = self._requestedProtocol
         self._transport.send(message)
@@ -213,12 +204,14 @@ class TPDU(LayerAutomata):
         next state is recvData
         @see : http://msdn.microsoft.com/en-us/library/cc240501.aspx
         '''
-        message = TPDUConnectMessage()
-        message.code = MessageType.X224_TPDU_CONNECTION_CONFIRM
+        message = TPDUConnectMessage(MessageType.X224_TPDU_CONNECTION_CONFIRM)
         message.protocolNeg.code = NegociationType.TYPE_RDP_NEG_REQ
         message.protocolNeg.selectedProtocol = self._selectedProtocol
         self._transport.send(message)
-        self.setNextState(self.recvConnectionConfirm)
+        #_transport is TPKT and transport is TCP layer of twisted
+        self._transport.transport.startTLS(ServerTLSContext())
+        #connection is done send to presentation
+        LayerAutomata.connect(self)
         
     def send(self, message):
         '''
@@ -241,4 +234,12 @@ class ClientTLSContext(ssl.ClientContextFactory):
         context.set_options(0x00020000)#SSL_OP_NO_COMPRESSION
         context.set_options(SSL.OP_DONT_INSERT_EMPTY_FRAGMENTS)
         context.set_options(SSL.OP_TLS_BLOCK_PADDING_BUG)
-        return context
+        return context
+    
+class ServerTLSContext(ssl.DefaultOpenSSLContextFactory):
+    '''
+    server context factory for open ssl
+    '''
+    def __init__(self, *args, **kw):
+        kw['sslmethod'] = SSL.TLSv1_METHOD
+        ssl.DefaultOpenSSLContextFactory.__init__(self, *args, **kw)

rdpy/protocol/rfb/message.py

@@ -1,147 +0,0 @@
-'''
-@author: sylvain
-'''
-
-from rdpy.network.type import UInt8, UInt16Be, UInt32Be, SInt32Be, String, CompositeType
-from rdpy.network.const import ConstAttributes, TypeAttributes
-
-@ConstAttributes
-@TypeAttributes(String)
-class ProtocolVersion(object):
-    '''
-    different ptotocol version
-    '''
-    UNKNOWN = ""
-    RFB003003 = "RFB 003.003\n"
-    RFB003007 = "RFB 003.007\n"
-    RFB003008 = "RFB 003.008\n"
-
-@ConstAttributes 
-@TypeAttributes(UInt8)
-class SecurityType(object):
-    '''
-    security type supported 
-    (or will be supported)
-    by rdpy
-    '''
-    INVALID = 0
-    NONE = 1
-    VNC = 2
-
-@ConstAttributes
-@TypeAttributes(UInt32Be)
-class Pointer(object):
-    '''
-    mouse event code (which button)
-    actually in RFB specification only$
-    three buttons are supported
-    '''
-    BUTTON1 = 0x1
-    BUTTON2 = 0x2
-    BUTTON3 = 0x4
-
-@ConstAttributes
-@TypeAttributes(SInt32Be)  
-class Encoding(object):
-    '''
-    encoding types
-    '''
-    RAW = 0
-
-@ConstAttributes
-@TypeAttributes(UInt8)
-class ClientToServerMessages(object):
-    '''
-    messages types
-    '''
-    PIXEL_FORMAT = 0
-    ENCODING = 2
-    FRAME_BUFFER_UPDATE_REQUEST = 3
-    KEY_EVENT = 4
-    POINTER_EVENT = 5
-    CUT_TEXT = 6
-    
-class PixelFormat(CompositeType):
-    '''
-    pixel format structure
-    '''
-    def __init__(self):
-        CompositeType.__init__(self)
-        self.BitsPerPixel = UInt8(32)
-        self.Depth = UInt8(24)
-        self.BigEndianFlag = UInt8(False)
-        self.TrueColorFlag = UInt8(True)
-        self.RedMax = UInt16Be(255)
-        self.GreenMax = UInt16Be(255)
-        self.BlueMax = UInt16Be(255)
-        self.RedShift = UInt8(16)
-        self.GreenShift = UInt8(8)
-        self.BlueShift = UInt8(0)
-        self.padding = (UInt16Be(), UInt8())
-    
-        
-class ServerInit(CompositeType):
-    '''
-    server init structure
-    framebuffer configuration
-    '''
-    def __init__(self):
-        CompositeType.__init__(self)
-        self.width = UInt16Be()
-        self.height = UInt16Be()
-        self.pixelFormat = PixelFormat()
-        
-class FrameBufferUpdateRequest(CompositeType):
-    '''
-    fb update request send from client to server
-    '''
-    def __init__(self, incremental = False, x = 0, y = 0, width = 0, height = 0):
-        CompositeType.__init__(self)
-        self.incremental = UInt8(incremental)
-        self.x = UInt16Be(x)
-        self.y = UInt16Be(y)
-        self.width = UInt16Be(width)
-        self.height = UInt16Be(height)
-
-    
-class Rectangle(CompositeType):
-    '''
-    header message of update rect
-    '''
-    def __init__(self):
-        CompositeType.__init__(self)
-        self.x = UInt16Be()
-        self.y = UInt16Be()
-        self.width = UInt16Be()
-        self.height = UInt16Be()
-        self.encoding = SInt32Be()
-        
-class KeyEvent(CompositeType):
-    '''
-    key event structure message
-    '''
-    def __init__(self, downFlag = False, key = 0):
-        CompositeType.__init__(self)
-        self.downFlag = UInt8(downFlag)
-        self.padding = UInt16Be()
-        self.key = UInt32Be(key)
-        
-class PointerEvent(CompositeType):
-    '''
-    pointer event structure message
-    '''
-    def __init__(self, mask = 0, x = 0, y = 0):
-        CompositeType.__init__(self)
-        self.mask = UInt8(mask)
-        self.x = UInt16Be(x)
-        self.y = UInt16Be(y)
-        
-class ClientCutText(CompositeType):
-    '''
-    client cut text message message
-    '''
-    def __init__(self, text = ""):
-        CompositeType.__init__(self)
-        self.padding = (UInt16Be(), UInt8())
-        self.size = UInt32Be(len(text))
-        self.message = String(text)

rdpy/protocol/rfb/rfb.py

@@ -1,10 +1,151 @@
 '''
-@author: sylvain
+@author: citronneur
 '''
 from twisted.internet import protocol
-from rdpy.network.type import String, UInt8, UInt16Be, UInt32Be
 from rdpy.network.layer import RawLayer, LayerMode
-from message import *
+from rdpy.network.type import UInt8, UInt16Be, UInt32Be, SInt32Be, String, CompositeType
+from rdpy.network.const import ConstAttributes, TypeAttributes
+
+@ConstAttributes
+@TypeAttributes(String)
+class ProtocolVersion(object):
+    '''
+    different ptotocol version
+    '''
+    UNKNOWN = ""
+    RFB003003 = "RFB 003.003\n"
+    RFB003007 = "RFB 003.007\n"
+    RFB003008 = "RFB 003.008\n"
+
+@ConstAttributes 
+@TypeAttributes(UInt8)
+class SecurityType(object):
+    '''
+    security type supported 
+    (or will be supported)
+    by rdpy
+    '''
+    INVALID = 0
+    NONE = 1
+    VNC = 2
+
+@ConstAttributes
+@TypeAttributes(UInt32Be)
+class Pointer(object):
+    '''
+    mouse event code (which button)
+    actually in RFB specification only$
+    three buttons are supported
+    '''
+    BUTTON1 = 0x1
+    BUTTON2 = 0x2
+    BUTTON3 = 0x4
+
+@ConstAttributes
+@TypeAttributes(SInt32Be)  
+class Encoding(object):
+    '''
+    encoding types
+    '''
+    RAW = 0
+
+@ConstAttributes
+@TypeAttributes(UInt8)
+class ClientToServerMessages(object):
+    '''
+    messages types
+    '''
+    PIXEL_FORMAT = 0
+    ENCODING = 2
+    FRAME_BUFFER_UPDATE_REQUEST = 3
+    KEY_EVENT = 4
+    POINTER_EVENT = 5
+    CUT_TEXT = 6
+    
+class PixelFormat(CompositeType):
+    '''
+    pixel format structure
+    '''
+    def __init__(self):
+        CompositeType.__init__(self)
+        self.BitsPerPixel = UInt8(32)
+        self.Depth = UInt8(24)
+        self.BigEndianFlag = UInt8(False)
+        self.TrueColorFlag = UInt8(True)
+        self.RedMax = UInt16Be(255)
+        self.GreenMax = UInt16Be(255)
+        self.BlueMax = UInt16Be(255)
+        self.RedShift = UInt8(16)
+        self.GreenShift = UInt8(8)
+        self.BlueShift = UInt8(0)
+        self.padding = (UInt16Be(), UInt8())
+    
+        
+class ServerInit(CompositeType):
+    '''
+    server init structure
+    framebuffer configuration
+    '''
+    def __init__(self):
+        CompositeType.__init__(self)
+        self.width = UInt16Be()
+        self.height = UInt16Be()
+        self.pixelFormat = PixelFormat()
+        
+class FrameBufferUpdateRequest(CompositeType):
+    '''
+    fb update request send from client to server
+    '''
+    def __init__(self, incremental = False, x = 0, y = 0, width = 0, height = 0):
+        CompositeType.__init__(self)
+        self.incremental = UInt8(incremental)
+        self.x = UInt16Be(x)
+        self.y = UInt16Be(y)
+        self.width = UInt16Be(width)
+        self.height = UInt16Be(height)
+
+    
+class Rectangle(CompositeType):
+    '''
+    header message of update rect
+    '''
+    def __init__(self):
+        CompositeType.__init__(self)
+        self.x = UInt16Be()
+        self.y = UInt16Be()
+        self.width = UInt16Be()
+        self.height = UInt16Be()
+        self.encoding = SInt32Be()
+        
+class KeyEvent(CompositeType):
+    '''
+    key event structure message
+    '''
+    def __init__(self, downFlag = False, key = 0):
+        CompositeType.__init__(self)
+        self.downFlag = UInt8(downFlag)
+        self.padding = UInt16Be()
+        self.key = UInt32Be(key)
+        
+class PointerEvent(CompositeType):
+    '''
+    pointer event structure message
+    '''
+    def __init__(self, mask = 0, x = 0, y = 0):
+        CompositeType.__init__(self)
+        self.mask = UInt8(mask)
+        self.x = UInt16Be(x)
+        self.y = UInt16Be(y)
+        
+class ClientCutText(CompositeType):
+    '''
+    client cut text message message
+    '''
+    def __init__(self, text = ""):
+        CompositeType.__init__(self)
+        self.padding = (UInt16Be(), UInt8())
+        self.size = UInt32Be(len(text))
+        self.message = String(text)
 
 class Rfb(RawLayer):
     '''

rdpy/tests/network/layer.py

@@ -35,4 +35,30 @@ class LayerCase(unittest.TestCase):
                 if s == "message":
                     raise LayerCase.LayerCaseException()
             
-        self.assertRaises(LayerCase.LayerCaseException, rdpy.network.layer.Layer(presentation = TestConnect()).recv, "message")
+        self.assertRaises(LayerCase.LayerCaseException, rdpy.network.layer.Layer(presentation = TestConnect()).recv, "message")
+        
+    def test_layer_automata_more_than_expected(self):
+        '''
+        test layer automata mechanism if data received is more than expected
+        '''
+        class TestAutomata(rdpy.network.layer.RawLayer):
+            def expectedCallBack(self, data):
+                if data.dataLen() == 4:
+                    raise LayerCase.LayerCaseException()
+            
+        t = TestAutomata(rdpy.network.layer.LayerMode.NONE)
+        t.expect(4, t.expectedCallBack)
+        self.assertRaises(LayerCase.LayerCaseException, t.dataReceived, "\x00\x00\x00\x00\x00")
+        
+    def test_layer_automata_less_than_expected(self):
+        '''
+        test layer automata mechanism
+        '''
+        class TestAutomata(rdpy.network.layer.RawLayer):
+            def expectedCallBack(self, data):
+                if data.dataLen() == 4:
+                    raise LayerCase.LayerCaseException()
+            
+        t = TestAutomata(rdpy.network.layer.LayerMode.NONE)
+        t.expect(4, t.expectedCallBack)
+        self.assertEqual(t.dataReceived("\x00\x00\x00"), None, "Not enough dada")

rdpy/tests/protocol/rfb/rfb.py

@@ -0,0 +1,11 @@
+'''
+@author: sylvain
+'''
+import unittest
+
+class RfbCase(unittest.TestCase):
+    '''
+    test casefor rfb layer (vnc)
+    '''
+    def testName(self):
+        pass