From d940d2897ced0c2cc8dd6247338ad05ba724ba8a Mon Sep 17 00:00:00 2001 From: Frank Villaro-Dixon Date: Wed, 28 Aug 2024 00:00:08 +0200 Subject: [PATCH] add pod cfg Signed-off-by: Frank Villaro-Dixon --- Pipfile | 2 ++ Pipfile.lock | 43 ++++++++++++++++++++++++++++++- client.py | 40 ++++++++++++++--------------- conn.py | 8 ++++++ pod.py | 72 ++++++++++++++++++++++++++++++++++++++++------------ 5 files changed, 128 insertions(+), 37 deletions(-) diff --git a/Pipfile b/Pipfile index ca5186a..c27ddc4 100644 --- a/Pipfile +++ b/Pipfile @@ -7,6 +7,8 @@ name = "pypi" websockets = "*" asyncio = "*" bencoder = "*" +rich = "*" +colorful = "*" [dev-packages] diff --git a/Pipfile.lock b/Pipfile.lock index 1b25a6e..761bd82 100644 --- a/Pipfile.lock +++ b/Pipfile.lock @@ -1,7 +1,7 @@ { "_meta": { "hash": { - "sha256": "32897fbdc87627e0a506eca3bb31e5b94fecf1d0b3ee5a3426c5def16fea1acc" + "sha256": "627eb531894ee5cbb395c8cd5e2802b75748a4e9d20257afade2167137bf9b1d" }, "pipfile-spec": 6, "requires": { @@ -33,6 +33,47 @@ "index": "pypi", "version": "==0.2.0" }, + "colorful": { + "hashes": [ + "sha256:b56d5c01db1dac4898308ea889edcb113fbee3e6ec5df4bacffd61d5241b5b8d", + "sha256:eab8c1c809f5025ad2b5238a50bd691e26850da8cac8f90d660ede6ea1af9f1e" + ], + "index": "pypi", + "version": "==0.5.6" + }, + "markdown-it-py": { + "hashes": [ + "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", + "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb" + ], + "markers": "python_version >= '3.8'", + "version": "==3.0.0" + }, + "mdurl": { + "hashes": [ + "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", + "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba" + ], + "markers": "python_version >= '3.7'", + "version": "==0.1.2" + }, + "pygments": { + "hashes": [ + "sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199", + "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a" + ], + "markers": "python_version >= '3.8'", + "version": "==2.18.0" + }, + "rich": { + "hashes": [ + "sha256:2e85306a063b9492dffc86278197a60cbece75bcb766022f3436f567cae11bdc", + "sha256:a5ac1f1cd448ade0d59cc3356f7db7a7ccda2c8cbae9c7a90c28ff463d3e91f4" + ], + "index": "pypi", + "markers": "python_full_version >= '3.7.0'", + "version": "==13.8.0" + }, "websockets": { "hashes": [ "sha256:02cc9bb1a887dac0e08bf657c5d00aa3fac0d03215d35a599130c2034ae6663a", diff --git a/client.py b/client.py index d47455a..3426328 100644 --- a/client.py +++ b/client.py @@ -1,6 +1,7 @@ import asyncio import websockets import hashlib +import random import conn @@ -10,18 +11,14 @@ LOCALHOST = '127.0.0.1' PORT = 8443 WEBSOCKET_URL = 'ws://localhost:9999/data' -ALREADY_DONE = False -socket_id = -1 async def handle_client(socket_reader, socket_writer): - global socket_id - socket_id += 1 + socket_id = random.randint(0, 9000) - try: - async with websockets.connect(WEBSOCKET_URL) as websocket: - - print(f"New client connected socket {socket_id}: {socket_reader} {socket_writer}") + async with websockets.connect(WEBSOCKET_URL) as websocket: + try: + print(f"New client connected socket {socket_id}") m = conn.WSMsg(socket_id, conn.MsgType.CONNECT) await websocket.send(m.to_bytes()) @@ -29,10 +26,14 @@ async def handle_client(socket_reader, socket_writer): async def client_to_websocket(): while True: data = await socket_reader.read(2024) - print(f'TCP{socket_id}>WS: ', hashlib.md5(data).hexdigest()) if not data: + c = conn.WSMsg(socket_id, conn.MsgType.DISCONNECT) + print(f"Client {socket_id} disconnected") + await websocket.send(c.to_bytes()) break + c = conn.WSMsg(socket_id, conn.MsgType.DATA, data) + print(f'TCP>WS: {c}') await websocket.send(c.to_bytes()) # Forwarding data from WebSocket to client @@ -46,8 +47,11 @@ async def handle_client(socket_reader, socket_writer): if c.msg == conn.MsgType.DISCONNECT: print(f"Client {socket_id} disconnected") break + elif c.msg == conn.MsgType.CONNECT: + ## This shouldn't really happen‽ + break else: - print(f'WS>TCP@{socket_id}: ', hashlib.md5(message).hexdigest()) + print(f'WS>TCP: {c}') socket_writer.write(c.payload) await socket_writer.drain() else: @@ -58,17 +62,13 @@ async def handle_client(socket_reader, socket_writer): await asyncio.gather(client_to_websocket(), websocket_to_client()) print(f">>>> Client {socket_id} disconnected") - except Exception as e: - print(f"Error: {e}") - import traceback - traceback.print_exc() - finally: - print(f">>>>>>>>>> Closing client {socket_id}") - m = conn.WSMsg(socket_id, conn.MsgType.DISCONNECT) - await websocket.send(m.to_bytes()) + finally: + print(f">>>>>>>>>> Closing client {socket_id}") + m = conn.WSMsg(socket_id, conn.MsgType.DISCONNECT) + await websocket.send(m.to_bytes()) - socket_writer.close() - await socket_writer.wait_closed() + socket_writer.close() + await socket_writer.wait_closed() async def main(): port = PORT diff --git a/conn.py b/conn.py index c6f32b7..09ed213 100644 --- a/conn.py +++ b/conn.py @@ -1,6 +1,7 @@ from typing import Optional import bencoder from enum import Enum +import hashlib @@ -39,3 +40,10 @@ class WSMsg: if self.payload: d[b"data"] = self.payload return bencoder.encode(d) + + def __repr__(self): + ret = f"MSG ({self.socketid}, {self.msg}" + if self.msg == MsgType.DATA: + ret += f", {hashlib.md5(self.payload).hexdigest()}" + ret += ")" + return ret diff --git a/pod.py b/pod.py index a782abe..dbc6295 100644 --- a/pod.py +++ b/pod.py @@ -1,13 +1,13 @@ import asyncio -import socket import websockets -import hashlib +import uuid +import os +import pathlib +from urllib.parse import urlparse +import colorful import conn -HOST = '192.168.21.30' -PORT = 6443 -WEBSOCKET_URL = 'ws://localhost:9999/data' async def handle_socket_read(socketid, tcpreader, ws): @@ -22,7 +22,7 @@ async def handle_socket_read(socketid, tcpreader, ws): break c = conn.WSMsg(socketid, conn.MsgType.DATA, data) - print(f'TCP@{socketid}>WS: ', hashlib.md5(data).hexdigest()) + print(f'TCP>WS: {c}') await ws.send(c.to_bytes()) except Exception as e: @@ -31,11 +31,10 @@ async def handle_socket_read(socketid, tcpreader, ws): traceback.print_exc() -async def handle_ws_incoming(ws, sockets): +async def handle_ws_incoming(cfg, ws, sockets): data = await ws.recv() c = conn.WSMsg.from_bytes(data) - print(f'NEW DATA: {c.socketid} {c.msg} {c.payload}') socketid = c.socketid if c.msg == conn.MsgType.CONNECT: @@ -44,7 +43,7 @@ async def handle_ws_incoming(ws, sockets): return else: print(f"New socket: {socketid}") - tcpreader, tcpwriter = await asyncio.open_connection(HOST, PORT) + tcpreader, tcpwriter = await asyncio.open_connection(cfg['kube_api_host'], cfg['kube_api_port']) sockets[socketid] = (tcpreader, tcpwriter) asyncio.create_task(handle_socket_read(socketid, tcpreader, ws)) @@ -54,29 +53,70 @@ async def handle_ws_incoming(ws, sockets): return else: print(f"Socket {socketid} disconnected") + tcpreader, tcpwriter = sockets[socketid] del sockets[socketid] tcpwriter.close() await tcpwriter.wait_closed() elif c.msg == conn.MsgType.DATA: tcpreader, tcpwriter = sockets[socketid] - print(f'WS@{socketid}>TCP: ', hashlib.md5(data).hexdigest()) + print(f'WS>TCP: {c}') tcpwriter.write(c.payload) +def get_config(): + KUBERNETES_PORT = os.environ.get('KUBERNETES_PORT', 'http://localhost:6443') + kube_api = urlparse(KUBERNETES_PORT) + + WEBSOCKET_ROOT_URL = os.environ.get('WEBSOCKET_ROOT_URL', 'ws://localhost:9999') + + + WS_ID = os.environ.get('WS_ID', None) + + ws_id = WS_ID + if ws_id is None: + HOSTNAME = os.environ.get('HOSTNAME') + # We want to name it as the replicaset, so it's somehow random, but doesn't + # change on every pod restart + if HOSTNAME.count('-') < 2: + ws_id = uuid.uuid4().hex + else: + ws_id = ''.join(HOSTNAME.split('-')[:2]) + + + + websocket_full_url = str(pathlib.Path(WEBSOCKET_ROOT_URL) / ws_id) + + # Pathlib replaces `ws://` into `ws:/`. This "fixes" it. lul + websocket_full_url = websocket_full_url.replace('s:/', 's://') + + return { + 'kube_api_host': kube_api.hostname, + 'kube_api_port': kube_api.port, + 'websocket_url': websocket_full_url, + } + async def main(): - tcpreader, tcpwriter = await asyncio.open_connection(HOST, PORT) - ws = await websockets.connect(WEBSOCKET_URL) + config = get_config() + print(\ +f""" === KUBE-ESCAPE === +A Kube API proxy over WebSockets + +Websocket URL: {colorful.bold_coral(config['websocket_url'])} +Kube API Host: {config['kube_api_host']} +Kube API Port: {config['kube_api_port']} + +Enjoy and fuck Citrix and VDIs!""") + + + ws = await websockets.connect(config['websocket_url']) sockets = {} while True: - await handle_ws_incoming(ws, sockets) - - - #await handle_client(tcpreader, tcpwriter, ws) + await handle_ws_incoming(config, ws, sockets) if __name__ == "__main__": asyncio.run(main())