128 lines
3.4 KiB
Python
Executable file
128 lines
3.4 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
import asyncio
|
|
import websockets
|
|
import uuid
|
|
import os
|
|
import pathlib
|
|
from urllib.parse import urlparse
|
|
import colorful
|
|
|
|
import conn
|
|
|
|
|
|
|
|
async def handle_socket_read(socketid, tcpreader, ws):
|
|
try:
|
|
print(f"New socket: {socketid}. Waiting on recv")
|
|
while True:
|
|
data = await tcpreader.read(2024)
|
|
print(f"TCP@{socketid} Received {len(data)} bytes")
|
|
if data == b'':
|
|
print(f"TCP@{socketid} Connection closed")
|
|
c = conn.WSMsg(socketid, conn.MsgType.DISCONNECT)
|
|
await ws.send(c.to_bytes())
|
|
break
|
|
|
|
c = conn.WSMsg(socketid, conn.MsgType.DATA, data)
|
|
print(f'TCP>WS: {c}')
|
|
await ws.send(c.to_bytes())
|
|
|
|
except Exception as e:
|
|
print(f"{socketid} Error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
async def handle_ws_incoming(cfg, ws, sockets):
|
|
data = await ws.recv()
|
|
|
|
c = conn.WSMsg.from_bytes(data)
|
|
socketid = c.socketid
|
|
|
|
if c.msg == conn.MsgType.CONNECT:
|
|
if socketid in sockets:
|
|
print(f"Socket {socketid} already connected")
|
|
return
|
|
else:
|
|
print(f"New socket: {socketid}")
|
|
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))
|
|
|
|
elif c.msg == conn.MsgType.DISCONNECT:
|
|
if socketid not in sockets:
|
|
print(f"Socket {socketid} not connected")
|
|
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>TCP: {c}')
|
|
tcpwriter.write(c.payload)
|
|
print('written')
|
|
|
|
|
|
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():
|
|
|
|
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(config, ws, sockets)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|
|
|
|
|
|
|