#!/usr/bin/python3 -I
"""pi-secret-broker: root broker of pi-sudo-prompt (docs/adr/0002-broker-root-para-segredos-duraveis.md).

Started by `pi-secret-run` through sudo (NOPASSWD rule for this file only):

    pi-secret-broker --label <description> -- <absolute consumer path> [args...]

1. asks the pi session (private socket + HMAC, env kept by sudoers) for a secret, sending the
   real argv; pi checks through /proc that the requester is this root process and shows a modal;
2. creates a pipe as root (0600), which a process of the user cannot reopen through /proc;
3. runs the consumer with the uid/gid of the user that called sudo (never root), with the
   pipe as its stdin, and writes the secret into it;
4. exits with the consumer's exit code.

The secret is never put in an environment variable. Runs as root: keep it small, stdlib only,
and self-contained. test/secret-interop.test.mjs keeps it in lockstep with the pi side.
"""

import hashlib
import hmac
import json
import os
import pwd
import secrets
import signal
import socket
import stat
import sys
import time

PROTOCOL_VERSION = 2
TICKET_DIR = "/run/pi-sudo-prompt"
ENV_SOCKET = "PI_ASKPASS_SOCKET"
ENV_SECRET = "PI_ASKPASS_SECRET"
MAX_LINE = 64 * 1024
CONSUMER_PATH = "/usr/local/bin:/usr/bin:/bin"
# Harmless variables sudo's env_reset keeps; everything else (PI_*, SUDO_*) stays out.
KEEP_ENV = ("LANG", "LANGUAGE", "TERM", "COLORTERM", "NO_COLOR", "TZ", "DISPLAY", "WAYLAND_DISPLAY")


class Refused(Exception):
    pass


def canonical(message):
    """Same bytes as canonical() in protocol.ts (JSON.stringify of sorted [key, value] pairs)."""
    pairs = [[key, message[key]] for key in sorted(message) if key != "signature"]
    return json.dumps(pairs, ensure_ascii=False, separators=(",", ":"))


def _mac(key, message):
    return hmac.new(key.encode(), canonical(message).encode(), hashlib.sha256).hexdigest()


def sign(key, message):
    return {**message, "signature": _mac(key, message)}


def verify(key, message):
    signature = message.get("signature") if isinstance(message, dict) else None
    return isinstance(signature, str) and hmac.compare_digest(_mac(key, message), signature)


def request_digest(message):
    """What the ticket file holds: binds the request to this root process."""
    return hashlib.sha256(canonical(message).encode()).hexdigest()


def parse_args(argv):
    if len(argv) < 4 or argv[0] != "--label" or argv[2] != "--":
        raise Refused("usage: pi-secret-broker --label <description> -- <absolute command> [args...]")
    label, command = argv[1], argv[3:]
    path = command[0]
    if not os.path.isabs(path):
        raise Refused(f"consumer must be an absolute path: {path}")
    try:
        info = os.stat(path)
    except OSError as error:
        raise Refused(f"cannot run {path}: {error.strerror}") from None
    if not stat.S_ISREG(info.st_mode) or not info.st_mode & 0o111:
        raise Refused(f"{path} is not an executable file")
    # Capture the real executable before consent. Never re-open a caller-controlled symlink
    # after the modal, or it could point to a different consumer when the secret is delivered.
    command[0] = os.path.realpath(path)
    return label, command


def caller():
    """The user that called sudo. The consumer runs as this user, never as root."""
    try:
        uid, gid = int(os.environ["SUDO_UID"]), int(os.environ["SUDO_GID"])
    except (KeyError, ValueError):
        raise Refused("must be started by sudo (through pi-secret-run)") from None
    if uid == 0 or gid == 0:
        raise Refused("the consumer never runs as root: call pi-secret-run as a regular user")
    return pwd.getpwuid(uid), uid, gid


def _write_ticket(name, digest):
    try:
        os.mkdir(TICKET_DIR, 0o711)
    except FileExistsError:
        pass
    info = os.lstat(TICKET_DIR)
    if not stat.S_ISDIR(info.st_mode) or info.st_uid != 0 or info.st_mode & 0o022:
        raise Refused(f"{TICKET_DIR} is not a root-owned directory")
    path = os.path.join(TICKET_DIR, name)
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o644)
    try:
        os.write(fd, digest.encode())
    finally:
        os.close(fd)
    return path


def ask_secret(socket_path, key, uid, label, command):
    try:
        if os.stat(socket_path).st_uid != uid:
            raise Refused("the pi socket does not belong to the calling user")
    except OSError:
        raise Refused("no pi session socket (is pi running in the TUI?)") from None

    request = {
        "type": "secret_request",
        "id": secrets.token_hex(16),
        "pid": os.getpid(),
        "version": PROTOCOL_VERSION,
        "argv": command,
        "label": label,
        "ticket": secrets.token_hex(32),
        "ts": time.time_ns() // 1_000_000,
    }
    ticket = _write_ticket(request["ticket"], request_digest(request))
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn:
            conn.connect(socket_path)
            conn.sendall((json.dumps(sign(key, request), ensure_ascii=False) + "\n").encode())
            buffer = b""
            while b"\n" not in buffer and len(buffer) <= MAX_LINE:
                chunk = conn.recv(4096)
                if not chunk:
                    break
                buffer += chunk
    except OSError as error:
        raise Refused(f"cannot reach the pi session: {error.strerror or error}") from None
    finally:
        os.unlink(ticket)

    if b"\n" not in buffer:
        raise Refused("pi closed the connection")
    try:
        answer = json.loads(buffer.split(b"\n", 1)[0])
    except ValueError:
        raise Refused("malformed answer from pi") from None
    if not verify(key, answer) or answer.get("id") != request["id"]:
        raise Refused("invalid answer signature")
    if answer.get("type") == "secret_response" and isinstance(answer.get("secret"), str):
        return bytearray(answer["secret"].encode())
    raise Refused(str(answer.get("reason") or "secret request denied"))


def consumer_env(user):
    env = {key: os.environ[key] for key in KEEP_ENV if key in os.environ}
    env.update({key: value for key, value in os.environ.items() if key.startswith("LC_")})
    env.update(HOME=user.pw_dir, USER=user.pw_name, LOGNAME=user.pw_name, SHELL=user.pw_shell, PATH=CONSUMER_PATH)
    return env


def run_consumer(command, env, user, uid, gid, secret):
    groups = os.getgrouplist(user.pw_name, gid)
    read_end, write_end = os.pipe()  # created by root: owner uid 0, mode 0600
    pid = os.fork()
    if pid == 0:
        try:
            os.close(write_end)
            os.dup2(read_end, 0)
            os.close(read_end)
            os.setgroups(groups)
            os.setresgid(gid, gid, gid)
            os.setresuid(uid, uid, uid)
            if 0 in os.getresuid() or 0 in os.getresgid():
                raise OSError("could not drop root")
            for sig in (signal.SIGPIPE, signal.SIGXFSZ, signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
                signal.signal(sig, signal.SIG_DFL)
            os.execve(command[0], command, env)
        except BaseException as error:  # noqa: BLE001 - must never return into the parent's code
            os.write(2, f"pi-secret-run: cannot run {command[0]}: {error}\n".encode(errors="replace"))
        os._exit(127)

    def forward(sig, _frame):
        try:
            os.kill(pid, sig)
        except ProcessLookupError:
            pass

    for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
        signal.signal(sig, forward)

    os.close(read_end)
    secret.extend(b"\n")
    try:
        with memoryview(secret) as view:
            while view:
                view = view[os.write(write_end, view) :]
    except BrokenPipeError:
        pass  # the consumer did not read it
    finally:
        for index in range(len(secret)):
            secret[index] = 0
        os.close(write_end)

    _, status = os.waitpid(pid, 0)
    code = os.waitstatus_to_exitcode(status)
    return 128 - code if code < 0 else code


def _abort(sig, _frame):
    sys.exit(128 + sig)  # unwinds through `finally`, so the ticket file is removed


def main(argv):
    os.umask(0o022)
    for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
        signal.signal(sig, _abort)
    try:
        if os.geteuid() != 0 or os.getuid() != 0:
            raise Refused("the broker must run as root (pi-secret-run starts it through sudo)")
        user, uid, gid = caller()
        label, command = parse_args(argv)
        socket_path, key = os.environ.get(ENV_SOCKET), os.environ.get(ENV_SECRET)
        if not socket_path or not key:
            raise Refused("no pi TUI session (missing socket; check the sudoers env_keep rule)")
        env = consumer_env(user)
        secret = ask_secret(socket_path, key, uid, label, command)
        return run_consumer(command, env, user, uid, gid, secret)
    except Refused as error:
        print(f"pi-secret-run: refusing: {error}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))