diff --git a/lib/bgr233_palette.py b/lib/bgr233_palette.py index 8fa16b2..cb2fa02 100644 --- a/lib/bgr233_palette.py +++ b/lib/bgr233_palette.py @@ -1,260 +1,15 @@ # BGR233 palette -palette = [ - 0, 0, 0, - 36, 0, 0, - 73, 0, 0, - 109, 0, 0, - 146, 0, 0, - 182, 0, 0, - 219, 0, 0, - 255, 0, 0, - 0, 36, 0, - 36, 36, 0, - 73, 36, 0, - 109, 36, 0, - 146, 36, 0, - 182, 36, 0, - 219, 36, 0, - 255, 36, 0, - 0, 73, 0, - 36, 73, 0, - 73, 73, 0, - 109, 73, 0, - 146, 73, 0, - 182, 73, 0, - 219, 73, 0, - 255, 73, 0, - 0, 109, 0, - 36, 109, 0, - 73, 109, 0, - 109, 109, 0, - 146, 109, 0, - 182, 109, 0, - 219, 109, 0, - 255, 109, 0, - 0, 146, 0, - 36, 146, 0, - 73, 146, 0, - 109, 146, 0, - 146, 146, 0, - 182, 146, 0, - 219, 146, 0, - 255, 146, 0, - 0, 182, 0, - 36, 182, 0, - 73, 182, 0, - 109, 182, 0, - 146, 182, 0, - 182, 182, 0, - 219, 182, 0, - 255, 182, 0, - 0, 219, 0, - 36, 219, 0, - 73, 219, 0, - 109, 219, 0, - 146, 219, 0, - 182, 219, 0, - 219, 219, 0, - 255, 219, 0, - 0, 255, 0, - 36, 255, 0, - 73, 255, 0, - 109, 255, 0, - 146, 255, 0, - 182, 255, 0, - 219, 255, 0, - 255, 255, 0, - 0, 0, 85, - 36, 0, 85, - 73, 0, 85, - 109, 0, 85, - 146, 0, 85, - 182, 0, 85, - 219, 0, 85, - 255, 0, 85, - 0, 36, 85, - 36, 36, 85, - 73, 36, 85, - 109, 36, 85, - 146, 36, 85, - 182, 36, 85, - 219, 36, 85, - 255, 36, 85, - 0, 73, 85, - 36, 73, 85, - 73, 73, 85, - 109, 73, 85, - 146, 73, 85, - 182, 73, 85, - 219, 73, 85, - 255, 73, 85, - 0, 109, 85, - 36, 109, 85, - 73, 109, 85, - 109, 109, 85, - 146, 109, 85, - 182, 109, 85, - 219, 109, 85, - 255, 109, 85, - 0, 146, 85, - 36, 146, 85, - 73, 146, 85, - 109, 146, 85, - 146, 146, 85, - 182, 146, 85, - 219, 146, 85, - 255, 146, 85, - 0, 182, 85, - 36, 182, 85, - 73, 182, 85, - 109, 182, 85, - 146, 182, 85, - 182, 182, 85, - 219, 182, 85, - 255, 182, 85, - 0, 219, 85, - 36, 219, 85, - 73, 219, 85, - 109, 219, 85, - 146, 219, 85, - 182, 219, 85, - 219, 219, 85, - 255, 219, 85, - 0, 255, 85, - 36, 255, 85, - 73, 255, 85, - 109, 255, 85, - 146, 255, 85, - 182, 255, 85, - 219, 255, 85, - 255, 255, 85, - 0, 0, 170, - 36, 0, 170, - 73, 0, 170, - 109, 0, 170, - 146, 0, 170, - 182, 0, 170, - 219, 0, 170, - 255, 0, 170, - 0, 36, 170, - 36, 36, 170, - 73, 36, 170, - 109, 36, 170, - 146, 36, 170, - 182, 36, 170, - 219, 36, 170, - 255, 36, 170, - 0, 73, 170, - 36, 73, 170, - 73, 73, 170, - 109, 73, 170, - 146, 73, 170, - 182, 73, 170, - 219, 73, 170, - 255, 73, 170, - 0, 109, 170, - 36, 109, 170, - 73, 109, 170, - 109, 109, 170, - 146, 109, 170, - 182, 109, 170, - 219, 109, 170, - 255, 109, 170, - 0, 146, 170, - 36, 146, 170, - 73, 146, 170, - 109, 146, 170, - 146, 146, 170, - 182, 146, 170, - 219, 146, 170, - 255, 146, 170, - 0, 182, 170, - 36, 182, 170, - 73, 182, 170, - 109, 182, 170, - 146, 182, 170, - 182, 182, 170, - 219, 182, 170, - 255, 182, 170, - 0, 219, 170, - 36, 219, 170, - 73, 219, 170, - 109, 219, 170, - 146, 219, 170, - 182, 219, 170, - 219, 219, 170, - 255, 219, 170, - 0, 255, 170, - 36, 255, 170, - 73, 255, 170, - 109, 255, 170, - 146, 255, 170, - 182, 255, 170, - 219, 255, 170, - 255, 255, 170, - 0, 0, 255, - 36, 0, 255, - 73, 0, 255, - 109, 0, 255, - 146, 0, 255, - 182, 0, 255, - 219, 0, 255, - 255, 0, 255, - 0, 36, 255, - 36, 36, 255, - 73, 36, 255, - 109, 36, 255, - 146, 36, 255, - 182, 36, 255, - 219, 36, 255, - 255, 36, 255, - 0, 73, 255, - 36, 73, 255, - 73, 73, 255, - 109, 73, 255, - 146, 73, 255, - 182, 73, 255, - 219, 73, 255, - 255, 73, 255, - 0, 109, 255, - 36, 109, 255, - 73, 109, 255, - 109, 109, 255, - 146, 109, 255, - 182, 109, 255, - 219, 109, 255, - 255, 109, 255, - 0, 146, 255, - 36, 146, 255, - 73, 146, 255, - 109, 146, 255, - 146, 146, 255, - 182, 146, 255, - 219, 146, 255, - 255, 146, 255, - 0, 182, 255, - 36, 182, 255, - 73, 182, 255, - 109, 182, 255, - 146, 182, 255, - 182, 182, 255, - 219, 182, 255, - 255, 182, 255, - 0, 219, 255, - 36, 219, 255, - 73, 219, 255, - 109, 219, 255, - 146, 219, 255, - 182, 219, 255, - 219, 219, 255, - 255, 219, 255, - 0, 255, 255, - 36, 255, 255, - 73, 255, 255, - 109, 255, 255, - 146, 255, 255, - 182, 255, 255, - 219, 255, 255, - 255, 255, 255 -] +def generate_bgr233_palette(): + palette = [] + for b in range(4): + for g in range(8): + for r in range(4): + red = int(r * 255 / 3) + green = int(g * 255 / 7) + blue = int(b * 255 / 3) + palette.extend([red, green, blue]) + return palette + + +palette = generate_bgr233_palette() diff --git a/lib/encodings/__init__.py b/lib/encodings/__init__.py index 2ee2e55..0a2254f 100644 --- a/lib/encodings/__init__.py +++ b/lib/encodings/__init__.py @@ -19,4 +19,6 @@ from . import common from . import raw from . import zlib +#from . import zrle +from . import tight from . import cursor diff --git a/lib/encodings/common.py b/lib/encodings/common.py index 4b421b6..6c94a8e 100644 --- a/lib/encodings/common.py +++ b/lib/encodings/common.py @@ -20,10 +20,14 @@ encodings = {} class ENCODINGS: raw = 0 zlib = 6 + tight = 7 + #zrle = 16 # supported pseudo-encodings cursor = -239 encodings_priority = [ + #ENCODINGS.zrle, + ENCODINGS.tight, ENCODINGS.zlib, ENCODINGS.raw ] diff --git a/lib/encodings/cursor.py b/lib/encodings/cursor.py index 6fbe3a6..c409921 100644 --- a/lib/encodings/cursor.py +++ b/lib/encodings/cursor.py @@ -15,10 +15,28 @@ # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . +import os +import struct +import ctypes +import ctypes.util +import platform +from PIL import Image, ImageChops, ImageDraw, ImagePalette +import base64 +from io import BytesIO + from . import common from struct import * from lib import log -import zlib + +OS_NAME = platform.system() + +if OS_NAME == 'Linux': + import lib.oshelpers.x11 as x11 + Xcursor = x11.XCursor + +if OS_NAME == 'Windows': + import lib.oshelpers.windows_cursor as windows_cursor + class Encoding: name = 'Cursor' @@ -31,6 +49,54 @@ class Encoding: def __init__(self): log.debug("Initialized", __name__) + self.default_cursor_data = 'iVBORw0KGgoAAAANSUhEUgAAABgAAAAYCAYAAADgdz34AAACc0lEQVR4nK2Wv0sbYRjHP/fDyNUGalKitBDI4FAojg5NF8M7BFIPUkvSopuQoYMgFrl/wbWbndrFuBocAroVHCIdJAaxzdJuNYXGBpIgsU8H7zSpGo3mhS/3Hvfe9/N837t7OI3zobXNhT4Nre2otZ3/7RdEbwOYSqk4MACYbdfuPDTXcKherzeUUklgyAX1BaK5ZkERkUql0lBKpYD7/YJogA94LCJi27YcHh42lVLpfkE0YBAIi4iEw2FJJpNSqVSaSqnX/YB0ACKRiEQiEZmenu4bpAMwNjZ2pnQ6fWfIhcWGYZxpd3eX+fn5wWw2+1Ep9cItxOgFcmGhaZodKhaLLCws3BrSNYGnYrHI4uKiB0n0Ark2gadSqcTS0tLg2traJxficyHaBddeE3gqlUo4juNB4m0proR4Dc4HjIjI92g0ekrWdQzDoNVqsbq6Sjgc7rix0Wg0bdt+tbW1ladLczQvS6DrOo7jsLe3Ry6XY319nUAg8GV2dvY90ASqwFfg51XG/6c4+w5isZjk83nZ39//XS6XZXJyUqampqRarR6HQqEo8Ah4CFj08AzEq8RxHCzL+jExMfGu1Wr9Gh8fp9FosL29PTA3N/cSqLkJmsDJdQnaAScAmqZ9i8fjb2u12ueVlZWcbduYpsnGxgaZTOYNp9vaterLtsgAApubmwXLsp4BIcA/PDz8vFqtHs/MzEgikZCjoyMJBoNPOG0ZPUN8lmWNct5zBoDR5eXl3MHBgezs7EihUCgDI7cBtCfxXl0duKfr+tNUKvUhk8lk/X5/DPDTQy/qVoUH8QEP3PkfoE4PPwU3iemBcI25qTnAP3ZG9LuVtmFhAAAAAElFTkSuQmCC' + self.default_cursor_img = Image.open(BytesIO(base64.b64decode(self.default_cursor_data))) + + def get_cursor_image(self): + if OS_NAME == 'Windows': + return windows_cursor.get_cursor_image() + + elif OS_NAME == 'Linux': + cursor = Xcursor() + return cursor.get_cursor_image() + + elif OS_NAME == 'Darwin': + if self.default_cursor_img.mode != 'RGBA': + self.default_cursor_img = self.default_cursor_img.convert('RGBA') + return self.default_cursor_img + + else: + return None + + + def send_cursor(self, x, y, cursor): + sendbuff = bytearray() + sendbuff.extend(pack("!B", 0)) # message type 0 == SetCursorPosition + sendbuff.extend(pack("!H", x)) + sendbuff.extend(pack("!H", y)) + self.cursor_sent = True + + if cursor is not None: + w, h = cursor.size + cursor_bytes = cursor.convert("RGBA").tobytes("raw", "BGRA") + + # Invert alpha channel if needed + pixels = bytearray(cursor_bytes) + for i in range(0, len(pixels), 4): + pixels[i + 3] = 255 - pixels[i + 3] + + sendbuff.extend(pack("!B", 1)) # message type 1 == SetCursorShape + sendbuff.extend(pack("!H", w)) # width + sendbuff.extend(pack("!H", h)) # height + sendbuff.extend(pixels) + + else: + sendbuff.extend(pack("!B", 0)) # message type 0 == SetCursorPosition + sendbuff.extend(pack("!H", x)) + sendbuff.extend(pack("!H", y)) + + return sendbuff + common.encodings[common.ENCODINGS.cursor] = Encoding log.debug("Loaded encoding: %s (%s)" % (__name__, Encoding.id)) diff --git a/lib/encodings/tight.py b/lib/encodings/tight.py new file mode 100644 index 0000000..94dbec7 --- /dev/null +++ b/lib/encodings/tight.py @@ -0,0 +1,101 @@ +from . import common +from lib import log +from struct import * +import zlib +from io import BytesIO +import numpy as np + +class Encoding: + name = 'tight' + id = 7 + description = 'Tight VNC encoding' + enabled = True + + def __init__(self): + self.compress_obj = zlib.compressobj( + zlib.Z_DEFAULT_COMPRESSION, + zlib.DEFLATED, + zlib.MAX_WBITS, + zlib.DEF_MEM_LEVEL, + zlib.Z_DEFAULT_STRATEGY + ) + self._last_compression_was_jpeg = False + + def send_image(self, x, y, w, h, image): + sendbuff = bytearray() + + if image.mode == 'RGBX' or image.mode == 'RGBA': + image = image.convert('RGB') + + rectangles = 1 + sendbuff.extend(pack("!BxH", 0, rectangles)) # FramebufferUpdate + sendbuff.extend(pack("!HHHH", x, y, w, h)) + sendbuff.extend(pack(">i", self.id)) + + if self._should_use_jpeg(image, 64): + self._last_compression_was_jpeg = True + compressed_data = self._compress_image_jpeg(image) + sendbuff.append(0x90) # 0x90 = 10010000 = JPEG subencoding + else: + compressed_data = self._compress_image_zlib(image) + sendbuff.append(0) # control byte + + # content lenght + sendbuff.extend(self._send_compact_length(len(compressed_data))) + + # Tight data + sendbuff.extend(compressed_data) + + return sendbuff + + def _send_compact_length(self, length): + sendbuff = bytearray() + while True: + # Toma los 7 bits más bajos del tamaño + byte = length & 0x7F + length >>= 7 + # Si aún hay más datos, establece el bit alto y continúa + if length > 0: + byte |= 0x80 + sendbuff.append(byte) + if length == 0: + break + return sendbuff + + def _should_use_jpeg(self, image, threshold=256): + if image.mode == 'P': + return False + + if image.mode == 'RGB': + width, height = image.size + sample_size = min(width * height, 1000) + sample = np.array(image).reshape(-1, 3)[:sample_size] + unique_colors = np.unique(sample, axis=0) + return len(unique_colors) > threshold + + return False + + def _compress_image_jpeg(self, image, quality=75): + buffer = BytesIO() + image.save(buffer, format="JPEG", quality=quality) + jpeg_data = buffer.getvalue() + buffer.close() + return jpeg_data + + def _compress_image_zlib(self, image): + if self.compress_obj is None or self._last_compression_was_jpeg: + self.compress_obj = zlib.compressobj( + zlib.Z_DEFAULT_COMPRESSION, + zlib.DEFLATED, + -zlib.MAX_WBITS, # El negativo omite la inclusión de un encabezado zlib. + zlib.DEF_MEM_LEVEL, + zlib.Z_DEFAULT_STRATEGY + ) + self._last_compression_was_jpeg = False + + zlib_data = self.compress_obj.compress(image.tobytes()) + self.compress_obj.flush(zlib.Z_SYNC_FLUSH) + return zlib_data + + +common.encodings[common.ENCODINGS.tight] = Encoding +log.debug("Loaded encoding: %s (%s)" % (__name__, Encoding.id)) diff --git a/lib/encodings/zlib.py b/lib/encodings/zlib.py index ec25348..c953083 100644 --- a/lib/encodings/zlib.py +++ b/lib/encodings/zlib.py @@ -1,6 +1,6 @@ from . import common -from struct import * from lib import log +from struct import * import zlib @@ -39,13 +39,13 @@ class Encoding: sendbuff.extend(pack(">i", self.id)) #log.debug("Compressing...") - zlibdata = self._compressObj.compress( image.tobytes() ) + zlibdata = self._compressObj.compress(image.tobytes()) zlibdata += self._compressObj.flush(zlib.Z_FULL_FLUSH) #log.debug("LEN", len(zlibdata)) - l = pack("!I", len(zlibdata) ) - sendbuff.extend( l ) # send length - sendbuff.extend( zlibdata ) # send compressed data + l = pack("!I", len(zlibdata)) + sendbuff.extend(l) # send length + sendbuff.extend(zlibdata) # send compressed data return sendbuff diff --git a/lib/encodings/zrle.py b/lib/encodings/zrle.py new file mode 100644 index 0000000..7a18aed --- /dev/null +++ b/lib/encodings/zrle.py @@ -0,0 +1,95 @@ +from . import common +from lib import log +import zlib +from struct import pack +from PIL import Image + +class Encoding: + name = 'zrle' + id = 16 + description = 'zrle VNC encoding' + enabled = True + + def __init__(self): + log.debug("Initialized", __name__) + self._compressObj = zlib.compressobj() + + + def send_image(self, x, y, w, h, image): + sendbuff = bytearray() + rectangles = 1 + sendbuff.extend(pack("!BxH", 0, rectangles)) # FramebufferUpdate + sendbuff.extend(pack("!HHHH", x, y, w, h)) + sendbuff.extend(pack(">i", self.id)) # ID de encoding ZRLE + + tmpbuf = bytearray() + + # Dividir la imagen en tiles y comprimirlas + for tile_y in range(0, h, 64): + for tile_x in range(0, w, 64): + tile = image.crop((tile_x, tile_y, min(tile_x + 64, w), min(tile_y + 64, h))) + encoded_tile = self.tile_encode(tile) + tmpbuf.extend(encoded_tile) + + compressed_data = self._compressObj.compress(tmpbuf) + compressed_data += self._compressObj.flush(zlib.Z_SYNC_FLUSH) + + sendbuff.extend(pack("!I", len(compressed_data))) + sendbuff.extend(compressed_data) + log.debug("zrle send_image", x, y, w, h, image) + return sendbuff + + def tile_encode(self, tile): + """Codifica una baldosa (tile) de la imagen usando ZRLE.""" + w, h = tile.size + pixels = list(tile.getdata()) + rle_data = bytearray() + + # Proceso RLE para la baldosa + prev_pixel = pixels[0] + count = 1 + for pixel in pixels[1:]: + if pixel == prev_pixel and count < 255: + count += 1 + else: + rle_data.extend(self._pack_pixel(prev_pixel, count)) + prev_pixel = pixel + count = 1 + rle_data.extend(self._pack_pixel(prev_pixel, count)) + + # Empaquetar la data RLE con el byte de subencoding + encoded_tile = bytearray() + encoded_tile.append(128) # Subencoding RLE + encoded_tile.extend(rle_data) + + return encoded_tile + + def _pack_pixel(self, pixel, count): + if isinstance(pixel, tuple): + # RGBA + r, g, b, a = pixel + pixel_data = bytes([r, g, b]) # Usar solo RGB para ZRLE. + else: + pixel_data = bytes([pixel, pixel, pixel]) + + count_data = self._encode_run_length(count) + return pixel_data + count_data + + def _encode_run_length(self, length): + """Codifica la longitud de una secuencia para RLE.""" + if length == 1: + return b"" + length -= 1 # La longitud se incrementa en 1 según el protocolo ZRLE + encoded_length = bytearray() + while length > 0: + byte = length % 255 + encoded_length.append(byte) + length //= 255 + if length > 0: + encoded_length.append(255) + return encoded_length + + +common.encodings[common.ENCODINGS.zrle] = Encoding + +log.debug("Loaded encoding: %s (%s)" % (__name__, Encoding.id)) diff --git a/lib/imagegrab.py b/lib/imagegrab.py index 7da093e..f86d6d7 100644 --- a/lib/imagegrab.py +++ b/lib/imagegrab.py @@ -2,14 +2,11 @@ import sys from PIL import Image from lib import log -if sys.platform == "linux" or sys.platform == "linux2": - log.debug("ImageGrab: running on Linux") - from Xlib import display, X - # take screen images, that's not the best way, so here - # we use directly use xlib to take the screenshot. - class ImageGrab(): - @staticmethod - def grab(): +class ImageGrab(): + @staticmethod + def grab(): + if sys.platform == "linux" or sys.platform == "linux2": + from Xlib import display, X dsp = display.Display() root = dsp.screen().root geom = root.get_geometry() @@ -19,12 +16,8 @@ if sys.platform == "linux" or sys.platform == "linux2": image = Image.frombytes("RGB", (w, h), raw.data, "raw", "BGRX") return image -elif sys.platform == "darwin": - log.debug("ImageGrab: running on darwin") - import Quartz.CoreGraphics as CG - class ImageGrab(): - @staticmethod - def grab(): + elif sys.platform == "darwin": + import Quartz.CoreGraphics as CG screenshot = CG.CGWindowListCreateImage(CG.CGRectInfinite, CG.kCGWindowListOptionOnScreenOnly, CG.kCGNullWindowID, CG.kCGWindowImageDefault) width = CG.CGImageGetWidth(screenshot) height = CG.CGImageGetHeight(screenshot) @@ -38,6 +31,10 @@ elif sys.platform == "darwin": return i -else: - log.debug("ImageGrab: running on Unknown!") - from PIL import ImageGrab + elif sys.platform == "win32": + from PIL import ImageGrab as WinImageGrab + return WinImageGrab.grab() + + else: + log.debug("ImageGrab: running on an unknown platform!") + raise EnvironmentError("Unsupported platform") diff --git a/lib/kbdctrl.py b/lib/kbdctrl.py index 5a9456a..23d1ea7 100644 --- a/lib/kbdctrl.py +++ b/lib/kbdctrl.py @@ -19,24 +19,18 @@ class KeyboardController: def process_event(self, data): # B = U8, L = U32 (self.downflag, self.key) = unpack("!BxxL", data) - log.debug("KeyEvent", self.downflag, hex(self.key)) + #log.debug("KeyEvent", self.downflag, hex(self.key)) # special key if self.key in self.kbdmap: self.kbdkey = self.kbdmap[self.key] - log.debug("SPECIAL KEY", self.kbdkey) + #log.debug("SPECIAL KEY", self.kbdkey) else: # normal key try: self.kbdkey = self.kbd.KeyCode.from_char(chr(self.key)) except: self.kbdkey = None - # debug keypress to stdout - try: - log.debug("KEY:", self.kbdkey) - except: - log.debug("KEY: (unprintable)") - # send the actual keyboard event try: if self.downflag: diff --git a/lib/mousectrl.py b/lib/mousectrl.py index 3fc9179..5aa0a78 100644 --- a/lib/mousectrl.py +++ b/lib/mousectrl.py @@ -23,38 +23,39 @@ class MouseController(): # process mouse button events if self.buttons[0] and not self.left_pressed: - log.debug("LEFT PRESSED") + #log.debug("LEFT PRESSED") mouse.Controller().press(mouse.Button.left) self.left_pressed = 1 elif not self.buttons[0] and self.left_pressed: - log.debug("LEFT RELEASED") + #log.debug("LEFT RELEASED") mouse.Controller().release(mouse.Button.left) self.left_pressed = 0 if self.buttons[1] and not self.middle_pressed: - log.debug("MIDDLE PRESSED") + #log.debug("MIDDLE PRESSED") mouse.Controller().press(mouse.Button.middle) self.middle_pressed = 1 elif not self.buttons[1] and self.middle_pressed: - log.debug("MIDDLE RELEASED") + #log.debug("MIDDLE RELEASED") mouse.Controller().release(mouse.Button.middle) self.middle_pressed = 0 if self.buttons[2] and not self.right_pressed: - log.debug("RIGHT PRESSED") + #log.debug("RIGHT PRESSED") mouse.Controller().press(mouse.Button.right) self.right_pressed = 1 elif not self.buttons[2] and self.right_pressed: - log.debug("RIGHT RELEASED") + #log.debug("RIGHT RELEASED") mouse.Controller().release(mouse.Button.right) self.right_pressed = 0 if self.buttons[3]: - log.debug("SCROLLUP PRESSED") + #log.debug("SCROLLUP PRESSED") mouse.Controller().scroll(0, 2) if self.buttons[4]: - log.debug("SCROLLDOWN PRESSED") + #log.debug("SCROLLDOWN PRESSED") mouse.Controller().scroll(0, -2) #log.debug("PointerEvent", buttonmask, x, y) + return x, y, self.buttonmask diff --git a/lib/oshelpers/windows_cursor.py b/lib/oshelpers/windows_cursor.py new file mode 100644 index 0000000..6e282e7 --- /dev/null +++ b/lib/oshelpers/windows_cursor.py @@ -0,0 +1,82 @@ +import ctypes +import ctypes.wintypes +from PIL import Image +import numpy as np + +class BITMAPINFOHEADER(ctypes.Structure): + _fields_ = [ + ("biSize", ctypes.wintypes.DWORD), + ("biWidth", ctypes.wintypes.LONG), + ("biHeight", ctypes.wintypes.LONG), + ("biPlanes", ctypes.wintypes.WORD), + ("biBitCount", ctypes.wintypes.WORD), + ("biCompression", ctypes.wintypes.DWORD), + ("biSizeImage", ctypes.wintypes.DWORD), + ("biXPelsPerMeter", ctypes.wintypes.LONG), + ("biYPelsPerMeter", ctypes.wintypes.LONG), + ("biClrUsed", ctypes.wintypes.DWORD), + ("biClrImportant", ctypes.wintypes.DWORD) + ] + +class RGBQUAD(ctypes.Structure): + _fields_ = [ + ("rgbBlue", ctypes.c_ubyte), + ("rgbGreen", ctypes.c_ubyte), + ("rgbRed", ctypes.c_ubyte), + ("rgbReserved", ctypes.c_ubyte) + ] + +class BITMAPINFO(ctypes.Structure): + _fields_ = [("bmiHeader", BITMAPINFOHEADER), ("bmiColors", RGBQUAD * 1)] + +class ICONINFO(ctypes.Structure): + _fields_ = [ + ("fIcon", ctypes.wintypes.BOOL), + ("xHotspot", ctypes.wintypes.DWORD), + ("yHotspot", ctypes.wintypes.DWORD), + ("hbmMask", ctypes.wintypes.HBITMAP), + ("hbmColor", ctypes.wintypes.HBITMAP), + ] + +class CURSORINFO(ctypes.Structure): + _fields_ = [ + ("cbSize", ctypes.wintypes.DWORD), + ("flags", ctypes.wintypes.DWORD), + ("hCursor", ctypes.wintypes.HANDLE), + ("ptScreenPos", ctypes.wintypes.POINT), + ] + +def get_cursor_image(): + ci = CURSORINFO() + ci.cbSize = ctypes.sizeof(CURSORINFO) + ctypes.windll.user32.GetCursorInfo(ctypes.byref(ci)) + + ii = ICONINFO() + ctypes.windll.user32.GetIconInfo(ci.hCursor, ctypes.byref(ii)) + + hdc = ctypes.windll.user32.GetDC(0) # Usar 0 en lugar de None + hbmp = ctypes.wintypes.HANDLE(ii.hbmColor) # Asegurarse de que hbmp es un HANDLE + bmpinfo = BITMAPINFO() + bmpinfo.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER) + ctypes.windll.gdi32.GetDIBits(hdc, hbmp, 0, 0, None, ctypes.byref(bmpinfo), 0) + + width, height = bmpinfo.bmiHeader.biWidth, bmpinfo.bmiHeader.biHeight + bmpinfo.bmiHeader.biCompression = 0 # BI_RGB + buffer = ctypes.create_string_buffer(width * height * 4) + ctypes.windll.gdi32.GetDIBits(hdc, hbmp, 0, height, buffer, ctypes.byref(bmpinfo), 0) + + img = np.frombuffer(buffer, dtype=np.uint8) + img = img.reshape((height, width, 4)) + img = np.flip(img, axis=0) # Las imágenes de bitmap en Windows están al revés + img = Image.fromarray(img, 'RGBA') + + # Free resources + try: + ctypes.windll.gdi32.DeleteObject(hbmp) + ctypes.windll.gdi32.DeleteObject(ii.hbmMask) + ctypes.windll.user32.ReleaseDC(None, hdc) + except: + pass + + return img + diff --git a/lib/oshelpers/x11.py b/lib/oshelpers/x11.py new file mode 100644 index 0000000..52f4643 --- /dev/null +++ b/lib/oshelpers/x11.py @@ -0,0 +1,73 @@ +import ctypes +from ctypes import POINTER, c_int, c_short, c_ushort, c_ulong, c_void_p, Structure, cast +from PIL import Image +import numpy as np + +# Definición de Atom para su uso en la estructura XFixesCursorImage +Atom = c_ulong + +# Definición de la estructura XFixesCursorImage +class XFixesCursorImage(Structure): + _fields_ = [ + ("x", c_short), + ("y", c_short), + ("width", c_ushort), + ("height", c_ushort), + ("xhot", c_ushort), + ("yhot", c_ushort), + ("cursor_serial", Atom), + ("pixels", POINTER(c_ulong)), # Suponiendo que 'pixels' es un puntero a c_ulong + ("atom", Atom), # Presente en la versión 2 y superiores de XFixes + ("name", ctypes.c_char_p) + ] + + +class XCursor: + def __init__(self): + # Cargar las bibliotecas X11 y Xfixes + self.xlib = ctypes.cdll.LoadLibrary("libX11.so") + self.xfixes = ctypes.cdll.LoadLibrary("libXfixes.so.3") + + # Configurar los tipos de retorno + self.xlib.XOpenDisplay.restype = POINTER(c_void_p) + self.xlib.XOpenDisplay.argtypes = [ctypes.c_char_p] + + self.xfixes.XFixesGetCursorImage.restype = POINTER(XFixesCursorImage) + self.xfixes.XFixesGetCursorImage.argtypes = [c_void_p] + + # Abrir la conexión con X + self.display = self.xlib.XOpenDisplay(None) + if not self.display: + raise Exception("No se pudo abrir el display") + + def __del__(self): + self.xlib.XCloseDisplay(self.display) + + def get_cursor_image(self): + # Llamar a XFixesGetCursorImage + cursor_image_ref = self.xfixes.XFixesGetCursorImage(self.display) + if not cursor_image_ref: + # return a 2x2 red image + return Image.fromarray(np.array([[[255, 0, 0, 255], [255, 0, 0, 255]], [[255, 0, 0, 255], [255, 0, 0, 255]]], dtype=np.uint8), 'RGBA') + + cursor_image = cursor_image_ref.contents + width, height = cursor_image.width, cursor_image.height + + # Leer los datos de píxeles + pixels_array_type = c_ulong * (cursor_image.width * cursor_image.height) + pixels_pointer = cast(cursor_image.pixels, POINTER(pixels_array_type)) + pixels_64bit = np.frombuffer(pixels_pointer.contents, dtype=np.uint64) + + # Convertir cada valor de 64 bits en un píxel RGBA + pixels_rgba = np.zeros((cursor_image.height, cursor_image.width, 4), dtype=np.uint8) + + for i in range(cursor_image.height): + for j in range(cursor_image.width): + pixel = int(pixels_64bit[i * cursor_image.width + j]) # Convertir a int para bit shifting + pixels_rgba[i, j, 0] = (pixel >> 16) & 0xFF # Rojo + pixels_rgba[i, j, 1] = (pixel >> 8) & 0xFF # Verde + pixels_rgba[i, j, 2] = pixel & 0xFF # Azul + pixels_rgba[i, j, 3] = (pixel >> 24) & 0xFF + + return Image.fromarray(pixels_rgba, 'RGBA') + diff --git a/lib/rfb_bitmap.py b/lib/rfb_bitmap.py index 3bdd827..7014ff2 100644 --- a/lib/rfb_bitmap.py +++ b/lib/rfb_bitmap.py @@ -15,29 +15,8 @@ class RfbBitmap(): self.red_shift = None self.green_shift = None self.blue_shift = None + self.bigendian = 0 - def __quantizetopalette(self, silf, palette, dither=False): - """Converts an RGB or L mode image to use a given P image's palette.""" - silf.load() - - # use palette from reference image - palette.load() - if palette.mode != "P": - raise ValueError("bad mode for palette image") - if silf.mode != "RGB" and silf.mode != "L": - raise ValueError( - "only RGB or L mode images can be quantized to a palette" - ) - im = silf.im.convert("P", 1 if dither else 0, palette.im) - # the 0 above means turn OFF dithering - - # Later versions of Pillow (4.x) rename _makeself to _new - try: - return silf._new(im) - except AttributeError: - return silf._makeself(im) - - def get_bitmap(self, rectangle): if self.bpp == 32: redBits = 8 @@ -55,53 +34,42 @@ class RfbBitmap(): a[..., 2] = ( a[..., 2] ) & blueMask >> self.blue_shift image = Image.fromarray(a) + if image.mode == "RGBA": + (r, g, b, a) = image.split() + image = Image.merge("RGB", (r, g, b)) + del r, g, b, a + if self.primaryOrder == "rgb": (b, g, r) = image.split() image = Image.merge("RGB", (r, g, b)) - del b,g,r + del b, g, r image = image.convert("RGBX") return image - elif self.bpp == 16: #BGR565 - greenBits = 5 - blueBits = 6 - redBits = 5 - image = rectangle - - if self.primaryOrder == "bgr": # FIXME: does not work - (b, g, r) = image.split() - image = Image.merge("RGB", (r, g, b)) - - if self.depth == 16: - image = image.convert('BGR;16') - if self.depth == 15: - image = image.convert('BGR;15') - + elif self.bpp == 16: + # BGR565 + a = np.array(rectangle) + r = (a[..., 0] >> 3) & 0x1F + g = (a[..., 1] >> 2) & 0x3F + b = (a[..., 2] >> 3) & 0x1F + bgr565 = (r << 11) | (g << 5) | b + bgr565 = bgr565.astype('uint16') + if self.bigendian == 0: + bgr565 = bgr565.byteswap().newbyteorder() + bgr565_bytes = bgr565.tobytes() + image = Image.frombytes('RGB', rectangle.size, bgr565_bytes, 'raw', 'BGR;16') return image - elif self.bpp == 8: #bgr233 - redBits = 3 - greenBits = 3 - blueBits = 2 - image = rectangle - - palette = bgr233_palette.palette - if self.primaryOrder == "rgb": - #(b, g, r) = image.split() - #image = Image.merge("RGB", (r, g, b)) - - palette = np.reshape(palette, (-3,3)) - palette[:,[0, 2]] = palette[:,[2, 0]] - palette = palette.flatten() - palette = list(palette) - - p = Image.new('P',(16,16)) - p.putpalette(palette) - - image = self.__quantizetopalette(image, p, dither=self.dither) - - #image = image.convert('RGB', colors=4).quantize(palette=p) - #log.debug(image) + elif self.bpp == 8: + # BGR233 + image = rectangle.convert('RGB') + a = np.array(image) + r = (a[..., 0] >> 6) & 0x03 + g = (a[..., 1] >> 5) & 0x07 + b = (a[..., 2] >> 6) & 0x03 + bgr233 = (b << 6) | (g << 3) | r + image = Image.fromarray(bgr233.astype('uint8'), 'P') + image.putpalette(bgr233_palette.palette) return image else: diff --git a/pyvncs/__init__.py b/pyvncs/__init__.py index 7b7400f..eace1cc 100644 --- a/pyvncs/__init__.py +++ b/pyvncs/__init__.py @@ -20,3 +20,4 @@ The main pyvncs module """ from . import server + diff --git a/pyvncs/server.py b/pyvncs/server.py index 2588cd3..7acb68c 100644 --- a/pyvncs/server.py +++ b/pyvncs/server.py @@ -22,11 +22,9 @@ from pynput import mouse, keyboard from PIL import Image, ImageChops, ImageDraw, ImagePalette import socket -import select -import os -import sys -import random +import errno import numpy as np +import time from lib import mousectrl from lib import kbdctrl @@ -38,12 +36,13 @@ from lib import log # encodings support import lib.encodings as encs from lib.encodings.common import ENCODINGS +from lib.encodings.cursor import Encoding as CursorEncoding # auth support from lib.auth.vnc_auth import VNCAuth from lib.auth.vencrypt import VeNCrypt -class VncServer(): +class VNCServer(): class RFB_SECTYPES: vncauth = 2 # plain VNC auth @@ -52,6 +51,8 @@ class VncServer(): encoding_object = None + last_cursor = None + def __init__(self, socket, password=None, auth_type=None, pem_file='', vnc_config = None): self.RFB_VERSION = '003.008' self.initmsg = ("RFB %s\n" % self.RFB_VERSION) @@ -62,6 +63,7 @@ class VncServer(): self.auth_type = auth_type self.pem_file = pem_file self.vnc_config = vnc_config + self.cursor_encoding = CursorEncoding() log.debug("Configured auth type:", self.auth_type) @@ -69,7 +71,7 @@ class VncServer(): def __del__(self): log.debug("VncServer died") - def sendmessage(self, message): + def send_message(self, message): ''' sends a RFB message, usually an error message ''' sock = self.socket message = bytes(message, 'iso8859-1') @@ -77,7 +79,7 @@ class VncServer(): buff = pack("I%ds" % (len(message),), len(message), message) sock.send(buff) - def getbuff(self, timeout): + def get_buffer(self, timeout): sock = self.socket sock.settimeout(timeout) @@ -94,7 +96,7 @@ class VncServer(): sock.send(self.initmsg.encode()) # RFB version handshake - data = self.getbuff(30) + data = self.get_buffer(30) log.debug("init received: '%s'" % data) server_version = float(self.RFB_VERSION) @@ -121,7 +123,7 @@ class VncServer(): sock.send(sendbuff) # get client choosen security type - data = self.getbuff(30) + data = self.get_buffer(30) try: sectype = unpack("B", data)[0] except: @@ -130,7 +132,7 @@ class VncServer(): if sectype not in sectypes: log.debug("Incompatible security type: %s" % data) sock.send(pack("B", 1)) # failed handshake - self.sendmessage("Incompatible security type") + self.send_message("Incompatible security type") sock.close() return False @@ -139,7 +141,7 @@ class VncServer(): # VNC Auth if sectype == self.RFB_SECTYPES.vncauth: auth = VNCAuth() - auth.getbuff = self.getbuff + auth.getbuff = self.get_buffer if not auth.auth(sock, self.password): msg = "Auth failed." sendbuff = pack("I", len(msg)) @@ -159,7 +161,7 @@ class VncServer(): return False auth = VeNCrypt(sock) - auth.getbuff = self.getbuff + auth.getbuff = self.get_buffer auth.send_subtypes() client_subtype = auth.client_subtype @@ -191,21 +193,21 @@ class VncServer(): return False # get ClientInit - data = self.getbuff(30) + data = self.get_buffer(30) log.debug("Clientinit (shared flag)", repr(data)) - self.ServerInit() + self.server_init() return True - def ServerInit(self): + def server_init(self): # ServerInit sock = self.socket screen = ImageGrab.grab() - log.debug("screen", repr(screen)) + #log.debug("screen", repr(screen)) size = screen.size - log.debug("size", repr(size)) + #log.debug("size", repr(size)) del screen width = size[0] @@ -249,36 +251,31 @@ class VncServer(): sock.send(sendbuff) - def protocol(self): + def handle_client(self): self.socket.settimeout(None) # set nonblocking socket - screen = ImageGrab.grab() - size = screen.size - width = size[0] - height = size[1] - del screen + last_rfbu = time.time() mousecontroller = mousectrl.MouseController() kbdcontroller = kbdctrl.KeyboardController() clipboardcontroller = clipboardctrl.ClipboardController() - self.primaryOrder = "rgb" + self.primaryOrder = "bgr" self.encoding = ENCODINGS.raw self.encoding_object = encs.common.encodings[self.encoding]() + sock = self.socket while True: - #log.debug(".", end='', flush=True) - r,_,_ = select.select([self.socket],[],[],0) - if r == []: - #no data - sleep(0.1) - continue - sock = r[0] try: data = sock.recv(1) # read first byte except socket.timeout: #log.debug("timeout") continue + except socket.error as e: + err = e.args[0] + # no data + if err == errno.EAGAIN or err == errno.EWOULDBLOCK: + continue except Exception as e: log.debug("exception '%s'" % e) sock.close() @@ -300,10 +297,27 @@ class VncServer(): log.debug("SHIFTS", self.red_shift, self.green_shift, self.blue_shift) log.debug("MAXS", self.red_maximum, self.green_maximum, self.blue_maximum) - if self.red_shift > self.blue_shift: - self.primaryOrder = "rgb" - else: - self.primaryOrder = "bgr" + # Configure primaryOrder + self.primaryOrder = "rgb" if self.red_shift > self.blue_shift else "bgr" + + # rfb_bitmap common config + self.rfb_bitmap.bpp = self.bpp + self.rfb_bitmap.depth = self.depth + self.rfb_bitmap.dither = self.vnc_config.eightbitdither + self.rfb_bitmap.primaryOrder = self.primaryOrder + self.rfb_bitmap.truecolor = self.truecolor + self.rfb_bitmap.red_shift = self.red_shift + self.rfb_bitmap.green_shift = self.green_shift + self.rfb_bitmap.blue_shift = self.blue_shift + self.rfb_bitmap.red_maximum = self.red_maximum + self.rfb_bitmap.green_maximum = self.green_maximum + self.rfb_bitmap.blue_maximum = self.blue_maximum + self.rfb_bitmap.bigendian = self.bigendian + + # fixed bpp for 8 bpp + if self.bpp == 8: + self.primaryOrder = "bgr" # assume BGR for 8 bpp + log.debug("Using order:", self.primaryOrder) continue @@ -319,20 +333,23 @@ class VncServer(): log.debug("client_encodings", repr(self.client_encodings), len(self.client_encodings)) # cursor support? + self.cursor_support = False if ENCODINGS.cursor in self.client_encodings: log.debug("client cursor support") + self.cursor_encoding = CursorEncoding() self.cursor_support = True # which pixel encoding to use? log.debug("encs.common.encodings_priority", encs.common.encodings_priority) for e in encs.common.encodings_priority: - log.debug("E", e) + #log.debug("E", e) if e in self.client_encodings: if self.encoding == e: # don't initialize same encoding again break self.encoding = e - log.debug("Using %s encoding" % self.encoding) + #log.debug("Using %s encoding" % self.encoding) + log.debug("Using %s encoding" % encs.common.encodings[self.encoding].name) self.encoding_object = encs.common.encodings[self.encoding]() break @@ -340,21 +357,33 @@ class VncServer(): if data[0] == 3: # FBUpdateRequest + # rate limit data2 = sock.recv(9, socket.MSG_WAITALL) - #log.debug("Client Message Type: FBUpdateRequest (3)") - #print(len(data2)) + if not data2: + log.debug("connection closed?") + break + if time.time() - last_rfbu < 0.1: + try: + sock.sendall(pack("!BxH", 0, 0)) + except: + log.debug("connection closed?") + break + continue + last_rfbu = time.time() (incremental, x, y, w, h) = unpack("!BHHHH", data2) #log.debug("RFBU:", incremental, x, y, w, h) - self.SendRectangles(sock, x, y, w, h, incremental) - + self.send_rectangles(sock, x, y, w, h, incremental) + if self.cursor_support: + self.send_cursor(x, y) continue + if data[0] == 4: # keyboard event kbdcontroller.process_event(sock.recv(7)) continue if data[0] == 5: # PointerEvent - mousecontroller.process_event(sock.recv(5, socket.MSG_WAITALL)) + x, y, _ = mousecontroller.process_event(sock.recv(5, socket.MSG_WAITALL)) continue if data[0] == 6: # ClientCutText @@ -365,8 +394,7 @@ class VncServer(): data2 = sock.recv(4096) log.debug("RAW Server received data:", repr(data[0]) , data+data2) - - def GetRectangle(self, x, y, w, h): + def get_rectangle(self, x, y, w, h): try: scr = ImageGrab.grab() except: @@ -385,11 +413,51 @@ class VncServer(): return crop - def SendRectangles(self, sock, x, y, w, h, incremental=0): - # send FramebufferUpdate to client + def send_cursor(self, x, y): + cursor_img = self.cursor_encoding.get_cursor_image() + if cursor_img is None: + return False - #log.debug("start SendRectangles") - rectangle = self.GetRectangle(x, y, w, h) + if self.last_cursor == cursor_img: + return True + + w, h = cursor_img.size + bitmap = self.rfb_bitmap + self.last_cursor = cursor_img + raw_pixels = bitmap.get_bitmap(cursor_img) + raw_pixels = raw_pixels.tobytes("raw", raw_pixels.mode) + + # Crear la máscara de forma (bitmask) + bitmask = bytearray() + for j in range(h): + row = 0 + for i in range(w): + if cursor_img.getpixel((i, j))[3]: # Verificar alfa del pixel + row |= (128 >> (i % 8)) + if (i % 8 == 7) or i == w - 1: + bitmask.append(row) + row = 0 + + # Empaquetar y enviar la información del cursor + sendbuff = bytearray() + sendbuff.extend(pack("!BxH", 0, 1)) # FramebufferUpdate, 1 rectangle + sendbuff.extend(pack("!HHHH", x, y, w, h)) # geometry + sendbuff.extend(pack("!i", -239)) # cursor pseudo encoding + sendbuff.extend(raw_pixels) + sendbuff.extend(bitmask) + + try: + self.socket.sendall(sendbuff) + except Exception as e: + print(f"Error al enviar el cursor: {e}") + return False + + return True + + + def send_rectangles(self, sock, x, y, w, h, incremental=0): + # send FramebufferUpdate to client + rectangle = self.get_rectangle(x, y, w, h) if not rectangle: rectangle = Image.new("RGB", [w, h], (0,0,0)) @@ -405,23 +473,26 @@ class VncServer(): # no changes... rectangles = 0 sendbuff.extend(pack("!BxH", 0, rectangles)) + # clear the incoming socket buffer + sleep(0.05) try: sock.sendall(sendbuff) except: + log.debug("connection closed?") return False - sleep(0.1) return - if diff.getbbox() is not None: + else: if hasattr(diff, "getbbox"): + #log.debug(f"RFB_REQ:", incremental, x, y, w, h) rectangle = rectangle.crop(diff.getbbox()) (x, y, _, _) = diff.getbbox() w = rectangle.width h = rectangle.height - #log.debug("XYWH:", x,y,w,h, "diff", repr(diff.getbbox())) + #log.debug(f"RFB_RES:", incremental, x, y, w, h) - stimeout = sock.gettimeout() - sock.settimeout(None) + #stimeout = sock.gettimeout() + #sock.settimeout(None) if self.bpp == 32 or self.bpp == 16 or self.bpp == 8: bitmap = self.rfb_bitmap @@ -447,5 +518,5 @@ class VncServer(): except: # connection closed? return False - sock.settimeout(stimeout) + #sock.settimeout(stimeout) #log.debug("end SendRectangles") diff --git a/requeriments.txt b/requirements.txt similarity index 72% rename from requeriments.txt rename to requirements.txt index 8d781eb..38bbd7e 100644 --- a/requeriments.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ pydes pynput -numpy==1.17 +numpy Pillow-PIL elevate diff --git a/requirements.win32.txt b/requirements.win32.txt new file mode 100644 index 0000000..bdf1f72 --- /dev/null +++ b/requirements.win32.txt @@ -0,0 +1,6 @@ +pydes +pynput +numpy +Pillow-PIL +elevate +pywin32 diff --git a/vncserver.py b/vncserver.py index b4357f8..40cf37a 100755 --- a/vncserver.py +++ b/vncserver.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- - import pyvncs from argparse import ArgumentParser from threading import Thread @@ -34,7 +33,7 @@ class ClientThread(Thread): def run(self): _debug("[+] New server socket thread started for " + self.ip + ":" + str(self.port)) - server = pyvncs.server.VncServer(self.sock, + server = pyvncs.server.VNCServer(self.sock, auth_type=self.vnc_config.auth_type, password=self.vnc_config.vnc_password, pem_file=self.vnc_config.pem_file, @@ -47,7 +46,7 @@ class ClientThread(Thread): _debug("Error negotiating client init") return False - server.protocol() + server.handle_client() def main(argv): @@ -100,19 +99,20 @@ def main(argv): _debug("Multithreaded Python server : Waiting for connections from TCP clients...") _debug("Runing on:", sys.platform) - if sys.platform in ['win32', 'win64']: - from lib.oshelpers import windows as win32 - if not win32.is_admin(): - ret = win32.run_as_admin() - if ret is None: - log.debug("Respawning with admin rights") - sys.exit(0) - elif ret is True: - # admin rights - log.debug("Running with admin rights!") - else: - print('Error(ret=%d): cannot elevate privilege.' % (ret)) - sys.exit(1) + # FIXME run_as_admin() is not working on windows + # if sys.platform in ['win32', 'win64']: + # from lib.oshelpers import windows as win32 + # if not win32.is_admin(): + # ret = win32.run_as_admin() + # if ret is None: + # log.debug("Respawning with admin rights") + # sys.exit(0) + # elif ret is True: + # # admin rights + # log.debug("Running with admin rights!") + # else: + # print('Error(ret=%d): cannot elevate privilege.' % (ret)) + # sys.exit(1) while True: sockServer.listen(4) (conn, (ip,port)) = sockServer.accept() @@ -121,11 +121,12 @@ def main(argv): newthread.start() -if __name__ == "__main__": +if __name__ == "__main__2": try: main(sys.argv) - except KeyboardInterrupt: - # quit + except KeyboardInterrupt as e: _debug("Exiting on ctrl+c...") sys.exit() - \ No newline at end of file + +if __name__ == "__main__": + main(sys.argv)