#!/usr/bin/env python3
"""One-shot local Owner authentication ingress for AIOS.

This module deliberately does not mint a reusable Owner token.

Authentication is based on a real local login session:
- a root-run one-shot server owns a Unix socket;
- the client must run as the configured Owner OS account;
- the server validates SO_PEERCRED, Linux audit loginuid and the controlling TTY;
- the Owner must interactively confirm the exact task/branch/base request;
- the server constructs a narrow OwnerPrincipal and invokes the existing
  OwnerAuthorizationBroker;
- one successful authorization consumes the service lifetime.

The resulting principal is fixed to one non-wildcard intake scope
(`commit:create` for the current deployment) and cannot carry production,
tag, baseline, main or history-rewrite capability.

This is a trust-root ingress, not a Git writer and not a production executor.
"""
from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
from pathlib import Path
import pwd
import re
import secrets
import socket
import stat
import struct
import sys
import tempfile
import time
from typing import Any, Mapping


BINDING_SCHEMA = "aios.local-owner-auth-ingress.v1"
CLIENT_BEGIN_SCHEMA = "aios.local-owner-auth-client-begin.v1"
CHALLENGE_SCHEMA = "aios.local-owner-auth-challenge.v1"
CLIENT_CONFIRM_SCHEMA = "aios.local-owner-auth-client-confirm.v1"
RESULT_SCHEMA = "aios.local-owner-auth-result.v1"

ACTION = "AUTHORIZE_TASK_WRITER_PROVISIONING"
OID_RE = re.compile(r"^[0-9a-f]{40}$")
IDENTITY_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._@:+/-]{1,254}$")
SCOPE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{1,127}$")
TASK_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{1,255}$")
BRANCH_RE = re.compile(r"^agent/[A-Za-z0-9._/-]{1,220}$")

MAX_FRAME = 64 * 1024
CHALLENGE_TTL_SECONDS = 120

EXPECTED_ALLOWED_OPERATIONS = frozenset(
    {
        "create_worktree",
        "create_target_branch_if_absent",
        "commit_to_target_branch",
        "push_target_branch",
    }
)
EXPECTED_DENIED_OPERATIONS = frozenset(
    {
        "direct_main_push",
        "force_push",
        "tag_create",
        "tag_update",
        "tag_delete",
        "branch_delete",
        "repository_admin",
        "baseline_write",
        "production_deploy",
        "production_marker_write",
        "history_rewrite",
    }
)

SENSITIVE_FRAGMENTS = (
    "authorization_header",
    "bearer",
    "client_secret",
    "password",
    "private_key",
    "raw_secret",
    "raw_token",
    "secret_value",
)


class BootstrapError(RuntimeError):
    def __init__(self, code: str, status: int = 1) -> None:
        super().__init__(code)
        self.code = code
        self.status = status


def canonical_json(value: Any) -> bytes:
    return (
        json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
        + "\n"
    ).encode("utf-8")


def sha256_json(value: Any) -> str:
    return hashlib.sha256(canonical_json(value)).hexdigest()


def atomic_json(path: Path, value: Mapping[str, Any], mode: int = 0o600) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.parent.is_symlink() or path.is_symlink():
        raise BootstrapError("STATE_SYMLINK_REJECTED")
    fd, name = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent))
    try:
        os.fchmod(fd, mode)
        with os.fdopen(fd, "wb", closefd=True) as stream:
            stream.write(canonical_json(dict(value)))
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(name, path)
        os.chmod(path, mode)
        dfd = os.open(path.parent, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
        try:
            os.fsync(dfd)
        finally:
            os.close(dfd)
    finally:
        try:
            os.unlink(name)
        except FileNotFoundError:
            pass


def secure_root_file(path: Path) -> dict[str, Any]:
    if path.is_symlink() or not path.is_file():
        raise BootstrapError("BINDING_FILE_INVALID")
    st = path.lstat()
    if st.st_uid != 0:
        raise BootstrapError("BINDING_FILE_ROOT_OWNER_REQUIRED")
    if stat.S_IMODE(st.st_mode) & 0o077:
        raise BootstrapError("BINDING_FILE_PERMISSIONS_INVALID")
    if st.st_size <= 0 or st.st_size > 1_048_576:
        raise BootstrapError("BINDING_FILE_SIZE_INVALID")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise BootstrapError("BINDING_JSON_INVALID") from exc
    if not isinstance(value, dict):
        raise BootstrapError("BINDING_JSON_INVALID")
    return value


def validate_binding(value: Mapping[str, Any]) -> dict[str, Any]:
    if set(value) != {
        "schema",
        "owner_identity",
        "owner_role",
        "authenticated_scope",
        "request",
    }:
        raise BootstrapError("BINDING_FIELDS_REJECTED")
    if value["schema"] != BINDING_SCHEMA:
        raise BootstrapError("BINDING_SCHEMA_REJECTED")
    identity = value["owner_identity"]
    if not isinstance(identity, str) or not IDENTITY_RE.fullmatch(identity):
        raise BootstrapError("OWNER_IDENTITY_INVALID")
    if value["owner_role"] != "owner":
        raise BootstrapError("OWNER_ROLE_INVALID")
    scope = value["authenticated_scope"]
    if (
        not isinstance(scope, str)
        or not SCOPE_RE.fullmatch(scope)
        or scope == "writer:provision"
        or scope == "*"
        or scope.endswith(":*")
    ):
        raise BootstrapError("OWNER_SCOPE_INVALID")

    request = value["request"]
    if not isinstance(request, dict):
        raise BootstrapError("REQUEST_INVALID")
    expected_fields = {
        "action",
        "task_id",
        "repository",
        "base_ref",
        "base_oid",
        "target_branch",
        "one_time_use",
        "expires_in_seconds",
        "production_change_authorized",
        "tag_change_authorized",
        "baseline_change_authorized",
        "main_write_authorized",
        "allowed_operations",
        "denied_operations",
    }
    if set(request) != expected_fields:
        raise BootstrapError("REQUEST_FIELDS_REJECTED")
    if request["action"] != ACTION:
        raise BootstrapError("REQUEST_ACTION_REJECTED")
    if request["repository"] != "aios/aios":
        raise BootstrapError("REPOSITORY_REJECTED")
    if request["base_ref"] != "refs/heads/main":
        raise BootstrapError("BASE_REF_REJECTED")
    if not isinstance(request["base_oid"], str) or not OID_RE.fullmatch(request["base_oid"]):
        raise BootstrapError("BASE_OID_REJECTED")
    if (
        not isinstance(request["target_branch"], str)
        or not BRANCH_RE.fullmatch(request["target_branch"])
    ):
        raise BootstrapError("TARGET_BRANCH_REJECTED")
    if (
        not isinstance(request["task_id"], str)
        or not TASK_RE.fullmatch(request["task_id"])
    ):
        raise BootstrapError("TASK_ID_REJECTED")
    if request["one_time_use"] is not True:
        raise BootstrapError("ONE_TIME_USE_REQUIRED")
    expiry = request["expires_in_seconds"]
    if (
        not isinstance(expiry, int)
        or isinstance(expiry, bool)
        or not 60 <= expiry <= 900
    ):
        raise BootstrapError("EXPIRY_REJECTED")
    for field in (
        "production_change_authorized",
        "tag_change_authorized",
        "baseline_change_authorized",
        "main_write_authorized",
    ):
        if request[field] is not False:
            raise BootstrapError("CONSEQUENTIAL_CAPABILITY_REJECTED")
    if (
        not isinstance(request["allowed_operations"], list)
        or frozenset(request["allowed_operations"]) != EXPECTED_ALLOWED_OPERATIONS
    ):
        raise BootstrapError("ALLOWED_OPERATIONS_MISMATCH")
    if (
        not isinstance(request["denied_operations"], list)
        or frozenset(request["denied_operations"]) != EXPECTED_DENIED_OPERATIONS
    ):
        raise BootstrapError("DENIED_OPERATIONS_MISMATCH")
    return {
        "schema": BINDING_SCHEMA,
        "owner_identity": identity,
        "owner_role": "owner",
        "authenticated_scope": scope,
        "request": dict(request),
    }


def _status_uids(pid: int) -> tuple[int, int, int, int]:
    try:
        lines = Path(f"/proc/{pid}/status").read_text(encoding="utf-8").splitlines()
    except OSError as exc:
        raise BootstrapError("PEER_STATUS_UNAVAILABLE") from exc
    for line in lines:
        if line.startswith("Uid:"):
            fields = line.split()[1:]
            if len(fields) != 4:
                break
            return tuple(int(item) for item in fields)  # type: ignore[return-value]
    raise BootstrapError("PEER_UID_STATUS_INVALID")


def _loginuid(pid: int) -> int:
    try:
        raw = Path(f"/proc/{pid}/loginuid").read_text(encoding="ascii").strip()
        value = int(raw)
    except (OSError, ValueError) as exc:
        raise BootstrapError("PEER_LOGINUID_UNAVAILABLE") from exc
    if value == 4294967295:
        raise BootstrapError("PEER_LOGIN_SESSION_REQUIRED")
    return value


def _tty_path(pid: int) -> Path:
    try:
        value = os.readlink(f"/proc/{pid}/fd/0")
    except OSError as exc:
        raise BootstrapError("PEER_INTERACTIVE_TTY_REQUIRED") from exc
    if not (value.startswith("/dev/pts/") or value.startswith("/dev/tty")):
        raise BootstrapError("PEER_INTERACTIVE_TTY_REQUIRED")
    path = Path(value)
    try:
        if not stat.S_ISCHR(path.stat().st_mode):
            raise BootstrapError("PEER_INTERACTIVE_TTY_REQUIRED")
    except OSError as exc:
        raise BootstrapError("PEER_INTERACTIVE_TTY_REQUIRED") from exc
    return path


def verify_peer_session(pid: int, uid: int, gid: int, owner_identity: str) -> dict[str, Any]:
    try:
        owner = pwd.getpwnam(owner_identity)
    except KeyError as exc:
        raise BootstrapError("OWNER_OS_ACCOUNT_MISSING") from exc
    if uid != owner.pw_uid:
        raise BootstrapError("PEER_OWNER_UID_MISMATCH")
    uids = _status_uids(pid)
    if any(item != owner.pw_uid for item in uids):
        raise BootstrapError("PEER_UID_TRANSITION_REJECTED")
    loginuid = _loginuid(pid)
    if loginuid != owner.pw_uid:
        # This deliberately rejects sudo/runuser from another authenticated
        # account; the session itself must have logged in as the Owner.
        raise BootstrapError("PEER_LOGINUID_OWNER_MISMATCH")
    tty = _tty_path(pid)
    if tty.stat().st_uid != owner.pw_uid:
        raise BootstrapError("PEER_TTY_OWNER_MISMATCH")
    return {
        "pid": pid,
        "uid": uid,
        "gid": gid,
        "loginuid": loginuid,
        "tty": str(tty),
        "owner_identity": owner_identity,
    }


def read_frame(connection: socket.socket) -> dict[str, Any]:
    data = bytearray()
    while len(data) <= MAX_FRAME:
        block = connection.recv(4096)
        if not block:
            break
        data.extend(block)
        if b"\n" in block:
            break
    if not data or len(data) > MAX_FRAME:
        raise BootstrapError("PROTOCOL_FRAME_INVALID")
    raw = bytes(data).split(b"\n", 1)[0]
    try:
        value = json.loads(raw)
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise BootstrapError("PROTOCOL_JSON_INVALID") from exc
    if not isinstance(value, dict):
        raise BootstrapError("PROTOCOL_JSON_INVALID")
    return value


def write_frame(connection: socket.socket, value: Mapping[str, Any]) -> None:
    connection.sendall(canonical_json(dict(value)))


def safe_result(value: Mapping[str, Any]) -> dict[str, Any]:
    for key in value:
        lowered = str(key).lower()
        if any(fragment in lowered for fragment in SENSITIVE_FRAGMENTS):
            raise BootstrapError("BROKER_RESPONSE_SENSITIVE_DATA_REJECTED")
    allowed = {
        "status",
        "approval_plan_id",
        "grant_id",
        "grant_digest_sha256",
        "branch_oid",
        "branch_ref_created",
        "grant_consumed",
        "owner_authorization_event_id",
        "owner_authorization_reference",
        "owner_authorization_digest_sha256",
        "owner_authorization_consumed",
        "owner_authorization_idempotent",
    }
    return {key: value.get(key) for key in sorted(allowed) if key in value}


def request_summary(binding: Mapping[str, Any]) -> dict[str, Any]:
    request = binding["request"]
    return {
        "owner_identity": binding["owner_identity"],
        "authenticated_scope": binding["authenticated_scope"],
        "task_id": request["task_id"],
        "repository": request["repository"],
        "base_ref": request["base_ref"],
        "base_oid": request["base_oid"],
        "target_branch": request["target_branch"],
        "expires_in_seconds": request["expires_in_seconds"],
        "allowed_operations": sorted(request["allowed_operations"]),
        "denied_operations": sorted(request["denied_operations"]),
        "production_change_authorized": False,
        "tag_change_authorized": False,
        "baseline_change_authorized": False,
        "main_write_authorized": False,
    }


def _ensure_runtime_paths(socket_path: Path, state_root: Path, owner_identity: str) -> None:
    if os.geteuid() != 0:
        raise BootstrapError("SERVER_ROOT_REQUIRED")
    owner = pwd.getpwnam(owner_identity)
    for directory, mode, uid, gid in (
        (socket_path.parent, 0o750, 0, owner.pw_gid),
        (state_root, 0o700, 0, 0),
    ):
        if directory.is_symlink():
            raise BootstrapError("RUNTIME_DIRECTORY_SYMLINK_REJECTED")
        directory.mkdir(parents=True, exist_ok=True)
        os.chown(directory, uid, gid)
        os.chmod(directory, mode)
    if socket_path.exists() or socket_path.is_symlink():
        raise BootstrapError("SOCKET_ALREADY_EXISTS")


def _lock(path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    fd = os.open(path, os.O_RDWR | os.O_CREAT | getattr(os, "O_NOFOLLOW", 0), 0o600)
    os.fchmod(fd, 0o600)
    fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    return fd


def invoke_broker(binding: Mapping[str, Any]) -> dict[str, Any]:
    # Import only on the root server side. The bootstrap package does not carry
    # a second broker/provisioner implementation; it reuses current main.
    from owner_authorization_broker import (  # type: ignore
        OwnerAuthorizationBroker,
        OwnerAuthorizationBrokerConfig,
        OwnerPrincipal,
    )

    config = OwnerAuthorizationBrokerConfig.from_environment()
    identity = str(binding["owner_identity"])
    scope = str(binding["authenticated_scope"])
    if identity not in config.allowed_owner_identities:
        raise BootstrapError("OWNER_IDENTITY_NOT_ALLOWED_BY_LIVE_BROKER")
    if scope not in config.allowed_intake_scopes:
        raise BootstrapError("OWNER_SCOPE_NOT_ALLOWED_BY_LIVE_BROKER")
    if config.max_expiry_seconds < int(binding["request"]["expires_in_seconds"]):
        raise BootstrapError("LIVE_BROKER_EXPIRY_TOO_SMALL")
    principal = OwnerPrincipal(
        identity=identity,
        role="owner",
        authenticated_scope=scope,
        token_scopes=frozenset({scope}),
        credential_id="",
    )
    return OwnerAuthorizationBroker(config=config).authorize_and_provision(
        binding["request"], principal
    )


def serve(binding_path: Path, socket_path: Path, state_root: Path, timeout: int) -> int:
    binding = validate_binding(secure_root_file(binding_path))
    identity = str(binding["owner_identity"])
    _ensure_runtime_paths(socket_path, state_root, identity)
    lock_fd = _lock(state_root / "bootstrap.lock")
    listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        listener.bind(str(socket_path))
        owner = pwd.getpwnam(identity)
        os.chown(socket_path, 0, owner.pw_gid)
        os.chmod(socket_path, 0o660)
        listener.listen(1)
        listener.settimeout(timeout)

        connection, _ = listener.accept()
        with connection:
            raw_peer = connection.getsockopt(
                socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i")
            )
            pid, uid, gid = struct.unpack("3i", raw_peer)
            session = verify_peer_session(pid, uid, gid, identity)

            begin = read_frame(connection)
            if begin != {"schema": CLIENT_BEGIN_SCHEMA, "action": "BEGIN"}:
                raise BootstrapError("CLIENT_BEGIN_REJECTED")

            challenge = secrets.token_hex(12)
            challenge_issued = time.monotonic()
            summary = request_summary(binding)
            write_frame(
                connection,
                {
                    "schema": CHALLENGE_SCHEMA,
                    "challenge": challenge,
                    "expires_in_seconds": CHALLENGE_TTL_SECONDS,
                    "request_digest_sha256": sha256_json(binding["request"]),
                    "summary": summary,
                },
            )

            confirm = read_frame(connection)
            expected_confirmation = f"AUTHORIZE {challenge}"
            if (
                confirm.get("schema") != CLIENT_CONFIRM_SCHEMA
                or confirm.get("challenge") != challenge
                or confirm.get("confirmation") != expected_confirmation
            ):
                raise BootstrapError("OWNER_INTERACTIVE_CONFIRMATION_REJECTED")
            if time.monotonic() - challenge_issued > CHALLENGE_TTL_SECONDS:
                raise BootstrapError("OWNER_CONFIRMATION_EXPIRED")

            result = safe_result(invoke_broker(binding))
            receipt = {
                "schema": "aios.local-owner-auth-ingress-receipt.v1",
                "status": "CONSUMED",
                "owner_identity": identity,
                "owner_role": "owner",
                "authenticated_scope": binding["authenticated_scope"],
                "peer": session,
                "request_digest_sha256": sha256_json(binding["request"]),
                "binding_digest_sha256": sha256_json(binding),
                "challenge_digest_sha256": hashlib.sha256(
                    challenge.encode("ascii")
                ).hexdigest(),
                "broker": result,
                "completed_at_unix": int(time.time()),
            }
            event_id = str(
                result.get("owner_authorization_event_id")
                or hashlib.sha256(
                    canonical_json(
                        {
                            "peer": session,
                            "request": sha256_json(binding["request"]),
                        }
                    )
                ).hexdigest()[:32]
            )
            atomic_json(state_root / "events" / f"{event_id}.json", receipt)
            write_frame(
                connection,
                {
                    "schema": RESULT_SCHEMA,
                    "ok": True,
                    "result": result,
                },
            )
            return 0
    except (BootstrapError, TimeoutError, socket.timeout) as exc:
        code = exc.code if isinstance(exc, BootstrapError) else "OWNER_BOOTSTRAP_TIMEOUT"
        try:
            if "connection" in locals():
                write_frame(
                    connection,
                    {
                        "schema": RESULT_SCHEMA,
                        "ok": False,
                        "error": code,
                    },
                )
        except Exception:
            pass
        atomic_json(
            state_root / "last-failure.json",
            {
                "schema": "aios.local-owner-auth-ingress-failure.v1",
                "status": "FAILED",
                "error": code,
                "binding_digest_sha256": sha256_json(binding),
                "failed_at_unix": int(time.time()),
            },
        )
        return 2
    finally:
        listener.close()
        try:
            socket_path.unlink()
        except FileNotFoundError:
            pass
        try:
            fcntl.flock(lock_fd, fcntl.LOCK_UN)
        finally:
            os.close(lock_fd)


def client(socket_path: Path) -> int:
    # The server performs the authoritative peer/session checks. The client
    # also refuses non-interactive execution to make accidental automation hard.
    if not sys.stdin.isatty() or not sys.stdout.isatty():
        raise BootstrapError("CLIENT_INTERACTIVE_TTY_REQUIRED")
    connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    connection.connect(str(socket_path))
    with connection:
        write_frame(connection, {"schema": CLIENT_BEGIN_SCHEMA, "action": "BEGIN"})
        challenge = read_frame(connection)
        if challenge.get("schema") != CHALLENGE_SCHEMA:
            raise BootstrapError("CHALLENGE_INVALID")
        summary = challenge.get("summary")
        if not isinstance(summary, dict):
            raise BootstrapError("CHALLENGE_INVALID")

        print("\nAIOS one-shot Owner authorization request:")
        print(json.dumps(summary, indent=2, sort_keys=True))
        print("\nThis authorizes only the task writer shown above.")
        phrase = f"AUTHORIZE {challenge['challenge']}"
        typed = input(f"Type exactly:\n{phrase}\n> ").strip()
        if typed != phrase:
            raise BootstrapError("OWNER_CONFIRMATION_MISMATCH")
        write_frame(
            connection,
            {
                "schema": CLIENT_CONFIRM_SCHEMA,
                "challenge": challenge["challenge"],
                "confirmation": typed,
            },
        )
        result = read_frame(connection)
        print(json.dumps(result, indent=2, sort_keys=True))
        return 0 if result.get("ok") is True else 3


def parser() -> argparse.ArgumentParser:
    value = argparse.ArgumentParser()
    mode = value.add_mutually_exclusive_group(required=True)
    mode.add_argument("--serve", action="store_true")
    mode.add_argument("--client", action="store_true")
    value.add_argument("--binding", type=Path)
    value.add_argument(
        "--socket",
        type=Path,
        default=Path("/run/aios-owner-bootstrap/owner.sock"),
    )
    value.add_argument(
        "--state-root",
        type=Path,
        default=Path("/var/lib/aios-gitops/owner-local-auth-ingress"),
    )
    value.add_argument("--timeout", type=int, default=300)
    return value


def main(argv: list[str] | None = None) -> int:
    args = parser().parse_args(argv)
    if args.serve:
        if args.binding is None:
            raise SystemExit("--binding is required with --serve")
        return serve(args.binding, args.socket, args.state_root, args.timeout)
    return client(args.socket)


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except BootstrapError as exc:
        print(json.dumps({"status": "BLOCKED", "error": exc.code}), file=sys.stderr)
        raise SystemExit(exc.status)