"""direct-ssh channel (PAD-12): OpenSSH from the dispatcher straight to the terminal over Tailscale, no Warpgate.

- Auth: one dedicated agentd key. Its private text arrives in the env var named by `agentd.ssh_key` (default
  AGENTD_SSH_KEY, injected by `agent-access.py exec --secret-env`), is popped from os.environ (children never see
  it) and loaded with `ssh-add -` into a per-process ssh-agent on a socket in a private 0700 temp dir. The key is
  never written to disk, argv, logs or exceptions. No key / agent failure -> the channel is unavailable
  (`Unreached`), so a terminal with `fallback: warpgate-ssh` uses Warpgate instead.
- Host key: pinned per terminal (`ssh.host_key`); a known_hosts file is generated from it (public data) and
  checked strictly. A mismatch raises HostKeyError: never a fallback.
- Reuse: ControlMaster=auto / ControlPersist=60 under `agentd.ssh_dir` (root-only 0700), so the second call of a
  tick (and the next tick within a minute) skips the handshake. The socket name carries a fingerprint of the pinned
  host key: a changed pin never rides an open master that was verified against the old one.
- Payload on stdin: Windows runs a fixed PowerShell bootstrap (-EncodedCommand) that reads the base64 of the real
  script from stdin; POSIX runs a fixed bash loader that reads the script from stdin. One call per operation, no
  argv size limit."""
import atexit
import base64
import hashlib
import os
import re
import shutil
import signal
import subprocess
import tempfile
import time

from . import NULL_LOG, HostKeyError, TransportError, Unreached
from .local import sh_quote

CONNECT_TIMEOUT = 8
CONTROL_PERSIST = 60
SOCKET_MAX = 104          # sun_path; ssh adds ".<16 random chars>" while it sets up the master socket
CTL_NAME = 1 + 9 + 40 + 17   # "/<pin8>-<%C: 40 hex>" + that suffix

# Fixed PowerShell bootstrap: the operation's script arrives on stdin as base64 of its UTF-8 text (ASCII on the
# wire, so the console input code page does not matter).
BOOTSTRAP = ("$ProgressPreference='SilentlyContinue'\n"
             "$ErrorActionPreference='Continue'\n"
             "[Console]::OutputEncoding=[Text.Encoding]::UTF8\n"
             "$s=[Text.Encoding]::UTF8.GetString([Convert]::FromBase64String([Console]::In.ReadToEnd().Trim()))\n"
             "& ([ScriptBlock]::Create($s))\n")


def encoded_command(script):
    """PowerShell -EncodedCommand value: base64 of the UTF-16LE text."""
    return base64.b64encode(script.encode("utf-16-le")).decode("ascii")


PS_COMMAND = "powershell.exe -NoProfile -NonInteractive -EncodedCommand " + encoded_command(BOOTSTRAP)
# POSIX loader: the script is read from stdin into a variable (no MAX_ARG_STRLEN), then runs in a subshell with
# stdin = /dev/null. Its exit 255 is reported as 254: 255 stays ssh's own "no session" code. (The PowerShell
# bootstrap exits 0/1 by $?; none of our PowerShell scripts calls `exit`.)
SH_COMMAND = "bash -c " + sh_quote('exec 3<&0 </dev/null; s=$(cat <&3); exec 3<&-; (eval "$s"); r=$?; '
                                   '[ $r = 255 ] && r=254; exit $r')


def ps_payload(script):
    return (base64.b64encode(script.encode("utf-8")).decode("ascii") + "\n").encode("ascii")


# ---------------------------------------------------------------- failure classification
# A host offering no ssh-ed25519 key at all counts as a host-key failure too (a downgrade is a security signal).
HOSTKEY = re.compile(r"Host key verification failed|REMOTE HOST IDENTIFICATION HAS CHANGED|"
                     r"No \S+ host key is known for|host key for \S+ has changed|no matching host key type", re.I)
# ssh exit 255 with one of these = no remote session ran the command. Only ssh's own messages, anchored at the start
# of a line (stderr also carries the remote command's stderr).
UNREACHED = re.compile(r"^(?:ssh: connect to host |ssh: Could not resolve hostname |banner exchange: |"
                       r"Connection timed out during banner exchange|kex_exchange_identification: |"
                       r"Connection (?:closed|reset) by \S+ port \d+\s*$|"
                       r"\S+: Permission denied \((?:publickey|password|keyboard-interactive|hostbased|gssapi)|"
                       r"Received disconnect from .*Too many authentication failures|"
                       r"sign_and_send_pubkey: |error fetching identities|Error connecting to agent|"
                       r"Unable to negotiate with |Session open refused by peer)", re.I | re.M)


def classify(rc, err, out=""):
    """-> 'ok' (the remote command ran, whatever its exit code), 'hostkey', 'unreached' (certainly not run) or
    'unknown' (ssh exit 255 without a recognisable pre-session error: the command may have run). Any stdout proves
    that a session ran the command, so it is never 'unreached'."""
    if rc != 255:
        return "ok"
    if (out or "").strip():
        return "unknown"
    if HOSTKEY.search(err or ""):
        return "hostkey"
    if UNREACHED.search(err or ""):
        return "unreached"
    return "unknown"


# ---------------------------------------------------------------- per-process key agent
def _die_with_parent():
    """preexec: the agent gets SIGTERM when agentd dies, even by SIGKILL (Linux PR_SET_PDEATHSIG)."""
    try:
        import ctypes
        ctypes.CDLL("libc.so.6", use_errno=True).prctl(1, signal.SIGTERM)
    except Exception:
        pass


def normalise_key(text):
    """Vault may hand the key with literal `\\n` escapes or without the final newline."""
    if "\n" not in text.strip() and "\\n" in text:
        text = text.replace("\\n", "\n")
    return text.strip() + "\n"


_KEYS = {}   # env var name -> key text, taken out of os.environ at CLI start until the agent loads it


def take_key(env_name):
    """Move the key out of os.environ right away (called by the CLI before anything else runs), so no child -- not
    even one started before the first direct call -- can inherit it."""
    if env_name not in _KEYS:
        _KEYS[env_name] = os.environ.pop(env_name, None)
    else:
        os.environ.pop(env_name, None)


def key_present(env_name):
    return bool((_KEYS.get(env_name) or os.environ.get(env_name) or "").strip())


class KeyAgent:
    """ssh-agent of this process holding the agentd key. `error` set = direct-ssh unavailable here."""

    def __init__(self, env_name):
        self.env_name = env_name
        self.dir = self.sock = self.pub = self.proc = None
        self.error = None

    def start(self):
        key = os.environ.pop(self.env_name, None)   # popped even if loading fails: children never inherit it
        key = _KEYS.pop(self.env_name, None) or key
        if not key or not key.strip():
            self.error = "%s is not set (agent-access.py exec --secret-env %s=<vault key>)" % (self.env_name, self.env_name)
            return self
        agent_bin, add_bin = shutil.which("ssh-agent"), shutil.which("ssh-add")
        if not (agent_bin and add_bin and shutil.which("ssh")):
            self.error = "ssh / ssh-agent / ssh-add not installed"
            return self
        env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin"), "LANG": "C"}
        try:
            self.dir = tempfile.mkdtemp(prefix="agentd-agent-")      # 0700
            self.sock = os.path.join(self.dir, "agent.sock")
            atexit.register(self.close)
            self.proc = subprocess.Popen([agent_bin, "-D", "-a", self.sock], stdin=subprocess.DEVNULL,
                                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env,
                                         close_fds=True, preexec_fn=_die_with_parent)
            end = time.monotonic() + 5
            while not os.path.exists(self.sock):
                if self.proc.poll() is not None or time.monotonic() > end:
                    self.error = "ssh-agent did not start"
                    return self
                time.sleep(0.02)
            env["SSH_AUTH_SOCK"] = self.sock
            add = subprocess.run([add_bin, "-q", "-"], input=normalise_key(key).encode("utf-8"), env=env,
                                 stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, timeout=15)
            key = None
            if add.returncode != 0:
                why = add.stderr.decode("utf-8", "replace").strip().splitlines()[:1]
                why = why[0][:120] if why and "-----" not in why[0] else ""
                self.error = "ssh-add failed (rc %d) %s" % (add.returncode, why)
                return self
            listed = subprocess.run([add_bin, "-L"], env=env, capture_output=True, timeout=15)
            pub = next((l for l in listed.stdout.decode("ascii", "replace").splitlines() if l.startswith("ssh-")), None)
            if not pub:
                self.error = "ssh-agent holds no key after ssh-add"
                return self
            self.pub = os.path.join(self.dir, "agentd.pub")   # public half: selects the agent key (IdentitiesOnly)
            with open(self.pub, "w", encoding="ascii") as fh:
                fh.write(pub + "\n")
        except (OSError, subprocess.SubprocessError) as exc:
            self.error = "ssh-agent setup failed: %s" % type(exc).__name__
        finally:
            key = None
            if self.error:
                self.close()
        return self

    def close(self):
        if self.proc and self.proc.poll() is None:
            self.proc.terminate()
            try:
                self.proc.wait(3)
            except subprocess.TimeoutExpired:
                self.proc.kill()
        if self.dir:
            shutil.rmtree(self.dir, ignore_errors=True)


_AGENTS = {}


def key_agent(cfg, log=NULL_LOG):
    """The process-wide agent for cfg's key env var, started on first use; raises Unreached when unavailable."""
    name = cfg.ssh_key_env
    agent = _AGENTS.get(name)
    if agent is None:
        agent = _AGENTS[name] = KeyAgent(name).start()
    if agent.error:
        log.event("direct-ssh unavailable in this process: %s" % agent.error, once="agent:" + name)
        raise Unreached("direct-ssh unavailable: %s" % agent.error)
    return agent


def control_dir(cfg):
    """agentd.ssh_dir: ControlMaster sockets + generated known_hosts; created 0700, must be ours."""
    d = cfg.ssh_dir
    try:
        d.mkdir(parents=True, exist_ok=True)
        os.chmod(d, 0o700)
        if os.stat(d).st_uid != os.geteuid():
            raise OSError("not owned by this user")
    except OSError as exc:
        raise Unreached("ssh_dir %s unusable: %s" % (d, exc)) from None
    return d


# ---------------------------------------------------------------- channel
class DirectChannel:
    """One terminal's direct OpenSSH path. ps() / sh() mirror the warpgate channel; payloads go on stdin."""
    stdin = True

    def __init__(self, cfg, terminal, log=None, label="direct"):
        self.cfg, self.terminal, self.log, self.via = cfg, terminal, log or NULL_LOG, label
        s = cfg.terminals[terminal]["ssh"]
        self.host, self.user, self.port = s["host"], s["user"], int(s.get("port", 22))
        self.host_key = " ".join(s["host_key"].split()[:2])
        self.pin = hashlib.sha256(self.host_key.encode("ascii")).hexdigest()[:8]
        self.alias = "agentd-" + terminal
        self.hostkey_failed = None
        self.name = "ssh %s@%s" % (self.user, self.host)

    def known_hosts(self, d):
        path = d / ("known_hosts." + self.terminal)
        line = "%s %s\n" % (self.alias, self.host_key)
        try:
            if not path.exists() or path.read_text(encoding="ascii") != line:
                tmp = path.with_name("%s.%d.tmp" % (path.name, os.getpid()))   # per terminal and process
                tmp.write_text(line, encoding="ascii")
                os.replace(tmp, path)
        except OSError as exc:
            raise Unreached("known_hosts %s: %s" % (path, exc)) from None
        return path

    def argv(self, agent, remote_command):
        d = control_dir(self.cfg)
        ctl = str(d / (self.pin + "-%C"))    # pin fingerprint: a new pin never reuses a master verified on the old one
        if len(str(d)) + CTL_NAME > SOCKET_MAX:
            self.log.event("agentd.ssh_dir %s is too long for a control socket; connection reuse off" % d, once="ctl")
            mux = ["-o", "ControlMaster=no", "-o", "ControlPath=none"]
        else:
            mux = ["-o", "ControlMaster=auto", "-o", "ControlPersist=%d" % CONTROL_PERSIST, "-o", "ControlPath=" + ctl]
        return (["ssh", "-F", "/dev/null", "-T", "-p", str(self.port), "-l", self.user,
                 "-o", "BatchMode=yes", "-o", "IdentitiesOnly=yes",
                 "-o", "IdentityAgent=" + agent.sock, "-o", "IdentityFile=" + agent.pub,
                 "-o", "PreferredAuthentications=publickey", "-o", "PasswordAuthentication=no",
                 "-o", "KbdInteractiveAuthentication=no",
                 "-o", "StrictHostKeyChecking=yes", "-o", "UserKnownHostsFile=" + str(self.known_hosts(d)),
                 "-o", "GlobalKnownHostsFile=/dev/null", "-o", "HostKeyAlias=" + self.alias,
                 "-o", "HostKeyAlgorithms=" + self.host_key.split()[0], "-o", "UpdateHostKeys=no",
                 "-o", "ConnectTimeout=%d" % CONNECT_TIMEOUT, "-o", "ServerAliveInterval=15",
                 "-o", "ServerAliveCountMax=3", "-o", "ForwardAgent=no", "-o", "ForwardX11=no",
                 "-o", "ClearAllForwardings=yes", "-o", "LogLevel=ERROR", "-o", "EscapeChar=none"]
                + mux + ["--", self.host, remote_command])

    def run(self, remote_command, payload, timeout, op):
        """-> (rc, stdout, stderr) of a call whose command ran; raises HostKeyError / Unreached / TransportError."""
        if self.hostkey_failed:
            raise HostKeyError(self.hostkey_failed)
        agent = key_agent(self.cfg, self.log)
        argv = self.argv(agent, remote_command)
        env = self.cfg.scrubbed_env()
        env.pop("SSH_AUTH_SOCK", None)
        t0 = time.monotonic()
        ms = lambda: (time.monotonic() - t0) * 1000  # noqa: E731
        # output goes to files, not pipes: a ControlPersist master forked by this call must never hold our pipes open
        with tempfile.TemporaryFile() as fo, tempfile.TemporaryFile() as fe:
            try:
                p = subprocess.run(argv, input=payload, stdout=fo, stderr=fe, env=env, timeout=timeout)
            except subprocess.TimeoutExpired:
                self.log.call(self.terminal, self.via, op, ms(), "timeout")
                raise TransportError("%s timeout after %ss" % (self.name, timeout)) from None   # may have run
            except OSError as exc:
                self.log.call(self.terminal, self.via, op, ms(), "launch-failed")
                raise Unreached("%s: ssh launch failed: %s" % (self.name, type(exc).__name__)) from None
            fo.seek(0)
            fe.seek(0)
            dec = lambda b: b.decode("utf-8", "replace").replace("\r", "")  # noqa: E731
            out, err = dec(fo.read()), dec(fe.read())
        kind = classify(p.returncode, err, out)
        self.log.call(self.terminal, self.via, op, ms(), p.returncode, "" if kind == "ok" else kind)
        if kind == "hostkey":
            self.hostkey_failed = ("%s:%d host key verification FAILED (pinned %s...; other key or no such key type "
                                   "offered): not falling back; check terminals.%s.ssh.host_key out of band"
                                   % (self.host, self.port, self.host_key[:40], self.terminal))
            self.log.event("SECURITY %s: %s" % (self.terminal, self.hostkey_failed), once="hostkey:" + self.terminal)
            raise HostKeyError(self.hostkey_failed)
        if kind == "unreached":
            raise Unreached("%s exit 255: %s" % (self.name, err.strip()[-300:]))
        if kind == "unknown":
            raise TransportError("%s exit 255: %s" % (self.name, err.strip()[-300:] or "(no message)"))
        return p.returncode, out, err

    def ps(self, script, timeout, op=None):
        rc, out, err = self.run(PS_COMMAND, ps_payload(script), timeout, op)
        if rc != 0 and not out.strip():
            raise TransportError("%s exit %s: %s" % (self.name, rc, err.strip()[-300:]))
        return out

    def sh(self, script, timeout, op=None):
        return self.run(SH_COMMAND, script.encode("utf-8"), timeout, op)