From adad5ed4fe37038442d247d9246795a82d31093c Mon Sep 17 00:00:00 2001 From: retoor Date: Tue, 13 May 2025 19:08:18 +0200 Subject: [PATCH] Update. --- src/snek/research/serpentarium.py | 70 -------------------------- src/snek/system/websocket.py | 81 +++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 70 deletions(-) create mode 100644 src/snek/system/websocket.py diff --git a/src/snek/research/serpentarium.py b/src/snek/research/serpentarium.py index c87b70b..e425a23 100644 --- a/src/snek/research/serpentarium.py +++ b/src/snek/research/serpentarium.py @@ -33,76 +33,6 @@ class DatasetTable: def __getattr__(self, name): return DatasetMethod(self, name) -class WebSocketClient: - def __init__(self): - self.buffer = b'' - self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.connect() - - def connect(self): - self.socket.connect(("127.0.0.1", 3131)) - key = base64.b64encode(b'1234123412341234').decode('utf-8') - handshake = ( - f"GET /db HTTP/1.1\r\n" - f"Host: localhost:3131\r\n" - f"Upgrade: websocket\r\n" - f"Connection: Upgrade\r\n" - f"Sec-WebSocket-Key: {key}\r\n" - f"Sec-WebSocket-Version: 13\r\n\r\n" - ) - self.socket.sendall(handshake.encode('utf-8')) - response = self.read_until(b'\r\n\r\n') - if b'101 Switching Protocols' not in response: - raise Exception("Failed to connect to WebSocket") - - def write(self, message): - message_bytes = message.encode('utf-8') - length = len(message_bytes) - if length <= 125: - self.socket.sendall(b'\x81' + bytes([length]) + message_bytes) - elif length >= 126 and length <= 65535: - self.socket.sendall(b'\x81' + bytes([126]) + length.to_bytes(2, 'big') + message_bytes) - else: - self.socket.sendall(b'\x81' + bytes([127]) + length.to_bytes(8, 'big') + message_bytes) - - - def read_until(self, delimiter): - while True: - find_pos = self.buffer.find(delimiter) - if find_pos != -1: - data = self.buffer[:find_pos+4] - self.buffer = self.buffer[find_pos+4:] - return data - - chunk = self.socket.recv(1024) - if not chunk: - return None - self.buffer += chunk - - def read_exactly(self, length): - while len(self.buffer) < length: - chunk = self.socket.recv(length - len(self.buffer)) - if not chunk: - return None - self.buffer += chunk - response = self.buffer[: length] - self.buffer = self.buffer[length:] - return response - - def read(self): - frame = None - frame = self.read_exactly(2) - length = frame[1] & 127 - if length == 126: - length = int.from_bytes(self.read_exactly(2), 'big') - elif length == 127: - length = int.from_bytes(self.read_exactly(8), 'big') - message = self.read_exactly(length) - return message - - def close(self): - self.socket.close() - diff --git a/src/snek/system/websocket.py b/src/snek/system/websocket.py new file mode 100644 index 0000000..c54b094 --- /dev/null +++ b/src/snek/system/websocket.py @@ -0,0 +1,81 @@ +class WebSocketClient: + def __init__(self, hostname, port): + self.buffer = b'' + self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.hostname = hostname + self.port = port + self.connect() + + def __getattr__(self, method, *args, **kwargs): + if method in self.__dict__.keys(): + return self.__dict__[method] + def call(*args, **kwargs): + self.write(json.dumps({'method': method, 'args': args, 'kwargs': kwargs})) + return json.loads(self.read()) + return call + + def connect(self): + self.socket.connect((self.hostname, self.port)) + key = base64.b64encode(b'1234123412341234').decode('utf-8') + handshake = ( + f"GET /db HTTP/1.1\r\n" + f"Host: localhost:3131\r\n" + f"Upgrade: websocket\r\n" + f"Connection: Upgrade\r\n" + f"Sec-WebSocket-Key: {key}\r\n" + f"Sec-WebSocket-Version: 13\r\n\r\n" + ) + self.socket.sendall(handshake.encode('utf-8')) + response = self.read_until(b'\r\n\r\n') + if b'101 Switching Protocols' not in response: + raise Exception("Failed to connect to WebSocket") + + def write(self, message): + message_bytes = message.encode('utf-8') + length = len(message_bytes) + if length <= 125: + self.socket.sendall(b'\x81' + bytes([length]) + message_bytes) + elif length >= 126 and length <= 65535: + self.socket.sendall(b'\x81' + bytes([126]) + length.to_bytes(2, 'big') + message_bytes) + else: + self.socket.sendall(b'\x81' + bytes([127]) + length.to_bytes(8, 'big') + message_bytes) + + + def read_until(self, delimiter): + while True: + find_pos = self.buffer.find(delimiter) + if find_pos != -1: + data = self.buffer[:find_pos+4] + self.buffer = self.buffer[find_pos+4:] + return data + + chunk = self.socket.recv(1024) + if not chunk: + return None + self.buffer += chunk + + def read_exactly(self, length): + while len(self.buffer) < length: + chunk = self.socket.recv(length - len(self.buffer)) + if not chunk: + return None + self.buffer += chunk + response = self.buffer[: length] + self.buffer = self.buffer[length:] + return response + + def read(self): + frame = None + frame = self.read_exactly(2) + length = frame[1] & 127 + if length == 126: + length = int.from_bytes(self.read_exactly(2), 'big') + elif length == 127: + length = int.from_bytes(self.read_exactly(8), 'big') + message = self.read_exactly(length) + return message + + def close(self): + self.socket.close() + +