"""eBPF 污点传播加载层（对齐 Muse P5-A，纯标准库 + 可选工具链）。

Muse 口径（见 docs/chv-muse-p5-gap.md §1 污点出口行）：用 eBPF
cgroup 程序做网络拦截与进程归因、及挂在 LSM hook 上的 eBPF
程序做污点传播。

本模块只含两样东西，绝不冒充已加载：
1. eBPF 程序源码（字符串常量，见 `EBPF_CGROUP_EGRESS_C` 与
   `EBPF_LSM_TAINT_C`）：cgroup egress 拦截 + 进程归因 map +
   LSM hook 污点传播骨架。源码本身不是“已在拦截流量”。
2. 加载器 `EbpfTaintLoader`：能力探测（内核版本、bpffs、
   CAP_BPF/CAP_SYS_ADMIN、bpftool/clang 是否在位）→ 状态
   loaded/degraded/blocked + 原因；加载失败一律明确降级为
   “策略引擎 + 代理层钩子”模式（`FALLBACK_MODE`），绝不把
   degraded 写成 loaded。

加载成功的硬条件（缺一即 blocked/degraded，不重试蒙混）：
- 内核 >= 5.8（cgroup_skb egress 挂载点；LSM hook 另需
  内核开 CONFIG_BPF_LSM，见 LSM 源码头注释）；
- /sys/fs/bpf 存在且可写（pin 落点）；
- 进程有效能力集含 CAP_BPF（bit 39）或 CAP_SYS_ADMIN
  （bit 21）——且 bpf() 系统调用本身不被 seccomp 挡；
- bpftool + clang 在位（编译与加载工具链）。

本机实测结论（2026-11-07，Muse VM）：不行。内核
7.0.0-43-generic、CapEff 含 CAP_BPF 与 CAP_SYS_ADMIN，但
bpf() 与 BPF_MAP_CREATE 均报 EPERM（Seccomp filter 在位，
4 条过滤器），且无 /sys/fs/bpf、无 bpftool/clang。能力位
在位不等于系统调用放行——探测只看 CapEff 会误判，所以
`probe()` 把“缺 bpffs/缺工具链”也记为 blocker，不以
CapEff 单项定生死。
"""

import os
import platform
import shutil
import subprocess
import tempfile
from collections.abc import Callable

# 加载失败后的降级模式：策略引擎（taint.py + sentinel.py 决策）
# + 代理层钩子（Tinyproxy 白名单等现有强制），不走 eBPF。
FALLBACK_MODE = "policy_engine_proxy_hooks"

# cgroup egress 拦截的最低内核版本。
MIN_KERNEL = (5, 8)

# CAP 位号（/proc/self/status 的 CapEff，按 Linux capability.h）。
_CAP_BPF_BIT = 39
_CAP_SYS_ADMIN_BIT = 21

# eBPF cgroup egress 拦截程序源码。
#
# 作用（按 section）：
# - `SEC("cgroup_skb/egress")`：挂在 cgroup 的 egress 路径上，
#   每个出站 skb 都过这里；返回 0 放行、返回非 0 丢弃。污点
#   进程的直连在这里被拦下，而不是等代理层才发现。
# - `taint_map`（BPF_MAP_TYPE_HASH，key=pid u32，value=污点
#   标志 u8）：进程归因 map。用户态追踪器（taint.py）与 LSM
#   程序共用这张表：LSM 侧写污点，cgroup 侧读污点。
# - `events`（BPF_MAP_TYPE_RINGBUF）：拦截/放行事件上报给
#   用户态审计，不含包内容，只含 pid 与 verdict。
# - 归因键：`bpf_get_current_pid_tgid() >> 32` 取当前进程 pid；
#   cgroup 归属由挂载时的 cgroup 决定，不靠包里的地址猜进程。
# 内核要求：>= 5.8，CAP_BPF（或 CAP_SYS_ADMIN），cgroup v2，
# bpffs 可写（pin 落点），clang -target bpf + bpftool 加载。
EBPF_CGROUP_EGRESS_C = r"""
#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>

/* 进程归因 map：pid -> 污点标志（1=污点，0=干净）。LSM 程序写，
   本程序读；用户态也可经 map fd 直接查。 */
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 65536);
    __type(key, __u32);
    __type(value, __u8);
} taint_map SEC(".maps");

/* 事件上报：只含 pid 与 verdict，不含包内容。 */
struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 1 << 16);
} events SEC(".maps");

struct event {
    __u32 pid;
    __u32 verdict; /* 0=放行，1=丢弃 */
};

/* cgroup egress 拦截：污点进程的出站包一律丢弃并上报，干净
   进程放行（放行不等于“自动放行通过”——窄范围自动放行仍由
   Sentinel 策略引擎在 L7 判定，这里只做强制层拦截）。 */
SEC("cgroup_skb/egress")
int egress_filter(struct __sk_buff *skb)
{
    __u32 pid = (__u32)(bpf_get_current_pid_tgid() >> 32);
    __u8 *tainted = bpf_map_lookup_elem(&taint_map, &pid);
    struct event *e;
    if (tainted && *tainted) {
        e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
        if (e) {
            e->pid = pid;
            e->verdict = 1;
            bpf_ringbuf_submit(e, 0);
        }
        return 0; /* 丢弃：skb 不出 cgroup */
    }
    return 1; /* 放行 */
}

char _license[] SEC("license") = "GPL";
"""

# LSM hook 污点传播程序源码（骨架）。
#
# 作用（按 section）：
# - `SEC("lsm/file_open")`：进程打开文件时触发。hook 内只判
#   “打开的是不是用户数据路径”（路径表由用户态经 map 下发，
#   这里不硬编码路径），是则把当前 pid 写进 `taint_map`
#   （与 cgroup 程序共用同一张 pin 住的 map）。
# - 污点传播方向：读用户数据的进程变污点；子进程的污点由
#   用户态 `fork_child` 继承（本骨架不替 fork 背书——fork
#   时子进程尚未读文件，LSM 看不到它，必须由追踪器继承）。
# 内核要求：>= 5.7 且内核开 CONFIG_BPF_LSM（否则 lsm/* 的
# section 挂不上）；CAP_BPF；LSM BPF 程序挂载还受系统 LSM
# 策略约束（与 SELinux/AppArmor 共存时以后者为准）。本骨架
# 只做“读事件 → 写 map”，不做文件内容读取、不记路径内容。
EBPF_LSM_TAINT_C = r"""
#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>

/* 与 cgroup 程序共用同一张 pin 住的 map（pin 路径
   /sys/fs/bpf/chv_taint_map，用户态建表后两程序同引）。 */
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 65536);
    __type(key, __u32);
    __type(value, __u8);
} taint_map SEC(".maps");

/* LSM hook 骨架：文件打开时触发。是否“用户数据路径”由用户
   态下发的路径表判定（本骨架只留判断位，不硬编码路径、不读
   文件内容）；命中则把当前 pid 记为污点。 */
SEC("lsm/file_open")
int taint_on_file_open(struct file *file)
{
    __u32 pid = (__u32)(bpf_get_current_pid_tgid() >> 32);
    __u8 flag = 1;
    /* TODO（真机接线时补）：查用户态下发的路径表，命中才写；
       未命中直接 return 0。骨架阶段不写 map，避免误标。 */
    (void)pid;
    (void)flag;
    return 0;
}

char _license[] SEC("license") = "GPL";
"""


def _parse_kernel(release: str) -> tuple[int, int, int]:
    """`7.0.0-43-generic` 这类 release 串解为三元组。"""
    parts = release.split("-")[0].split(".")
    nums = []
    for part in parts[:3]:
        nums.append(int(part) if part.isdigit() else 0)
    while len(nums) < 3:
        nums.append(0)
    return (nums[0], nums[1], nums[2])


def _cap_bit(cap_eff_hex: str, bit: int) -> bool:
    """CapEff 十六进制串里某一位是否在位。"""
    try:
        value = int(cap_eff_hex.strip(), 16)
    except ValueError:
        return False
    return bool((value >> bit) & 1)


class EbpfTaintLoader:
    """eBPF 加载器：探测 → 加载（或明确降级）→ 状态上报。

    状态机（`load_status()` 读到的 `status`）：
    - `blocked`：探测到硬 blocker（内核太旧 / 无 bpffs /
      无能力位 / 缺工具链），加载不会尝试；
    - `degraded`：未加载成功，工作在 `FALLBACK_MODE`（策略
      引擎 + 代理层钩子），这是失败时的唯一出口；
    - `loaded`：bpftool 加载 + pin 成功，且 pin 文件在位。
    只有 `loaded` 才表示 eBPF 在拦截；其余一律不冒充。
    """

    def __init__(self, bpffs_path: str = "/sys/fs/bpf",
                 status_path: str = "/proc/self/status",
                 which: Callable[[str], str | None] | None = None
                 ) -> None:
        self._bpffs_path = bpffs_path
        self._status_path = status_path
        self._which = which or shutil.which
        self._status: dict = {
            "status": "degraded",
            "mode": FALLBACK_MODE,
            "reason": "not_probed_yet",
            "blockers": [],
        }

    def _read_caps(self) -> tuple[bool, bool]:
        """读 CapEff：返回 (CAP_BPF 在位, CAP_SYS_ADMIN 在位)。"""
        try:
            with open(self._status_path) as fh:
                for line in fh:
                    if line.startswith("CapEff:"):
                        hex_value = line.split(":", 1)[1]
                        return (_cap_bit(hex_value, _CAP_BPF_BIT),
                                _cap_bit(hex_value, _CAP_SYS_ADMIN_BIT))
        except OSError:
            return (False, False)
        return (False, False)

    def _check_tools(self) -> dict:
        """bpftool / clang / llvm-objdump 是否在位。"""
        found = {}
        for tool in ("bpftool", "clang", "llvm-objdump"):
            found[tool] = self._which(tool) is not None
        return found

    def _blockers(self, kernel: tuple[int, int, int],
                  cap_bpf: bool, cap_sys_admin: bool,
                  tools: dict) -> list:
        """按硬条件列 blocker（空表 = 探测关全过，可试加载）。"""
        blockers = []
        if kernel < MIN_KERNEL:
            blockers.append("kernel_too_old:%d.%d.%d" % kernel)
        if not os.path.isdir(self._bpffs_path):
            blockers.append("bpffs_missing:" + self._bpffs_path)
        elif not os.access(self._bpffs_path, os.W_OK):
            blockers.append("bpffs_not_writable:" + self._bpffs_path)
        if not cap_bpf and not cap_sys_admin:
            blockers.append("cap_bpf_and_cap_sys_admin_missing")
        for tool, present in sorted(tools.items()):
            if not present:
                blockers.append("tool_missing:" + tool)
        return blockers

    def probe(self) -> dict:
        """能力探测：只读检查，不加载、不改系统状态。

        返回探测结果副本；`status` 为 blocked（有 blocker）
        或 degraded（探测关全过、但尚未加载——仍在降级模式）。
        注意：CapEff 在位不等于 bpf() 放行（本机实测：能力位
        全在，bpf() 仍 EPERM，见模块头注释），所以探测结果
        从不写“eBPF 可用”，只写“探测关是否全过”。"""
        kernel = _parse_kernel(platform.uname().release)
        cap_bpf, cap_sys_admin = self._read_caps()
        tools = self._check_tools()
        bpffs_exists = os.path.isdir(self._bpffs_path)
        blockers = self._blockers(kernel, cap_bpf, cap_sys_admin,
                                  tools)
        status = "blocked" if blockers else "degraded"
        reason = ("; ".join(blockers) if blockers
                  else "probe_passed_not_loaded")
        self._status = {
            "status": status, "mode": FALLBACK_MODE,
            "reason": reason, "blockers": list(blockers),
            "kernel": "%d.%d.%d" % kernel,
            "kernel_ok": kernel >= MIN_KERNEL,
            "bpffs_exists": bpffs_exists,
            "bpffs_writable": (bpffs_exists and os.access(
                self._bpffs_path, os.W_OK)),
            "cap_bpf": cap_bpf, "cap_sys_admin": cap_sys_admin,
            "tools": dict(tools),
        }
        return dict(self._status)

    def load_status(self) -> dict:
        """当前加载状态上报（未探测过时先跑一次探测）。"""
        if self._status.get("reason") == "not_probed_yet":
            return self.probe()
        return dict(self._status)

    def try_load(self) -> dict:
        """尝试加载：探测关不过即降级，不尝试；过了才编译加载。

        编译/加载任一步失败也降级，`status`  stays degraded，
        `reason` 写清失败步。成功须 bpftool 加载 + pin 文件
        在位双重确认，缺一律不报 loaded。"""
        result = self.probe()
        if result["blockers"]:
            return dict(self._status)
        loaded = self._compile_and_load()
        if loaded:
            self._status = dict(result)
            self._status["status"] = "loaded"
            self._status["mode"] = "ebpf"
            self._status["reason"] = "loaded_and_pinned"
            self._status["blockers"] = []
        else:
            self._status["reason"] = "load_failed_fallback"
        return dict(self._status)

    def _compile_and_load(self) -> bool:
        """clang 编译 + bpftool 加载 + pin（只在探测关全过时跑）。

        任一步失败返回 False（调用方降级）。本机永远走不到
        这里（探测关不过）；真机上失败原因由 bpftool 的
        stderr 决定，这里只记成败、不猜原因。"""
        with tempfile.TemporaryDirectory() as tmp:
            src = os.path.join(tmp, "egress.c")
            obj = os.path.join(tmp, "egress.o")
            pin = os.path.join(self._bpffs_path,
                               "chv_egress_filter")
            with open(src, "w") as fh:
                fh.write(EBPF_CGROUP_EGRESS_C)
            compiled = subprocess.run(
                ["clang", "-target", "bpf", "-O2", "-c", src,
                 "-o", obj],
                capture_output=True, timeout=120)
            if compiled.returncode != 0:
                return False
            loaded = subprocess.run(
                ["bpftool", "prog", "load", obj, pin],
                capture_output=True, timeout=120)
            if loaded.returncode != 0:
                return False
            return os.path.exists(pin)