Gan, Jimmy
74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
import asyncio
|
|
import os
|
|
import pty
|
|
import signal
|
|
import struct
|
|
import fcntl
|
|
import termios
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
|
|
async def ws_endpoint(ws: WebSocket):
|
|
# Block non-Tailscale IPs
|
|
client_ip = ws.client.host if ws.client else ""
|
|
if not client_ip.startswith("100."):
|
|
await ws.close(code=1008, reason="Tailscale only")
|
|
return
|
|
|
|
await ws.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)
|
|
loop = asyncio.get_event_loop()
|
|
|
|
async def read_pty():
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(0.01)
|
|
if select_readable(master_fd):
|
|
data = os.read(master_fd, 4096)
|
|
if not data:
|
|
break
|
|
await ws.send_bytes(data)
|
|
except (OSError, WebSocketDisconnect):
|
|
pass
|
|
|
|
def select_readable(fd):
|
|
import select
|
|
r, _, _ = select.select([fd], [], [], 0)
|
|
return bool(r)
|
|
|
|
read_task = asyncio.create_task(read_pty())
|
|
|
|
try:
|
|
while True:
|
|
msg = await ws.receive()
|
|
if "bytes" in msg and msg["bytes"]:
|
|
data = msg["bytes"]
|
|
# Handle resize: first byte 0x01 means resize message
|
|
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)
|