#!/usr/bin/env python3

"""SSH keys, the agent holding them, and the hosts this machine has met.

Nothing here ever reads a private key. Fingerprints and comments come from the
matching .pub file, and whether a key is encrypted is answered by asking
ssh-keygen to derive the PUBLIC key with an empty passphrase: it succeeds for an
unencrypted key and fails for an encrypted one, and either way the only thing it
can print is public material.

No passphrase passes through this tool at all. Adding an encrypted key to the
agent lets ssh-add prompt through the system's own askpass, which is where that
belongs -- a settings page collecting a passphrase and handing it on would be a
worse place for it to live, and putting one in argv would publish it to every
process on the machine.

    panama-ssh-keys snapshot
    panama-ssh-keys agent-add PATH | agent-remove PATH
    panama-ssh-keys forget-host HOST
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
from pathlib import Path

SSH_DIR = Path.home() / ".ssh"
KNOWN_HOSTS = SSH_DIR / "known_hosts"

# A host as it may appear in known_hosts, including [host]:port forms.
HOST = re.compile(r"^[A-Za-z0-9._:\[\]-]{1,253}$")

# gnome-keyring's agent, which is what runs on this desktop. Only used when the
# environment has not already named one, so an ssh-agent started by hand wins.
KEYRING_SOCKET = Path(os.environ.get("XDG_RUNTIME_DIR", "/run/user/1000")) / "keyring" / "ssh"


class BoundaryError(RuntimeError):
    """A user-visible validation or ssh failure."""


def run(command: list[str], timeout: float = 15.0, env: dict | None = None):
    try:
        return subprocess.run(command, capture_output=True, text=True,
                              timeout=timeout, env=env)
    except FileNotFoundError as error:
        raise BoundaryError(f"{command[0]} is not installed.") from error
    except subprocess.TimeoutExpired as error:
        raise BoundaryError(f"{command[0]} did not respond.") from error


def agent_environment() -> dict:
    """The environment an ssh-add call should run in.

    A settings window inherits whatever the shell was started with, which on
    this desktop does not include SSH_AUTH_SOCK -- so without this the page
    would report "no agent" while one is plainly running.
    """
    environment = dict(os.environ)
    if not environment.get("SSH_AUTH_SOCK") and KEYRING_SOCKET.is_socket():
        environment["SSH_AUTH_SOCK"] = str(KEYRING_SOCKET)
    return environment


def agent_state() -> dict:
    environment = agent_environment()
    socket = environment.get("SSH_AUTH_SOCK", "")
    if not socket:
        return {"available": False, "socket": "", "kind": "", "durableRemoval": False,
                "fingerprints": [], "detail": "No SSH agent is running."}

    result = run(["ssh-add", "-l"], env=environment)
    # ssh-add exits 1 for "no identities" and 2 for "cannot connect", which are
    # very different things to report.
    if result.returncode == 2:
        return {"available": False, "socket": socket, "kind": "", "durableRemoval": False,
                "fingerprints": [], "detail": "An agent socket exists but could not be reached."}

    fingerprints = []
    for line in (result.stdout or "").splitlines():
        parts = line.split()
        if len(parts) >= 2 and parts[1].startswith("SHA256:"):
            fingerprints.append(parts[1])
    # gnome-keyring's agent enumerates whatever keys it finds in ~/.ssh, so
    # `ssh-add -d` reports "Identity removed" and the key is still listed a
    # second later -- it comes straight back from disk. A plain ssh-agent
    # removes durably. Measured on this machine rather than assumed, because an
    # Unload button that reports success and changes nothing is worse than no
    # button at all.
    keyring = "/keyring/" in socket
    return {
        "available": True,
        "socket": socket,
        "kind": "gnome-keyring" if keyring else "ssh-agent",
        "durableRemoval": not keyring,
        "fingerprints": fingerprints,
        "detail": "" if fingerprints else "The agent is running but holds no keys.",
    }


def encrypted(private: Path) -> bool | None:
    """Whether a private key needs a passphrase.

    Asked by deriving the public key with an empty passphrase. That reads the
    file, but the only thing it can ever emit is the public half, and the answer
    is not obtainable any other way without parsing key material directly.
    """
    result = run(["ssh-keygen", "-y", "-P", "", "-f", str(private)], timeout=10.0)
    if result.returncode == 0:
        return False
    detail = (result.stderr or "").lower()
    if "incorrect passphrase" in detail or "load failed" in detail:
        return True
    return None


def keys(agent: dict) -> list[dict]:
    if not SSH_DIR.is_dir():
        return []

    held = set(agent.get("fingerprints") or [])
    found = []
    for public in sorted(SSH_DIR.glob("*.pub")):
        private = public.with_suffix("")
        described = run(["ssh-keygen", "-l", "-f", str(public)], timeout=10.0)
        if described.returncode != 0:
            continue
        parts = (described.stdout or "").split()
        if len(parts) < 3:
            continue
        bits, fingerprint = parts[0], parts[1]
        kind = parts[-1].strip("()")
        comment = " ".join(parts[2:-1]).strip()

        found.append({
            "name": private.name,
            "path": str(private),
            "publicPath": str(public),
            "type": kind,
            "bits": int(bits) if bits.isdigit() else 0,
            "fingerprint": fingerprint,
            "comment": comment if comment != "no" else "",
            "hasPrivate": private.is_file(),
            "encrypted": encrypted(private) if private.is_file() else None,
            "loaded": fingerprint in held,
            # Read so the page can say when a key is readable by other people;
            # a private key must be 0600.
            "mode": oct(private.stat().st_mode & 0o777)[2:] if private.is_file() else "",
        })
    return found


def hosts() -> list[dict]:
    """Hosts in known_hosts, grouped by name.

    A hashed known_hosts cannot be listed -- that is the entire point of hashing
    it -- so that is reported rather than shown as an empty list.
    """
    if not KNOWN_HOSTS.is_file():
        return []

    grouped: dict[str, dict] = {}
    try:
        lines = KNOWN_HOSTS.read_text(encoding="utf-8", errors="replace").splitlines()
    except OSError as error:
        raise BoundaryError("known_hosts could not be read.") from error

    for line in lines:
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split()
        if len(parts) < 3:
            continue
        names, kind = parts[0], parts[1]
        if names.startswith("|1|"):
            entry = grouped.setdefault("", {"host": "", "hashed": True, "types": [], "count": 0})
            entry["count"] += 1
            if kind not in entry["types"]:
                entry["types"].append(kind)
            continue
        for name in names.split(","):
            entry = grouped.setdefault(name, {"host": name, "hashed": False, "types": [], "count": 0})
            entry["count"] += 1
            if kind not in entry["types"]:
                entry["types"].append(kind)

    ordered = [entry for key, entry in sorted(grouped.items()) if key != ""]
    if "" in grouped:
        ordered.append(grouped[""])
    return ordered


def snapshot() -> dict:
    agent = agent_state()
    return {
        "available": SSH_DIR.is_dir(),
        "directory": str(SSH_DIR),
        "agent": agent,
        "keys": keys(agent),
        "hosts": hosts(),
        "error": "",
    }


def resolve_key(path: str) -> Path:
    """A key path, confined to ~/.ssh.

    Resolved and compared against the directory so that a name cannot walk out
    of it, and refused if it is not a file this tool put there.
    """
    candidate = Path(path)
    try:
        resolved = candidate.resolve(strict=True)
    except OSError as error:
        raise BoundaryError("That key no longer exists.") from error
    if resolved.parent != SSH_DIR.resolve(strict=False):
        raise BoundaryError("That key is not in the SSH directory.")
    if not resolved.is_file():
        raise BoundaryError("That is not a key file.")
    return resolved


def agent_add(path: str) -> None:
    key = resolve_key(path)
    environment = agent_environment()
    if not environment.get("SSH_AUTH_SOCK"):
        raise BoundaryError("No SSH agent is running.")
    # No passphrase is supplied here on purpose. An encrypted key makes ssh-add
    # prompt through the system's askpass, which is the right place for it.
    result = run(["ssh-add", str(key)], timeout=120.0, env=environment)
    if result.returncode != 0:
        detail = (result.stderr or "").strip().splitlines()
        raise BoundaryError(detail[-1] if detail else "That key could not be added.")


def agent_remove(path: str) -> None:
    key = resolve_key(path)
    environment = agent_environment()
    if not environment.get("SSH_AUTH_SOCK"):
        raise BoundaryError("No SSH agent is running.")
    if not agent_state().get("durableRemoval", True):
        raise BoundaryError(
            "This desktop's agent lists every key in ~/.ssh, so removing one "
            "does not stick. Move the key out of ~/.ssh to stop it being offered.")
    result = run(["ssh-add", "-d", str(key)], env=environment)
    if result.returncode != 0:
        detail = (result.stderr or "").strip().splitlines()
        raise BoundaryError(detail[-1] if detail else "That key could not be removed.")


def forget_host(host: str) -> None:
    """Drop a host's keys from known_hosts.

    The reason anyone reaches for this is a host key that changed, which is
    either a rebuilt machine or something worth being alarmed about -- so the
    page says which before offering the button. ssh-keygen -R rewrites the file
    and keeps a .old copy itself.
    """
    if not HOST.match(host or ""):
        raise BoundaryError("That is not a host name.")
    if not KNOWN_HOSTS.is_file():
        raise BoundaryError("There is no known_hosts file.")
    result = run(["ssh-keygen", "-R", host, "-f", str(KNOWN_HOSTS)], timeout=20.0)
    if result.returncode != 0:
        detail = (result.stderr or "").strip().splitlines()
        raise BoundaryError(detail[-1] if detail else "That host could not be removed.")


def main(arguments: list[str]) -> int:
    try:
        if arguments == ["snapshot"]:
            print(json.dumps(snapshot(), separators=(",", ":")))
            return 0

        if len(arguments) == 2 and arguments[0] == "agent-add":
            agent_add(arguments[1])
        elif len(arguments) == 2 and arguments[0] == "agent-remove":
            agent_remove(arguments[1])
        elif len(arguments) == 2 and arguments[0] == "forget-host":
            forget_host(arguments[1])
        else:
            raise BoundaryError(
                "Usage: panama-ssh-keys snapshot | agent-add PATH | agent-remove PATH | "
                "forget-host HOST")
    except BoundaryError as error:
        try:
            state = snapshot()
        except BoundaryError:
            state = {"available": False, "directory": str(SSH_DIR),
                     "agent": {"available": False, "socket": "", "fingerprints": [], "detail": ""},
                     "keys": [], "hosts": []}
        state["error"] = str(error)
        print(json.dumps(state, separators=(",", ":")))
        return 0

    print(json.dumps(snapshot(), separators=(",", ":")))
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))