From 68441864349aaf266cc043914c4463707781841e Mon Sep 17 00:00:00 2001 From: citronneur Date: Mon, 16 Jun 2014 22:25:13 +0200 Subject: [PATCH] add test + ssl server connection --- rdpy/protocol/rdp/tpdu.py | 45 ++++----- rdpy/protocol/rfb/message.py | 147 ---------------------------- rdpy/protocol/rfb/rfb.py | 147 +++++++++++++++++++++++++++- rdpy/tests/network/layer.py | 28 +++++- rdpy/tests/protocol/__init__.py | 0 rdpy/tests/protocol/rfb/__init__.py | 0 rdpy/tests/protocol/rfb/rfb.py | 11 +++ 7 files changed, 205 insertions(+), 173 deletions(-) delete mode 100644 rdpy/protocol/rfb/message.py create mode 100644 rdpy/tests/protocol/__init__.py create mode 100644 rdpy/tests/protocol/rfb/__init__.py create mode 100644 rdpy/tests/protocol/rfb/rfb.py diff --git a/rdpy/protocol/rdp/tpdu.py b/rdpy/protocol/rdp/tpdu.py index ab22d50..331b143 100644 --- a/rdpy/protocol/rdp/tpdu.py +++ b/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 \ No newline at end of file + 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) \ No newline at end of file diff --git a/rdpy/protocol/rfb/message.py b/rdpy/protocol/rfb/message.py deleted file mode 100644 index 3baf936..0000000 --- a/rdpy/protocol/rfb/message.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/rdpy/protocol/rfb/rfb.py b/rdpy/protocol/rfb/rfb.py index bfad6c5..f7e1acf 100644 --- a/rdpy/protocol/rfb/rfb.py +++ b/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): ''' diff --git a/rdpy/tests/network/layer.py b/rdpy/tests/network/layer.py index cad5904..bfa35da 100644 --- a/rdpy/tests/network/layer.py +++ b/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") \ No newline at end of file + 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") \ No newline at end of file diff --git a/rdpy/tests/protocol/__init__.py b/rdpy/tests/protocol/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rdpy/tests/protocol/rfb/__init__.py b/rdpy/tests/protocol/rfb/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rdpy/tests/protocol/rfb/rfb.py b/rdpy/tests/protocol/rfb/rfb.py new file mode 100644 index 0000000..786e7b2 --- /dev/null +++ b/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