diff --git a/rdpy/protocol/common/error.py b/rdpy/protocol/common/error.py index 4b4f791..b171008 100644 --- a/rdpy/protocol/common/error.py +++ b/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) + diff --git a/rdpy/protocol/common/layer.py b/rdpy/protocol/common/layer.py index 68b3202..128f523 100644 --- a/rdpy/protocol/common/layer.py +++ b/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()) \ No newline at end of file + 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") \ No newline at end of file diff --git a/rdpy/protocol/common/network.py b/rdpy/protocol/common/network.py index 41b90a6..ec56896 100644 --- a/rdpy/protocol/common/network.py +++ b/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): diff --git a/rdpy/protocol/rdp/tpkt.py b/rdpy/protocol/rdp/tpkt.py index 41fa1e1..5d5346e 100644 --- a/rdpy/protocol/rdp/tpkt.py +++ b/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): ''' diff --git a/rdpy/protocol/rfb/rfb.py b/rdpy/protocol/rfb/rfb.py index fcf588d..bcc3824 100644 --- a/rdpy/protocol/rfb/rfb.py +++ b/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 - self._sharedFlag = UInt8 - #framebuffer width - self._width = 0 - #framebuffer height - self._height = 0 - #pixel format structure - self._pixelFormat = PixelFormat() + #shared framebuffer client init message + self._sharedFlag = UInt8(False) + #server init message + #that contain framebuffer dim and pixel format + self._serverInit = ServerInit() + #client pixel format + self._clientPixelFormat = 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() - s.writeType(self._securityLevel) - self.send(s) + self.send(self._securityLevel) 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() - self._height = data.read_beuint16() - serverPixelFomat = PixelFormat() - serverPixelFomat.read(data) + data.readType(self._serverInit) self.expectWithHeader(4, self.readServerName) def readServerName(self, data): ''' read server name from server init packet ''' - self._serverName = data.getvalue() - print "Server name %s"%self._serverName + data.readType(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() - s.writeType(self._sharedFlag) - self.send(s) + self.send(self._sharedFlag) 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 diff --git a/rdpy/protocol/rfb/types.py b/rdpy/protocol/rfb/types.py index af8415c..3b27c69 100644 --- a/rdpy/protocol/rfb/types.py +++ b/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 -class Pointer: +@ConstAttributes +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 - -class Encoding: - RAW = 0 \ No newline at end of file + self.Encoding = 0 \ No newline at end of file