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