ubuntu-bare-metal/stacks/server-ui/cli_pty.py

"""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
import uuid
from dataclasses import dataclass, field
from typing import Any

from starlette.websockets import WebSocket, WebSocketDisconnect

log = logging.getLogger(__name__)

TIOCSWINSZ = termios.TIOCSWINSZ
OUTPUT_BUFFER_MAX = 256 * 1024


class CliSessionLimitError(Exception):
    pass


class CliSessionNotFoundError(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 _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()


@dataclass
class CliSession:
    id: str
    master_fd: int
    process: asyncio.subprocess.Process
    cols: int = 80
    rows: int = 24
    ws: WebSocket | None = None
    output_buffer: bytearray = field(default_factory=bytearray)
    running: asyncio.Event = field(default_factory=asyncio.Event)
    read_task: asyncio.Task[None] | None = None
    detach_timer: asyncio.Task[None] | None = None
    closed: bool = False

    def append_output(self, data: bytes) -> None:
        self.output_buffer.extend(data)
        overflow = len(self.output_buffer) - OUTPUT_BUFFER_MAX
        if overflow > 0:
            del self.output_buffer[:overflow]

    async def send_to_ws(self, data: bytes) -> bool:
        if not self.ws:
            return False
        try:
            await self.ws.send_bytes(data)
            return True
        except Exception:
            return False

    async def send_json(self, obj: dict[str, Any]) -> bool:
        if not self.ws:
            return False
        try:
            await self.ws.send_text(json.dumps(obj))
            return True
        except Exception:
            return False


class CliSessionManager:
    def __init__(self, *, session_ttl: int = 300, max_sessions: int = 5) -> None:
        self.session_ttl = session_ttl
        self.max_sessions = max_sessions
        self._sessions: dict[str, CliSession] = {}
        self._lock = asyncio.Lock()

    def sessions_available(self) -> bool:
        return len(self._sessions) < self.max_sessions

    def get(self, session_id: str) -> CliSession | None:
        return self._sessions.get(session_id)

    async def create_session(
        self,
        *,
        shell: str,
        cwd: str,
        env: dict[str, str],
    ) -> CliSession:
        async with self._lock:
            if len(self._sessions) >= self.max_sessions:
                raise CliSessionLimitError(
                    f"Limit sesji CLI ({self.max_sessions}) — spróbuj później"
                )

        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)

        session = CliSession(
            id=str(uuid.uuid4()),
            master_fd=master_fd,
            process=process,
        )
        session.running.set()

        async with self._lock:
            self._sessions[session.id] = session

        session.read_task = asyncio.create_task(self._read_pty_loop(session))
        log.info("CLI session created: %s", session.id)
        return session

    async def _read_pty_loop(self, session: CliSession) -> None:
        loop = asyncio.get_running_loop()
        while session.running.is_set() and not session.closed:
            try:
                data = await loop.run_in_executor(None, os.read, session.master_fd, 4096)
            except OSError:
                break
            if not data:
                break
            session.append_output(data)
            await session.send_to_ws(data)
            if session.process.returncode is not None:
                break

        if not session.closed:
            log.info("CLI session ended (shell exit): %s", session.id)
            await self.close_session(session.id)

    async def attach_session(self, session_id: str, websocket: WebSocket) -> CliSession:
        session = self.get(session_id)
        if session is None or session.closed:
            raise CliSessionNotFoundError(session_id)

        await self._cancel_detach_timer(session)

        if session.ws is not None and session.ws is not websocket:
            try:
                await session.ws.close()
            except Exception:
                pass

        session.ws = websocket

        if session.output_buffer:
            await session.send_to_ws(bytes(session.output_buffer))

        await session.send_json({"type": "replay", "complete": True})
        _set_winsize(session.master_fd, session.rows, session.cols)
        log.info("CLI session reattached: %s", session.id)
        return session

    async def detach_session(self, session_id: str) -> None:
        session = self.get(session_id)
        if session is None or session.closed:
            return
        session.ws = None
        await self._schedule_detach_timer(session)
        log.info("CLI session detached (TTL %ds): %s", self.session_ttl, session.id)

    async def close_session(self, session_id: str) -> None:
        async with self._lock:
            session = self._sessions.pop(session_id, None)
        if session is None or session.closed:
            return

        session.closed = True
        session.running.clear()
        await self._cancel_detach_timer(session)

        if session.read_task is not None:
            session.read_task.cancel()
            try:
                await session.read_task
            except asyncio.CancelledError:
                pass

        await _kill_process(session.process)
        try:
            os.close(session.master_fd)
        except OSError:
            pass

        if session.ws is not None:
            try:
                await session.ws.close()
            except Exception:
                pass

        log.info("CLI session closed: %s", session_id)

    async def _cancel_detach_timer(self, session: CliSession) -> None:
        if session.detach_timer is not None:
            session.detach_timer.cancel()
            try:
                await session.detach_timer
            except asyncio.CancelledError:
                pass
            session.detach_timer = None

    async def _schedule_detach_timer(self, session: CliSession) -> None:
        await self._cancel_detach_timer(session)

        async def _expire() -> None:
            try:
                await asyncio.sleep(self.session_ttl)
            except asyncio.CancelledError:
                return
            log.info("CLI session expired after detach TTL: %s", session.id)
            await self.close_session(session.id)

        session.detach_timer = asyncio.create_task(_expire())

    async def handle_resize(self, session: CliSession, rows: int, cols: int) -> None:
        session.rows = rows
        session.cols = cols
        _set_winsize(session.master_fd, rows, cols)

    async def write_input(self, session: CliSession, payload: bytes) -> None:
        try:
            os.write(session.master_fd, payload)
        except OSError:
            await self.close_session(session.id)


_manager: CliSessionManager | None = None


def init_session_manager(*, session_ttl: int = 300, max_sessions: int = 5) -> CliSessionManager:
    global _manager
    _manager = CliSessionManager(session_ttl=session_ttl, max_sessions=max_sessions)
    return _manager


def get_session_manager() -> CliSessionManager:
    if _manager is None:
        raise RuntimeError("CLI session manager not initialized")
    return _manager


def sessions_available(max_sessions: int) -> bool:
    if _manager is not None:
        return _manager.sessions_available()
    return True


async def _handle_ws_input(session: CliSession, message: dict[str, Any], mgr: CliSessionManager) -> bool:
    """Process one WS message. Returns False to end the connection loop."""
    msg_type = message.get("type")
    if msg_type == "websocket.disconnect":
        return False

    if message.get("bytes"):
        await mgr.write_input(session, message["bytes"])
        return True

    if message.get("text"):
        text = message["text"]
        if text.startswith("{"):
            try:
                obj: dict[str, Any] = json.loads(text)
                cmd = obj.get("type")
                if cmd == "resize":
                    rows = int(obj.get("rows", 24))
                    cols = int(obj.get("cols", 80))
                    await mgr.handle_resize(session, rows, cols)
                    return True
                if cmd == "close":
                    await mgr.close_session(session.id)
                    return False
            except (json.JSONDecodeError, TypeError, ValueError):
                pass
        await mgr.write_input(session, text.encode("utf-8"))

    if session.closed or session.process.returncode is not None:
        return False
    return True


async def run_pty_session(
    websocket: WebSocket,
    *,
    shell: str,
    cwd: str,
    env: dict[str, str],
    max_sessions: int,
    session_id: str | None = None,
    session_ttl: int = 300,
) -> None:
    if _manager is None:
        init_session_manager(session_ttl=session_ttl, max_sessions=max_sessions)
    mgr = get_session_manager()

    session: CliSession

    if session_id:
        try:
            session = await mgr.attach_session(session_id, websocket)
        except CliSessionNotFoundError:
            await websocket.close(code=1008, reason="Session expired")
            return
    else:
        session = await mgr.create_session(shell=shell, cwd=cwd, env=env)
        session.ws = websocket
        await websocket.send_text(json.dumps({"type": "session", "id": session.id}))
        _set_winsize(session.master_fd, session.rows, session.cols)

    try:
        while not session.closed:
            message = await websocket.receive()
            if not await _handle_ws_input(session, message, mgr):
                break
    except WebSocketDisconnect:
        pass
    finally:
        if not session.closed:
            if session.ws is websocket:
                session.ws = None
            await mgr.detach_session(session.id)