#!/usr/bin/env python3
"""chv-broker：单元外 IPC broker（客机侧、单元外服务）。

对齐 Muse P2（见 docs/chv-muse-p2-gap.md）：单元内外走 Unix socket +
SO_PEERCRED 的最小权限 IPC。broker 只认白名单操作，凭证与浏览器
都在单元外——单元看不见 token，只能请 broker 代办。

校验（双重，缺一不可才放行；cell.pid 不存在时只做 uid 检查并打日志）：
  1. SO_PEERCRED 的 uid 须在单元映射范围 [1000, 1000+65536) 内；
  2. 对端 pid 的 pid ns 须与 /run/chv/cell.pid 记的单元 init 同 ns
     （读不到 cell.pid 时跳过第 2 项并警告；cell.pid 存在但对端 ns
     读不到时拒绝，不降级）。

操作白名单：
  ping   -> {"ok": true}
  sync   {action: pull|push|status, repo?} -> 单元外以 chv 身份跑
            chv-sync（list argv，不进 shell；仓目录只允许
            ~/workspace 下，realpath + commonpath 校验）。
  browser {…} -> 原样转发给 chv-browser.sock（浏览器子代理），
            回包原样返回；浏览器服务不在时报错，不替它跑 playwright。

只用标准库。单连接一请求一响应（JSON 行，≤64KB）；异常只断本
连接，绝不杀死 broker 循环。
"""
import grp
import json
import os
import socket
import struct
import subprocess
import sys

BROKER_SOCK = "/run/chv/broker.sock"
BROWSER_SOCK = "/run/chv/browser.sock"
CELL_PID_FILE = "/run/chv/cell.pid"
UID_BASE = 1000
UID_COUNT = 65536
WORKSPACE = "/home/chv/workspace"
SYNC_USER = "chv"
SYNC_BIN = "/usr/local/bin/chv-sync"
MAX_REQUEST_BYTES = 65536
MAX_OUTPUT_BYTES = 8192

SYNC_ACTIONS = ("pull", "push", "status")


def _read_ns_link(pid, proc_root="/proc"):
    """读 /proc/<pid>/ns/pid 的链接目标；读不到返回 None。"""
    try:
        return os.readlink(os.path.join(proc_root, str(pid), "ns", "pid"))
    except OSError:
        return None


def peer_allowed(peer_uid, peer_pid, cell_pid=None, proc_root="/proc"):
    """对端是否允许用 broker。返回 (allowed, reason)。
    cell_pid 为 None 时只做 uid 检查（broker 打日志警告，调用方负责）。"""
    if not (UID_BASE <= peer_uid < UID_BASE + UID_COUNT):
        return False, "uid %d 不在单元映射范围 [%d, %d)" % (
            peer_uid, UID_BASE, UID_BASE + UID_COUNT)
    if cell_pid is None:
        return True, "uid 在范围内（cell.pid 缺失，未做 pid ns 对照）"
    cell_ns = _read_ns_link(cell_pid, proc_root)
    if cell_ns is None:
        return False, "cell.pid 的 pid ns 读不到，拒绝（不降级）"
    peer_ns = _read_ns_link(peer_pid, proc_root)
    if peer_ns is None:
        return False, "对端 pid ns 读不到，拒绝"
    if peer_ns != cell_ns:
        return False, "对端 pid ns 与单元不一致"
    return True, "ok"


def _read_cell_pid():
    try:
        with open(CELL_PID_FILE) as fh:
            return int(fh.read().strip())
    except (OSError, ValueError):
        return None


def _resolve_repo(repo):
    """仓目录只允许 ~/workspace 下。repo 为空时返回 None（由
    chv-sync 自己取唯一仓）。返回解析后的绝对路径或抛 ValueError。"""
    if not repo:
        return None
    base = os.path.realpath(WORKSPACE)
    target = os.path.realpath(repo)
    if os.path.commonpath([base, target]) != base or target == base:
        raise ValueError("仓目录只允许 %s 下：%s" % (WORKSPACE, repo))
    return target


def _truncate(text, limit=MAX_OUTPUT_BYTES):
    if text is None:
        return ""
    if len(text) > limit:
        return text[:limit] + "\n…（已截断）"
    return text


def _run_sync(action, repo):
    if action not in SYNC_ACTIONS:
        return {"ok": False,
                "error": "sync action 只允许 %s：%s"
                         % ("/".join(SYNC_ACTIONS), action)}
    try:
        repo_path = _resolve_repo(repo)
    except ValueError as e:
        return {"ok": False, "error": str(e)}
    cmd = ["runuser", "-u", SYNC_USER, "--", "env", "HOME=/home/chv",
           SYNC_BIN, action]
    if repo_path:
        cmd.append(repo_path)
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True,
                              timeout=300)
    except subprocess.TimeoutExpired:
        return {"ok": False, "error": "chv-sync 超时（300 秒）"}
    except OSError as e:
        return {"ok": False, "error": "chv-sync 起不来：%s" % e}
    return {"ok": proc.returncode == 0,
            "stdout": _truncate(proc.stdout),
            "stderr": _truncate(proc.stderr),
            "returncode": proc.returncode}


def _forward_browser(req, sock_path=BROWSER_SOCK, timeout=120):
    """转发给浏览器子代理；回包须是带 ok 的 dict，否则报错。"""
    try:
        conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        conn.settimeout(timeout)
        conn.connect(sock_path)
    except OSError as e:
        return {"ok": False,
                "error": "浏览器子代理未就绪：%s" % e}
    try:
        conn.sendall((json.dumps(req, ensure_ascii=False) + "\n").encode())
        data = b""
        while len(data) < MAX_REQUEST_BYTES:
            chunk = conn.recv(4096)
            if not chunk:
                break
            data += chunk
            if data.endswith(b"\n"):
                break
    except OSError as e:
        return {"ok": False, "error": "浏览器子代理通信失败：%s" % e}
    finally:
        conn.close()
    try:
        resp = json.loads(data.decode())
    except (ValueError, UnicodeDecodeError):
        return {"ok": False, "error": "浏览器子代理回包不是 JSON"}
    if not isinstance(resp, dict) or "ok" not in resp:
        return {"ok": False, "error": "浏览器子代理回包格式不对"}
    return resp


def handle_request(req, browser_sock=BROWSER_SOCK):
    """处理一个请求 dict，返回响应 dict（纯逻辑，可单测）。"""
    if not isinstance(req, dict):
        return {"ok": False, "error": "请求须是 JSON 对象"}
    op = req.get("op")
    if op == "ping":
        return {"ok": True, "cell": True}
    if op == "sync":
        return _run_sync(req.get("action"), req.get("repo"))
    if op == "browser":
        return _forward_browser(req, sock_path=browser_sock)
    return {"ok": False, "error": "未知操作：%r（白名单：ping/sync/browser）"
                                  % (op,)}


def _peer_creds(conn):
    """SO_PEERCRED：返回 (pid, uid, gid)。Linux struct ucred = 3 个 int。"""
    raw = conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED,
                          struct.calcsize("3i"))
    pid, uid, gid = struct.unpack("3i", raw)
    return pid, uid, gid


def serve(sock_path=BROKER_SOCK, browser_sock=BROWSER_SOCK):
    """监听 Unix socket，逐连接处理。返回码 0 只能是被信号杀死；
    socket 建不起来即抛错（systemd 会重启服务，不要静默带病跑）。"""
    if os.path.exists(sock_path):
        os.remove(sock_path)
    os.makedirs(os.path.dirname(sock_path), exist_ok=True)
    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(sock_path)
    srv.listen(16)
    # 0660 + 组 chv：单元内 root 在客机侧是 uid 1000（chv），
    # 组 chv 可连；其它用户连不上（再叠加 SO_PEERCRED 校验）。
    os.chmod(sock_path, 0o660)
    try:
        gid = grp.getgrnam("chv").gr_gid
    except KeyError:
        print("chv-broker: 组 chv 不存在，socket 权限置不全，拒绝启动",
              file=sys.stderr)
        raise SystemExit(1)
    os.chown(sock_path, -1, gid)
    print("chv-broker: 监听 %s" % sock_path, file=sys.stderr, flush=True)
    while True:
        try:
            conn, _ = srv.accept()
        except OSError:
            continue
        try:
            pid, uid, _gid = _peer_creds(conn)
            cell_pid = _read_cell_pid()
            allowed, reason = peer_allowed(uid, pid, cell_pid)
            if not allowed:
                print("chv-broker: 拒绝连接 pid=%d uid=%d：%s"
                      % (pid, uid, reason), file=sys.stderr, flush=True)
                conn.close()
                continue
            if cell_pid is None:
                print("chv-broker: 警告：cell.pid 缺失，只做了 uid 检查 "
                      "(pid=%d uid=%d)" % (pid, uid), file=sys.stderr,
                      flush=True)
            conn.settimeout(310)
            data = b""
            while len(data) < MAX_REQUEST_BYTES:
                chunk = conn.recv(4096)
                if not chunk:
                    break
                data += chunk
                if data.endswith(b"\n"):
                    break
            try:
                req = json.loads(data.decode())
            except (ValueError, UnicodeDecodeError):
                resp = {"ok": False, "error": "请求不是合法 JSON"}
            else:
                resp = handle_request(req, browser_sock=browser_sock)
            conn.sendall((json.dumps(resp, ensure_ascii=False)
                          + "\n").encode())
        except OSError as e:
            print("chv-broker: 连接处理异常：%s" % e, file=sys.stderr,
                  flush=True)
        finally:
            try:
                conn.close()
            except OSError:
                pass


def main(argv):
    sock_path = BROKER_SOCK
    browser_sock = BROWSER_SOCK
    i = 0
    while i < len(argv):
        if argv[i] == "--socket" and i + 1 < len(argv):
            sock_path = argv[i + 1]
            i += 2
        elif argv[i] == "--browser-socket" and i + 1 < len(argv):
            browser_sock = argv[i + 1]
            i += 2
        else:
            print("chv-broker: 未知参数：%s" % argv[i], file=sys.stderr)
            return 2
    serve(sock_path, browser_sock)
    return 0


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