#!/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.

A passphrase enters this tool in exactly one place -- creating a key -- and it
enters on standard input, which no other process can read. From there it is
typed at ssh-keygen over a pseudo-terminal, the same way a person would type it,
because the two obvious alternatives are both worse: a passphrase in argv is
published to every process on the machine, and a passphrase in a temporary file
is written to disk. It is never logged, never echoed back, and never included in
an error message.

Adding an existing encrypted key to the agent is different: no passphrase is
collected for that at all, because ssh-add prompts through the system's own
askpass, which is where that belongs.

    panama-ssh-keys snapshot
    panama-ssh-keys generate NAME COMMENT   (passphrase on stdin)
    panama-ssh-keys fix-permissions NAME
    panama-ssh-keys agent-add PATH | agent-remove PATH
    panama-ssh-keys forget-host HOST
"""

from __future__ import annotations

import json
import os
import pty
import re
import select
import shutil
import signal
import subprocess
import sys
import time
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}$")

# A key file name, and nothing that could be a path. No slash is in the class,
# so a name cannot describe another directory at all -- the resolve-and-compare
# below is the second lock on the same door rather than the only one.
KEY_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")

# A key comment. Free text, but nothing that could break out of a terminal line
# or be mistaken for one of ssh-keygen's own prompts.
KEY_COMMENT = re.compile(r"^[^\x00-\x1f\x7f]{0,128}$")

# ssh-keygen's own floor. Checked here so the refusal arrives before a terminal
# is opened, rather than as a re-prompt nobody is there to answer.
MINIMUM_PASSPHRASE = 5

# Generating an ed25519 key takes milliseconds. The budget is this large only so
# that a machine starved of entropy fails with a message rather than a hang.
KEYGEN_TIMEOUT_SECONDS = 120.0

# 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 resolve_new_key(name: str) -> Path:
    """Where a key by this name would go, or a refusal.

    Refuses anything that already exists -- both halves, because a stray .pub
    beside no private key still means ssh-keygen would be asked to overwrite,
    and this tool does not overwrite keys. Losing a private key is not
    recoverable and a settings page is the wrong place to learn that.
    """
    if not KEY_NAME.match(name or "") or name in (".", ".."):
        raise BoundaryError(
            "A key name can use letters, numbers, dots, dashes and underscores.")
    if name.endswith(".pub"):
        raise BoundaryError("Name the key itself, not its public half.")

    try:
        SSH_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
        directory = SSH_DIR.resolve(strict=True)
    except OSError as error:
        raise BoundaryError("The SSH directory could not be opened.") from error
    if not directory.is_dir():
        raise BoundaryError("The SSH directory is not a directory.")

    target = directory / name
    if target.parent.resolve(strict=True) != directory:
        raise BoundaryError("That key is not in the SSH directory.")

    public = Path(str(target) + ".pub")
    for candidate in (target, public):
        if candidate.exists() or candidate.is_symlink():
            raise BoundaryError(f"{candidate.name} already exists, so nothing was written.")
    return target


def terminal_environment() -> dict:
    """The environment ssh-keygen must run in to ask its question at the terminal.

    This desktop sets SSH_ASKPASS_REQUIRE=prefer, which makes ssh-keygen open a
    graphical passphrase dialog even when it has a perfectly good terminal in
    front of it -- so the first version of this hung, waiting for a prompt that
    had been drawn on somebody's screen instead. The terminal is supplied
    deliberately here, so the askpass route is switched off just as deliberately.

    LC_ALL is pinned so the prompts read below are the ones OpenSSH ships.
    """
    environment = dict(os.environ)
    environment["SSH_ASKPASS_REQUIRE"] = "never"
    environment["LC_ALL"] = "C"
    for name in ("SSH_ASKPASS", "DISPLAY", "WAYLAND_DISPLAY"):
        environment.pop(name, None)
    return environment


def type_at_keygen(command: list[str], passphrase: str) -> None:
    """Run ssh-keygen on a pseudo-terminal and answer its prompts.

    ssh-keygen reads a passphrase through readpassphrase(), which opens
    /dev/tty: a pipe on standard input is not read at all, which is why this
    needs a terminal rather than a simpler subprocess call. The passphrase is
    written to the terminal's master side, exactly as typing it would, and
    ssh-keygen asks twice, so it is typed twice.

    Nothing about the passphrase is kept. It is not written to disk, does not
    appear in the command, and is scrubbed out of anything reported back in case
    a future ssh-keygen ever echoes it.
    """
    try:
        pid, master = pty.fork()
    except OSError as error:
        raise BoundaryError("A terminal could not be opened for ssh-keygen.") from error

    if pid == 0:
        # The child. Nothing may return from here into the parent's code holding
        # the parent's file descriptors, so a failed exec exits outright.
        try:
            os.execvpe(command[0], command, terminal_environment())
        except OSError:
            pass
        os._exit(127)

    typed = 0
    pending = ""
    transcript = ""
    problem = ""
    deadline = time.monotonic() + KEYGEN_TIMEOUT_SECONDS

    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            problem = "ssh-keygen did not finish."
            break
        try:
            ready, _, _ = select.select([master], [], [], min(remaining, 1.0))
        except OSError:
            break
        if not ready:
            continue
        try:
            chunk = os.read(master, 4096)
        except OSError:
            # EIO on Linux: the child closed the terminal, which is how a pty
            # reports end of output.
            break
        if not chunk:
            break

        text = chunk.decode("utf-8", errors="replace")
        pending += text
        transcript += text
        lowered = pending.lower()

        # Matched on ssh-keygen's whole prompt rather than one word of it: the
        # passphrase prompt quotes the key's path back, and "overwrite" is a
        # perfectly legal key name.
        if "overwrite (y/n)" in lowered:
            # Unreachable in practice -- an existing key is refused before this
            # runs -- but answering anything other than "no" here would destroy
            # a key, so it answers no and stops.
            os.write(master, b"n\n")
            problem = "That key already exists, so nothing was written."
            break
        if "passphrase is too short" in lowered:
            problem = (f"ssh-keygen wants a passphrase of at least "
                       f"{MINIMUM_PASSPHRASE} characters.")
            break
        if "passphrases do not match" in lowered:
            problem = "Those passphrases did not match."
            break
        if typed < 2 and "passphrase" in lowered and pending.rstrip().endswith(":"):
            os.write(master, passphrase.encode("utf-8") + b"\n")
            typed += 1
            pending = ""

    if problem:
        try:
            os.kill(pid, signal.SIGKILL)
        except OSError:
            pass

    try:
        os.close(master)
    except OSError:
        pass
    try:
        _, status = os.waitpid(pid, 0)
    except OSError:
        status = 0

    if problem:
        raise BoundaryError(problem)
    if not (os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0):
        detail = scrubbed(transcript, passphrase)
        last = [line.strip() for line in detail.splitlines() if line.strip()]
        raise BoundaryError(last[-1] if last else "ssh-keygen could not create that key.")


def scrubbed(text: str, secret: str) -> str:
    return text.replace(secret, "********") if secret else text


def generate(name: str, comment: str, passphrase: str, allow_empty: bool) -> None:
    """Create an ed25519 key.

    ed25519 and nothing else: it is the key everything current accepts, the
    choice between it and RSA is not one a settings page should make someone
    make, and offering a size field for a curve that has one size would be
    theatre.
    """
    if not KEY_COMMENT.match(comment or ""):
        raise BoundaryError("A comment cannot contain control characters.")
    if "\n" in passphrase or "\r" in passphrase:
        raise BoundaryError("A passphrase cannot contain a line break.")
    if passphrase == "":
        if not allow_empty:
            raise BoundaryError(
                "A passphrase is required. A key with none is usable by anyone "
                "who reads the file.")
    elif len(passphrase) < MINIMUM_PASSPHRASE:
        raise BoundaryError(
            f"ssh-keygen wants a passphrase of at least {MINIMUM_PASSPHRASE} characters.")

    if shutil.which("ssh-keygen") is None:
        raise BoundaryError("ssh-keygen is not installed.")

    target = resolve_new_key(name)
    command = ["ssh-keygen", "-t", "ed25519", "-f", str(target)]
    if comment:
        command += ["-C", comment]
    type_at_keygen(command, passphrase)

    if not target.is_file() or not Path(str(target) + ".pub").is_file():
        raise BoundaryError("ssh-keygen finished but the key is not there.")


def fix_permissions(name: str) -> None:
    """Make a private key readable only by its owner.

    ssh refuses to use a key other people can read, and says so in a message
    most people meet for the first time at the worst moment. The path is
    resolved and compared against the SSH directory first, so a name that is a
    link to something elsewhere is refused rather than followed -- this changes
    a file's mode, and that is not a thing to do to a file you have not checked.
    """
    if not KEY_NAME.match(name or "") or name in (".", ".."):
        raise BoundaryError("That is not a key name.")
    key = resolve_key(str(SSH_DIR / name))
    try:
        os.chmod(key, 0o600)
    except OSError as error:
        raise BoundaryError("That key's permissions could not be changed.") from error


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 read_passphrase() -> str:
    """The passphrase, from standard input, and only from there.

    One trailing newline is dropped because the caller writes one to end the
    line; anything else is taken literally, including spaces, because a
    passphrase is allowed to end in one.
    """
    try:
        raw = sys.stdin.buffer.read().decode("utf-8")
    except (OSError, UnicodeDecodeError) as error:
        raise BoundaryError("The passphrase could not be read.") from error
    if raw.endswith("\n"):
        raw = raw[:-1]
    if raw.endswith("\r"):
        raw = raw[:-1]
    return raw


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

        if arguments and arguments[0] == "generate":
            # The flag exists so a machine-shaped caller can ask for a key with
            # no passphrase deliberately. Panama's own page never passes it: it
            # requires a passphrase and validates that both fields match.
            allow_empty = "--no-passphrase" in arguments[1:]
            rest = [value for value in arguments[1:] if value != "--no-passphrase"]
            if len(rest) != 2:
                raise BoundaryError("Usage: panama-ssh-keys generate NAME COMMENT")
            generate(rest[0], rest[1], read_passphrase(), allow_empty)
        elif len(arguments) == 2 and arguments[0] == "fix-permissions":
            fix_permissions(arguments[1])
        elif 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 | generate NAME COMMENT | "
                "fix-permissions NAME | 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:]))
