LLeMbas/src/lembas/services/agent/ssh.py

"""Acting on a machine over SSH.

Connections are made per call, for the reason MCP sessions are, plus one more: a
live `SSHClientConnection` is exactly the kind of state `ToolContext` exists so
that nothing holds. A command is already a network round trip inside a reply
that takes seconds, so a second one to open the channel is not the cost worth
optimising.

**Four asyncssh defaults are actively wrong here, and all four are passed
explicitly on every connection.** Every LLeMbas user shares one unix account, so
"whatever the account has lying around" is never the right answer:

* `known_hosts` unset reads that shared `~/.ssh/known_hosts` -- one trust store
  for everybody. Set to `None` it disables host key checking altogether, which
  is never correct and is the single easiest way to make this insecure.
* `client_keys` unset loads `~/.ssh/id_*`, so one person's chat could
  authenticate with a key another person left there, or with the server's own.
* `config` unset reads `~/.ssh/config`, where a `Hostname` or `ProxyCommand`
  can send the connection somewhere else entirely.
* `agent_path` unset silently uses `$SSH_AUTH_SOCK`.

`asyncssh` is an optional dependency, imported inside the functions that need it
so an instance with agents switched off never pays for it and an instance that
forgot to install it gets a sentence rather than an ImportError at startup.
"""

from __future__ import annotations

import contextlib
import logging
import time
from typing import Any

from lembas.db.models import AUTH_PASSWORD, SshProfile
from lembas.services.agent.base import (
    Conflict,
    ExecError,
    ExecRequest,
    ExecResult,
    RemoteEntry,
    RemoteFile,
    clean_output,
    revision_of,
)
from lembas.services.crypto import decrypt

log = logging.getLogger(__name__)

# A file read into a model's context, and one written out of it. Both bounded:
# the first because a 40 MB log would fill the window, the second because
# nothing a model writes in one call should be larger than this.
MAX_READ_BYTES = 256 * 1024
MAX_WRITE_BYTES = 1024 * 1024

# How many entries a directory listing returns before it is cut short.
MAX_ENTRIES = 500

INSTALL_HINT = (
    "SSH support is not installed. Run `pip install -e \".[ssh]\"` in the "
    "LLeMbas virtual environment and restart."
)


def available() -> str:
    """Empty when SSH can be used, else why it cannot.

    Shaped like `search.availability`, and used the same way: the feature stays
    visible in the UI with an install hint rather than silently missing.
    """
    try:
        import asyncssh  # noqa: F401
    except ImportError:
        return INSTALL_HINT
    return ""


def spec_from(profile: SshProfile) -> dict[str, Any]:
    """A session-free snapshot of one profile, credential decrypted.

    Called while the session is open. The plaintext lives in the returned dict
    and nowhere else; `generation` drops it when the reply ends.
    """
    return {
        "id": profile.id,
        "label": profile.label,
        "host": profile.host,
        "port": int(profile.port or 22),
        "username": profile.username,
        "auth": profile.auth,
        "password": decrypt(profile.password_encrypted),
        "private_key": decrypt(profile.private_key_encrypted),
        "key_passphrase": decrypt(profile.key_passphrase_encrypted),
        "host_key": profile.host_key,
        "connect_timeout": int(profile.connect_timeout or 15),
    }


def connect_kwargs(spec: dict[str, Any]) -> dict[str, Any]:
    """Everything asyncssh must be told rather than left to discover.

    See the module docstring: every one of these has a default that is wrong
    when one unix account is shared by every user of the instance.
    """
    if not spec.get("host_key"):
        raise ExecError(
            "This connection's host key has not been confirmed yet. Open it "
            "under Agents and press Check, then accept the fingerprint."
        )

    keys: list = []
    if spec.get("auth") != AUTH_PASSWORD and spec.get("private_key"):
        import asyncssh

        try:
            keys = [
                asyncssh.import_private_key(
                    spec["private_key"], passphrase=spec.get("key_passphrase") or None
                )
            ]
        except Exception as exc:  # noqa: BLE001 - any failure here is one message
            raise ExecError(f"That private key could not be read: {exc}") from exc

    timeout = int(spec.get("connect_timeout") or 15)
    return {
        "username": spec["username"],
        "port": int(spec.get("port") or 22),
        # Bytes, never None. None turns host key checking off entirely.
        "known_hosts": spec["host_key"].encode(),
        "client_keys": keys,
        "password": (spec.get("password") or None) if spec.get("auth") == AUTH_PASSWORD else None,
        "config": None,
        "agent_path": None,
        "connect_timeout": timeout,
        "login_timeout": timeout,
    }


async def capture_host_key(host: str, port: int, *, timeout: int = 15) -> tuple[str, str]:
    """The host's key as a known_hosts line, and its SHA256 fingerprint.

    `get_server_host_key` completes the key exchange and stops, so nothing is
    offered to a host that has not been accepted yet -- no username, no
    password, no key. That is what makes trust-on-first-use safe to do from a
    button rather than only from a terminal.
    """
    if problem := available():
        raise ExecError(problem)
    import asyncio

    import asyncssh

    try:
        key = await asyncio.wait_for(
            asyncssh.get_server_host_key(host, port=port), timeout=timeout
        )
    except TimeoutError as exc:
        raise ExecError(f"{host} did not answer within {timeout}s.") from exc
    except (OSError, asyncssh.Error) as exc:
        raise ExecError(f"Could not reach {host}: {exc}") from exc

    if key is None:
        raise ExecError(f"{host} offered no host key.")

    algorithm = key.get_algorithm()
    encoded = key.export_public_key("openssh").decode().split()[1]
    where = f"[{host}]:{port}" if port != 22 else host
    return f"{where} {algorithm} {encoded}\n", key.get_fingerprint("sha256")


class SshExecutor:
    """One target, reached over SSH. A connection per call."""

    def __init__(self, spec: dict[str, Any], project_dir: str = "") -> None:
        self.spec = spec
        self.project_dir = project_dir or ""
        self.label = str(spec.get("label") or spec.get("host") or "the remote host")

    def _connect(self):
        if problem := available():
            raise ExecError(problem)
        import asyncssh

        return asyncssh.connect(self.spec["host"], **connect_kwargs(self.spec))

    def _wrap(self, exc: Exception) -> ExecError:
        import asyncssh

        if isinstance(exc, asyncssh.HostKeyNotVerifiable):
            return ExecError(
                f"{self.label} presented a different host key than the one that "
                "was confirmed. Nothing was sent. If the host was rebuilt, open "
                "it under Agents and confirm the new fingerprint."
            )
        if isinstance(exc, asyncssh.PermissionDenied):
            return ExecError(f"{self.label} refused the credential.")
        return ExecError(f"Could not reach {self.label}: {exc}")

    async def run(self, request: ExecRequest) -> ExecResult:
        """Run one command and read back what it said.

        Every command is a fresh shell, so `cd` does not carry between calls --
        the working directory is set here, from `cwd` or the chat's project
        directory, and never spliced into the command string.
        """
        import asyncssh

        started = time.monotonic()
        directory = request.cwd or self.project_dir
        # A single-quoted path, with any embedded quote escaped. `cd` needs a
        # shell, so this is the one place a path meets one -- and it is a path
        # from the chat's own configuration, not from the model, except when the
        # model passed `cwd`, which is why it is quoted rather than trusted.
        command = request.command
        if directory:
            command = f"cd {_quote(directory)} && {command}"

        try:
            async with self._connect() as conn:
                result = await conn.run(
                    command,
                    check=False,
                    timeout=request.timeout,
                    # Interleaved, because a shell transcript is what the model
                    # has to read and separating them loses the ordering.
                    stderr=asyncssh.STDOUT,
                    # A command that waits for input fails at once instead of
                    # sitting out its whole timeout in silence.
                    stdin=asyncssh.DEVNULL,
                )
        except TimeoutError:
            elapsed = int((time.monotonic() - started) * 1000)
            return ExecResult(
                exit_status=-1,
                output=f"The command was still running after {request.timeout:g}s and was stopped.",
                timed_out=True,
                duration_ms=elapsed,
            )
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc

        output, truncated = clean_output(result.stdout or "", limit=request.max_bytes)
        return ExecResult(
            exit_status=result.exit_status if result.exit_status is not None else -1,
            output=output,
            truncated=truncated,
            duration_ms=int((time.monotonic() - started) * 1000),
        )

    # --- Files go over SFTP, never through a shell ---------------------------
    # The SSH exec protocol carries one command *string* that the far side's
    # shell parses; there is no argv form. So a path in a command line is
    # unavoidably a quoting problem, and a model-supplied path is exactly the
    # input that must not become one. Over SFTP a path is a path.
    async def read_file(self, path: str, *, max_bytes: int = MAX_READ_BYTES) -> str:
        import asyncssh

        try:
            async with (
                self._connect() as conn,
                conn.start_sftp_client() as sftp,
                sftp.open(self._resolve(path), "rb") as handle,
            ):
                data = await handle.read(max_bytes + 1)
        except asyncssh.SFTPNoSuchFile as exc:
            raise ExecError(f"There is no file at {path}.") from exc
        except asyncssh.SFTPPermissionDenied as exc:
            raise ExecError(f"Not allowed to read {path}.") from exc
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc

        text, _truncated = clean_output(data[:max_bytes], limit=max_bytes)
        return text

    async def write_file(self, path: str, text: str) -> int:
        import asyncssh

        payload = text.encode("utf-8")[:MAX_WRITE_BYTES]
        try:
            async with (
                self._connect() as conn,
                conn.start_sftp_client() as sftp,
                sftp.open(self._resolve(path), "wb") as handle,
            ):
                await handle.write(payload)
        except asyncssh.SFTPPermissionDenied as exc:
            raise ExecError(f"Not allowed to write {path}.") from exc
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc
        return len(payload)

    # --- The same files, for somebody about to edit them ---------------------
    # Deliberately not `read_file`/`write_file`, and those two are deliberately
    # left exactly as they are: what they return is a contract a model has been
    # shown, and it is the right contract for a model.
    #
    # It is the wrong one for an editor. `read_file` ends in `clean_output`,
    # which strips ANSI escape sequences and decodes with errors="replace" --
    # correct for the output of a command, and for a file it means that opening
    # one containing an escape byte and pressing Save rewrites it with the
    # escapes gone and every undecodable byte replaced by U+FFFD. `write_file`
    # truncates at MAX_WRITE_BYTES, which a model is told about and a person
    # pressing Save is not.
    async def read_text(self, path: str, *, max_bytes: int = MAX_READ_BYTES) -> RemoteFile:
        """A file as somebody is about to edit it.

        Strict decoding, so a file this cannot represent faithfully is reported
        as binary rather than silently mangled into something that would be
        saved back. The stat and the read share one connection: connections are
        per call, so doing it in two is two handshakes and two authentications
        to open one file.
        """
        import asyncssh

        try:
            async with (
                self._connect() as conn,
                conn.start_sftp_client() as sftp,
                sftp.open(self._resolve(path), "rb") as handle,
            ):
                attrs = await handle.stat()
                data = await handle.read(max_bytes + 1)
        except asyncssh.SFTPNoSuchFile as exc:
            raise ExecError(f"There is no file at {path}.") from exc
        except asyncssh.SFTPPermissionDenied as exc:
            raise ExecError(f"Not allowed to read {path}.") from exc
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc

        truncated = len(data) > max_bytes
        data = data[:max_bytes]
        size = int(getattr(attrs, "size", None) or len(data))
        mtime = int(getattr(attrs, "mtime", None) or 0)

        # A NUL in the first few kilobytes, or anything that will not decode.
        # Either way there is nothing safe to put in a textarea.
        if b"\0" in data[:8192]:
            return RemoteFile("", size, mtime, truncated, binary=True)
        try:
            text = data.decode("utf-8")
        except UnicodeDecodeError:
            return RemoteFile("", size, mtime, truncated, binary=True)
        return RemoteFile(text, size, mtime, truncated, binary=False)

    async def write_text(self, path: str, text: str, *, if_unchanged: str = "") -> RemoteFile:
        """Write a file, refusing if it moved under the editor.

        `if_unchanged` is the token `read_text` handed out. The re-stat and the
        write happen on one connection, which is the narrowest window SFTP
        allows; there is no compare-and-swap here and this does not pretend to
        be atomic. It catches what it exists for -- another editor, a build, a
        checkout between opening a tab and pressing Save -- and not a race
        measured in milliseconds.

        Oversize is refused rather than truncated. `write_file` truncates
        because a model is told how many bytes it wrote; somebody pressing Save
        would lose the tail of their file with nothing said.
        """
        import asyncssh

        payload = text.encode("utf-8")
        if len(payload) > MAX_WRITE_BYTES:
            raise ExecError(
                f"That is {len(payload) // 1024}KB and the limit is "
                f"{MAX_WRITE_BYTES // 1024}KB. Nothing was written."
            )

        target = self._resolve(path)
        try:
            async with self._connect() as conn, conn.start_sftp_client() as sftp:
                if if_unchanged:
                    current = ""
                    with contextlib.suppress(asyncssh.SFTPNoSuchFile):
                        attrs = await sftp.stat(target)
                        current = revision_of(
                            int(getattr(attrs, "mtime", None) or 0),
                            int(getattr(attrs, "size", None) or 0),
                        )
                    if current and current != if_unchanged:
                        raise Conflict(current)
                async with sftp.open(target, "wb") as handle:
                    await handle.write(payload)
                attrs = await sftp.stat(target)
        except asyncssh.SFTPPermissionDenied as exc:
            raise ExecError(f"Not allowed to write {path}.") from exc
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc

        return RemoteFile(
            text,
            len(payload),
            int(getattr(attrs, "mtime", None) or 0),
            truncated=False,
            binary=False,
        )

    async def list_dir(self, path: str = "") -> list[str]:
        import asyncssh

        try:
            async with self._connect() as conn, conn.start_sftp_client() as sftp:
                target = self._resolve(path) if path else (self.project_dir or ".")
                names = await sftp.listdir(target)
        except asyncssh.SFTPNoSuchFile as exc:
            raise ExecError(f"There is no directory at {path or self.project_dir}.") from exc
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc

        visible = sorted(n for n in names if n not in (".", ".."))
        return visible[:MAX_ENTRIES]

    async def scan_dir(self, path: str = "") -> list[RemoteEntry]:
        """A listing with types, for a picker rather than for a model.

        `readdir` rather than `listdir`: the latter returns bare names, and a
        browser has to know which rows can be walked into before it can draw
        them. Directories sort first and then by name, because that is the
        order somebody navigating expects -- `list_dir` keeps its plain
        lexicographic sort, since changing what a tool returns is changing a
        contract a model has already been shown.
        """
        import stat

        import asyncssh

        try:
            async with self._connect() as conn, conn.start_sftp_client() as sftp:
                target = self._resolve(path) if path else (self.project_dir or ".")
                names = await sftp.readdir(target)
        except asyncssh.SFTPNoSuchFile as exc:
            raise ExecError(f"There is no directory at {path or self.project_dir}.") from exc
        except asyncssh.SFTPPermissionDenied as exc:
            raise ExecError(f"Not allowed to read {path or self.project_dir}.") from exc
        except (OSError, asyncssh.Error) as exc:
            raise self._wrap(exc) from exc

        entries: list[RemoteEntry] = []
        for item in names:
            name = item.filename
            if name in (".", ".."):
                continue
            attrs = item.attrs
            permissions = getattr(attrs, "permissions", None) or 0
            entries.append(
                RemoteEntry(
                    name=name,
                    is_dir=stat.S_ISDIR(permissions),
                    size=getattr(attrs, "size", None) or 0,
                    modified=int(getattr(attrs, "mtime", None) or 0),
                )
            )

        entries.sort(key=lambda entry: (not entry.is_dir, entry.name.lower()))
        return entries[:MAX_ENTRIES]

    def _resolve(self, path: str) -> str:
        """A path relative to the project directory, unless it is absolute.

        Deliberately *not* a containment check. The account on the far side is
        the boundary -- a profile whose user can only see /srv/project can only
        reach things under it -- and pretending otherwise here would be a
        comfort rather than a control, since `shell_run` could walk out of it in
        one line anyway.
        """
        if not path:
            return self.project_dir or "."
        if path.startswith("/") or not self.project_dir:
            return path
        return f"{self.project_dir.rstrip('/')}/{path.lstrip('/')}"


def _quote(value: str) -> str:
    return "'" + value.replace("'", "'\\''") + "'"


async def check(spec: dict[str, Any], project_dir: str = "") -> dict[str, Any]:
    """Connect, confirm the project directory, and report what was found.

    Used by the Check button on a profile. Runs one harmless command rather than
    only opening a connection, because "the credential works" and "the directory
    is there" are the two things somebody is actually asking about.
    """
    executor = SshExecutor(spec, project_dir)
    result = await executor.run(
        ExecRequest(command="uname -sr 2>/dev/null; pwd", timeout=15, max_bytes=4096)
    )
    lines = [line for line in result.output.splitlines() if line.strip()]
    return {
        "ok": result.ok,
        "system": lines[0] if lines else "",
        "cwd": lines[-1] if len(lines) > 1 else "",
        "output": result.output,
    }


__all__ = [
    "INSTALL_HINT",
    "MAX_ENTRIES",
    "MAX_READ_BYTES",
    "SshExecutor",
    "available",
    "capture_host_key",
    "check",
    "connect_kwargs",
    "spec_from",
]