#!/usr/bin/env python3
"""
Safe DirtyFrag (CVE-2026-43284 / CVE-2026-43500) mitigation verification script
仅检测，不做任何写入、不触发利用

检测目标：
1. install esp4 /bin/false 是否存在
2. install esp6 /bin/false 是否存在
3. install rxrpc /bin/false 是否存在
4. esp4 / esp6 / rxrpc 当前是否已加载
5. modprobe 是否还能手动加载这些模块
6. /proc/sys/vm/drop_caches 基础检查（仅确认存在）
7. 简单判断当前内核版本

参考：
[DirtyFrag GitHub](https://github.com/V4bel/dirtyfrag)
"""

import os
import subprocess
import platform


TARGET_CONF = "/etc/modprobe.d/dirtyfrag.conf"


# =====================
# Color Helpers
# =====================

def red(msg):
    return f"\033[31m{msg}\033[0m"


def green(msg):
    return f"\033[32m{msg}\033[0m"


def yellow(msg):
    return f"\033[33m{msg}\033[0m"


def blue(msg):
    return f"\033[34m{msg}\033[0m"


def cyan(msg):
    return f"\033[36m{msg}\033[0m"


# =====================
# Command Runner
# =====================

def run(cmd):
    try:
        r = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        return r.returncode, r.stdout.strip(), r.stderr.strip()
    except Exception as e:
        return -1, "", str(e)


# =====================
# Check 1
# =====================

def check_modprobe_config():
    print(blue("[*] Checking modprobe disable rules..."))

    expected_rules = [
        "install esp4 /bin/false",
        "install esp6 /bin/false",
        "install rxrpc /bin/false",
    ]

    if not os.path.exists(TARGET_CONF):
        print(red(f"[-] Config not found: {TARGET_CONF}"))
        return False

    try:
        with open(TARGET_CONF, "r") as f:
            content = f.read()

        all_found = True

        for rule in expected_rules:
            if rule in content:
                print(green(f"[+] Found: {rule}"))
            else:
                print(red(f"[!] Missing: {rule}"))
                all_found = False

        return all_found

    except Exception as e:
        print(red(f"[!] Failed reading config: {e}"))
        return False


# =====================
# Check 2
# =====================

def check_loaded_modules():
    print(blue("\n[*] Checking loaded kernel modules..."))

    targets = [
        "esp4",
        "esp6",
        "rxrpc",
    ]

    loaded_risk = False

    try:
        with open("/proc/modules", "r") as f:
            modules = f.read()

        for mod in targets:
            if mod in modules:
                print(red(f"[!] Loaded (RISK): {mod}"))
                loaded_risk = True
            else:
                print(green(f"[+] Not loaded (GOOD): {mod}"))

        return loaded_risk

    except Exception as e:
        print(red(f"[!] Failed reading /proc/modules: {e}"))
        return None


# =====================
# Check 3
# =====================

def check_modprobe_behavior():
    print(blue("\n[*] Verifying modprobe behavior..."))

    targets = [
        "esp4",
        "esp6",
        "rxrpc",
    ]

    all_blocked = True

    for mod in targets:
        print(cyan(f"\n--- Checking: {mod} ---"))

        code, out, err = run(["modprobe", "-n", "-v", mod])

        print(f"[+] Return code: {code}")

        if out:
            print(cyan("[stdout]"))
            print(out)

        if err:
            print(cyan("[stderr]"))
            print(err)

        if "/bin/false" in out or "/bin/false" in err:
            print(green(f"[+] {mod} blocked successfully"))
        else:
            print(red(f"[!] {mod} may still be loadable"))
            all_blocked = False

    return all_blocked


# =====================
# Check 4
# =====================

def check_drop_caches():
    print(blue("\n[*] Checking drop_caches availability..."))

    path = "/proc/sys/vm/drop_caches"

    if os.path.exists(path):
        print(green("[+] drop_caches exists"))
        print(green("    cache cleanup can be performed if needed"))
        return True
    else:
        print(red("[-] drop_caches not found"))
        return False


# =====================
# Check 5
# =====================

def check_kernel_version():
    print(blue("\n[*] Checking kernel version..."))

    try:
        version = platform.release()
        print(f"[+] Kernel: {version}")

        print(yellow(
            "[*] DirtyFrag affects a wide range of kernels "
            "(roughly 2017+ / many mainstream distros)"
        ))
        print(yellow(
            "[*] Please also verify whether your distro has "
            "released an official patched kernel"
        ))

    except Exception as e:
        print(red(f"[!] Failed reading kernel version: {e}"))


# =====================
# Final Assessment
# =====================

def final_judgement(conf_ok, loaded_risk, modprobe_blocked):
    print(cyan("\n=============================="))
    print(cyan("Final Assessment"))
    print(cyan("=============================="))

    if conf_ok and (loaded_risk is False) and modprobe_blocked:
        print(green("[SAFE] mitigation is likely effective"))
        print(green("       esp4/esp6/rxrpc blocked + unloaded"))

    elif conf_ok and modprobe_blocked:
        print(yellow("[PARTIAL]"))
        print(yellow(
            "       future loading blocked, "
            "but risky modules may still be loaded"
        ))
        print(yellow(
            "       reboot or manual rmmod may still be needed"
        ))

    else:
        print(red("[RISK]"))
        print(red("       mitigation not fully effective"))


# =====================
# Main
# =====================

if __name__ == "__main__":
    print(cyan(
        "=== DirtyFrag Safe Mitigation Verification ===\n"
    ))

    conf_ok = check_modprobe_config()
    loaded_risk = check_loaded_modules()
    modprobe_blocked = check_modprobe_behavior()

    check_drop_caches()
    check_kernel_version()

    final_judgement(
        conf_ok,
        loaded_risk,
        modprobe_blocked
    )