import struct from array import array import PyD3DES from AppKit import * from twisted.internet.protocol import Protocol, ClientFactory from twisted.internet import reactor class NotEnoughDataException(Exception): pass def tocard32(i): try: return struct.pack('!i', i) except struct.error, e: raise NotEnoughDataException(e) def card32(i): if len(i) != 4: raise NotEnoughDataException(i) return array('i', i.tostring())[0] def tocard16(i): try: return struct.pack('!h', i) except struct.error, e: raise NotEnoughDataException(e) def card16(i): if len(i) != 2: raise NotEnoughDataException(i) return array('h', i.tostring())[0] def tocard8(i): return chr(i) def card8(i): return i class PixelFormat(object): def __init__(self, **args): for k, v in args.items(): setattr(self, k, v) def parse(klass, data): parts = list(data[:4]) parts.append(card16(data[4:6])) parts.append(card16(data[6:8])) parts.append(card16(data[8:10])) parts += map(None, data[10:13]) args = {} for attr, value in zip(('bpp', 'depth', 'bigendian', 'truecolor', 'red_max', 'green_max', 'blue_max', 'red_shift', 'green_shift', 'blue_shift'), parts): args[attr] = value args['size'] = args['bpp'] / 8 return klass(**args) parse = classmethod(parse) def __str__(self): parts = [x for x in dir(self) if not x.startswith('_')] parts.sort() s = '\n'.join([('%12s: %s' % (x, getattr(self, x))) for x in parts]) return s def pack(self): parts = [tocard8(self.bpp), tocard8(self.depth), tocard8(self.bigendian), tocard8(self.truecolor), tocard16(self.red_max), tocard16(self.green_max), tocard16(self.blue_max), tocard8(self.red_shift), tocard8(self.green_shift), tocard8(self.blue_shift), '\0' * 3] return ''.join(parts) def default(klass): return klass(bpp=32, depth=32, bigendian=0, truecolor=1, red_max=255, green_max=255, blue_max=255, red_shift=0, green_shift=8, blue_shift=16, size=4) default = classmethod(default) class VNCClient(Protocol): RAW_ENCODING, COPYRECT_ENCODING, RRE_ENCODING, CORRE_ENCODING, HEXTILE_ENCODING = 0, 1, 2, 4, 5 def __init__(self, passwd=''): self.read_buffer = array('B') self.state = 'ProtocolVersion' self.states = { 'ProtocolVersion': self.do_ProtocolVersion, 'Authentication': self.do_Auth, 'Initialization': self.do_Init, 'Connected': self.do_ConnectedHandler } self.auth_types = [self.do_ConnectionFailed, self.do_NoAuth, self.do_VNCAuth, ] self.vnc_auth_state = 0 self.vnc_auth_states = [self.do_VNCAuth1, self.do_VNCAuth2, ] self.connected_msgs = {0: self.do_FramebufferUpdate, 2: self.do_Bell, 3: self.do_ServerCutText, } def dataReceived(self, data): self.read_buffer.extend(array('B', data)) self.states[self.state]() def setFrameBufferRect(self, x, y, w, h, data): destRect = ((x, y), (w, h)) srcRect = ((0, 0), (w, h)) self.frameBuffer.lockFocus() data.drawInRect_fromRect_operation_fraction_(destRect, srcRect, NSCompositeCopy, 1.0) self.frameBuffer.unlockFocus() def setFrameBufferRectToRawData(self, x, y, w, h, rawData): imgRep = NSBitmapImageRep.alloc().initWithBitmapDataPlanes_pixelsWide_pixelsHigh_bitsPerSample_samplesPerPixel_hasAlpha_isPlanar_colorSpaceName_bytesPerRow_bitsPerPixel_((rawData, None, None, None, None), w, h, 8, 4, True, False, NSCalibratedRGBColorSpace, 0, 4 * 8) imgRep.setAlpha_(False) img = NSImage.alloc().initWithSize_((w, h)) img.setFlipped_(True) img.lockFocus() imgRep.drawInRect_(((0, 0), (w, h))) img.unlockFocus() self.setFrameBufferRect(x, y, w, h, img) def dumpFramebufferToFile(self): self.frameBuffer.TIFFRepresentationUsingCompression_factor_(NSTIFFCompressionLZW, 1.0).writeToFile_atomically_('VNC - %s.tiff' % (self.name, ), False) def do_ProtocolVersion(self): # Remove the servers version and return our own print('Protocol version recieved "%s"' % (''.join(self.read_buffer[:11].tostring()), )) self.read_buffer = self.read_buffer[12:] self.transport.write('RFB 003.003\n') self.state = 'Authentication' def do_ConnectionFailed(self): try: reason_length = card32(self.read_buffer[4:8]) reason = self.read_buffer[8:8 + reason_length] self.read_buffer = self.read_buffer[8 + reason_length:] self.transport.loseConnection() print('Connection failed: %s' % (reason,)) return 1 except Exception, e: print repr(e) def do_Auth(self): t = self.read_buffer[:4] t = card32(t) print('Authentication type is %d (%r)' % (t, self.auth_types[t])) if self.auth_types[t](): pass def do_NoAuth(self): pass def do_VNCAuth(self): self.vnc_auth_states[self.vnc_auth_state]() def do_VNCAuth1(self): try: if len(self.read_buffer) < 20: raise IndexError challenge = self.read_buffer[4:20] self.read_buffer = self.read_buffer[20:] self.read_buffer.extend(array('B', '\x00\x00\x00\x02')) PyD3DES.setkey(self.factory.password) self.transport.write(PyD3DES.encrypt(challenge[:8]) + PyD3DES.encrypt(challenge[8:])) self.vnc_auth_state += 1 except Exception, e: print repr(e), e return def do_VNCAuth2(self): try: if len(self.read_buffer) < 8: raise IndexError response = card32(self.read_buffer[4:8]) if response == 0: print('Authentication succeeded') self.state = 'Initialization' self.read_buffer = self.read_buffer[8:] # Send the ClientInit message that consists # of a card32 that is 1 if the session is shared try: self.transport.write(chr(self.shared)) except: self.transport.write(chr(0)) elif response == 1: print('Invalid password') self.transport.loseConnection() elif response == 2: print('Too many connections!') self.transport.loseConnection() except Exception, e: print(e) def do_Init(self): """ Handle a ServerInitialisation message. """ if len(self.read_buffer) >= 24: print('ServerInitialisation buffer length: %d' % (len(self.read_buffer), )) self.fbsize = (card16(self.read_buffer[:2]), card16(self.read_buffer[2:4])) self.pixel_format = PixelFormat.parse(self.read_buffer[4:20]) name_length = card32(self.read_buffer[20:24]) else: return if len(self.read_buffer) < 24 + name_length: return self.name = ''.join(self.read_buffer[24:24 + name_length].tostring()) print('Connected to %s' % (self.name, )) print('Servers pixel format: \n%s' % (self.pixel_format, )) self.frameBuffer = NSImage.alloc().initWithSize_(self.fbsize) self.frameBuffer.setFlipped_(True) self.alreadyProcessedRectangles = (0, 4) self.read_buffer = self.read_buffer[24 + name_length:] self.state = 'Connected' self.send_SetPixelFormat(PixelFormat.default()) self.send_SetEncoding((self.CORRE_ENCODING, )) self.send_FramebufferUpdateRequest((0, 0, self.fbsize[0], self.fbsize[1]), 0) def do_ConnectedHandler(self): """ """ p = NSAutoreleasePool.alloc().init() self.connected_msgs[self.read_buffer[0]]() del p def do_FramebufferUpdate(self): buf = self.read_buffer try: messageType = buf[0] numberOfRectangles = card16(buf[2:4]) firstRectangle, currentRectangleStart = self.alreadyProcessedRectangles for currentRectangle in range(firstRectangle, numberOfRectangles): currentRectangleLength = 0 x, y, width, height = card16(buf[currentRectangleStart :currentRectangleStart + 2]), \ card16(buf[currentRectangleStart + 2:currentRectangleStart + 4]), \ card16(buf[currentRectangleStart + 4:currentRectangleStart + 6]), \ card16(buf[currentRectangleStart + 6:currentRectangleStart + 8]) encodingType = card32(buf[currentRectangleStart + 8:currentRectangleStart + 12]) currentRectangleLength += 12 currentEncodingStart = currentRectangleStart + 12 currentEncodingLength = 0 if encodingType == self.RAW_ENCODING: currentEncodingLength = VNCClient._Raw(self, currentEncodingStart, x, y, width, height, NotEnoughDataException) elif encodingType == self.CORRE_ENCODING: currentEncodingLength = VNCClient._CoRRE(self, currentEncodingStart, x, y, width, height, NotEnoughDataException) else: print '\t\tUnknown encoding!', encodingType currentRectangleLength += currentEncodingLength currentRectangleStart += currentRectangleLength self.alreadyProcessedRectangles = (currentRectangle + 1, currentRectangleStart) self.read_buffer = self.read_buffer[currentRectangleStart:] self.alreadyProcessedRectangles = (0, 4) # self.dumpFramebufferToFile() self.send_FramebufferUpdateRequest((0, 0, self.fbsize[0], self.fbsize[1]), 0) # self.transport.loseConnection() except NotEnoughDataException, e: pass def _Raw(self, currentEncodingStart, x, y, width, height, NotEnoughDataException): rawDataLength = width * height * self.pixel_format.size rawData = self.read_buffer[currentEncodingStart:currentEncodingStart + rawDataLength] if len(rawData) < rawDataLength: raise NotEnoughDataException self.setFrameBufferRectToRawData(x, y, width, height, rawData) return rawDataLength def _CoRRE(self, currentEncodingStart, x, y, width, height, NotEnoughDataException): buf = self.read_buffer if len(buf) < currentEncodingStart + 8: raise NotEnoughDataException numberOfSubRectangles = card32(buf[currentEncodingStart:currentEncodingStart + 4]) r, g, b, a = buf[currentEncodingStart + 4:currentEncodingStart + 8] self.frameBuffer.lockFocus() NSColor.colorWithCalibratedRed_green_blue_alpha_(r / 255.0, g / 255.0, b / 255.0, 1.0).set() NSRectFill(((x, y), (width, height))) currentSubRectangleStart = currentEncodingStart + 8 if len(buf) < currentSubRectangleStart + (8 * numberOfSubRectangles): raise NotEnoughDataException subRectangleRects = [] subRectangleColors = [] if numberOfSubRectangles: for currentSubRectangle in range(numberOfSubRectangles): r, g, b, a, srX, srY, srWidth, srHeight = buf[currentSubRectangleStart:currentSubRectangleStart + 8] subRectangleRects.append(((x + srX, y + srY), (srWidth, srHeight))) subRectangleColors.append((r / 255.0, g / 255.0, b / 255.0, 1.0)) currentSubRectangleStart += 8 subRectangleColors = map(NSColor.colorWithCalibratedRed_green_blue_alpha_, *zip(*subRectangleColors)) NSRectFillListWithColors(subRectangleRects, subRectangleColors) self.frameBuffer.unlockFocus() currentEncodingLength = 8 * (numberOfSubRectangles + 1) return currentEncodingLength try: from _CoRRE import _CoRRE except ImportError: pass def do_Bell(self): self.read_buffer = self.read_buffer[1:] def do_ServerCutText(self): length = card32(self.read_buffer[4:8]) text = self.read_buffer[8:8 + length] print text self.read_buffer = self.read_buffer[8 + length:] def send_SetPixelFormat(self, pixelFormat): self.pixel_format = pixelFormat msg = tocard8(0) msg += '\0' * 3 msg += pixelFormat.pack() self.transport.write(msg) def send_SetEncoding(self, encodings): msg = tocard8(2) msg = msg + '\0' msg = msg + tocard16(len(encodings)) encs = map(tocard32, encodings) msg = msg + ''.join(encs) self.transport.write(msg) def send_FramebufferUpdateRequest(self, pos, incremental): msg = tocard8(3) msg = msg + tocard8(incremental) msg = msg + ''.join(map(tocard16, pos)) self.transport.write(msg) def send_KeyEvent(self, down, key): msg = tocard8(4) msg = msg + tocard8(down) + '\0' msg = msg + tocard32(key) self.transport.write(msg) def send_PointerEvent(self, pos, buttons): msg = tocard8(5) msg = msg + tocard8(buttons) msg = msg + ''.join(map(tocard16, pos)) self.transport.write(msg) def send_ClientCutText(self, text): msg = tocard8(6) msg = msg + ' ' + card32(len(text)) msg = msg + text self.transport.write(msg) class VNCClientFactory(ClientFactory): protocol = VNCClient def __init__(self, password): self.password = password def startedConnecting(self, connector): print 'startedConnecting' def clientConnectionFailed(self, connector, reason): print 'Connection failed. Reason:', reason def clientConnectionLost(self, connector, reason): print 'Connection lost. Reason:', reason if reactor.running: reactor.stop() import hotshot, hotshot.stats def main(): import sys if len(sys.argv) < 4: print 'Usage: %s hostname port password' % (sys.argv[0], ) return NSApplicationLoad() reactor.connectTCP(sys.argv[1], int(sys.argv[2]), VNCClientFactory(sys.argv[3])) reactor.run() import profile, pstats if __name__ == '__main__': main()