#!/usr/bin/env python3
"""Gate0 P4: exercise Linux SO_PEERCRED and correlate it to /proc."""

from __future__ import annotations

import json
import os
import socket
import stat
import tempfile
from pathlib import Path


def proc_identity(pid: int) -> dict[str, int | str]:
    stat_text = Path(f"/proc/{pid}/stat").read_text()
    close = stat_text.rfind(")")
    fields = stat_text[close + 2 :].split()
    # fields[0] is field 3 (state); ppid is field 4 and starttime is field 22.
    return {
        "pid": pid,
        "ppid": int(fields[1]),
        "starttime_ticks": int(fields[19]),
        "uid": int(Path(f"/proc/{pid}/status").read_text().split("Uid:", 1)[1].split()[0]),
        "exe": os.readlink(f"/proc/{pid}/exe"),
    }


def main() -> None:
    with tempfile.TemporaryDirectory(prefix="gate0-p4-") as tmp:
        root = Path(tmp)
        os.chmod(root, 0o700)
        socket_path = root / "broker.sock"
        server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server.bind(str(socket_path))
        os.chmod(socket_path, 0o600)
        server.listen(1)

        child = os.fork()
        if child == 0:
            client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            client.connect(str(socket_path))
            identity = proc_identity(os.getpid())
            client.sendall((json.dumps(identity, sort_keys=True) + "\n").encode())
            # Keep /proc/<pid> alive until the server has correlated peercred.
            if client.recv(2) != b"OK":
                os._exit(2)
            client.close()
            os._exit(0)

        conn, _ = server.accept()
        raw = conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
        peer_pid = int.from_bytes(raw[0:4], byteorder="little", signed=True)
        peer_uid = int.from_bytes(raw[4:8], byteorder="little", signed=True)
        peer_gid = int.from_bytes(raw[8:12], byteorder="little", signed=True)
        claimed = json.loads(conn.makefile("r", encoding="utf-8").readline())
        observed = proc_identity(peer_pid)
        conn.sendall(b"OK")
        _, status = os.waitpid(child, 0)

        root_mode = stat.S_IMODE(root.stat().st_mode)
        socket_mode = stat.S_IMODE(socket_path.stat().st_mode)
        if not (
            peer_pid == claimed["pid"] == observed["pid"]
            and peer_uid == claimed["uid"] == observed["uid"]
            and claimed["starttime_ticks"] == observed["starttime_ticks"]
            and root_mode == 0o700
            and socket_mode == 0o600
            and os.waitstatus_to_exitcode(status) == 0
        ):
            raise AssertionError("SO_PEERCRED, /proc identity, or socket-mode correlation failed")
        print("machine_assertions=PASS")
        print(f"server_pid={os.getpid()} server_uid={os.getuid()} server_gid={os.getgid()}")
        print(f"socket_path={socket_path}")
        print(f"directory_mode={root_mode:04o} socket_mode={socket_mode:04o}")
        print(f"SO_PEERCRED pid={peer_pid} uid={peer_uid} gid={peer_gid}")
        print("client_claim=" + json.dumps(claimed, sort_keys=True))
        print("proc_observed=" + json.dumps(observed, sort_keys=True))
        print(f"pid_match={peer_pid == claimed['pid'] == observed['pid']}")
        print(f"uid_match={peer_uid == claimed['uid'] == observed['uid']}")
        print(
            "starttime_match="
            + str(claimed["starttime_ticks"] == observed["starttime_ticks"])
        )
        print(f"client_exit_status={os.waitstatus_to_exitcode(status)}")
        print("same_principal_socket=true")
        print(
            "posture=0700 parent + 0600 socket excludes other UIDs, but does not prevent "
            "the same UID from unlinking/rebinding; distinct-principal system service remains "
            "required for a claim stronger than T-C against same-UID counterfeit replacement"
        )

        conn.close()
        server.close()


if __name__ == "__main__":
    main()