"""PTY shell sessions for browser CLI (WebSocket bridge).""" from __future__ import annotations import asyncio import fcntl import json import logging import os import pty import shlex import signal import struct import termios from typing import Any from starlette.websockets import WebSocket, WebSocketDisconnect log = logging.getLogger(__name__) TIOCSWINSZ = termios.TIOCSWINSZ _active_sessions = 0 _session_lock = asyncio.Lock() class CliSessionLimitError(Exception): pass def _set_winsize(fd: int, rows: int, cols: int) -> None: winsize = struct.pack("HHHH", rows, cols, 0, 0) fcntl.ioctl(fd, TIOCSWINSZ, winsize) def _parse_shell(shell: str) -> list[str]: parts = shlex.split(shell.strip()) return parts if parts else ["/bin/bash"] async def _read_pty_loop( master_fd: int, websocket: WebSocket, running: asyncio.Event, ) -> None: loop = asyncio.get_running_loop() while running.is_set(): try: data = await loop.run_in_executor(None, os.read, master_fd, 4096) except OSError: break if not data: break try: await websocket.send_bytes(data) except Exception: break async def _kill_process(process: asyncio.subprocess.Process) -> None: if process.returncode is not None: return try: os.killpg(process.pid, signal.SIGTERM) except ProcessLookupError: return try: await asyncio.wait_for(process.wait(), timeout=2.0) except asyncio.TimeoutError: try: os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: pass await process.wait() def sessions_available(max_sessions: int) -> bool: return _active_sessions < max_sessions async def run_pty_session( websocket: WebSocket, *, shell: str, cwd: str, env: dict[str, str], max_sessions: int, ) -> None: global _active_sessions async with _session_lock: if _active_sessions >= max_sessions: raise CliSessionLimitError( f"Limit sesji CLI ({max_sessions}) — spróbuj później" ) _active_sessions += 1 master_fd: int | None = None process: asyncio.subprocess.Process | None = None read_task: asyncio.Task[None] | None = None running = asyncio.Event() running.set() try: master_fd, slave_fd = pty.openpty() shell_cmd = _parse_shell(shell) proc_env = {**env, "TERM": "xterm-256color"} process = await asyncio.create_subprocess_exec( *shell_cmd, stdin=slave_fd, stdout=slave_fd, stderr=slave_fd, cwd=cwd, env=proc_env, preexec_fn=os.setsid, ) os.close(slave_fd) slave_fd = -1 read_task = asyncio.create_task(_read_pty_loop(master_fd, websocket, running)) while True: message = await websocket.receive() msg_type = message.get("type") if msg_type == "websocket.disconnect": break payload: bytes | None = None if message.get("bytes"): payload = message["bytes"] elif message.get("text"): text = message["text"] if text.startswith("{"): try: obj: dict[str, Any] = json.loads(text) if obj.get("type") == "resize": rows = int(obj.get("rows", 24)) cols = int(obj.get("cols", 80)) if master_fd is not None: _set_winsize(master_fd, rows, cols) continue except (json.JSONDecodeError, TypeError, ValueError): pass payload = text.encode("utf-8") if payload and master_fd is not None: try: os.write(master_fd, payload) except OSError: break if process.returncode is not None: break except WebSocketDisconnect: pass finally: running.clear() if read_task is not None: read_task.cancel() try: await read_task except asyncio.CancelledError: pass if process is not None: await _kill_process(process) if master_fd is not None: try: os.close(master_fd) except OSError: pass async with _session_lock: _active_sessions = max(0, _active_sessions - 1)