"""污点追踪器（对齐 Muse P5-A，叶子模块，纯标准库）。

Muse 口径（见 docs/chv-muse-p5-gap.md §1 污点出口行）：每个工具
执行进程初始干净，读取用户数据后变“污点”；干净且符合窄范围
自动放行策略并通过出站 URL 其他检查的请求可免打扰放行，污点
或不可验证进程失去自动放行、回到正常审批流。

本模块只是用户态追踪器，不是强制层：
- “读用户数据”事件靠调用方上报（工具执行层在读用户数据前
  调 `mark_tainted`，或 eBPF 通了之后由 LSM hook 回调上报）。
  没接 LSM hook 之前漏报是可能的，本模块不宣称“所有读用户
  数据的行为都被捕获”。
- 污点是单向的：没有清除入口，只能随进程退出消失。
- `source` 只记“哪类用户数据”（见 `_SOURCE_CATEGORIES`），
  绝不记内容；不在分类表里的上报值一律归一化为 `unknown`，
  原始字符串不会落进任何记录。审计记录里只有事件名、pid、
  原因与分类，绝无用户数据内容、无凭证值。
- 未知 pid（未注册）一律不可自动放行：`is_tainted` 对它报
  False（没证据说它污），`auto_allow_eligible` 对它也报
  False（没证据说它干净可信）。
"""

import time
from collections.abc import Callable

# 允许的 source 分类：只记“哪类用户数据”，不记内容。
_SOURCE_CATEGORIES = frozenset({
    "gitea_repo",
    "user_file",
    "user_memory",
    "chat_history",
    "credential_vault",
    "tool_output",
    "unknown",
})

# 审计事件名：沿用 P3 的 `sentinel_*` 契约（见 gap §1 审计行）。
_EVENT_MARKED = "sentinel_taint_marked"
_EVENT_INHERITED = "sentinel_taint_inherited"
_EVENT_REGISTERED = "sentinel_taint_registered"
_EVENT_RELEASED = "sentinel_taint_released"


def _normalize_source(source: str) -> str:
    """source 归一化为分类表里的值，不在表里的记 `unknown`。

    调用方误把内容当 source 传进来时，原始字符串在这里就被
    丢掉，绝不会进入进程记录或审计记录。"""
    value = (source or "").strip().lower()
    if value in _SOURCE_CATEGORIES:
        return value
    return "unknown"


class TaintTracker:
    """进程污点两态追踪器（内存态，进程退出即清）。

    每个进程一条记录：pid、cgroup、连接器、用例、是否污点。
    污点只能由 `mark_tainted` 置位、由 `fork_child` 继承、由
    `process_exit` 随进程消失；没有反向清除入口。
    """

    def __init__(self, clock: Callable[[], float] | None = None
                 ) -> None:
        self._records: dict[int, dict] = {}
        self._audit: list[dict] = []
        self._clock = clock or time.time

    def _emit(self, event: str, pid: int, reason: str,
              source: str = "") -> None:
        """记一条审计（事件名 + pid + 原因 + 分类，无内容）"""
        record = {"event": event, "pid": pid, "reason": reason,
                  "ts": self._clock()}
        if source:
            record["source"] = source
        self._audit.append(record)

    def register(self, pid: int, cgroup: str = "",
                 connector: str = "", use_case: str = "") -> bool:
        """登记一个干净进程。pid 非法或已登记返回 False。"""
        if pid <= 0 or pid in self._records:
            return False
        self._records[pid] = {
            "pid": pid, "cgroup": cgroup,
            "connector": connector, "use_case": use_case,
            "tainted": False,
        }
        self._emit(_EVENT_REGISTERED, pid, "process_registered")
        return True

    def mark_tainted(self, pid: int, source: str) -> bool:
        """上报“读用户数据”事件：该进程变污点，不可逆。

        未登记的 pid 也接受上报（fail-closed：读了用户数据的
        进程先记为污点，不等它补登记）；source 只存分类。"""
        category = _normalize_source(source)
        record = self._records.get(pid)
        if record is None:
            self._records[pid] = {
                "pid": pid, "cgroup": "", "connector": "",
                "use_case": "", "tainted": True,
            }
        else:
            record["tainted"] = True
        reason = "read_user_data:" + category
        self._emit(_EVENT_MARKED, pid, reason, source=category)
        return True

    def fork_child(self, parent_pid: int, child_pid: int) -> bool:
        """子进程继承污点：父污则子污，父干净则子干净。

        父未登记或子 pid 非法返回 False，不建记录；子已登记
        时只改它的污点位，不覆盖它自己的 cgroup/连接器/用例。"""
        if child_pid <= 0 or child_pid == parent_pid:
            return False
        parent = self._records.get(parent_pid)
        if parent is None:
            return False
        child = self._records.get(child_pid)
        if child is None:
            self._records[child_pid] = {
                "pid": child_pid, "cgroup": parent["cgroup"],
                "connector": parent["connector"],
                "use_case": parent["use_case"],
                "tainted": parent["tainted"],
            }
        else:
            child["tainted"] = parent["tainted"]
        if parent["tainted"]:
            self._emit(_EVENT_INHERITED, child_pid,
                       "inherited_from_parent:%d" % parent_pid)
        return True

    def process_exit(self, pid: int) -> bool:
        """进程退出：记录连同污点一起消失。未知 pid 返回 False。"""
        if pid not in self._records:
            return False
        del self._records[pid]
        self._emit(_EVENT_RELEASED, pid, "process_exit")
        return True

    def is_tainted(self, pid: int) -> bool:
        """该进程是否污点。未知 pid 报 False（无证据说它污）。"""
        record = self._records.get(pid)
        return bool(record and record["tainted"])

    def is_registered(self, pid: int) -> bool:
        """该 pid 是否已登记。"""
        return pid in self._records

    def state(self, pid: int) -> dict | None:
        """该进程记录的副本（无记录返回 None）。"""
        record = self._records.get(pid)
        if record is None:
            return None
        return dict(record)

    def auto_allow_eligible(self, pid: int) -> bool:
        """出口自动放行资格：已登记且干净才 True。

        污点进程一律 False；未知 pid（不可验证）也一律 False。
        这只是决策输入，不放行任何请求——放行仍须走 Sentinel
        的出站 URL 其他检查（见 gap §1）。"""
        record = self._records.get(pid)
        if record is None:
            return False
        return not record["tainted"]

    def audit(self) -> list:
        """全部审计记录的副本（事件名 + pid + 原因 + 分类）。"""
        return [dict(record) for record in self._audit]