Gan, Jimmy
ee2e5dcb3c
- Rewrite terminal.py to use asyncssh instead of pty/fork - Add SSH key pair for dashboard container auth - Mount SSH key and add SSH config vars in docker-compose - Install openssh-client in Dockerfile - Add asyncssh to requirements.txt
64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import asyncio
|
|
import struct
|
|
from fastapi import WebSocket, WebSocketDisconnect, Query
|
|
import asyncssh
|
|
import jwt
|
|
import config
|
|
|
|
|
|
async def ws_endpoint(websocket: WebSocket, token: str = Query("")):
|
|
try:
|
|
payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
|
|
if payload.get("sub") != config.ADMIN_USER:
|
|
await websocket.close(code=1008, reason="Unauthorized")
|
|
return
|
|
except Exception:
|
|
await websocket.close(code=1008, reason="Unauthorized")
|
|
return
|
|
|
|
await websocket.accept()
|
|
|
|
try:
|
|
conn = await asyncssh.connect(
|
|
config.SSH_HOST, username=config.SSH_USER,
|
|
client_keys=[config.SSH_KEY_PATH], known_hosts=None,
|
|
)
|
|
except Exception:
|
|
await websocket.close(code=1011, reason="SSH connection failed")
|
|
return
|
|
|
|
try:
|
|
chan, session = await conn.create_session(
|
|
asyncssh.SSHClientProcess, term_type="xterm-256color",
|
|
term_size=(80, 24),
|
|
)
|
|
|
|
async def read_ssh():
|
|
try:
|
|
async for data in chan.stdout:
|
|
await websocket.send_bytes(data.encode() if isinstance(data, str) else data)
|
|
except (asyncssh.Error, WebSocketDisconnect):
|
|
pass
|
|
|
|
read_task = asyncio.create_task(read_ssh())
|
|
|
|
try:
|
|
while True:
|
|
msg = await websocket.receive()
|
|
if "bytes" in msg and msg["bytes"]:
|
|
data = msg["bytes"]
|
|
if data[0:1] == b"\x01" and len(data) == 5:
|
|
cols = struct.unpack("!H", data[1:3])[0]
|
|
rows = struct.unpack("!H", data[3:5])[0]
|
|
chan.change_terminal_size(cols, rows)
|
|
else:
|
|
chan.write(data.decode("utf-8", errors="replace"))
|
|
elif "text" in msg and msg["text"]:
|
|
chan.write(msg["text"])
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
read_task.cancel()
|
|
finally:
|
|
conn.close()
|