"""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 # Stands up something real -- see the `slow` marker in pyproject.toml. pytestmark = pytest.mark.slow 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")