Gan, Jimmy
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import asyncio
|
|
import os
|
|
import pty
|
|
import signal
|
|
import struct
|
|
import fcntl
|
|
import termios
|
|
import select
|
|
from fastapi import WebSocket, WebSocketDisconnect, Query
|
|
import jwt
|
|
import config
|
|
|
|
|
|
async def ws_endpoint(websocket: WebSocket, token: str = Query("")):
|
|
# Verify auth token
|
|
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()
|
|
master_fd, slave_fd = pty.openpty()
|
|
pid = os.fork()
|
|
|
|
if pid == 0:
|
|
os.close(master_fd)
|
|
os.setsid()
|
|
os.dup2(slave_fd, 0)
|
|
os.dup2(slave_fd, 1)
|
|
os.dup2(slave_fd, 2)
|
|
os.close(slave_fd)
|
|
os.execvp("/bin/bash", ["/bin/bash"])
|
|
|
|
os.close(slave_fd)
|
|
|
|
async def read_pty():
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(0.01)
|
|
r, _, _ = select.select([master_fd], [], [], 0)
|
|
if r:
|
|
data = os.read(master_fd, 4096)
|
|
if not data:
|
|
break
|
|
await websocket.send_bytes(data)
|
|
except (OSError, WebSocketDisconnect):
|
|
pass
|
|
|
|
read_task = asyncio.create_task(read_pty())
|
|
|
|
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]
|
|
fcntl.ioctl(master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
|
|
else:
|
|
os.write(master_fd, data)
|
|
elif "text" in msg and msg["text"]:
|
|
os.write(master_fd, msg["text"].encode())
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
read_task.cancel()
|
|
os.close(master_fd)
|
|
os.kill(pid, signal.SIGTERM)
|
|
os.waitpid(pid, 0)
|