101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
import asyncio
|
|
import os
|
|
import pty
|
|
import fcntl
|
|
import tty, termios
|
|
import uvicorn
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
import json, struct
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
class PTYHandler:
|
|
def __init__(self):
|
|
self.master_fd, self.slave_fd = pty.openpty()
|
|
self.shell_pid = os.fork()
|
|
if self.shell_pid == 0: # child
|
|
os.close(self.master_fd)
|
|
os.setsid()
|
|
# tty.setcbreak(self.slave_fd)
|
|
tty.setraw(self.slave_fd, when=tty.TCSANOW)
|
|
os.dup2(self.slave_fd, 0)
|
|
os.dup2(self.slave_fd, 1)
|
|
os.dup2(self.slave_fd, 2)
|
|
os.close(self.slave_fd)
|
|
# cmd = ["/bin/bash"]
|
|
cmd = ["./shell"]
|
|
os.execvp(cmd[0], cmd)
|
|
else: # parent
|
|
os.close(self.slave_fd)
|
|
tty.setraw(self.master_fd, when=tty.TCSANOW)
|
|
# tty.setcbreak(self.master_fd)
|
|
flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
|
|
fcntl.fcntl(self.master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
|
|
|
|
async def proxy_websocket(self, websocket: WebSocket):
|
|
read_task = None
|
|
write_task = None
|
|
try:
|
|
await websocket.accept()
|
|
read_task = asyncio.create_task(self.read_from_master(websocket))
|
|
write_task = asyncio.create_task(self.write_to_master(websocket))
|
|
await asyncio.gather(read_task, write_task)
|
|
except asyncio.CancelledError:
|
|
pass # Task cancellation should not be treated as an error
|
|
finally:
|
|
if read_task: read_task.cancel()
|
|
if write_task: write_task.cancel()
|
|
|
|
async def read_from_master(self, websocket: WebSocket):
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(0.01) # Wait 10ms to prevent busy loop while checking for data
|
|
try:
|
|
data = os.read(self.master_fd, 1024)
|
|
if data:
|
|
await websocket.send_bytes(data)
|
|
except BlockingIOError:
|
|
# There is no data available to read; move on
|
|
pass
|
|
except WebSocketDisconnect:
|
|
pass
|
|
|
|
async def write_to_master(self, websocket: WebSocket):
|
|
try:
|
|
while True:
|
|
message = await websocket.receive()
|
|
print(message)
|
|
if "bytes" in message:
|
|
os.write(self.master_fd, message["bytes"])
|
|
elif "text" in message:
|
|
if message["text"].startswith('{'):
|
|
# Handle the JSON message that sets the terminal size
|
|
resize_message = json.loads(message["text"])
|
|
if resize_message["action"] == "resize":
|
|
cols = resize_message["cols"]
|
|
rows = resize_message["rows"]
|
|
self.resize_pty(cols, rows)
|
|
else:
|
|
# This part is for regular text input
|
|
os.write(self.master_fd, message["text"].encode('utf-8'))
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
|
|
def resize_pty(self, cols, rows):
|
|
# Set the terminal size of the PTY
|
|
winsize = struct.pack("HHHH", rows, cols, 0, 0)
|
|
fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, winsize)
|
|
|
|
pty_handler = PTYHandler()
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
await pty_handler.proxy_websocket(websocket)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="127.0.0.1", port=38000)
|