LLeMbas/tests/test_agent_ssh.py

"""Acting on a machine over SSH.

Driven against a real asyncssh server on 127.0.0.1 with a generated host key, so
nothing here touches an outside network and the host-key path is exercised for
real rather than mocked.
"""

from __future__ import annotations

import asyncio

import pytest

from lembas.services.agent import ssh
from lembas.services.agent.base import ExecError, ExecRequest

asyncssh = pytest.importorskip("asyncssh")


# --- A server to talk to -------------------------------------------------------
def _generate_key():
    return asyncssh.generate_private_key("ssh-ed25519")


class _Server(asyncssh.SSHServer):
    def begin_auth(self, username: str) -> bool:
        return False  # no authentication wanted; every login succeeds


async def _handler(process):
    """A shell that understands just enough to be tested against."""
    command = process.command or ""
    if "sleep" in command:
        await asyncio.sleep(5)
        process.exit(0)
        return
    if "flood" in command:
        process.stdout.write("x" * 200_000)
        process.exit(0)
        return
    if "fail" in command:
        process.stderr.write("it went wrong\n")
        process.exit(3)
        return
    if "colour" in command:
        process.stdout.write("\x1b[31mred\x1b[0m\n")
        process.exit(0)
        return
    process.stdout.write(f"ran: {command}\n")
    process.exit(0)


async def _start(host_key=None):
    server = await asyncssh.create_server(
        _Server,
        "127.0.0.1",
        0,
        server_host_keys=[host_key or _generate_key()],
        process_factory=_handler,
        sftp_factory=True,
    )
    port = next(iter(server.sockets)).getsockname()[1]
    return server, port


def _spec(port: int, host_key: str, **overrides) -> dict:
    base = {
        "id": "p1",
        "label": "test box",
        "host": "127.0.0.1",
        "port": port,
        "username": "tester",
        "auth": "key",
        "password": "",
        "private_key": "",
        "key_passphrase": "",
        "host_key": host_key,
        "connect_timeout": 5,
    }
    base.update(overrides)
    return base


@pytest.fixture
async def project(tmp_path):
    """A directory of its own.

    Not `tmp_path` itself: the autouse database fixture points the data
    directory there, so a listing would find lembas.db and the uploads folder.
    """
    path = tmp_path / "project"
    path.mkdir()
    return path


@pytest.fixture
async def box(project):
    """A running server, its pinned host key line, and a project directory."""
    key = _generate_key()
    server, port = await _start(key)
    line, fingerprint = await ssh.capture_host_key("127.0.0.1", port)
    try:
        yield {
            "port": port,
            "host_key": line,
            "fingerprint": fingerprint,
            "dir": str(project),
            "server": server,
        }
    finally:
        server.close()
        await server.wait_closed()


# --- The four defaults that must never be left to asyncssh --------------------
def test_the_dangerous_defaults_are_all_passed_explicitly():
    """Every LLeMbas user shares one unix account, so "whatever the account has
    lying around" is never the right answer. This is the most important test in
    the file and it needs no server at all."""
    kwargs = ssh.connect_kwargs(_spec(22, "127.0.0.1 ssh-ed25519 AAAA\n"))

    # Bytes, never None: None turns host key checking off entirely.
    assert isinstance(kwargs["known_hosts"], bytes)
    assert kwargs["known_hosts"] != b""
    # Not left to load ~/.ssh/id_*, which could be another person's key.
    assert kwargs["client_keys"] == []
    # Not left to read ~/.ssh/config, where ProxyCommand could redirect us.
    assert kwargs["config"] is None
    # Not left to use $SSH_AUTH_SOCK.
    assert kwargs["agent_path"] is None


def test_a_profile_with_no_confirmed_host_key_refuses_to_connect():
    with pytest.raises(ExecError, match="host key has not been confirmed"):
        ssh.connect_kwargs(_spec(22, ""))


def test_a_password_profile_sends_a_password_and_no_keys():
    kwargs = ssh.connect_kwargs(
        _spec(22, "h ssh-ed25519 AAAA\n", auth="password", password="hunter2")
    )
    assert kwargs["password"] == "hunter2"
    assert kwargs["client_keys"] == []


def test_a_key_profile_sends_no_password():
    kwargs = ssh.connect_kwargs(_spec(22, "h ssh-ed25519 AAAA\n", password="stale"))
    assert kwargs["password"] is None


# --- Trust on first use --------------------------------------------------------
async def test_the_first_look_captures_a_key_and_a_fingerprint(box):
    assert box["host_key"].startswith(f"[127.0.0.1]:{box['port']} ssh-ed25519 ")
    assert box["fingerprint"].startswith("SHA256:")


async def test_capturing_a_key_offers_no_credential():
    """get_server_host_key completes the key exchange and stops, which is what
    makes accepting a fingerprint from a button safe: nothing is sent to a host
    that has not been accepted yet.

    The server here records every authentication attempt, so an empty list is
    evidence rather than an absence of it.
    """
    attempts: list[str] = []

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

    server = await asyncssh.create_server(
        Watchful, "127.0.0.1", 0, server_host_keys=[_generate_key()]
    )
    port = next(iter(server.sockets)).getsockname()[1]
    try:
        line, fingerprint = await ssh.capture_host_key("127.0.0.1", port)
        assert line and fingerprint.startswith("SHA256:")
        assert attempts == [], "a host not yet accepted was offered a username"
    finally:
        server.close()
        await server.wait_closed()


async def test_a_host_that_answers_with_a_different_key_is_refused(box):
    """The pinned key is the whole of the protection. A host presenting another
    one is either rebuilt or is not the host."""
    other_line, _ = await ssh.capture_host_key("127.0.0.1", box["port"])
    wrong = other_line.rsplit(" ", 1)[0] + " " + "A" * 68 + "\n"

    executor = ssh.SshExecutor(_spec(box["port"], wrong), box["dir"])
    with pytest.raises(ExecError):
        await executor.run(ExecRequest(command="echo hi", timeout=5))


async def test_an_unreachable_host_says_so():
    with pytest.raises(ExecError, match="Could not reach"):
        await ssh.capture_host_key("127.0.0.1", 1, timeout=3)


# --- Running commands ----------------------------------------------------------
async def test_a_command_runs_and_comes_back(box):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    result = await executor.run(ExecRequest(command="echo hi", timeout=5))

    assert result.ok
    assert result.exit_status == 0
    assert "echo hi" in result.output


async def test_the_working_directory_is_set_and_never_spliced_in(box):
    """Every command is a fresh shell, so `cd` cannot carry between calls. The
    directory is quoted, because `cwd` may come from the model."""
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), "/tmp/a dir")
    result = await executor.run(ExecRequest(command="echo hi", timeout=5))
    assert "cd '/tmp/a dir' && echo hi" in result.output


async def test_a_directory_with_a_quote_in_it_cannot_end_the_quoting(box):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), "/tmp/it's; rm -rf /")
    result = await executor.run(ExecRequest(command="echo hi", timeout=5))
    assert "'/tmp/it'\\''s; rm -rf /'" in result.output


async def test_a_failing_command_is_a_result_not_an_error(box):
    """A command that ran and failed is something the model should read and
    react to. Only being unable to act at all is an ExecError."""
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    result = await executor.run(ExecRequest(command="fail", timeout=5))

    assert result.exit_status == 3
    assert result.ok is False
    assert "it went wrong" in result.output, "stderr is interleaved, not dropped"


async def test_a_slow_command_times_out_and_says_so(box):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    result = await executor.run(ExecRequest(command="sleep", timeout=1))

    assert result.timed_out is True
    assert result.ok is False
    assert "was stopped" in result.output


async def test_a_flood_of_output_is_capped(box):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    result = await executor.run(ExecRequest(command="flood", timeout=10, max_bytes=2000))

    assert result.truncated is True
    assert len(result.output) < 2200
    assert result.output.endswith("(truncated)")


async def test_terminal_escapes_are_stripped(box):
    """Inert in escaped HTML, but this text re-enters the model's context, where
    they are a known way to hide instructions -- and a log a person later cats."""
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    result = await executor.run(ExecRequest(command="colour", timeout=5))

    assert "red" in result.output
    assert "\x1b" not in result.output


# --- Files, over SFTP ----------------------------------------------------------
async def test_a_file_is_written_and_read_back(box):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])

    written = await executor.write_file("hello.txt", "a mallorn tree\n")
    assert written == len("a mallorn tree\n")
    assert await executor.read_file("hello.txt") == "a mallorn tree\n"


async def test_a_relative_path_resolves_against_the_project_directory(box, project):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    await executor.write_file("nested.txt", "here")

    assert (project / "nested.txt").read_text() == "here"


async def test_a_missing_file_is_reported_plainly(box):
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    with pytest.raises(ExecError, match="no file at"):
        await executor.read_file("nope.txt")


async def test_a_directory_lists(box, project):
    (project / "one.txt").write_text("1")
    (project / "two.txt").write_text("2")

    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])
    assert await executor.list_dir() == ["one.txt", "two.txt"]


async def test_reading_a_file_is_capped(box, project):
    (project / "big.txt").write_text("y" * 50_000)
    executor = ssh.SshExecutor(_spec(box["port"], box["host_key"]), box["dir"])

    text = await executor.read_file("big.txt", max_bytes=1000)
    assert len(text) <= 1000


# --- The check a person presses ------------------------------------------------
async def test_check_reports_what_it_found(box):
    found = await ssh.check(_spec(box["port"], box["host_key"]), box["dir"])
    assert found["ok"] is True
    assert "uname" in found["output"]


# --- The snapshot --------------------------------------------------------------
def test_spec_from_decrypts_the_credential_and_the_row_does_not_hold_it(db, user_id):
    from lembas.db.models import SshProfile
    from lembas.services.crypto import encrypt

    profile = SshProfile(
        owner_id=user_id,
        name="box",
        host="10.0.0.5",
        username="root",
        private_key_encrypted=encrypt("PRIVATE KEY MATERIAL"),
        host_key="h ssh-ed25519 AAAA\n",
    )
    db.add(profile)
    db.commit()

    spec = ssh.spec_from(profile)
    assert spec["private_key"] == "PRIVATE KEY MATERIAL"
    assert "PRIVATE KEY MATERIAL" not in profile.private_key_encrypted