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
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)