#!/usr/bin/env python3
"""Build a local, public-only Rspamd DKIM cache from a remote SSH host.

The helper never downloads private key bytes. It reads only Rspamd mapping/config
metadata and derives the public key on the remote host with OpenSSL.
"""
from __future__ import annotations

import argparse
import pathlib
import posixpath
import re
import shlex
import subprocess
import sys
from dataclasses import dataclass


@dataclass
class Entry:
    domain: str
    selector: str
    path: str


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--ssh", default="ssh")
    p.add_argument("--host", required=True)
    p.add_argument("--user", required=True)
    p.add_argument("--port", type=int, default=22)
    p.add_argument("--identity", default="")
    p.add_argument("--known-hosts", default="")
    p.add_argument("--signing-config", required=True)
    p.add_argument("--selector-map", required=True)
    p.add_argument("--path-map", required=True)
    p.add_argument("--key-dir", required=True)
    p.add_argument("--default-selector", default="mail")
    p.add_argument("--output-dir", required=True)
    return p.parse_args()


def ssh_base(a: argparse.Namespace) -> list[str]:
    out = [a.ssh, "-T", "-p", str(a.port), "-o", "BatchMode=yes", "-o", "ConnectTimeout=8"]
    if a.identity:
        out += ["-i", a.identity]
    if a.known_hosts:
        out += ["-o", f"UserKnownHostsFile={a.known_hosts}"]
    out.append(f"{a.user}@{a.host}")
    return out


def remote(a: argparse.Namespace, command: str, *, check: bool = False) -> subprocess.CompletedProcess[bytes]:
    return subprocess.run(ssh_base(a) + [command], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=check)


def read_remote_text(a: argparse.Namespace, path: str) -> str:
    cmd = f"test -r {shlex.quote(path)} && cat -- {shlex.quote(path)} || true"
    cp = remote(a, cmd)
    return cp.stdout.decode("utf-8", "replace")


def parse_map(text: str) -> dict[str, str]:
    result: dict[str, str] = {}
    for raw in text.splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        parts = line.split()
        if len(parts) >= 2:
            result[parts[0].lower()] = parts[1]
    return result


def unquote_assignment(line: str, key: str) -> str | None:
    m = re.match(rf"^\s*{re.escape(key)}\s*=\s*\"((?:\\.|[^\"])*)\"\s*;?\s*,?\s*$", line)
    if not m:
        return None
    return bytes(m.group(1), "utf-8").decode("unicode_escape")


def parse_legacy(text: str, default_selector: str) -> list[Entry]:
    entries: list[Entry] = []
    in_domain = False
    depth = 0
    domain: str | None = None
    selector = ""
    path = ""
    for raw in text.splitlines():
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        if not in_domain:
            if line.startswith("domain") and "{" in line:
                in_domain = True
                depth = 1
            continue
        if depth == 1 and domain is None and "{" in line:
            name = line.split("{", 1)[0].strip().rstrip(",").strip('"')
            if name and "=" not in name:
                domain = name.lower()
                selector = ""
                path = ""
        if domain is not None and depth >= 2:
            selector = unquote_assignment(line, "selector") or selector
            path = unquote_assignment(line, "path") or path
        depth += line.count("{") - line.count("}")
        if domain is not None and depth == 1 and "}" in line:
            if path:
                entries.append(Entry(domain, selector or default_selector, path))
            domain = None
            selector = ""
            path = ""
        if depth <= 0:
            break
    return entries


def under_key_dir(path: str, key_dir: str) -> bool:
    normalized = posixpath.normpath(path)
    normalized_root = posixpath.normpath(key_dir)
    if not normalized.startswith("/") or not normalized_root.startswith("/"):
        return False
    p = pathlib.PurePosixPath(normalized)
    root = pathlib.PurePosixPath(normalized_root)
    return p == root or root in p.parents


def remote_sidecar(a: argparse.Namespace, private_path: str) -> str:
    sidecar = str(pathlib.PurePosixPath(private_path).with_suffix(".txt"))
    text = read_remote_text(a, sidecar).strip()
    return text if "DKIM1" in text.upper() and "p=" in text else ""


def remote_public_record(a: argparse.Namespace, private_path: str, domain: str, selector: str) -> str:
    existing = remote_sidecar(a, private_path)
    if existing:
        return existing
    q = shlex.quote(private_path)
    cmd = (
        "pub=$("
        f"(openssl pkey -in {q} -pubout -outform DER 2>/dev/null "
        f"|| openssl rsa -in {q} -pubout -outform DER 2>/dev/null) "
        "| openssl base64 -A 2>/dev/null"
        ") && printf '__PA_DKIM__%s\\n' \"$pub\""
    )
    cp = remote(a, cmd)
    if cp.returncode != 0:
        return ""
    text = cp.stdout.decode("ascii", "ignore")
    marker = "__PA_DKIM__"
    line = next((line for line in text.splitlines() if line.startswith(marker)), "")
    public = line[len(marker):].strip()
    if not public or not re.fullmatch(r"[A-Za-z0-9+/=]+", public):
        return ""
    return f'{selector}._domainkey.{domain} IN TXT "v=DKIM1; k=rsa; p={public}"'


def main() -> int:
    a = parse_args()
    out = pathlib.Path(a.output_dir).resolve()
    keys = out / "keys"
    out.mkdir(parents=True, exist_ok=True)
    keys.mkdir(parents=True, exist_ok=True)

    signing = read_remote_text(a, a.signing_config)
    selectors = parse_map(read_remote_text(a, a.selector_map))
    paths = parse_map(read_remote_text(a, a.path_map))

    for entry in parse_legacy(signing, a.default_selector):
        selectors.setdefault(entry.domain, entry.selector)
        paths.setdefault(entry.domain, entry.path)

    local_selectors: dict[str, str] = {}
    local_paths: dict[str, str] = {}
    synced = 0
    skipped = 0

    for domain in sorted(set(selectors) | set(paths)):
        selector = selectors.get(domain, a.default_selector)
        raw_path = paths.get(domain)
        if not raw_path:
            skipped += 1
            continue
        remote_path = raw_path.replace("$domain", domain).replace("$selector", selector)
        if not under_key_dir(remote_path, a.key_dir):
            print(f"warning: skip {domain}: key path is outside {a.key_dir}", file=sys.stderr)
            skipped += 1
            continue
        record = remote_public_record(a, remote_path, domain, selector)
        if not record:
            print(f"warning: skip {domain}: unable to derive public DKIM record", file=sys.stderr)
            skipped += 1
            continue
        domain_dir = keys / domain
        domain_dir.mkdir(parents=True, exist_ok=True)
        local_private = domain_dir / f"{selector}.private"
        local_private.with_suffix(".txt").write_text(record.strip() + "\n", encoding="utf-8")
        local_selectors[domain] = selector
        local_paths[domain] = str(local_private)
        synced += 1

    (out / "selectors.map").write_text(
        "# Public-only development cache generated over SSH.\n" + "".join(f"{k} {v}\n" for k, v in sorted(local_selectors.items())),
        encoding="utf-8",
    )
    (out / "paths.map").write_text(
        "# Public-only development cache generated over SSH.\n" + "".join(f"{k} {v}\n" for k, v in sorted(local_paths.items())),
        encoding="utf-8",
    )
    (out / "dkim_signing.conf").write_text("# Public-only development cache; no private keys are stored here.\n", encoding="utf-8")
    print(f"Rspamd DKIM public cache: {synced} domain(s), {skipped} skipped")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())