LLeMbas/tests/test_terminal_socket.py

"""The WebSocket behind the terminal panel: who gets one, and what it carries.

The refusals are most of this file on purpose. It is the one endpoint where
getting past the door means a shell on somebody's machine, so every "no" is
pinned -- including *where* it happens, since a browser can read a frame and
cannot read a rejected handshake.
"""

from __future__ import annotations

import json

import pytest
from fastapi.testclient import TestClient

from lembas.db.models import KIND_AGENT, KIND_CHAT, Chat, Connection, Model, SshProfile, User
from lembas.services import settings_store
from lembas.services.agent import policy
from lembas.services.agent import terminal as terminal_service

asyncssh = pytest.importorskip("asyncssh")


class _Server(asyncssh.SSHServer):
    def begin_auth(self, username: str) -> bool:
        return False


@pytest.fixture
def shell_host():
    """A real SSH server with a shell-ish session, on its own thread and loop.

    Its own loop matters: `websocket_connect` is synchronous and blocks the
    test's loop, so a server sharing it could never accept the connection the
    endpoint is trying to make.
    """
    import asyncio
    import threading

    loop = asyncio.new_event_loop()
    thread = threading.Thread(target=loop.run_forever, daemon=True)
    thread.start()
    seen: dict = {}

    async def handler(process):
        seen["command"] = process.command
        seen["term_size"] = process.get_terminal_size()
        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, and so must this
                # one, or every resize would look like the shell exiting.
                seen["term_size"] = process.get_terminal_size()
                continue
            except Exception:  # noqa: BLE001 - the session ending is not a failure
                break
            if not line:
                break
            text = line.rstrip("\n")
            if text == "exit":
                break
            seen.setdefault("typed", []).append(text)
            process.stdout.write(f"echo:{text}\n")
        process.exit(0)

    async def start():
        from lembas.services.agent import ssh as ssh_service

        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]
        line, fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
        return server, port, line, fingerprint

    server, port, host_key, fingerprint = asyncio.run_coroutine_threadsafe(start(), loop).result(10)
    try:
        yield {"port": port, "host_key": host_key, "fingerprint": fingerprint, "seen": seen}
    finally:

        async def stop():
            server.close()
            await server.wait_closed()

        asyncio.run_coroutine_threadsafe(stop(), loop).result(10)
        loop.call_soon_threadsafe(loop.stop)
        thread.join(timeout=5)


@pytest.fixture(autouse=True)
def close_terminals(client, shell_host):
    """End every shell before the server goes.

    A session outlives the socket on purpose, so a test that opens one leaves it
    open -- and `wait_closed` then waits on a client connection that nothing is
    going to close. It runs on the TestClient's portal because that is the loop
    the asyncssh connection belongs to.
    """
    yield
    client.portal.call(terminal_service.shutdown)


def _agent_chat(db, user_id, shell_host, *, kind=KIND_AGENT, project_dir="") -> Chat:
    """A chat pointed at the test server, with everything switched on."""
    settings_store.update(
        db, {"enabled": True, "terminal_enabled": True}, key=settings_store.AGENTS
    )

    profile = SshProfile(
        owner_id=user_id,
        name="Test box",
        host="127.0.0.1",
        port=shell_host["port"],
        username="tester",
        host_key=shell_host["host_key"],
        host_fingerprint=shell_host["fingerprint"],
    )
    db.add(profile)
    db.commit()

    connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted="")
    db.add(connection)
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
    db.commit()

    chat = Chat(
        user_id=user_id,
        model_id="m",
        connection_id=connection.id,
        kind=kind,
        ssh_profile_id=profile.id,
        project_dir=project_dir,
        agent_mode=policy.MODE_MANUAL,
    )
    db.add(chat)
    db.commit()
    return chat


def _url(chat: Chat, **params) -> str:
    query = "".join(f"&{k}={v}" for k, v in params.items())
    return f"/api/chats/{chat.id}/terminal/ws?{query.lstrip('&')}"


def _headers(client: TestClient) -> dict:
    """What a browser on this origin sends. The endpoint requires both."""
    return {"origin": str(client.base_url).rstrip("/"), "host": client.base_url.host}


def _first_json(socket) -> dict:
    return json.loads(socket.receive_text())


def _read_until(socket, needle: bytes, *, frames: int = 40) -> bytes:
    seen = b""
    for _ in range(frames):
        message = socket.receive()
        if message.get("bytes") is not None:
            seen += message["bytes"]
            if needle in seen:
                return seen
        elif message.get("text"):
            payload = json.loads(message["text"])
            raise AssertionError(f"the socket said {payload} before {needle!r} arrived")
    raise AssertionError(f"{needle!r} never arrived; saw {seen!r}")


# --- Getting in -------------------------------------------------------------------
def test_a_shell_opens_and_says_hello(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        hello = _first_json(socket)
        assert hello["t"] == "ready"
        assert hello["label"] == "Test box"
        assert b"READY" in _read_until(socket, b"READY")


def test_what_is_typed_reaches_the_shell(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        assert _first_json(socket)["t"] == "ready"
        _read_until(socket, b"READY")
        socket.send_bytes(b"hello\n")
        assert b"echo:hello" in _read_until(socket, b"echo:hello")


def test_the_requested_size_reaches_the_pty(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(
        _url(chat, cols=120, rows=40), headers=_headers(client)
    ) as socket:
        hello = _first_json(socket)
        _read_until(socket, b"READY")

    assert hello["cols"] == 120
    assert shell_host["seen"]["term_size"][:2] == (120, 40)


def test_the_project_directory_is_where_the_shell_starts(
    client, db, registered, user_id, shell_host
):
    chat = _agent_chat(db, user_id, shell_host, project_dir="/srv/work")
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        assert _first_json(socket)["dir"] == "/srv/work"
        _read_until(socket, b"READY")

    assert "cd '/srv/work'" in shell_host["seen"]["command"]


def test_the_mode_does_not_apply_to_what_a_person_types(
    client, db, registered, user_id, shell_host
):
    """Plan mode stops the *model* running anything. It is not a keyboard lock.

    Pinned because it looks like a bug to anybody reading `policy.py` next to
    this, and "fixing" it would make the panel useless in the mode people spend
    the most time in.
    """
    chat = _agent_chat(db, user_id, shell_host)
    chat.agent_mode = policy.MODE_PLAN
    db.commit()

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        assert _first_json(socket)["t"] == "ready"
        _read_until(socket, b"READY")
        socket.send_bytes(b"rm -rf /tmp/nothing\n")
        _read_until(socket, b"echo:rm -rf /tmp/nothing")

    assert "rm -rf /tmp/nothing" in shell_host["seen"]["typed"]


# --- Staying out ------------------------------------------------------------------
def test_a_stranger_is_refused_before_the_socket_is_accepted(client, db, user_id, shell_host):
    """No cookie, no handshake. Accepting first would mean an unauthenticated
    socket existed at all, however briefly."""
    chat = _agent_chat(db, user_id, shell_host)
    client.cookies.clear()
    # noqa: B017 - starlette raises on a rejected handshake, and the class differs
    with pytest.raises(Exception), client.websocket_connect(  # noqa: B017
        _url(chat), headers=_headers(client)
    ):
        pass


def test_a_cross_site_origin_is_refused_before_the_socket_is_accepted(
    client, db, registered, user_id, shell_host
):
    chat = _agent_chat(db, user_id, shell_host)
    headers = _headers(client) | {"origin": "https://evil.example"}
    with pytest.raises(Exception), client.websocket_connect(  # noqa: B017
        _url(chat), headers=headers
    ):
        pass


def test_an_absent_origin_is_refused(client, db, registered, user_id, shell_host):
    """Required rather than checked-when-present: nothing without an Origin is
    a browser, and a non-browser client has no business here."""
    chat = _agent_chat(db, user_id, shell_host)
    with pytest.raises(Exception), client.websocket_connect(  # noqa: B017
        _url(chat), headers={"host": client.base_url.host}
    ):
        pass


def test_someone_elses_chat_is_refused_with_a_reason(
    client, db, registered, user_id, shell_host
):
    chat = _agent_chat(db, user_id, shell_host)
    stranger = User(name="Sam", email="sam@shire.test", password_hash="x")
    db.add(stranger)
    db.commit()
    chat.user_id = stranger.id
    db.commit()

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        payload = _first_json(socket)
    assert payload["t"] == "error"
    assert "no longer exists" in payload["message"]


def test_an_ordinary_chat_has_no_machine_to_open(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host, kind=KIND_CHAT)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        payload = _first_json(socket)
    assert payload["t"] == "error"
    assert "ordinary chat" in payload["message"]


def test_the_instance_switch_refuses_with_a_reason(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    settings_store.update(db, {"terminal_enabled": False}, key=settings_store.AGENTS)

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        payload = _first_json(socket)
    assert payload["t"] == "error"
    assert "switched off" in payload["message"]


def test_a_disabled_connection_refuses_with_a_reason(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    profile = db.get(SshProfile, chat.ssh_profile_id)
    profile.enabled = False
    db.commit()

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        payload = _first_json(socket)
    assert payload["t"] == "error"
    assert "not usable" in payload["message"]


def test_an_unconfirmed_host_key_says_so_rather_than_failing_silently(
    client, db, registered, user_id, shell_host
):
    """The one refusal that is genuinely actionable: press Check and accept."""
    chat = _agent_chat(db, user_id, shell_host)
    profile = db.get(SshProfile, chat.ssh_profile_id)
    profile.host_key = ""
    db.commit()

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        payload = _first_json(socket)
    assert payload["t"] == "error"
    assert "host key has not been confirmed" in payload["message"]


def test_without_the_permission_there_is_no_terminal(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    user = db.get(User, user_id)
    user.role = "user"  # an admin bypasses every permission, deliberately
    db.commit()
    settings_store.update(
        db,
        {"default_permissions": {"agent.terminal": False, "tools.agent": True, "agent.ssh": True}},
    )

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        payload = _first_json(socket)
    assert payload["t"] == "error"
    assert "permission" in payload["message"]


# --- Lifetime ---------------------------------------------------------------------
def test_closing_the_socket_leaves_the_shell_running(client, db, registered, user_id, shell_host):
    """The whole reason the session is not owned by the socket."""
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        _first_json(socket)
        _read_until(socket, b"READY")
        socket.send_bytes(b"remember\n")
        _read_until(socket, b"echo:remember")

    session = terminal_service.get(chat.id)
    assert session is not None
    assert not session.closed


def test_reconnecting_replays_what_was_missed(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        _first_json(socket)
        _read_until(socket, b"READY")
        socket.send_bytes(b"before\n")
        _read_until(socket, b"echo:before")

    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        assert _first_json(socket)["t"] == "ready"
        # Arrives as the scrollback, in one frame, before anything new.
        assert b"echo:before" in _read_until(socket, b"echo:before")


def test_the_shell_exiting_is_announced(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        _first_json(socket)
        _read_until(socket, b"READY")
        socket.send_bytes(b"exit\n")

        payload = None
        for _ in range(40):
            message = socket.receive()
            if message.get("text"):
                payload = json.loads(message["text"])
                break
    assert payload is not None
    assert payload["t"] == "closed"
    assert payload["reason"] == terminal_service.CLOSED_EXITED


def test_a_resize_frame_reaches_the_far_side(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        _first_json(socket)
        _read_until(socket, b"READY")
        socket.send_text(json.dumps({"t": "resize", "cols": 132, "rows": 43}))
        # Round-tripped through the shell so the resize has certainly been read.
        socket.send_bytes(b"after\n")
        _read_until(socket, b"echo:after")

    assert shell_host["seen"]["term_size"][:2] == (132, 43)


# --- The panel on the page --------------------------------------------------------
def test_the_panel_and_its_vendored_terminal_are_on_an_agent_chat(
    client, db, registered, user_id, shell_host
):
    chat = _agent_chat(db, user_id, shell_host)
    body = client.get(f"/chat/{chat.id}").text

    assert 'id="terminal"' in body
    assert "vendor/xterm.js" in body
    assert 'data-toggle-group="side"' in body


def test_an_ordinary_chat_loads_none_of_it(client, db, registered, user_id, shell_host):
    """280KB of terminal on a page that could never use it."""
    chat = _agent_chat(db, user_id, shell_host, kind=KIND_CHAT)
    body = client.get(f"/chat/{chat.id}").text

    assert 'id="terminal"' not in body
    assert "xterm" not in body


def test_the_panel_is_absent_when_the_instance_switch_is_off(
    client, db, registered, user_id, shell_host
):
    chat = _agent_chat(db, user_id, shell_host)
    settings_store.update(db, {"terminal_enabled": False}, key=settings_store.AGENTS)

    assert 'id="terminal"' not in client.get(f"/chat/{chat.id}").text


def test_the_panel_is_absent_without_ssh_installed(
    client, db, registered, user_id, shell_host, monkeypatch
):
    """A button whose only outcome is an error frame is worse than no button."""
    from lembas.services.agent import ssh as ssh_service

    chat = _agent_chat(db, user_id, shell_host)
    monkeypatch.setattr(ssh_service, "available", lambda: ssh_service.INSTALL_HINT)

    assert 'id="terminal"' not in client.get(f"/chat/{chat.id}").text


# --- Things that must reach a shell already open ----------------------------------
def _open_and_leave(client, chat) -> None:
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        _first_json(socket)
        _read_until(socket, b"READY")
    assert terminal_service.get(chat.id) is not None


def test_deleting_the_chat_closes_its_terminal(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    _open_and_leave(client, chat)

    assert client.delete(f"/api/chats/{chat.id}").status_code == 204
    assert terminal_service.get(chat.id) is None


def test_disabling_the_connection_closes_its_terminal(
    client, db, registered, user_id, shell_host
):
    """The model stops on the next reply anyway. A shell already open would not."""
    chat = _agent_chat(db, user_id, shell_host)
    _open_and_leave(client, chat)

    profile = db.get(SshProfile, chat.ssh_profile_id)
    client.post(
        f"/api/agents/{profile.id}",
        data={
            "name": profile.name,
            "host": profile.host,
            "port": str(profile.port),
            "username": profile.username,
            "auth": profile.auth,
            "connect_timeout": "15",
        },
        follow_redirects=False,
    )
    db.expire_all()
    assert db.get(SshProfile, profile.id).enabled is False
    assert terminal_service.get(chat.id) is None


def test_deleting_the_connection_closes_its_terminal(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    _open_and_leave(client, chat)

    profile_id = chat.ssh_profile_id
    client.post(f"/api/agents/{profile_id}/delete", follow_redirects=False)
    assert terminal_service.get(chat.id) is None


def test_forgetting_the_host_key_closes_its_terminal(client, db, registered, user_id, shell_host):
    chat = _agent_chat(db, user_id, shell_host)
    _open_and_leave(client, chat)

    client.post(f"/api/agents/{chat.ssh_profile_id}/forget")
    assert terminal_service.get(chat.id) is None


def test_nonsense_control_frames_are_ignored(client, db, registered, user_id, shell_host):
    """A frame is text somebody could forge; none of it may crash the socket."""
    chat = _agent_chat(db, user_id, shell_host)
    with client.websocket_connect(_url(chat), headers=_headers(client)) as socket:
        _first_json(socket)
        _read_until(socket, b"READY")
        for frame in ("not json", "[]", '{"t":"unknown"}', '{"t":"resize"}'):
            socket.send_text(frame)
        socket.send_bytes(b"alive\n")
        assert b"echo:alive" in _read_until(socket, b"echo:alive")