add pod cfg

Signed-off-by: Frank Villaro-Dixon <frank@villaro-dixon.eu>

Pipfile

@@ -7,6 +7,8 @@ name = "pypi"
 websockets = "*"
 asyncio = "*"
 bencoder = "*"
+rich = "*"
+colorful = "*"
 
 [dev-packages]
 

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",

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 = random.randint(0, 9000)
-    socket_id += 1
 
 
-    try:
     async with websockets.connect(WEBSOCKET_URL) as websocket:
-
+        try:
-            print(f"New client connected socket {socket_id}: {socket_reader} {socket_writer}")
+            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,10 +62,6 @@ 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)

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

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_ws_incoming(config, ws, sockets)
-    
-
-    #await handle_client(tcpreader, tcpwriter, ws)
 
 if __name__ == "__main__":
     asyncio.run(main())