diff --git a/rdpy/display/qt/adaptor.py b/rdpy/display/qt/adaptor.py index 7cfca41..409a865 100644 --- a/rdpy/display/qt/adaptor.py +++ b/rdpy/display/qt/adaptor.py @@ -4,6 +4,7 @@ Created on 4 sept. 2013 ''' from PyQt4 import QtGui from rdpy.protocol.rfb.observer import RfbObserver +from rdpy.protocol.common.network import UInt8, UInt16Be class QAdaptor(object): ''' @@ -41,14 +42,14 @@ class RfbAdaptor(RfbObserver, QAdaptor): ''' imageFormat = None - if pixelFormat.BitsPerPixel == 32 and pixelFormat.RedShift == 16: + if pixelFormat.BitsPerPixel == UInt8(32) and pixelFormat.RedShift == UInt8(16): imageFormat = QtGui.QImage.Format_RGB32 else: print "Receive image in bad format" return - image = QtGui.QImage(data, width, height, imageFormat) - self.notifyImage(x, y, image) + image = QtGui.QImage(data, width.value, height.value, imageFormat) + self.notifyImage(x.value, y.value, image) def sendMouseEvent(self, e): ''' diff --git a/rdpy/protocol/common/layer.py b/rdpy/protocol/common/layer.py index 128f523..210a146 100644 --- a/rdpy/protocol/common/layer.py +++ b/rdpy/protocol/common/layer.py @@ -121,15 +121,10 @@ class RawLayer(protocol.Protocol, LayerAutomata): #default callback is recv from LayerAutomata self.setNextState(callback) - def send(self, message): + def sendMessage(self, message): ''' send stream on tcp layer ''' - 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 + s = Stream() + s.writeType(message) + self.transport.write(s.getvalue()) \ No newline at end of file diff --git a/rdpy/protocol/common/network.py b/rdpy/protocol/common/network.py index ec56896..ad7c7fd 100644 --- a/rdpy/protocol/common/network.py +++ b/rdpy/protocol/common/network.py @@ -83,30 +83,23 @@ class CompositeType(Type): ''' magic function to update type list ''' - if name[0] != '_' and isinstance(value, Type): + if name[0] != '_' and (isinstance(value, Type) or isinstance(value, tuple)): self._type.append(value) self.__dict__[name] = value - - def __iter__(self): - ''' - iteration over object - ''' - 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) + s.readType(i) def write(self, s): ''' call write on each ordered subtype ''' for i in self._type: - i.write(s) + s.writeType(i) def sizeof(self): ''' @@ -299,11 +292,16 @@ class Stream(StringIO): ''' return self.len - self.pos - def readType(self, t): + def readType(self, value): ''' call specific read on type object ''' - t.read(self) + #read each tuple + if isinstance(value, tuple): + for element in value: + self.readType(element) + return + value.read(self) def readNextType(self, t): ''' @@ -312,11 +310,16 @@ class Stream(StringIO): self.readType(t) self.pos -= t.sizeof() - def writeType(self, t): + def writeType(self, value): ''' call specific write on type object ''' - t.write(self) + #write each element of tuple + if isinstance(value, tuple): + for element in value: + self.writeType(element) + return + value.write(self) def write_unistr(self, value): for c in value: diff --git a/rdpy/protocol/rdp/tpkt.py b/rdpy/protocol/rdp/tpkt.py index 5d5346e..4577dae 100644 --- a/rdpy/protocol/rdp/tpkt.py +++ b/rdpy/protocol/rdp/tpkt.py @@ -20,7 +20,7 @@ class TPKT(RawLayer): #length may be coded on more than 1 bytes self._lastShortLength = 0 - def connectionMade(self): + def connect(self): ''' call when transport layer connection is made (inherit from RawLayer) @@ -28,7 +28,8 @@ class TPKT(RawLayer): #header is on two bytes self.expect(2, self.readHeader) #no connection automata on this layer - self.connect() + if not self._presentation is None: + self._presentation.connect() def readHeader(self, data): ''' diff --git a/rdpy/protocol/rfb/rfb.py b/rdpy/protocol/rfb/rfb.py index bcc3824..13a8e40 100644 --- a/rdpy/protocol/rfb/rfb.py +++ b/rdpy/protocol/rfb/rfb.py @@ -6,8 +6,7 @@ Created on 12 aout 2013 from rdpy.protocol.common.network import Stream, String, UInt8, UInt16Be, UInt32Be from rdpy.protocol.common.layer import RawLayer -from types import ServerInit, PixelFormat, ProtocolVersion, SecurityType, Rectangle, Encoding -from types import PixelFormatMessage, SetEncodingMessage +from types import ServerInit, PixelFormat, FrameBufferUpdateRequest, ProtocolVersion, SecurityType, Rectangle, Encoding class Rfb(RawLayer): ''' @@ -37,7 +36,7 @@ class Rfb(RawLayer): #that contain framebuffer dim and pixel format self._serverInit = ServerInit() #client pixel format - self._clientPixelFormat = PixelFormat() + self._pixelFormat = PixelFormat() #server name self._serverName = String() #nb rectangle @@ -86,7 +85,7 @@ class Rfb(RawLayer): if self._mode == Rfb.CLIENT: self.expect(12, self.readProtocolVersion) else: - self.send(self._version) + self.sendMessage(self._version) def readProtocolVersionFormat(self, data): ''' @@ -107,7 +106,7 @@ class Rfb(RawLayer): #protocol version is unknow try best version we can handle self._version = ProtocolVersion.RFB003008 #send same version of - self.send(self._version) + self.sendMessage(self._version) #next state read security if self._version == ProtocolVersion.RFB003003: @@ -138,7 +137,7 @@ class Rfb(RawLayer): self._securityLevel = s break #send back security level choosen - self.send(self._securityLevel) + self.sendMessage(self._securityLevel) self.expect(4, self.readSecurityResult) def readSecurityResult(self, data): @@ -173,11 +172,11 @@ class Rfb(RawLayer): print "Server name %s"%str(self._serverName) #end of handshake #send pixel format - self.send(PixelFormatMessage(self._clientPixelFormat)) + self.sendPixelFormat(self._pixelFormat) #write encoding - self.send(SetEncodingMessage()) + self.sendSetEncoding() #request entire zone - self.sendFramebufferUpdateRequest(False, 0, 0, self._width, self._height) + self.sendMessage(FrameBufferUpdateRequest(False, 0, 0, self._serverInit.width.value, self._serverInit.height.value)) self.expect(1, self.readServerOrder) def readServerOrder(self, data): @@ -194,34 +193,30 @@ class Rfb(RawLayer): read frame buffer update packet header ''' #padding - data.readType(UInt8()) - self._nbRect = data.read_beuint16(); + nbRect = UInt16Be() + self._nbRect = data.readType((UInt8(), nbRect)) + self._nbRect = nbRect.value self.expect(12, self.readRectHeader) def readRectHeader(self, data): ''' read rectangle header ''' - self._currentRect.X = data.read_beuint16() - self._currentRect.Y = data.read_beuint16() - self._currentRect.Width = data.read_beuint16() - self._currentRect.Height = data.read_beuint16() - self._currentRect.Encoding = data.read_besint32() - - if self._currentRect.Encoding == Encoding.RAW: - self.expect(self._currentRect.Width * self._currentRect.Height * (self._pixelFormat.BitsPerPixel / 8), self.readRectBody) + data.readType(self._currentRect) + if self._currentRect.encoding == Encoding.RAW: + self.expect(self._currentRect.width.value * self._currentRect.height.value * (self._pixelFormat.BitsPerPixel.value / 8), self.readRectBody) def readRectBody(self, data): ''' read body of rect ''' for observer in self._observer: - observer.notifyFramebufferUpdate(self._currentRect.Width, self._currentRect.Height, self._currentRect.X, self._currentRect.Y, self._pixelFormat, self._currentRect.Encoding, data.getvalue()) + observer.notifyFramebufferUpdate(self._currentRect.width, self._currentRect.height, self._currentRect.x, self._currentRect.y, self._pixelFormat, self._currentRect.encoding, data.getvalue()) self._nbRect = self._nbRect - 1 #if there is another rect to read if self._nbRect == 0: #job is finish send a request - self.sendFramebufferUpdateRequest(True, 0, 0, self._width, self._height) + self.sendMessage(FrameBufferUpdateRequest(True, 0, 0, self._serverInit.width.value, self._serverInit.height.value)) self.expect(1, self.readServerOrder) else: self.expect(12, self.readRectHeader) @@ -230,37 +225,27 @@ class Rfb(RawLayer): ''' write client init packet ''' - self.send(self._sharedFlag) + self.sendMessage(self._sharedFlag) self.expect(20, self.readServerInit) + def sendPixelFormat(self, pixelFormat): + ''' + send pixel format structure + ''' + self.sendMessage((UInt8(0), UInt16Be(), UInt8(), pixelFormat)) + def sendSetEncoding(self): ''' - write set encoding packet + send set encoding packet ''' - s = Stream() - #message type - s.write_uint8(2) - #padding - s.write_uint8(0) - #nb encoding - s.write_beuint16(1) - #raw encoding - s.write_besint32(0) - self.transport.write(s.getvalue()) + self.sendMessage((UInt8(2), UInt8(), UInt16Be(1), Encoding.RAW)) def sendFramebufferUpdateRequest(self, incremental, x, y, width, height): ''' request server the specified zone incremental means request only change before last update ''' - s = Stream() - s.write_uint8(3) - s.write_uint8(incremental) - s.write_beuint16(x) - s.write_beuint16(y) - s.write_beuint16(width) - s.write_beuint16(height) - self.transport.write(s.getvalue()) + self.sendMessage(FrameBufferUpdateRequest(incremental, x, y, width, height)) def sendKeyEvent(self, downFlag, key): ''' diff --git a/rdpy/protocol/rfb/types.py b/rdpy/protocol/rfb/types.py index 3b27c69..16283e3 100644 --- a/rdpy/protocol/rfb/types.py +++ b/rdpy/protocol/rfb/types.py @@ -61,38 +61,12 @@ class PixelFormat(CompositeType): self.RedShift = UInt8(16) self.GreenShift = UInt8(8) self.BlueShift = UInt8(0) - 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 + self.padding = (UInt16Be(), UInt8()) class ServerInit(CompositeType): ''' - message send by server to indicate + server init structure framebuffer configuration ''' def __init__(self): @@ -100,12 +74,26 @@ class ServerInit(CompositeType): 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.type = UInt8(3) + self.incremental = UInt8(incremental) + self.x = UInt16Be(x) + self.y = UInt16Be(y) + self.width = UInt16Be(width) + self.height = UInt16Be(height) -class Rectangle(object): +class Rectangle(CompositeType): def __init__(self): - self.X = 0 - self.Y = 0 - self.Width = 0 - self.Height = 0 - self.Encoding = 0 \ No newline at end of file + CompositeType.__init__(self) + self.x = UInt16Be() + self.y = UInt16Be() + self.width = UInt16Be() + self.height = UInt16Be() + self.encoding = SInt32Be() \ No newline at end of file