安全攻防中的WASM安全威胁建模与防护策略

wasm-sandbox-security-demo/
├── core/
│   ├── __init__.py
│   ├── runtime.py              # 核心安全运行时
│   ├── threats.py              # 威胁模型定义与攻击模块
│   └── defenses.py             # 防护策略实现
├── modules/                    # WASM模块（恶意/良性）
│   ├── benign/
│   │   ├── simple_loop.wat
│   │   └── memory_access.wat
│   └── malicious/
│       ├── cfi_hijack.wat
│       ├── memory_oob.wat
│       └── syscall_abuse.wat
├── config/
│   └── security_policy.yaml    # 安全策略配置文件
├── tools/
│   ├── wat2wasm.py             # 辅助工具：.wat转.wasm
│   └── module_analyzer.py      # 模块静态分析
├── tests/
│   ├── test_runtime_security.py
│   └── test_defenses.py
├── requirements.txt
├── main.py                     # 主演示入口
└── README.md                   # （根据要求，此文件在结构树中不展示）

#!/usr/bin/env python3
"""
核心安全运行时。
提供带有安全检查的WASM模块加载与执行环境。
"""
import logging
from typing import Any, Dict, Optional, Callable
from wasmer import Instance, Module, Store, ImportObject, Memory, Function
from wasmer import Type as WasmType
import wasmer_compiler_cranelift

from .defenses import (
    ControlFlowIntegrityValidator,
    ImportFilter,
    MemoryAccessMonitor,
    SyscallSandbox
)
from .threats import ThreatDetector

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class SecuritySandbox:
    """安全增强的WASM沙箱运行时。"""

    def __init__(self, wasm_bytes: bytes, security_policy: Optional[Dict] = None):
        """
        初始化安全沙箱。
        :param wasm_bytes: 编译好的WASM模块字节码。
        :param security_policy: 安全策略配置字典。
        """
        self.wasm_bytes = wasm_bytes
        self.policy = security_policy or {}
        self.store = Store(wasmer_compiler_cranelift.Compiler)
        self.module = Module(self.store, wasm_bytes)
        self.instance = None
        self.memory = None

        # 初始化防护组件
        self.import_filter = ImportFilter(self.policy.get('allowed_imports', []))
        self.cfi_validator = ControlFlowIntegrityValidator(self.module)
        self.memory_monitor = MemoryAccessMonitor()
        self.syscall_sandbox = SyscallSandbox(self.policy.get('syscall_restrictions', {}))
        self.threat_detector = ThreatDetector()

        # 准备安全的导入对象
        self.import_object = self._create_secure_imports()

    def _create_secure_imports(self) -> ImportObject:
        """创建受安全策略约束的导入对象。"""
        imports = ImportObject()

        # 1. 处理内存导入（如果模块需要）
        # 本项目假设模块创建自己的内存。

        # 2. 处理函数导入：这是攻击面关键点。
        def create_wrapped_import(namespace: str, func_name: str, original_func: Callable) -> Function:
            """为导入函数创建安全包装。"""
            def safe_wrapper(*args):
                # 钩子：调用前检查
                self.threat_detector.log_import_call(namespace, func_name, args)
                if not self.syscall_sandbox.is_allowed(namespace, func_name):
                    logger.error(f"🚨 违规系统调用: {namespace}::{func_name}")
                    raise RuntimeError(f"Syscall {namespace}::{func_name} is forbidden by policy.")

                # 执行原始函数（在真实场景中，可能是宿主提供的功能）
                try:
                    result = original_func(*args)
                    # 钩子：调用后检查
                    return result
                except Exception as e:
                    logger.error(f"导入函数执行失败: {e}")
                    raise
            # 需要将Python函数转换为WASM可调用的Function对象，这里简化处理。
            # 实际使用中，需要根据original_func的签名来创建。
            return Function(self.store, lambda *args: safe_wrapper(*args))

        # 示例：提供一个安全的"env"命名空间导入。
        # 假设模块期望一个 `env::log` 函数。
        def host_log(ptr: int, len: int):
            if self.instance:
                # 使用memory_monitor检查读取是否越界
                if self.memory_monitor.check_read(ptr, len, self.instance):
                    data = self.instance.memory.uint8_view()[ptr:ptr+len].tobytes().decode('utf-8')
                    logger.info(f"📝 WASM日志: {data}")
                else:
                    logger.error("🚨 阻止越界内存读取（在host_log中）")
        # 将host_log包装为WASM函数并加入导入对象
        # 此处省略wasmer具体的ImportObject构建细节，以下为概念代码。
        # imports.register("env", {"log": Function(self.store, host_log)})

        # 注意：实际项目中需要根据模块的导入段(import section)动态构建。
        # 这里返回一个空的导入对象作为示例。
        return imports

    def instantiate(self):
        """实例化WASM模块，应用安全防护。"""
        logger.info("开始实例化WASM模块，应用安全策略...")

        # 前置防护：静态分析（模拟）
        # 在实际中，可以在这里调用静态分析工具分析self.module
        self.cfi_validator.analyze_module()

        # 实例化模块
        try:
            self.instance = Instance(self.module, self.import_object)
        except Exception as e:
            logger.error(f"实例化失败: {e}")
            raise

        # 获取模块内存
        self.memory = self.instance.memory if hasattr(self.instance, 'memory') else None
        if self.memory:
            self.memory_monitor.register_memory(self.memory)

        # 后置防护：动态校验点注入（示例）
        # 我们可以劫持或包装导出函数，加入CFI检查。
        self._wrap_exported_functions()

        logger.info("WASM模块安全实例化完成。")

    def _wrap_exported_functions(self):
        """包装导出函数，插入动态安全检查。"""
        if not self.instance:
            return
        exports = self.instance.exports
        # 遍历所有导出函数（简化示例，仅包装第一个）
        for name, export in exports.__dict__.items():
            if callable(export):
                # 保存原始函数
                original_func = export
                # 创建包装函数
                def make_wrapper(func, func_name):
                    def wrapper(*args, **kwargs):
                        # 动态CFI检查：确保调用目标有效
                        if not self.cfi_validator.validate_call_target(func_name, args):
                            logger.error(f"🚨 CFI违规：非法调用目标 {func_name}")
                            raise RuntimeError("Control Flow Integrity violation")
                        # 内存访问监控（如果函数涉及指针参数）
                        self.memory_monitor.flag_operation('call', func_name, args)
                        # 执行原始函数
                        return func(*args, **kwargs)
                    return wrapper
                # 替换导出（注意：wasmer的导出可能不允许直接替换，此为概念演示）
                # 更实际的做法是在调用入口（如call_export）进行拦截。
                # setattr(exports, name, make_wrapper(original_func, name))
                break # 只演示一个

    def call_export(self, func_name: str, *args) -> Any:
        """调用导出函数，并经过安全检查。"""
        if not self.instance:
            raise RuntimeError("实例未初始化")
        if not hasattr(self.instance.exports, func_name):
            raise ValueError(f"导出函数 '{func_name}' 不存在")

        func = getattr(self.instance.exports, func_name)

        # 调用前安全检查
        self.threat_detector.before_export_call(func_name, args)
        self.memory_monitor.before_call(func_name, args)

        try:
            result = func(*args)
            # 调用后检查
            self.threat_detector.after_export_call(func_name, result)
            return result
        except Exception as e:
            logger.error(f"执行导出函数 {func_name} 时出错: {e}")
            self.threat_detector.on_execution_error(func_name, e)
            raise

    def get_memory_view(self):
        """获取内存视图用于检查。"""
        if self.memory:
            return self.memory.uint8_view()
        return None

#!/usr/bin/env python3
"""
防护策略实现模块。
"""
import struct
from typing import List, Dict, Any, Optional, Set
from wasmer import Module, Memory, FunctionType

class ControlFlowIntegrityValidator:
    """控制流完整性验证器（基础静态分析）。"""

    def __init__(self, module: Module):
        self.module = module
        # 存储合法的函数索引（从类型段和函数段分析得出）
        self.valid_func_indices: Set[int] = set()
        # 间接调用表（Table）的合法目标（从元素段分析得出）
        self.valid_table_targets: Dict[int, Set[int]] = {}

    def analyze_module(self):
        """分析模块，建立合法的控制流图（简化版）。"""
        # 注意：wasmer的Module对象对WASM内部结构的直接暴露有限。
        # 此处为演示逻辑，实际需要解析原始wasm字节码或使用其他库（如wasm-parser）。
        logger.info("[CFI] 开始模块静态分析...")
        # 模拟分析过程：假设我们通过某种方式获取了函数索引范围。
        # 例如，遍历所有可能的导出和已知的内部索引。
        self.valid_func_indices = set(range(0, 100))  # 示例：假设有100个内部函数
        # 记录合法的间接调用目标（例如，从元素段初始化数据推导）
        self.valid_table_targets[0] = {10, 11, 12, 20}  # 表索引0的合法目标
        logger.info(f"[CFI] 分析完成。合法函数索引数: {len(self.valid_func_indices)}")

    def validate_call_target(self, target_name: str, args: tuple) -> bool:
        """验证动态调用目标是否合法。"""
        # 此函数在实际中应在间接调用（call_indirect）发生时被调用。
        # 我们简化处理：检查目标是否在预定义的合法集合中。
        # 这里target_name可能是函数索引或导出名。
        # 示例逻辑：假设第一个参数是函数索引
        if args and isinstance(args[0], int):
            func_idx = args[0]
            if func_idx not in self.valid_func_indices:
                logger.warning(f"[CFI] 检测到非法函数索引调用: {func_idx}")
                return False
        return True

class ImportFilter:
    """导入函数过滤器，基于白名单策略。"""

    def __init__(self, allowed_imports: List[str]):
        # allowed_imports 格式: ["env::log", "wasi_snapshot_preview1::fd_write"]
        self.allowed = set(allowed_imports)

    def is_allowed(self, namespace: str, name: str) -> bool:
        key = f"{namespace}::{name}"
        return key in self.allowed

class MemoryAccessMonitor:
    """内存访问监视器，检测越界访问。"""

    def __init__(self):
        self.memory: Optional[Memory] = None
        self.access_log = []

    def register_memory(self, memory: Memory):
        self.memory = memory

    def check_read(self, offset: int, size: int, instance: Any) -> bool:
        """检查读取操作是否在内存边界内。"""
        if not self.memory:
            return True
        memory_size = self.memory.data_size
        if offset < 0 or offset + size > memory_size:
            logger.error(f"[MemoryMonitor] 越界读取尝试: offset={offset}, size={size}, mem_size={memory_size}")
            return False
        return True

    def check_write(self, offset: int, size: int, instance: Any) -> bool:
        """检查写入操作是否在内存边界内。"""
        # 类似 check_read
        return self.check_read(offset, size, instance)

    def before_call(self, func_name: str, args: tuple):
        """在调用前记录上下文。"""
        self.access_log.append(f"即将调用: {func_name}, 参数: {args}")

class SyscallSandbox:
    """系统调用沙箱，限制WASM模块对宿主功能的访问。"""

    def __init__(self, restrictions: Dict[str, List[str]]):
        # restrictions: {"namespace": ["allowed_func1", "allowed_func2"]}
        self.restrictions = restrictions

    def is_allowed(self, namespace: str, func_name: str) -> bool:
        allowed_funcs = self.restrictions.get(namespace, [])
        return func_name in allowed_funcs

#!/usr/bin/env python3
"""
威胁模型与攻击检测。
"""
import logging
import struct

logger = logging.getLogger(__name__)

class ThreatDetector:
    """威胁检测器，记录和识别可疑行为。"""

    def __init__(self):
        self.suspicious_activities = []

    def log_import_call(self, namespace: str, func_name: str, args: tuple):
        msg = f"[ThreatDetector] 导入调用: {namespace}::{func_name}({args})"
        self.suspicious_activities.append(msg)
        # 示例检测：频繁调用特定敏感函数
        if namespace == "env" and func_name == "get_secret":
            logger.warning(f"⚠️  检测到对敏感函数 get_secret 的调用")

    def before_export_call(self, func_name: str, args: tuple):
        msg = f"[ThreatDetector] 准备调用导出: {func_name}({args})"
        self.suspicious_activities.append(msg)

    def after_export_call(self, func_name: str, result: Any):
        # 检测异常结果，如返回一个可疑的内存地址
        if isinstance(result, int) and result > 0x80000000:  # 假设的高位地址
            logger.warning(f"⚠️  导出函数 {func_name} 返回了可疑指针: {hex(result)}")

    def on_execution_error(self, func_name: str, error: Exception):
        logger.error(f"[ThreatDetector] 执行 {func_name} 时发生错误: {error}")
        # 可以将特定错误（如陷阱、越界）视为攻击指标

    def generate_report(self) -> str:
        return "\n".join(self.suspicious_activities)


# 恶意模块生成助手（用于创建攻击测试用例）
class MaliciousModuleGenerator:
    """生成演示特定威胁的WASM模块（.wat格式）。"""

    @staticmethod
    def generate_cfi_hijack_wat() -> str:
        """生成一个尝试通过篡改间接函数表进行控制流劫持的WASM模块。"""
        wat = """
        (module
          (type $FUNC_TYPE (func (result i32)))
          (table $table 10 funcref)
          (elem $table (i32.const 0) $legit_func $legit_func2) ;; 初始化表，索引0和1是合法的

          ;; 合法函数1
          (func $legit_func (type $FUNC_TYPE)
            i32.const 42
          )
          ;; 合法函数2
          (func $legit_func2 (type $FUNC_TYPE)
            i32.const 100
          )
          ;; 攻击者注入的恶意函数（理论上不应在表中）
          (func $malicious_func (type $FUNC_TYPE)
            i32.const 9999 ;; 恶意返回值
          )
          ;; 导出函数：演示通过越界写入将 $malicious_func 的引用写入表
          (func $attack (export "attack")
            ;; 假设我们通过其他漏洞（如内存越界）能够修改表数据
            ;; 这里我们直接模拟：将 $malicious_func 写入表索引5（一个未初始化的槽位）
            ;; 注意：标准WASM文本格式不能直接"写入"表，这通常通过初始化(elem)或table.set指令完成。
            ;; 以下使用一个伪指令表示攻击意图。
            ;; (table.set $table (i32.const 5) $malicious_func) ; 如果支持
            ;; 然后通过 call_indirect 调用索引5，试图执行恶意代码。
            (call_indirect $table (type $FUNC_TYPE)
              (i32.const 5) ;; 被篡改的索引
            )
            drop ;; 丢弃结果
          )
        )
        """
        return wat

    @staticmethod
    def generate_memory_oob_wat() -> str:
        """生成一个尝试越界内存读写的WASM模块。"""
        wat = """
        (module
          (memory $mem 1) ;; 1页 = 64KB
          (export "memory" (memory $mem))
          (func $read_oob (export "read_oob") (param $addr i32) (result i32)
            ;; 尝试读取任意地址的内存
            (i32.load (local.get $addr))
          )
          (func $write_oob (export "write_oob") (param $addr i32) (param $value i32)
            ;; 尝试向任意地址写入
            (i32.store (local.get $addr) (local.get $value))
          )
        )
        """
        return wat

(module
  (type $t0 (func (result i32)))
  (table $T0 5 funcref)
  ;; 初始化表：索引0和1是合法的函数
  (elem (i32.const 0) $func_legit_1 $func_legit_2)
  ;; 合法函数1
  (func $func_legit_1 (type $t0)
    i32.const 100
    return)
  ;; 合法函数2
  (func $func_legit_2 (type $t0)
    i32.const 200
    return)
  ;; 攻击者想执行的恶意函数（未在elem中初始化，但可通过漏洞引入）
  (func $func_malicious (type $t0)
    i32.const 666  ;; 恶意标记
    return)
  ;; 导出一个"脆弱"函数，它使用一个来自参数的表索引进行间接调用
  (func $vulnerable_indirect_call (export "vulnerable_call") (param $table_idx i32) (result i32)
    ;; 没有验证$table_idx是否在合法范围内（0,1）
    (call_indirect (type $t0)
      (local.get $table_idx)
    )
  )
  ;; 导出一个"攻击"函数，模拟通过漏洞（如相邻内存覆盖）将恶意函数引用写入表索引3
  ;; 注意：在纯文本模块中无法动态写入表，此函数仅为演示逻辑。
  ;; 实际攻击可能通过 memory->table 的越界写入或利用初始化漏洞实现。
  (func (export "dummy_attack")
    ;; 假设执行后，表索引3被覆盖为 $func_malicious
    nop
  )
)

# 安全策略配置文件
runtime:
  # 启用的防护模块
  enabled_defenses:

    - cfi
    - import_filter
    - memory_monitor
    - syscall_sandbox

# 导入函数白名单
allowed_imports:

  - "env::log"          # 只允许日志函数
  # - "wasi_snapshot_preview1::fd_write" # 不允许文件写入

# 系统调用限制（按命名空间）
syscall_restrictions:
  env:

    - log
  wasi_snapshot_preview1:

    - proc_exit
    # 明确禁用其他功能
  # 另一个常见的命名空间
  # 禁止任何未列出的命名空间

# 内存安全策略
memory:
  max_pages: 4          # 最大内存页数 (4 * 64KB = 256KB)
  guard_region_pages: 1 # 警戒区域页数（运行时添加，用于检测越界）
  zero_on_free: true    # 释放时清零内存（如支持）

# 控制流完整性策略
control_flow_integrity:
  enforce_indirect_call_signature: true
  validate_return_address: false # 栈保护，实现较复杂

logging:
  level: INFO
  log_suspicious_activity: true

#!/usr/bin/env python3
"""
将 .wat 文本文件编译为 .wasm 字节码。
依赖外部工具 `wat2wasm` (WABT工具集的一部分)。
"""
import subprocess
import sys
import os

def compile_wat_to_wasm(wat_path: str, wasm_path: str):
    """调用系统安装的 wat2wasm 进行编译。"""
    if not os.path.exists(wat_path):
        print(f"错误: 文件 {wat_path} 不存在")
        sys.exit(1)
    try:
        cmd = ["wat2wasm", wat_path, "-o", wasm_path]
        subprocess.run(cmd, check=True)
        print(f"成功: {wat_path} -> {wasm_path}")
    except subprocess.CalledProcessError as e:
        print(f"编译失败: {e}")
        sys.exit(1)
    except FileNotFoundError:
        print("错误: 未找到 'wat2wasm' 命令。请安装 WABT (WebAssembly Binary Toolkit)。")
        print(" macOS: brew install wabt")
        print(" Ubuntu/Debian: sudo apt install wabt")
        sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("用法: python wat2wasm.py <input.wat> <output.wasm>")
        sys.exit(1)
    compile_wat_to_wasm(sys.argv[1], sys.argv[2])

#!/usr/bin/env python3
"""
WASM安全攻防演示主入口。
"""
import sys
import os
import yaml
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from core.runtime import SecuritySandbox
from core.threats import MaliciousModuleGenerator, ThreatDetector
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def load_policy() -> dict:
    """加载安全策略配置文件。"""
    policy_path = os.path.join('config', 'security_policy.yaml')
    try:
        with open(policy_path, 'r') as f:
            policy = yaml.safe_load(f)
            logger.info(f"安全策略加载成功。")
            return policy
    except FileNotFoundError:
        logger.warning(f"策略文件 {policy_path} 未找到，使用默认策略。")
        return {}

def demonstrate_cfi_attack():
    """演示控制流完整性攻击与防护。"""
    print("\n" + "="*60)
    print("演示 1: 控制流完整性 (CFI) 攻击")
    print("="*60)

    # 1. 生成恶意模块的WAT代码并编译（模拟）
    wat_code = MaliciousModuleGenerator.generate_cfi_hijack_wat()
    # 在实际运行前，需要将wat_code编译为.wasm。此处简化，假设已存在编译好的文件。
    # 我们使用项目内预置的恶意模块。
    malicious_wasm_path = os.path.join('modules', 'malicious', 'cfi_hijack.wasm')
    if not os.path.exists(malicious_wasm_path):
        logger.error(f"预编译的恶意模块不存在: {malicious_wasm_path}")
        logger.error("请先运行 'python tools/wat2wasm.py modules/malicious/cfi_hijack.wat modules/malicious/cfi_hijack.wasm'")
        return

    with open(malicious_wasm_path, 'rb') as f:
        wasm_bytes = f.read()

    # 2. 使用安全沙箱加载
    policy = load_policy()
    sandbox = SecuritySandbox(wasm_bytes, policy)
    try:
        sandbox.instantiate()
        logger.info("✅ 恶意模块加载成功（防护已加载）。")

        # 3. 尝试合法调用 (table index 0 或 1)
        result = sandbox.call_export("vulnerable_call", 0)
        logger.info(f"   合法间接调用(索引0) 结果: {result}")
        result = sandbox.call_export("vulnerable_call", 1)
        logger.info(f"   合法间接调用(索引1) 结果: {result}")

        # 4. 尝试攻击性调用 (table index 3, 假设被篡改)
        # 在我们的演示模块中，索引3未被初始化，调用应导致陷阱(trap)或CFI检查失败。
        try:
            result = sandbox.call_export("vulnerable_call", 3)
            logger.info(f"   ⚠️  非法间接调用(索引3) 结果: {result} -- CFI防护可能未生效！")
        except RuntimeError as e:
            logger.info(f"   ✅ 非法间接调用被阻断: {e}")
        except Exception as e:
            logger.info(f"   ✅ 非法间接调用导致陷阱或错误: {type(e).__name__}: {e}")

    except Exception as e:
        logger.error(f"演示过程中出错: {e}")

    # 5. 输出威胁检测报告
    report = sandbox.threat_detector.generate_report()
    if report:
        print("\n威胁检测报告:")
        print(report)

def demonstrate_memory_oob():
    """演示内存越界访问攻击与防护。"""
    print("\n" + "="*60)
    print("演示 2: 内存越界 (OOB) 访问攻击")
    print("="*60)

    malicious_wasm_path = os.path.join('modules', 'malicious', 'memory_oob.wasm')
    if not os.path.exists(malicious_wasm_path):
        logger.error(f"预编译的恶意模块不存在: {malicious_wasm_path}")
        return

    with open(malicious_wasm_path, 'rb') as f:
        wasm_bytes = f.read()

    policy = load_policy()
    sandbox = SecuritySandbox(wasm_bytes, policy)
    try:
        sandbox.instantiate()
        logger.info("✅ 内存OOB模块加载成功。")

        memory = sandbox.get_memory_view()
        if memory is None:
            logger.error("模块未导出内存。")
            return
        memory_size = len(memory)
        logger.info(f"   模块内存大小: {memory_size} 字节")

        # 尝试合法访问（前4字节）
        try:
            result = sandbox.call_export("read_oob", 0)
            logger.info(f"   合法读取地址0: {result}")
        except Exception as e:
            logger.info(f"   合法读取出错（可能未初始化）: {e}")

        # 尝试越界访问（超过内存大小）
        oob_address = memory_size + 100
        try:
            result = sandbox.call_export("read_oob", oob_address)
            logger.info(f"   ⚠️  越界读取地址{oob_address} 结果: {result} -- 内存防护可能未生效！")
        except RuntimeError as e:
            logger.info(f"   ✅ 越界读取被内存监视器阻断: {e}")
        except Exception as e:
            # wasmer运行时本身可能会抛出异常
            logger.info(f"   ✅ 越界读取导致运行时错误: {type(e).__name__}")

    except Exception as e:
        logger.error(f"演示过程中出错: {e}")

def main():
    print("🚀 WASM 安全攻防演示项目启动")
    print("本演示将展示常见WASM安全威胁及运行时防护策略。\n")

    # 演示1: CFI攻击与防护
    demonstrate_cfi_attack()

    # 演示2: 内存越界攻击与防护
    demonstrate_memory_oob()

    print("\n" + "="*60)
    print("演示结束。")
    print("="*60)

if __name__ == "__main__":
    main()

mkdir wasm-sandbox-security-demo
    cd wasm-sandbox-security-demo

python3 -m venv venv
    source venv/bin/activate  # Linux/macOS
    # 或 venv\Scripts\activate  # Windows

pip install wasmer wasmer-compiler-cranelift pyyaml

# 编译恶意演示模块
python tools/wat2wasm.py modules/malicious/cfi_hijack.wat modules/malicious/cfi_hijack.wasm
python tools/wat2wasm.py modules/malicious/memory_oob.wat modules/malicious/memory_oob.wasm

# 编译良性演示模块（可选）
python tools/wat2wasm.py modules/benign/simple_loop.wat modules/benign/simple_loop.wasm

python main.py

#!/usr/bin/env python3
"""
防护模块单元测试。
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import unittest
from core.defenses import ImportFilter, MemoryAccessMonitor

class TestImportFilter(unittest.TestCase):
    def setUp(self):
        self.allowed = ["env::log", "wasi::fd_write"]
        self.filter = ImportFilter(self.allowed)

    def test_allowed_import(self):
        self.assertTrue(self.filter.is_allowed("env", "log"))
        self.assertTrue(self.filter.is_allowed("wasi", "fd_write"))

    def test_denied_import(self):
        self.assertFalse(self.filter.is_allowed("env", "get_secret"))
        self.assertFalse(self.filter.is_allowed("malicious", "exploit"))

class TestMemoryAccessMonitor(unittest.TestCase):
    def setUp(self):
        self.monitor = MemoryAccessMonitor()
        # 模拟一个内存对象（此处简化）
        class MockMemory:
            def __init__(self):
                self.size = 65536  # 64KB
        self.mock_mem = MockMemory()

    def test_check_within_bounds(self):
        self.monitor.register_memory(self.mock_mem)
        # 假设check_read需要一个instance参数，这里传入None
        self.assertTrue(self.monitor.check_read(0, 100, None))
        self.assertTrue(self.monitor.check_read(65436, 100, None)) # 边界

    def test_check_out_of_bounds(self):
        self.monitor.register_memory(self.mock_mem)
        self.assertFalse(self.monitor.check_read(65500, 100, None)) # 越界
        self.assertFalse(self.monitor.check_read(-10, 20, None))    # 负偏移

if __name__ == '__main__':
    unittest.main()

python -m pytest tests/test_defenses.py -v

# 在 security_policy.yaml 中修改
runtime:
  enabled_defenses:

    - cfi
    - import_filter
    # - memory_monitor  # 注释掉此行

    - syscall_sandbox

allowed_imports:

  - "env::*"  # 允许所有env下的导入（危险！）