add message handle

This commit is contained in:
speyrefitte
2013-10-16 18:03:57 +02:00
parent 0f09d2909c
commit 9a862a401f
6 changed files with 145 additions and 71 deletions

View File

@@ -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)

View File

@@ -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")

View File

@@ -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):

View File

@@ -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):
'''

View File

@@ -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

View File

@@ -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
self.Encoding = 0