LLeMbas/tests/test_terminal.py

"""The shell behind the terminal panel.

Everything here runs against a real SSH server with a real PTY, because the
whole point of this module is what a pseudo-terminal does and a stub would agree
with any design at all.
"""

from __future__ import annotations

import asyncio

import pytest

from lembas.services.agent import terminal as terminal_service
from lembas.services.agent.base import ExecError
from lembas.services.agent.terminal import Session

# Stands up something real -- see the `slow` marker in pyproject.toml.
pytestmark = pytest.mark.slow

asyncssh = pytest.importorskip("asyncssh")


# --- A machine with a shell on it ----------------------------------------------
class _Server(asyncssh.SSHServer):
    def begin_auth(self, username: str) -> bool:
        return False


@pytest.fixture
async def shell_host():
    """A server whose sessions behave enough like a shell to be tested against.

    It announces itself, echoes what is typed at it, and remembers what the PTY
    was asked for -- which is the part the client side has to get right.
    """
    seen: dict = {}

    async def handler(process):
        seen["command"] = process.command
        seen["term_type"] = process.get_terminal_type()
        seen["term_size"] = process.get_terminal_size()
        seen["process"] = process
        process.stdout.write("READY\n")
        while True:
            try:
                line = await process.stdin.readline()
            except asyncssh.TerminalSizeChanged:
                # A window change interrupts the read rather than arriving as
                # data. A real shell redraws and carries on.
                continue
            except Exception:  # noqa: BLE001 - the session ending is not a failure
                break
            if not line or line.rstrip("\n") == "exit":
                break
            process.stdout.write(f"echo:{line.rstrip()}\n")
        process.exit(0)

    server = await asyncssh.create_server(
        _Server,
        "127.0.0.1",
        0,
        server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")],
        process_factory=handler,
    )
    port = next(iter(server.sockets)).getsockname()[1]
    from lembas.services.agent import ssh as ssh_service

    line, _fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
    try:
        yield {"port": port, "host_key": line, "seen": seen}
    finally:
        # Before the server, always: a client connection still open makes
        # `wait_closed` wait for it, and a test that failed mid-way has one.
        await terminal_service.shutdown()
        server.close()
        await server.wait_closed()


def _spec(host) -> dict:
    return {
        "host": "127.0.0.1",
        "port": host["port"],
        "username": "tester",
        "auth": "password",
        "password": "",
        "private_key": "",
        "key_passphrase": "",
        "host_key": host["host_key"],
        "connect_timeout": 10,
    }


async def _open(host, *, chat_id="chat-1", owner="user-1", project_dir="", **kwargs) -> Session:
    return await terminal_service.open_session(
        chat_id,
        owner_id=owner,
        profile_id="profile-1",
        label="Test box",
        spec=_spec(host),
        project_dir=project_dir,
        **kwargs,
    )


async def _read(viewer, *, timeout: float = 5.0) -> bytes:
    return await asyncio.wait_for(viewer.queue.get(), timeout=timeout)


async def _read_until(viewer, needle: bytes, *, timeout: float = 5.0) -> bytes:
    """Frames are whatever the far side wrote, so a line can arrive in pieces.

    Starts from the snapshot, because a viewer that attached after the shell
    had already said something finds it there rather than on the queue -- which
    is the whole design, and would otherwise make this hang.
    """
    seen = viewer.snapshot
    while needle not in seen:
        chunk = await _read(viewer, timeout=timeout)
        assert chunk is not None, f"the session closed before {needle!r} arrived"
        seen += chunk
    return seen


# --- The PTY --------------------------------------------------------------------
async def test_the_pty_is_asked_for_with_a_term_type_and_a_size(shell_host):
    session = await _open(shell_host, cols=100, rows=30)
    viewer = session.attach(cols=100, rows=30)
    await _read_until(viewer, b"READY")

    assert shell_host["seen"]["term_type"] == "xterm-256color"
    assert shell_host["seen"]["term_size"][:2] == (100, 30)
    await session.close()


async def test_a_project_directory_becomes_a_cd_before_the_shell(shell_host):
    session = await _open(shell_host, project_dir="/srv/work")
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    command = shell_host["seen"]["command"]
    assert command is not None
    assert "cd '/srv/work'" in command
    assert "exec" in command
    await session.close()


async def test_no_project_directory_and_no_integration_is_the_plain_login_shell(shell_host):
    """The original behaviour, kept reachable and kept tested.

    Switching integration off has to give back *exactly* what was there before,
    byte for byte -- including the None that means "whatever the account logs
    in with". A fallback that is nearly the old behaviour is not a fallback.
    """
    session = await _open(shell_host, integrate=False)
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    assert shell_host["seen"]["command"] is None
    await session.close()


async def test_an_unknown_shell_falls_through_to_the_plain_login_shell(shell_host):
    """With integration on the command is no longer None -- but every branch
    that does not recognise the shell ends in the same `exec` that was there
    before, because a terminal that works without markers is worth more than
    markers that break a terminal."""
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    command = shell_host["seen"]["command"]
    assert "case ${SHELL##*/} in" in command
    assert command.rstrip().endswith("exec ${SHELL:-/bin/sh} -l")
    # Every step is silenced, so a full /tmp or a read-only home costs the
    # markers and nothing else.
    assert "2>/dev/null" in command
    await session.close()


async def test_a_single_quote_in_the_directory_cannot_end_the_quoting(shell_host):
    session = await _open(shell_host, project_dir="/tmp/it's here; rm -rf /")
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    command = shell_host["seen"]["command"]
    assert command.startswith("cd '/tmp/it'\\''s here; rm -rf /'")
    await session.close()


async def test_what_is_typed_reaches_the_shell_and_comes_back(shell_host):
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    await session.send(b"hello\n")
    assert b"echo:hello" in await _read_until(viewer, b"echo:hello")
    await session.close()


async def test_a_resize_reaches_the_far_side(shell_host):
    session = await _open(shell_host)
    viewer = session.attach(cols=80, rows=24)
    await _read_until(viewer, b"READY")

    session.resize(viewer, 120, 40)
    process = shell_host["seen"]["process"]
    for _ in range(50):
        if process.get_terminal_size()[:2] == (120, 40):
            break
        await asyncio.sleep(0.02)
    assert process.get_terminal_size()[:2] == (120, 40)
    await session.close()


async def test_a_silly_size_is_clamped_rather_than_forwarded(shell_host):
    session = await _open(shell_host)
    viewer = session.attach(cols=100_000, rows=100_000)
    await _read_until(viewer, b"READY")

    assert viewer.cols == terminal_service.MAX_COLS
    assert viewer.rows == terminal_service.MAX_ROWS
    await session.close()


# --- Two tabs, one shell ---------------------------------------------------------
async def test_two_attachments_share_one_shell(shell_host):
    session = await _open(shell_host)
    first = session.attach()
    await _read_until(first, b"READY")
    second = session.attach()

    await session.send(b"both\n")
    assert b"echo:both" in await _read_until(first, b"echo:both")
    assert b"echo:both" in await _read_until(second, b"echo:both")
    await session.close()


async def test_opening_twice_returns_the_same_session(shell_host):
    first = await _open(shell_host)
    second = await _open(shell_host)
    assert first is second
    await first.close()


async def test_the_smaller_viewer_decides_the_size(shell_host):
    session = await _open(shell_host)
    wide = session.attach(cols=200, rows=60)
    await _read_until(wide, b"READY")
    session.attach(cols=90, rows=25)

    assert (session.cols, session.rows) == (90, 25)
    await session.close()


async def test_the_size_lifts_again_when_the_smaller_window_goes(shell_host):
    session = await _open(shell_host)
    wide = session.attach(cols=200, rows=60)
    await _read_until(wide, b"READY")
    narrow = session.attach(cols=90, rows=25)

    session.detach(narrow)
    assert (session.cols, session.rows) == (200, 60)
    await session.close()


# --- Lifetime --------------------------------------------------------------------
async def test_detaching_leaves_the_session_running(shell_host):
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")
    session.detach(viewer)

    assert terminal_service.get("chat-1") is session
    assert not session.closed
    await session.close()


async def test_reattaching_replays_the_scrollback(shell_host):
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")
    await session.send(b"before\n")
    await _read_until(viewer, b"echo:before")
    session.detach(viewer)

    returning = session.attach()
    assert b"READY" in returning.snapshot
    assert b"echo:before" in returning.snapshot
    await session.close()


async def test_scrollback_is_bounded(shell_host, monkeypatch):
    monkeypatch.setattr(terminal_service, "SCROLLBACK_BYTES", 1024)
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    for index in range(200):
        session._remember(f"line {index} ".encode() + b"y" * 100)
    assert session._scrollback_bytes <= 1024 + 110
    assert len(session.attach().snapshot) <= 1024 + 110
    await session.close()


async def test_a_slow_reader_is_dropped_rather_than_buffered(monkeypatch):
    """No server needed: fanning out is a plain function over the viewers.

    The alternative design -- a pump that waits for a full queue -- would stall
    every other viewer behind the one that stopped reading, and buffer without
    bound while it did.
    """
    monkeypatch.setattr(terminal_service, "VIEWER_QUEUE", 4)
    session = Session("chat-1", owner_id="user-1", profile_id="p", label="box")
    keeping_up = session.attach()
    stalled = session.attach()

    for index in range(20):
        session._fan_out(f"chunk {index}\n".encode())
        while not keeping_up.queue.empty():
            keeping_up.queue.get_nowait()

    assert stalled.dropped
    assert stalled.id not in session.viewers
    assert keeping_up.id in session.viewers
    assert not session.closed
    # Woken rather than left waiting: the sentinel is how the socket learns to
    # reconnect, and the scrollback is what makes that free.
    assert stalled.queue.get_nowait() is None


async def test_the_shell_exiting_closes_the_session(shell_host):
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    await session.send(b"exit\n")
    for _ in range(200):
        if session.closed:
            break
        await asyncio.sleep(0.02)
    assert session.closed
    assert session.closed_reason == terminal_service.CLOSED_EXITED
    # The viewer is woken with the sentinel rather than left waiting.
    frame = await _read(viewer)
    while frame is not None:
        frame = await _read(viewer)


async def test_a_closed_session_lingers_with_its_reason(shell_host):
    session = await _open(shell_host)
    viewer = session.attach()
    await _read_until(viewer, b"READY")
    await session.close(terminal_service.CLOSED_IDLE)

    assert terminal_service.get("chat-1") is None
    lingering = terminal_service.peek("chat-1")
    assert lingering is not None
    assert lingering.closed_reason == terminal_service.CLOSED_IDLE


async def test_an_idle_session_is_reaped(shell_host, monkeypatch):
    monkeypatch.setattr(terminal_service, "REAP_INTERVAL", 0.05)
    session = await _open(shell_host, idle_timeout=0.1)
    viewer = session.attach()
    await _read_until(viewer, b"READY")
    session.detach(viewer)

    for _ in range(200):
        if session.closed:
            break
        await asyncio.sleep(0.02)
    assert session.closed
    assert session.closed_reason == terminal_service.CLOSED_IDLE


async def test_a_watched_session_is_never_idle(shell_host, monkeypatch):
    monkeypatch.setattr(terminal_service, "REAP_INTERVAL", 0.05)
    session = await _open(shell_host, idle_timeout=0.1)
    viewer = session.attach()
    await _read_until(viewer, b"READY")

    await asyncio.sleep(0.4)
    assert not session.closed
    assert session.idle_for == 0.0
    await session.close()


async def test_closing_a_profile_ends_its_shells(shell_host):
    first = await _open(shell_host, chat_id="chat-1")
    second = await _open(shell_host, chat_id="chat-2")
    for session in (first, second):
        await _read_until(session.attach(), b"READY")

    assert await terminal_service.close_for_profile("profile-1") == 2
    assert first.closed and second.closed
    assert first.closed_reason == terminal_service.CLOSED_REVOKED


async def test_closing_a_chat_ends_only_its_shell(shell_host):
    first = await _open(shell_host, chat_id="chat-1")
    second = await _open(shell_host, chat_id="chat-2")

    assert await terminal_service.close_chat("chat-1")
    assert first.closed
    assert not second.closed
    await second.close()


async def test_shutdown_closes_every_shell(shell_host):
    first = await _open(shell_host, chat_id="chat-1")
    second = await _open(shell_host, chat_id="chat-2")

    await terminal_service.shutdown()
    assert first.closed and second.closed
    assert first.closed_reason == terminal_service.CLOSED_SHUTDOWN
    assert terminal_service.count() == 0


# --- Caps -------------------------------------------------------------------------
async def test_the_instance_cap_refuses_a_new_shell(shell_host):
    await _open(shell_host, chat_id="chat-1", max_sessions=1)
    with pytest.raises(ExecError, match="as many terminals open as it allows"):
        await _open(shell_host, chat_id="chat-2", max_sessions=1)


async def test_the_per_person_cap_counts_only_that_person(shell_host):
    await _open(shell_host, chat_id="chat-1", owner="user-1", max_per_user=1)
    with pytest.raises(ExecError, match="already have 1 terminal"):
        await _open(shell_host, chat_id="chat-2", owner="user-1", max_per_user=1)

    # Somebody else is unaffected.
    other = await _open(shell_host, chat_id="chat-3", owner="user-2", max_per_user=1)
    assert other is terminal_service.get("chat-3")
    await terminal_service.shutdown()


async def test_a_closed_session_does_not_count_against_the_cap(shell_host):
    session = await _open(shell_host, chat_id="chat-1", max_per_user=1)
    await terminal_service.close_chat("chat-1")
    assert session.closed

    replacement = await _open(shell_host, chat_id="chat-2", max_per_user=1)
    assert not replacement.closed
    await replacement.close()


# --- Refusals ---------------------------------------------------------------------
async def test_an_unconfirmed_host_key_is_refused_before_anything_is_sent(shell_host):
    spec = _spec(shell_host)
    spec["host_key"] = ""
    session = Session("chat-1", owner_id="user-1", profile_id="p", label="Test box")

    with pytest.raises(ExecError, match="host key has not been confirmed"):
        await session.start(spec)


async def test_a_different_host_key_is_refused(shell_host):
    spec = _spec(shell_host)
    spec["host_key"] = f"[127.0.0.1]:{shell_host['port']} ssh-ed25519 {'A' * 68}\n"
    session = Session("chat-1", owner_id="user-1", profile_id="p", label="Test box")

    with pytest.raises(ExecError):
        await session.start(spec)