mirror of
https://github.com/xcat2/confluent.git
synced 2026-08-29 01:56:46 +00:00
Markus Hilger
2af402b13c
Apply ruff's safe autofixes. The changes are mechanical and behaviour-preserving. Issues fixed: - F401: remove unused imports. - F841: drop unused local variables and assignments, including discarded await/return values, unused "except ... as e" bindings, and unused "with ... as name" targets. - F541: remove the f prefix from f-strings that contain no placeholders. - E711: compare against None with "is"/"is not" instead of "=="/"!=". - E712: test truthiness directly instead of comparing to True. - E713: use "x not in y" instead of "not x in y". - E714: use "is not" instead of "not ... is". - E731: convert lambdas bound to a name into def statements. - W291/W293: trim trailing whitespace on touched lines.
277 lines
11 KiB
Python
277 lines
11 KiB
Python
import asyncio
|
|
from PIL import Image
|
|
import io
|
|
import numpy as np
|
|
import zlib
|
|
|
|
# This results in an RGBA organization of pixels
|
|
MYPIXFORMAT = bytearray([
|
|
32, # bits per pixel
|
|
24, # depth
|
|
0, # big endian
|
|
1, # true color
|
|
0, 255, # red max
|
|
0, 255, # green max
|
|
0, 255, # blue max
|
|
0, 8, 16, # red shift, green shift, blue shift
|
|
0, 0, 0 # padding
|
|
])
|
|
|
|
class ByteStream:
|
|
def __init__(self):
|
|
self.buffer = b''
|
|
|
|
def add_number(self, number, num_bytes):
|
|
data = number.to_bytes(num_bytes, byteorder='big', signed=True)
|
|
self.buffer += data
|
|
|
|
def extend(self, data):
|
|
self.buffer += data
|
|
|
|
def get_bytes(self):
|
|
return self.buffer
|
|
|
|
def clear(self):
|
|
self.buffer = b''
|
|
|
|
def flush(self, writer):
|
|
writer.write(self.buffer)
|
|
self.clear()
|
|
|
|
class VNCClient:
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
await self.close()
|
|
return False
|
|
|
|
@classmethod
|
|
async def create(cls, url):
|
|
self = cls()
|
|
if url.startswith('unix://'):
|
|
url = url.replace('unix://', '')
|
|
if url.startswith('/'):
|
|
self.reader, self.writer = await asyncio.open_unix_connection(url)
|
|
elif url.startswith('@'):
|
|
url = '\0' + url[1:]
|
|
self.reader, self.writer = await asyncio.open_unix_connection(url)
|
|
elif url.startswith('tcp://'):
|
|
url = url.replace('tcp://', '')
|
|
host, port = url.split(':')
|
|
self.reader, self.writer = await asyncio.open_connection(host, int(port))
|
|
else:
|
|
raise ValueError('Unsupported URL: {}'.format(url))
|
|
self.receiver = None
|
|
self.framebuffer = None
|
|
self.copytext = None
|
|
self._updating = True
|
|
self.decompressor = zlib.decompressobj()
|
|
self._input_queue = asyncio.Queue()
|
|
self._input_task = asyncio.create_task(self._input_worker())
|
|
await self._do_vnc_handshake()
|
|
return self
|
|
|
|
async def _input_worker(self):
|
|
while True:
|
|
keys, modifierkeys = await self._input_queue.get()
|
|
payload = ByteStream()
|
|
for modkey in (modifierkeys or []):
|
|
payload.add_number(4, 1) # Key event
|
|
payload.add_number(1, 1) # Down
|
|
payload.add_number(0, 2) # Padding
|
|
payload.add_number(modkey.value, 4)
|
|
payload.flush(self.writer)
|
|
await self.writer.drain()
|
|
for key in keys:
|
|
keynumber = key.value if hasattr(key, 'value') else key
|
|
payload.add_number(4, 1) # Key event
|
|
payload.add_number(1, 1) # Down
|
|
payload.add_number(0, 2) # Padding
|
|
payload.add_number(keynumber, 4)
|
|
payload.add_number(4, 1) # Key event
|
|
payload.add_number(0, 1) # Up
|
|
payload.add_number(0, 2) # Padding
|
|
payload.add_number(keynumber, 4)
|
|
payload.flush(self.writer)
|
|
await self.writer.drain()
|
|
for modkey in (modifierkeys or []):
|
|
payload.add_number(4, 1) # Key event
|
|
payload.add_number(0, 1) # Up
|
|
payload.add_number(0, 2) # Padding
|
|
payload.add_number(modkey.value, 4)
|
|
payload.flush(self.writer)
|
|
await self.writer.drain()
|
|
await asyncio.sleep(0.01) # Have to slow down keypresses for some servers
|
|
# Still shouldn't be noticable interactively, but does slow down paste to a fast typist...
|
|
self._input_queue.task_done()
|
|
|
|
async def send_keypresses(self, keys, modifierkeys=None):
|
|
await self._input_queue.put((keys, modifierkeys))
|
|
|
|
async def _read_number(self, num_bytes):
|
|
data = await self.reader.readexactly(num_bytes)
|
|
return int.from_bytes(data, byteorder='big', signed=True)
|
|
|
|
def _write_number(self, number, num_bytes):
|
|
data = number.to_bytes(num_bytes, byteorder='big', signed=True)
|
|
self.writer.write(data)
|
|
return data
|
|
|
|
async def get_screenshot(self):
|
|
while self._updating:
|
|
await asyncio.sleep(0.1)
|
|
await asyncio.sleep(0)
|
|
if self.framebuffer is None:
|
|
raise Exception('No framebuffer data available')
|
|
self._updating = True
|
|
return self.framebuffer.copy()
|
|
|
|
async def _do_vnc_handshake(self):
|
|
rfbver = await self.reader.readline()
|
|
if not rfbver.startswith(b'RFB 003.008'):
|
|
self.writer.close()
|
|
await self.writer.wait_closed()
|
|
raise Exception('Unsupported RFB version')
|
|
self.writer.write(b'RFB 003.008\n')
|
|
numsectypes = await self._read_number(1)
|
|
if not numsectypes:
|
|
self.writer.close()
|
|
await self.writer.wait_closed()
|
|
raise Exception('No security types supported by the server')
|
|
sectypes = await self.reader.readexactly(numsectypes)
|
|
sectypes = bytearray(sectypes)
|
|
secresult = 1
|
|
if 1 in sectypes:
|
|
self.writer.write(b'\x01')
|
|
await self.writer.drain()
|
|
secresult = await self._read_number(4) # Security result
|
|
if secresult != 0:
|
|
self.writer.close()
|
|
await self.writer.wait_closed()
|
|
raise Exception('VNC authentication failed')
|
|
self.writer.write(b'\x01') # Share display
|
|
self.width = await self._read_number(2)
|
|
self.height = await self._read_number(2)
|
|
pixformat = await self.reader.readexactly(16)
|
|
name_length = await self._read_number(4)
|
|
self.name = await self.reader.readexactly(name_length)
|
|
payload = ByteStream()
|
|
if pixformat != MYPIXFORMAT:
|
|
payload.add_number(0, 1) # Set pixel format
|
|
payload.add_number(0, 3) # Padding
|
|
payload.extend(MYPIXFORMAT)
|
|
payload.flush(self.writer)
|
|
self.receiver = asyncio.create_task(self._receive_loop())
|
|
payload.add_number(2, 1) # Set encodings
|
|
payload.add_number(0, 1) # Padding
|
|
payload.add_number(4, 2) # Number of encodings
|
|
payload.add_number(6, 4) # zlib
|
|
payload.add_number(7, 4) # tight
|
|
payload.add_number(-223, 4) # desktopsize
|
|
payload.add_number(-308, 4) # extended desktopsize
|
|
payload.flush(self.writer)
|
|
self._request_screen_update(incremental=False)
|
|
|
|
def _request_screen_update(self, incremental=True):
|
|
incremental = 1 if incremental else 0
|
|
payload = ByteStream()
|
|
payload.add_number(3, 1) # Framebuffer update request
|
|
payload.add_number(incremental, 1) # Incremental
|
|
payload.add_number(0, 2) # x position
|
|
payload.add_number(0, 2) # y position
|
|
payload.add_number(self.width, 2) # width
|
|
payload.add_number(self.height, 2) # height
|
|
payload.flush(self.writer)
|
|
|
|
async def _receive_loop(self):
|
|
while True:
|
|
try:
|
|
message_type = await self._read_number(1)
|
|
if message_type == 0: # Framebuffer update
|
|
await self._handle_framebuffer_update()
|
|
elif message_type == 1: # Set color map entries
|
|
raise NotImplementedError('Set color map entries not implemented')
|
|
elif message_type == 2: # Bell
|
|
pass
|
|
elif message_type == 3: # Server cut text
|
|
padding = await self._read_number(3)
|
|
length = await self._read_number(4)
|
|
self.copytext = await self.reader.readexactly(length)
|
|
else:
|
|
raise Exception(f'Unknown message type: {message_type}')
|
|
except Exception as e:
|
|
print(f"Error in receive loop: {e}")
|
|
break
|
|
|
|
async def _handle_framebuffer_update(self):
|
|
_ = await self._read_number(1) # Padding
|
|
num_rects = await self._read_number(2)
|
|
self._updating = True
|
|
for _ in range(num_rects):
|
|
await self._handle_rectangle()
|
|
self._updating = False
|
|
self._request_screen_update(incremental=True)
|
|
|
|
async def _handle_rectangle(self):
|
|
if self.framebuffer is None:
|
|
self.framebuffer = Image.new('RGBA', (self.width, self.height))
|
|
x = await self._read_number(2)
|
|
y = await self._read_number(2)
|
|
width = await self._read_number(2)
|
|
height = await self._read_number(2)
|
|
encoding_type = await self._read_number(4)
|
|
pixel_data = None
|
|
if encoding_type == 6:
|
|
compressed_data_length = await self._read_number(4)
|
|
compressed_data = await self.reader.readexactly(compressed_data_length)
|
|
# Decompress the data using zlib and store it in the framebuffer
|
|
pixel_data = self.decompressor.decompress(compressed_data)
|
|
elif encoding_type == 0:
|
|
pixel_data = await self.reader.readexactly(width * height * 4) # Assuming 32 bits per pixel
|
|
if encoding_type in (-223, -308): # desktopsize
|
|
self.width = width
|
|
self.height = height
|
|
self.framebuffer = Image.new('RGBA', (self.width, self.height))
|
|
if encoding_type == -308:
|
|
nscreens = await self._read_number(1)
|
|
_ = await self._read_number(3) # padding
|
|
for _ in range(nscreens):
|
|
_ = await self.reader.readexactly(16) # screen info
|
|
elif pixel_data:
|
|
pixel_data = np.frombuffer(pixel_data, dtype=np.uint8).reshape((height, width, 4)).copy()
|
|
pixel_data[:, :, 3] = 0xff
|
|
img = Image.fromarray(pixel_data, 'RGBA')
|
|
self.framebuffer.paste(img, (x, y))
|
|
elif encoding_type == 7: # tight
|
|
# Best document I could see was:
|
|
# https://github.com/TurboVNC/tightvnc/blob/main/vnc_winsrc/rfb/rfbproto.h
|
|
tightheader = await self._read_number(1)
|
|
streamid = tightheader & 0x0F
|
|
if streamid:
|
|
raise NotImplementedError('tight encoding with streamid not implemented')
|
|
comptype = (tightheader >> 4) & 0x0F
|
|
if comptype != 9:
|
|
raise NotImplementedError(f'tight encoding with comptype {comptype} not implemented')
|
|
compressed_data_length = await self._read_tight_length()
|
|
compressed_data = await self.reader.readexactly(compressed_data_length)
|
|
with io.BytesIO(compressed_data) as jpgimg:
|
|
img = Image.open(jpgimg)
|
|
img.load()
|
|
self.framebuffer.paste(img, (x, y))
|
|
else:
|
|
raise Exception(f'Unsupported encoding type: {encoding_type}')
|
|
|
|
async def _read_tight_length(self):
|
|
length = 0
|
|
for i in range(3):
|
|
byte = await self._read_number(1)
|
|
length |= ((byte & 0x7F) << (i * 7))
|
|
if not (byte & 0x80):
|
|
break
|
|
return length
|
|
async def close(self):
|
|
self.writer.close()
|
|
await self.writer.wait_closed()
|