add pod cfg

Signed-off-by: Frank Villaro-Dixon <frank@villaro-dixon.eu>
This commit is contained in:
Frank Villaro-Dixon 2024-08-28 00:00:08 +02:00
parent 80184c4b55
commit d940d2897c
5 changed files with 128 additions and 37 deletions

View file

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

43
Pipfile.lock generated
View file

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

View file

@ -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

View file

@ -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

72
pod.py
View file

@ -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())