add message handle

rdpy/protocol/common/error.py

@@ -32,3 +32,13 @@ class NegotiationFailure(Exception):
         '''
         Exception.__init__(self, message)
         
+class InvalidType(Exception):
+    '''
+    raise when invalid value type occured
+    '''
+    def __init__(self, message):
+        '''
+        constructor with message
+        '''
+        Exception.__init__(self, message)
+        

rdpy/protocol/common/layer.py

@@ -70,7 +70,8 @@ class LayerAutomata(Layer):
 #twitsed layer concept
 from twisted.internet import protocol
 #first that handle stream     
-from network import Stream
+from network import Stream, Type
+from error import InvalidType
 
 class RawLayer(protocol.Protocol, LayerAutomata):
     '''
@@ -120,8 +121,15 @@ class RawLayer(protocol.Protocol, LayerAutomata):
         #default callback is recv from LayerAutomata
         self.setNextState(callback)
         
-    def send(self, s):
+    def send(self, message):
         '''
         send stream on tcp layer
         '''
-        self.transport.write(s.getvalue())
+        if isinstance(message, Type):
+            s = Stream()
+            s.writeType(message)
+            self.transport.write(s.getvalue())
+        elif isinstance(message, Stream):
+            self.transport.write(message.getvalue())
+        else:
+            raise InvalidType("expected Stream or Type")

rdpy/protocol/common/network.py

@@ -21,6 +21,12 @@ class Type(object):
         interface definition of read value
         '''
         pass
+    
+    def sizeof(self):
+        '''
+        return size of type
+        '''
+        pass
 
 class SimpleType(Type):
     '''
@@ -56,6 +62,12 @@ class SimpleType(Type):
         '''
         self._value = struct.unpack(self._structFormat,s.read(self._typeSize))[0]
         
+    def sizeof(self):
+        '''
+        return size of type
+        '''
+        return self._typeSize
+        
 class CompositeType(Type):
     '''
     keep ordering declaration of simple type
@@ -71,7 +83,7 @@ class CompositeType(Type):
         '''
         magic function to update type list
         '''
-        if isinstance(value, Type):
+        if name[0] != '_' and isinstance(value, Type):
             self._type.append(value)
         self.__dict__[name] = value
         
@@ -82,12 +94,27 @@ class CompositeType(Type):
         for i in self._type:
             yield i
             
+    def read(self, s):
+        '''
+        call read on each ordered subtype 
+        '''
+        for i in self._type:
+            i.read(s)
+            
     def write(self, s):
         '''
-        call format on each ordered subtype
+        call write on each ordered subtype
         '''
         for i in self._type:
             i.write(s)
+            
+    def sizeof(self):
+        '''
+        call sizeof on each subtype
+        '''
+        result = 0
+        for i in self._type:
+            result += i.sizeof()
 
 class UInt8(SimpleType):
     '''
@@ -237,6 +264,12 @@ class String(Type):
         '''
         return self._value == other._value
     
+    def __str__(self):
+        '''
+        call when str function is call
+        '''
+        return self._value
+    
     def write(self, s):
         '''
         write the entire raw value
@@ -248,6 +281,12 @@ class String(Type):
         read all stream
         '''
         self._value = s.getvalue()
+        
+    def sizeof(self):
+        '''
+        return len of string
+        '''
+        return len(self._value)
     
 
 class Stream(StringIO):
@@ -261,9 +300,22 @@ class Stream(StringIO):
         return self.len - self.pos
     
     def readType(self, t):
+        '''
+        call specific read on type object
+        '''
         t.read(self)
+        
+    def readNextType(self, t):
+        '''
+        read next type but didn't consume it
+        '''
+        self.readType(t)
+        self.pos -= t.sizeof()
     
     def writeType(self, t):
+        '''
+        call specific write on type object
+        '''
         t.write(self)
         
     def write_unistr(self, value):

rdpy/protocol/rdp/tpkt.py

@@ -2,7 +2,7 @@
 @author: sylvain
 '''
 from rdpy.protocol.common.layer import RawLayer
-from rdpy.protocol.common.stream import Stream
+from rdpy.protocol.common.network import Stream
 
 class TPKT(RawLayer):
     '''

rdpy/protocol/rfb/rfb.py

@@ -4,9 +4,10 @@ Created on 12 aout 2013
 @author: sylvain
 '''
 
-from rdpy.protocol.common.network import Stream, UInt8, UInt16Be, UInt32Be
+from rdpy.protocol.common.network import Stream, String, UInt8, UInt16Be, UInt32Be
 from rdpy.protocol.common.layer import RawLayer
-from types import PixelFormat,ProtocolVersion,SecurityType, Rectangle, Encoding
+from types import ServerInit, PixelFormat, ProtocolVersion, SecurityType, Rectangle, Encoding
+from types import PixelFormatMessage, SetEncodingMessage
 
 class Rfb(RawLayer):
     '''
@@ -30,16 +31,15 @@ class Rfb(RawLayer):
         self._version = ProtocolVersion.RFB003008
         #nb security launch by server
         self._securityLevel = SecurityType.INVALID
-        #shared framebuffer
+        #shared framebuffer client init message
-        self._sharedFlag = UInt8
+        self._sharedFlag = UInt8(False)
-        #framebuffer width
+        #server init message
-        self._width = 0
+        #that contain framebuffer dim and pixel format
-        #framebuffer height
+        self._serverInit = ServerInit()
-        self._height = 0
+        #client pixel format
-        #pixel format structure
+        self._clientPixelFormat = PixelFormat()
-        self._pixelFormat = PixelFormat()
         #server name
-        self._serverName = None
+        self._serverName = String()
         #nb rectangle
         self._nbRect = 0
         #current rectangle header
@@ -78,19 +78,6 @@ class Rfb(RawLayer):
         data.readType(bodyLen)
         self.expect(bodyLen.value, self._callbackBody)
         
-    def readProtocolVersionFormat(self, data):
-        '''
-        read protocol version
-        '''
-        data.readType(self._version)
-        if not self._version in [ProtocolVersion.RFB003003, ProtocolVersion.RFB003007, ProtocolVersion.RFB003008]:
-            self._version = ProtocolVersion.UNKNOWN
-    
-    def sendProtocolVersionFormat(self):
-        s = Stream()
-        s.writeType(self._version)
-        self.send(s)
-        
     def connect(self):
         '''
         call when transport layer connection
@@ -99,7 +86,15 @@ class Rfb(RawLayer):
         if self._mode == Rfb.CLIENT:
             self.expect(12, self.readProtocolVersion)
         else:
-            self.sendProtocolVersionFormat()
+            self.send(self._version)
+        
+    def readProtocolVersionFormat(self, data):
+        '''
+        read protocol version
+        '''
+        data.readType(self._version)
+        if not self._version in [ProtocolVersion.RFB003003, ProtocolVersion.RFB003007, ProtocolVersion.RFB003008]:
+            self._version = ProtocolVersion.UNKNOWN
     
     def readProtocolVersion(self, data):
         '''
@@ -112,7 +107,7 @@ class Rfb(RawLayer):
             #protocol version is unknow try best version we can handle
             self._version = ProtocolVersion.RFB003008
         #send same version of 
-        self.sendProtocolVersionFormat()
+        self.send(self._version)
         
         #next state read security
         if self._version == ProtocolVersion.RFB003003:
@@ -143,9 +138,7 @@ class Rfb(RawLayer):
                 self._securityLevel = s
                 break
         #send back security level choosen
-        s = Stream()
+        self.send(self._securityLevel)
-        s.writeType(self._securityLevel)
-        self.send(s)
         self.expect(4, self.readSecurityResult)
         
     def readSecurityResult(self, data):
@@ -154,7 +147,7 @@ class Rfb(RawLayer):
         '''
         result = UInt32Be()
         data.readType(result)
-        if result == 1:
+        if result == UInt32Be(1):
             print "Authentification failed"
             if self._version == ProtocolVersion.RFB003008:
                 self.expectWithHeader(4, self.readSecurityFailed)
@@ -169,23 +162,20 @@ class Rfb(RawLayer):
         '''
         read server init packet
         '''
-        self._width = data.read_beuint16()
+        data.readType(self._serverInit)
-        self._height = data.read_beuint16()
-        serverPixelFomat = PixelFormat()
-        serverPixelFomat.read(data)
         self.expectWithHeader(4, self.readServerName)
     
     def readServerName(self, data):
         '''
         read server name from server init packet
         '''
-        self._serverName = data.getvalue()
+        data.readType(self._serverName)
-        print "Server name %s"%self._serverName
+        print "Server name %s"%str(self._serverName)
         #end of handshake
         #send pixel format
-        self.sendSetPixelFormat(self._pixelFormat)
+        self.send(PixelFormatMessage(self._clientPixelFormat))
         #write encoding
-        self.sendSetEncoding()
+        self.send(SetEncodingMessage())
         #request entire zone
         self.sendFramebufferUpdateRequest(False, 0, 0, self._width, self._height)
         self.expect(1, self.readServerOrder)
@@ -195,7 +185,7 @@ class Rfb(RawLayer):
         read order receive from server
         '''
         packet_type = UInt8()
-        data.readType(packet_type)
+        data.readNextType(packet_type)
         if packet_type == UInt8(0):
             self.expect(3, self.readFrameBufferUpdateHeader)
         
@@ -240,25 +230,9 @@ class Rfb(RawLayer):
         '''
         write client init packet
         '''
-        s = Stream()
+        self.send(self._sharedFlag)
-        s.writeType(self._sharedFlag)
-        self.send(s)
         self.expect(20, self.readServerInit)
         
-    def sendSetPixelFormat(self, pixelFormat):
-        '''
-        write set pixel format packet
-        '''
-        s = Stream()
-        #message type
-        s.write_uint8(0)
-        #padding
-        s.write_uint8(0)
-        s.write_uint8(0)
-        s.write_uint8(0)
-        pixelFormat.write(s)
-        self.transport.write(s.getvalue())
-        
     def sendSetEncoding(self):
         '''
         write set encoding packet

rdpy/protocol/rfb/types.py

@@ -2,7 +2,7 @@
 @author: sylvain
 '''
 
-from rdpy.protocol.common.network import UInt8, UInt16Be, UInt32Be, String
+from rdpy.protocol.common.network import UInt8, UInt16Be, UInt32Be, SInt32Be, String
 from rdpy.protocol.common.network import CompositeType
 from rdpy.utils.const import ConstAttributes
 
@@ -17,7 +17,7 @@ class ProtocolVersion(object):
     RFB003008 = String("RFB 003.008\n")
 
 @ConstAttributes 
-class SecurityType:
+class SecurityType(object):
     '''
     security type supported 
     (or will be supported)
@@ -27,8 +27,8 @@ class SecurityType:
     NONE = UInt8(1)
     VNC = UInt8(2)
 
-@ConstAttributes 
+@ConstAttributes
-class Pointer:
+class Pointer(object):
     '''
     mouse event code (which button)
     actually in RFB specification only$
@@ -37,6 +37,13 @@ class Pointer:
     BUTTON1 = UInt32Be(0x1)
     BUTTON2 = UInt32Be(0x2)
     BUTTON3 = UInt32Be(0x4)
+
+@ConstAttributes  
+class Encoding(object):
+    '''
+    encoding types
+    '''
+    RAW = SInt32Be(0)
     
 class PixelFormat(CompositeType):
     '''
@@ -57,6 +64,32 @@ class PixelFormat(CompositeType):
         self.padding1 = UInt16Be()
         self.padding2 = UInt8()
         
+class PixelFormatMessage(CompositeType):
+    '''
+    message structure used in rfb
+    to send pixel format structure
+    '''
+    def __init__(self, pixelFormat):
+        CompositeType.__init__(self)
+        self.type = UInt8(0)
+        self.padding1 = UInt16Be()
+        self.padding2 = UInt8()
+        self.pixelFormat = pixelFormat
+        
+class SetEncodingMessage(CompositeType):
+    '''
+    message structure used in rfb
+    to send set encoding
+    Actually basic message that only send
+    raw encoding
+    '''
+    def __init__(self):
+        self.type = UInt8(2)
+        self.padding = UInt8()
+        self.nbEncoding = UInt16Be(1)
+        self.raw = Encoding.RAW
+    
+        
 class ServerInit(CompositeType):
     '''
     message send by server to indicate
@@ -69,13 +102,10 @@ class ServerInit(CompositeType):
         self.pixelFormat = PixelFormat()
 
     
-class Rectangle:
+class Rectangle(object):
     def __init__(self):
         self.X = 0
         self.Y = 0
         self.Width = 0
         self.Height = 0
-        self.Encoding = 0
+        self.Encoding = 0
-        
-class Encoding:
-    RAW = 0