"""privsep worker 协议与客户端（feat-chv-muse-p6，叶子模块）。

对齐 Muse 官方安全文（security.muse.ai 2026-09-08，见
docs/chv-muse-p6-gap.md）：单元内的连接器 CLI 只解析参数、通过
Unix socket 传递类型化参数与文件描述符；真正业务逻辑由 systemd
沙箱化的 privsep worker 执行；每个 worker 由 cgroup 标识、有明确
凭证白名单（改参数也索取不到别的凭证）。

本模块是协议 + 客户端 + 注册表（只用标准库，不依赖 cloud.*，
客内 worker 守护进程可直接 import 它）：
- `WORKER_SPECS`：本期只认一个 worker（gitea），凭证白名单只含
  GITEA_TOKEN，操作只含 pull/push/status/ping，仓库只限
  ~/workspace 下；
- `WorkerRequest` / `WorkerResponse`：类型化 JSON 行（≤64KB），
  请求里绝不带任何凭证字段——凭证只由 git 经 P3 的 credential
  helper 向 authd 按次取，worker 与客户端都碰不到 token 值；
- `PrivsepClient`：发请求前先按白名单校验（越权请求不出门），
  worker 崩溃/超时只返回类型化失败（WorkerRefusal），不抛原始
  异常、不拖垮调用方；返回的外部数据先过安全扫描再交给调用方，
  block 即阻断（内容不往下流）、flag 即标记复核。
"""

import json
import os
import socket
import time

from cloud.chv.safety import (
    InjectionScanner,
    check_context_input,
    mark_untrusted,
)

# 客内工作区根（worker 的仓库 confinement 基准）与默认 socket。
GUEST_WORKSPACE = "/home/chv/workspace"
DEFAULT_SOCKET_PATH = "/run/chv-gitea/worker.sock"
MAX_MESSAGE_BYTES = 65536
DEFAULT_TIMEOUT_SECONDS = 300


class WorkerRefusal(Exception):
    """privsep 拒绝：kind + 原因，不含任何凭证值。调用方按类型化
    失败处理，不把它当成 agent 崩溃。"""

    def __init__(self, kind: str, message: str) -> None:
        super().__init__("privsep_refused:%s: %s" % (kind, message))
        self.kind = kind


class WorkerSpec:
    """一个 privsep worker 的注册信息：连接器 + 凭证白名单 +
    操作白名单 + socket + systemd 单元。"""

    def __init__(self, connector: str, allowed_credentials: frozenset,
                 operations: frozenset, socket_path: str,
                 unit: str) -> None:
        self.connector = connector
        self.allowed_credentials = frozenset(allowed_credentials)
        self.operations = frozenset(operations)
        self.socket_path = socket_path
        self.unit = unit


WORKER_SPECS: dict[str, WorkerSpec] = {
    "gitea": WorkerSpec(
        connector="gitea",
        allowed_credentials=frozenset({"GITEA_TOKEN"}),
        operations=frozenset({"pull", "push", "status", "ping"}),
        socket_path=DEFAULT_SOCKET_PATH,
        unit="chv-gitea-worker.service",
    ),
}


def spec_for(connector: str) -> WorkerSpec:
    """取某连接器的 worker 注册信息；未知连接器即拒绝（不设默认
    放行，调用方不得以“没注册”绕过白名单）。"""
    spec = WORKER_SPECS.get(connector)
    if spec is None:
        raise WorkerRefusal(
            "unknown_connector",
            "连接器 %s 没有注册的 privsep worker" % connector)
    return spec


def worker_cgroup_for(pid: int) -> str:
    """读 /proc/<pid>/cgroup 的首行（worker 的 cgroup 标识，审计
    用）。读不到即空串，不抛错。"""
    try:
        with open("/proc/%d/cgroup" % pid) as fh:
            first = fh.readline().strip()
        return first
    except OSError:
        return ""


def _repo_within_workspace(repo: str, workspace: str) -> str:
    """仓库路径 confinement：realpath 后必须在 workspace 之下，
    且不能就是 workspace 本身。返回解析后的绝对路径，否则
    WorkerRefusal。"""
    base = os.path.realpath(workspace)
    target = os.path.realpath(repo)
    if os.path.commonpath([base, target]) != base or target == base:
        raise WorkerRefusal(
            "repo_outside_workspace",
            "仓目录只允许 %s 下" % workspace)
    return target


class WorkerRequest:
    """一次类型化 worker 请求：操作 + 连接器 + 仓库 + 请求 id。
    请求里绝无凭证字段（凭证只走 authd 替身，见差距表 §1）。"""

    def __init__(self, op: str, connector: str,
                 repo: str | None, request_id: str) -> None:
        self.op = op
        self.connector = connector
        self.repo = repo
        self.request_id = request_id

    def to_dict(self) -> dict:
        """序列化（JSON 行用）。"""
        return {"op": self.op, "connector": self.connector,
                "repo": self.repo, "request_id": self.request_id}

    @classmethod
    def from_dict(cls: type["WorkerRequest"], data: dict,
                  workspace: str = GUEST_WORKSPACE) -> "WorkerRequest":
        """反序列化 + 校验：未知连接器/操作、仓库越界、请求里带
        凭证字段一律拒绝。"""
        if not isinstance(data, dict):
            raise WorkerRefusal(
                "bad_request", "请求须是 JSON 对象")
        if "credential" in data or "token" in data \
                or "password" in data:
            raise WorkerRefusal(
                "credential_in_request",
                "请求里绝不带凭证字段")
        connector = data.get("connector")
        if not isinstance(connector, str):
            raise WorkerRefusal(
                "bad_request", "connector 缺失或不是字符串")
        spec = spec_for(connector)
        op = data.get("op")
        if not isinstance(op, str):
            raise WorkerRefusal(
                "bad_request", "op 缺失或不是字符串")
        if op not in spec.operations:
            raise WorkerRefusal(
                "operation_not_allowed",
                "操作 %r 不在 %s 的白名单内" % (op, connector))
        repo = data.get("repo")
        if repo is not None and repo != "":
            if not isinstance(repo, str):
                raise WorkerRefusal(
                    "bad_request", "repo 不是字符串")
            repo = _repo_within_workspace(repo, workspace)
        else:
            repo = None
        request_id = data.get("request_id")
        if not isinstance(request_id, str) or not request_id:
            raise WorkerRefusal(
                "bad_request", "request_id 缺失或不是字符串")
        return cls(op, connector, repo, request_id)


class WorkerResponse:
    """一次类型化 worker 响应：ok + git 输出 + 返回码 + 安全扫描
    摘要 + 请求 id。ok 为 False 时 error 说明原因，不含值。"""

    def __init__(self, ok: bool, stdout: str, stderr: str,
                 returncode: int, safety: dict, request_id: str,
                 error: str | None = None) -> None:
        self.ok = ok
        self.stdout = stdout
        self.stderr = stderr
        self.returncode = returncode
        self.safety = dict(safety)
        self.request_id = request_id
        self.error = error

    def to_dict(self) -> dict:
        """序列化（JSON 行用）。"""
        return {"ok": self.ok, "stdout": self.stdout,
                "stderr": self.stderr, "returncode": self.returncode,
                "safety": self.safety, "request_id": self.request_id,
                "error": self.error}

    @classmethod
    def from_dict(cls: type["WorkerResponse"], data: dict
                  ) -> "WorkerResponse":
        """反序列化 + 校验：形状不对即拒绝，不把脏回包当成成功。"""
        if not isinstance(data, dict) or "ok" not in data:
            raise WorkerRefusal(
                "bad_response", "worker 回包格式不对")
        if not isinstance(data["ok"], bool):
            raise WorkerRefusal(
                "bad_response", "worker 回包 ok 不是布尔值")
        request_id = data.get("request_id")
        if not isinstance(request_id, str):
            raise WorkerRefusal(
                "bad_response", "worker 回包 request_id 不对")
        return cls(bool(data["ok"]), str(data.get("stdout") or ""),
                   str(data.get("stderr") or ""),
                   int(data.get("returncode") or 0),
                   dict(data.get("safety") or {}), request_id,
                   data.get("error"))


class PrivsepClient:
    """privsep 客户端（agent 侧薄 CLI 的宿主版）：发请求前按白名单
    校验，worker 崩溃/超时只给类型化失败；返回的外部数据先过安全
    扫描再交给调用方——block 即阻断（内容不往下流，只剩原因），
    flag 即原样返回、附扫描摘要标记复核。"""

    def __init__(self, socket_path: str = DEFAULT_SOCKET_PATH,
                 connector: str = "gitea",
                 timeout: int = DEFAULT_TIMEOUT_SECONDS,
                 scanner: InjectionScanner | None = None,
                 workspace: str = GUEST_WORKSPACE) -> None:
        self.socket_path = socket_path
        self.connector = connector
        self.timeout = timeout
        self.scanner = scanner or InjectionScanner()
        self.workspace = workspace

    def _check_request(self, request: WorkerRequest) -> WorkerSpec:
        """客户端白名单校验：连接器须注册、操作须在白名单内；
        请求的 repo 须在 workspace 下。越权请求不出门。"""
        spec = spec_for(request.connector)
        if request.op not in spec.operations:
            raise WorkerRefusal(
                "operation_not_allowed",
                "操作 %r 不在 %s 的白名单内"
                % (request.op, request.connector))
        if request.repo is not None:
            _repo_within_workspace(request.repo, self.workspace)
        return spec

    def _send(self, payload: dict) -> dict:
        """一请求一响应（JSON 行，≤64KB）。socket 层面的失败
        全部转成 WorkerRefusal，不抛原始异常。"""
        try:
            conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            conn.settimeout(self.timeout)
            conn.connect(self.socket_path)
        except OSError as exc:
            raise WorkerRefusal(
                "worker_unavailable",
                "worker socket 连不上：%s" % exc)
        try:
            conn.sendall((json.dumps(payload) + "\n").encode("utf-8"))
            data = b""
            while len(data) < MAX_MESSAGE_BYTES:
                chunk = conn.recv(4096)
                if not chunk:
                    break
                data += chunk
                if data.endswith(b"\n"):
                    break
        except OSError as exc:
            raise WorkerRefusal(
                "worker_io_failed",
                "worker 通信失败：%s" % exc)
        finally:
            conn.close()
        try:
            parsed = json.loads(data.decode("utf-8"))
        except (ValueError, UnicodeDecodeError):
            raise WorkerRefusal(
                "bad_response", "worker 回包不是合法 JSON")
        if not isinstance(parsed, dict):
            raise WorkerRefusal(
                "bad_response", "worker 回包不是 JSON 对象")
        return parsed

    def call(self, request: WorkerRequest) -> WorkerResponse:
        """发一次类型化请求，返回扫描过的类型化响应。"""
        self._check_request(request)
        raw = self._send(request.to_dict())
        response = WorkerResponse.from_dict(raw)
        if response.request_id != request.request_id:
            raise WorkerRefusal(
                "bad_response",
                "worker 回包 request_id 与请求不一致")
        source = "gitea:%s" % (request.repo or "unknown")
        checked = check_context_input(
            mark_untrusted(response.stdout + "\n" + response.stderr,
                           source),
            scanner=self.scanner)
        worker_verdict = response.safety.get("verdict")
        if worker_verdict in ("block", "flag"):
            effective = worker_verdict
        else:
            effective = checked.verdict
            response.safety = checked.describe()
        if effective == "block":
            rules = response.safety.get("rules", [])
            if not rules:
                rules = [f.rule for f in checked.findings]
            return WorkerResponse(
                False, "", "", response.returncode,
                response.safety, response.request_id,
                "safety_blocked:%s" % ",".join(rules))
        return response

    def ping(self) -> WorkerResponse:
        """worker 存活探针（不跑 git，只走一遍协议与校验）。"""
        request = WorkerRequest(
            "ping", self.connector, None,
            "ping-%d" % int(time.time() * 1000))
        return self.call(request)