#!/usr/bin/env python3
"""Own managed MCP authority outside the agent's writable filesystem.

The runtime manager remains unprivileged. Only that exact child process may
publish authority it received through the authenticated control plane. Tool
processes running under the same Unix user cannot publish, replace, or clear it.
"""

import argparse
import fcntl
import json
import os
from pathlib import Path
import pwd
import re
import signal
import socket
import stat
import struct
import subprocess
import sys
import time
import uuid
from urllib.parse import urlsplit

AUTH_ROOT = Path("/run/sellable-managed-auth")
SCHEMA = "sellable-managed-auth/v1"
MAX_REQUEST = 32768
SAFE_ID = re.compile(r"[A-Za-z0-9_-]{1,128}\Z")
RUNTIME_ID = re.compile(r"fly-runtime_[a-f0-9]{32}\Z")
COMMANDS = {
    "ADMIN": "/usr/local/lib/sellable-agent/admin-runtime/admin-runtime.mjs",
    "CUSTOMER": "/usr/local/lib/sellable-agent/fly-customer-image/customer-runtime.mjs",
}


class Rejected(Exception):
    pass


def canonical(value):
    return (json.dumps(value, sort_keys=True, separators=(",", ":")) + "\n").encode()


def secure_directory(path, gid):
    try:
        path.mkdir(mode=0o750)
        os.chown(path, 0, gid)
    except FileExistsError:
        pass
    entry = path.lstat()
    if not stat.S_ISDIR(entry.st_mode) or entry.st_uid != 0 or entry.st_gid != gid or entry.st_mode & 0o027:
        raise Rejected("managed_auth_directory_rejected")


class AuthorityStore:
    def __init__(self, *, gid, runtime_id, runtime_generation, agent_kind, origin, allow_http_loopback=False):
        if os.geteuid() != 0:
            raise Rejected("managed_auth_supervisor_requires_root")
        url = urlsplit(origin)
        test_http = allow_http_loopback and url.scheme == "http" and url.hostname in ("127.0.0.1", "localhost", "host.docker.internal")
        if (not RUNTIME_ID.fullmatch(runtime_id) or type(runtime_generation) is not int
                or runtime_generation < 1 or agent_kind not in COMMANDS
                or (url.scheme != "https" and not test_http) or not url.hostname or url.username or url.password
                or url.path not in ("", "/") or url.query or url.fragment):
            raise Rejected("managed_auth_runtime_identity_rejected")
        parent = AUTH_ROOT.parent.lstat()
        if not stat.S_ISDIR(parent.st_mode) or parent.st_uid != 0 or parent.st_mode & 0o022:
            raise Rejected("managed_auth_parent_rejected")
        self.gid = gid
        self.runtime_id = runtime_id
        self.runtime_generation = runtime_generation
        self.agent_kind = agent_kind
        self.origin = origin.rstrip("/")
        self.bindings = {}
        self.fences = {}
        secure_directory(AUTH_ROOT, gid)

    def validate(self, value):
        fields = {"schemaVersion", "profileId", "runtimeId", "runtimeGeneration", "agentKind", "apiUrl", "token", "agent"}
        if not isinstance(value, dict) or set(value) != fields:
            raise Rejected("managed_auth_binding_rejected")
        agent = value["agent"]
        if (value["schemaVersion"] != SCHEMA or not isinstance(value["profileId"], str)
                or not SAFE_ID.fullmatch(value["profileId"]) or value["runtimeId"] != self.runtime_id
                or type(value["runtimeGeneration"]) is not int or value["runtimeGeneration"] != self.runtime_generation
                or value["agentKind"] != self.agent_kind or value["apiUrl"] != self.origin
                or not isinstance(value["token"], str) or not re.fullmatch(r"sat_[^_\s]+_\S+", value["token"])
                or len(value["token"]) > 4096 or not isinstance(agent, dict)
                or set(agent) != {"workspaceId", "agentId", "generation", "policyHash", "selectedChannelIds"}):
            raise Rejected("managed_auth_binding_rejected")
        channels = agent["selectedChannelIds"]
        if (any(not isinstance(agent[key], str) or not SAFE_ID.fullmatch(agent[key]) for key in ("workspaceId", "agentId"))
                or type(agent["generation"]) is not int or agent["generation"] < 1
                or not isinstance(agent["policyHash"], str) or not re.fullmatch(r"[a-f0-9]{64}", agent["policyHash"])
                or not isinstance(channels, list) or not 1 <= len(channels) <= 100
                or any(not isinstance(channel, str) or not SAFE_ID.fullmatch(channel) for channel in channels)
                or len(set(channels)) != len(channels)):
            raise Rejected("managed_auth_binding_rejected")
        prior = self.fences.get(value["profileId"])
        if prior:
            old_agent = prior["agent"]
            if any(agent[key] != old_agent[key] for key in ("workspaceId", "agentId")):
                raise Rejected("managed_auth_identity_change_rejected")
            if agent["generation"] < old_agent["generation"]:
                raise Rejected("managed_auth_stale_generation")
            if agent["generation"] == old_agent["generation"] and value != prior:
                raise Rejected("managed_auth_generation_conflict")
        return value

    def write(self, binding):
        directory = AUTH_ROOT / binding["profileId"]
        secure_directory(directory, self.gid)
        destination = directory / "identity.json"
        temporary = directory / (".identity-" + uuid.uuid4().hex)
        fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
        try:
            with os.fdopen(fd, "wb") as stream:
                stream.write(canonical(binding))
                stream.flush()
                os.fchown(stream.fileno(), 0, self.gid)
                os.fchmod(stream.fileno(), 0o440)
                os.fsync(stream.fileno())
            os.replace(temporary, destination)
        finally:
            temporary.unlink(missing_ok=True)
        return str(destination)

    def publish(self, binding):
        binding = self.validate(binding)
        path = self.write(binding)
        self.bindings[binding["profileId"]] = binding
        self.fences[binding["profileId"]] = binding
        return {"ok": True, "path": path, "generation": binding["agent"]["generation"]}

    def clear(self, profile_id):
        if not isinstance(profile_id, str) or not SAFE_ID.fullmatch(profile_id):
            raise Rejected("managed_auth_profile_rejected")
        directory = AUTH_ROOT / profile_id
        if directory.exists():
            secure_directory(directory, self.gid)
            (directory / "identity.json").unlink(missing_ok=True)
        self.bindings.pop(profile_id, None)
        return {"ok": True}

    def reconcile(self):
        # In-memory bindings originate only from the authenticated manager.
        # Repair filesystem damage without replaying any integration operation.
        for profile_id, binding in self.bindings.items():
            path = AUTH_ROOT / profile_id / "identity.json"
            try:
                entry = path.lstat()
                valid = (stat.S_ISREG(entry.st_mode) and entry.st_uid == 0
                         and entry.st_gid == self.gid and stat.S_IMODE(entry.st_mode) == 0o440
                         and entry.st_size <= MAX_REQUEST and path.read_bytes() == canonical(binding))
            except OSError:
                valid = False
            if not valid:
                self.write(binding)
                print(json.dumps({"code": "managed_auth_binding_restored", "profileId": profile_id}), file=sys.stderr, flush=True)


def serve(command, *, store, runtime_uid):
    lock_fd = os.open(AUTH_ROOT / ".supervisor.lock", os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
    try:
        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        os.close(lock_fd)
        raise Rejected("managed_auth_supervisor_already_running")
    # /run survives a process restart. Never adopt authority left by a killed
    # predecessor: the new authenticated manager must publish its own binding.
    for directory in AUTH_ROOT.iterdir():
        if SAFE_ID.fullmatch(directory.name) and directory.name != "control.sock":
            secure_directory(directory, store.gid)
            (directory / "identity.json").unlink(missing_ok=True)
    socket_path = AUTH_ROOT / "control.sock"
    if socket_path.exists():
        entry = socket_path.lstat()
        if not stat.S_ISSOCK(entry.st_mode) or entry.st_uid != 0:
            raise Rejected("managed_auth_socket_rejected")
        socket_path.unlink()
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(str(socket_path))
    os.chown(socket_path, 0, store.gid)
    os.chmod(socket_path, 0o660)
    server.listen(8)
    server.settimeout(0.25)
    child = subprocess.Popen(command, stdin=subprocess.DEVNULL, start_new_session=True)
    stopping_at = None

    def stop(signum, _frame):
        nonlocal stopping_at
        if stopping_at is None:
            stopping_at = time.monotonic()
            if child.poll() is None:
                child.send_signal(signum)

    old_handlers = {sig: signal.signal(sig, stop) for sig in (signal.SIGTERM, signal.SIGINT)}
    try:
        while child.poll() is None:
            if stopping_at is not None and time.monotonic() - stopping_at > 15:
                child.kill()
                break
            try:
                connection, _ = server.accept()
            except socket.timeout:
                store.reconcile()
                continue
            with connection:
                connection.settimeout(1)
                try:
                    pid, uid, _gid = struct.unpack("3i", connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
                    if pid != child.pid or uid != runtime_uid:
                        raise Rejected("managed_auth_peer_rejected")
                    data = b""
                    while b"\n" not in data and len(data) <= MAX_REQUEST:
                        chunk = connection.recv(min(4096, MAX_REQUEST + 1 - len(data)))
                        if not chunk:
                            break
                        data += chunk
                    if len(data) > MAX_REQUEST or not data.endswith(b"\n"):
                        raise Rejected("managed_auth_request_rejected")
                    request = json.loads(data)
                    if isinstance(request, dict) and set(request) == {"action", "binding"} and request["action"] == "publish":
                        result = store.publish(request["binding"])
                    elif isinstance(request, dict) and set(request) == {"action", "profileId"} and request["action"] == "clear":
                        result = store.clear(request["profileId"])
                    else:
                        raise Rejected("managed_auth_request_rejected")
                except Rejected as error:
                    result = {"ok": False, "code": str(error)}
                except (OSError, ValueError, TypeError, KeyError):
                    result = {"ok": False, "code": "managed_auth_request_failed"}
                try:
                    connection.sendall(canonical(result))
                except OSError:
                    pass
            store.reconcile()
        return child.wait()
    finally:
        for sig, handler in old_handlers.items():
            signal.signal(sig, handler)
        if child.poll() is None:
            child.terminate()
            try:
                child.wait(timeout=15)
            except subprocess.TimeoutExpired:
                child.kill()
                child.wait()
        for profile_id in list(store.bindings):
            store.clear(profile_id)
        server.close()
        socket_path.unlink(missing_ok=True)
        os.close(lock_fd)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--kind", required=True, choices=list(COMMANDS))
    args = parser.parse_args()
    account = pwd.getpwnam("hermes")
    store = AuthorityStore(gid=account.pw_gid, runtime_id=os.environ.get("SELLABLE_RUNTIME_ID", ""),
                           runtime_generation=int(os.environ.get("SELLABLE_AGENT_RUNTIME_GENERATION", os.environ.get("SELLABLE_RUNTIME_GENERATION", "0"))),
                           agent_kind=args.kind, origin=os.environ.get("SELLABLE_AGENT_CONTROL_ORIGIN", ""),
                           allow_http_loopback=os.environ.get("SELLABLE_AGENT_ALLOW_TEST_OIDC") == "1")
    return serve(["/usr/local/bin/node", COMMANDS[args.kind]], store=store, runtime_uid=account.pw_uid)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (Rejected, ValueError, KeyError):
        print('{"ok":false,"code":"managed_auth_supervisor_failed"}', file=sys.stderr)
        sys.exit(1)