#!/usr/bin/python3
"""Prepare, but never activate, the canonical Secure DNS runtime candidate."""

from __future__ import annotations

import grp
import hashlib
import json
import os
import pwd
import re
import shutil
import stat
import subprocess
import sys
from pathlib import Path
from typing import Callable

from octava_policy_runtime import PolicyRuntimeError, compile_policy_runtime
from octava_runtime_transaction import RuntimeArtifact, RuntimeTransactionError, prepare_release


STATE_DIRECTORY = Path("/var/lib/octava-secure-dns")
POLICY_PATH = STATE_DIRECTORY / "policy.json"
CLIENT_NETWORKS_PATH = STATE_DIRECTORY / "client-networks.json"
RELEASES_DIRECTORY = STATE_DIRECTORY / "runtime" / "releases"
RESOLVER_STATE_DIRECTORY = STATE_DIRECTORY / "resolvers"
TRUST_ANCHOR_SOURCE = Path("/usr/share/dns/root.key")
PLAN_NAME = "plan.json"


class PolicyPrepareError(RuntimeError):
    """Stored state, validator, or prepared release is unsafe."""


def _load_fixed_json(path: Path, *, default=None, require_root: bool = True):
    if not path.exists():
        if default is not None:
            return default
        raise PolicyPrepareError(f"required state is missing: {path.name}")
    info = path.lstat()
    if not stat.S_ISREG(info.st_mode) or path.is_symlink():
        raise PolicyPrepareError(f"unsafe state object: {path.name}")
    if require_root and (info.st_uid != 0 or info.st_mode & 0o022):
        raise PolicyPrepareError(f"unsafe state ownership or mode: {path.name}")
    if info.st_size > 1024 * 1024:
        raise PolicyPrepareError(f"state object is too large: {path.name}")
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise PolicyPrepareError(f"invalid state JSON: {path.name}") from exc


def _run_validator(argv: list[str], runner: Callable = subprocess.run) -> None:
    result = runner(
        argv,
        text=True,
        capture_output=True,
        timeout=20,
        check=False,
        env={"PATH": "/usr/sbin:/usr/bin:/sbin:/bin"},
    )
    if result.returncode != 0:
        raise PolicyPrepareError("runtime configuration syntax validation failed")


def validate_candidate(candidate: Path, *, runner: Callable = subprocess.run) -> None:
    unbound_files = sorted((candidate / "unbound").glob("*.conf"))
    dnsmasq_files = sorted((candidate / "dnsmasq").glob("*.conf"))
    if not unbound_files or not dnsmasq_files:
        raise PolicyPrepareError("runtime candidate is incomplete")
    for path in unbound_files:
        if path.is_symlink() or not path.is_file():
            raise PolicyPrepareError("unsafe Unbound candidate object")
        _run_validator(["/usr/sbin/unbound-checkconf", str(path)], runner)
    for path in dnsmasq_files:
        if path.is_symlink() or not path.is_file():
            raise PolicyPrepareError("unsafe dnsmasq candidate object")
        _run_validator(["/usr/sbin/dnsmasq", "--test", f"--conf-file={path}"], runner)


def ensure_resolver_state(
    plan: dict,
    *,
    state_directory: Path = RESOLVER_STATE_DIRECTORY,
    trust_anchor_source: Path = TRUST_ANCHOR_SOURCE,
) -> None:
    if not trust_anchor_source.is_file() or trust_anchor_source.is_symlink():
        raise PolicyPrepareError("DNSSEC trust anchor source is missing or unsafe")
    unbound = pwd.getpwnam("unbound")
    state_directory.mkdir(parents=True, exist_ok=True)
    if state_directory.is_symlink() or not state_directory.is_dir():
        raise PolicyPrepareError("resolver state directory is unsafe")
    os.chown(state_directory, 0, unbound.pw_gid)
    os.chmod(state_directory, 0o750)
    for resolver in plan["resolvers"]:
        namespace = resolver["cache_namespace"]
        if not isinstance(namespace, str) or not re.fullmatch(r"[A-Za-z0-9_.-]{1,64}", namespace):
            raise PolicyPrepareError("resolver state namespace is unsafe")
        directory = state_directory / namespace
        directory.mkdir(mode=0o770, exist_ok=True)
        if directory.is_symlink() or not directory.is_dir():
            raise PolicyPrepareError("resolver state object is unsafe")
        os.chown(directory, 0, unbound.pw_gid)
        os.chmod(directory, 0o770)
        anchor = directory / "root.key"
        if not anchor.exists():
            descriptor = os.open(anchor, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o640)
            with os.fdopen(descriptor, "wb") as output, trust_anchor_source.open("rb") as source:
                shutil.copyfileobj(source, output)
                output.flush()
                os.fsync(output.fileno())
        if anchor.is_symlink() or not anchor.is_file():
            raise PolicyPrepareError("resolver trust anchor object is unsafe")
        os.chown(anchor, unbound.pw_uid, unbound.pw_gid)
        os.chmod(anchor, 0o640)


def prepare_saved_policy(
    *,
    policy_path: Path = POLICY_PATH,
    client_networks_path: Path = CLIENT_NETWORKS_PATH,
    releases_directory: Path = RELEASES_DIRECTORY,
    runner: Callable = subprocess.run,
    require_root: bool = True,
    prepare_resolver_state: bool | None = None,
) -> dict:
    if require_root and os.geteuid() != 0:
        raise PolicyPrepareError("Secure DNS policy preparer must run as root")
    policy = _load_fixed_json(policy_path, require_root=require_root)
    networks = _load_fixed_json(client_networks_path, default=[], require_root=require_root)
    try:
        compiled = compile_policy_runtime(policy, networks)
    except PolicyRuntimeError as exc:
        raise PolicyPrepareError("saved Secure DNS state was rejected") from exc
    if prepare_resolver_state is None:
        prepare_resolver_state = require_root
    if prepare_resolver_state:
        ensure_resolver_state(compiled["plan"])
    plan_content = json.dumps(compiled["plan"], sort_keys=True, indent=2) + "\n"
    artifacts = [*compiled["artifacts"], RuntimeArtifact(PLAN_NAME, plan_content)]

    def validator(candidate: Path) -> None:
        validate_candidate(candidate, runner=runner)

    try:
        result = prepare_release(
            artifacts,
            releases_directory=releases_directory,
            validator=validator,
        )
    except RuntimeTransactionError as exc:
        raise PolicyPrepareError("Secure DNS runtime candidate preparation failed") from exc
    result.update(
        {
            "prepared": True,
            "applied": False,
            "dns_changed": False,
            "routing_changed": False,
            "dhcp_changed": False,
            "dpi_changed": False,
            "runtime_plan_sha256": compiled["plan"]["runtime_plan_sha256"],
            "requires_explicit_connectivity_approval": True,
        }
    )
    return result


def inspect_saved_policy(
    *,
    policy_path: Path = POLICY_PATH,
    client_networks_path: Path = CLIENT_NETWORKS_PATH,
    require_root: bool = True,
) -> dict:
    policy = _load_fixed_json(policy_path, require_root=require_root)
    networks = _load_fixed_json(client_networks_path, default=[], require_root=require_root)
    compiled = compile_policy_runtime(policy, networks)
    return {
        "runtime_plan_sha256": compiled["plan"]["runtime_plan_sha256"],
        "artifacts": [
            {"path": item.path, "sha256": hashlib.sha256(item.content.encode()).hexdigest()}
            for item in compiled["artifacts"]
        ],
        "prepared": False,
        "applied": False,
        "requires_explicit_connectivity_approval": True,
    }


def _secure_release_ownership(path: Path) -> None:
    unbound_gid = grp.getgrnam("unbound").gr_gid
    for item in [path, *path.rglob("*")]:
        if item.is_symlink():
            raise PolicyPrepareError("prepared release contains a symlink")
        os.chown(item, 0, unbound_gid)
        os.chmod(item, 0o750 if item.is_dir() else 0o640)


def main() -> int:
    if len(sys.argv) != 2 or sys.argv[1] not in {"plan", "prepare"}:
        raise SystemExit("usage: octava-secure-dns-policy-prepare plan|prepare")
    if os.geteuid() != 0:
        raise SystemExit("octava-secure-dns-policy-prepare must run as root")
    try:
        if sys.argv[1] == "plan":
            result = inspect_saved_policy()
        else:
            result = prepare_saved_policy()
            _secure_release_ownership(Path(result["path"]))
        print(json.dumps(result, sort_keys=True))
        return 0
    except (OSError, KeyError, PolicyPrepareError, PolicyRuntimeError) as exc:
        print(str(exc), file=sys.stderr)
        return 1


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