import os
import json
import importlib.util
import sys


class MiniLLM:
    """
    一个简易的 LLM 加载器，模拟了特定的逻辑漏洞：
    在检查 trust_remote_code 之前，就处理了 auto_map 里的 import。
    """

    def __init__(self, model_path, trust_remote_code=False):
        self.model_path = model_path
        self.trust_remote_code = trust_remote_code
        self.config = self._load_config()

        # 模拟加载流程
        print(f"[*] 正在加载模型: {model_path}")
        self._resolve_model_class()

        # 这一步本来应该在最前面，但现在放到了后面
        if not self.trust_remote_code:
            print("[-] 安全检查: trust_remote_code 为 False，拒绝加载远程代码。")
            raise RuntimeError("Aborted: Remote code not allowed!")

        print("[+] 安全检查通过，模型加载完成。")

    def _load_config(self):
        config_path = os.path.join(self.model_path, "config.json")
        if not os.path.exists(config_path):
            raise FileNotFoundError("Config not found")
        with open(config_path, 'r') as f:
            return json.load(f)

    def _resolve_model_class(self):
        """
        [漏洞核心]
        试图解析 config 中的 auto_map 以确定要使用哪个类。
        这个函数会在 trust_remote_code 检查之前运行。
        """
        auto_map = self.config.get("auto_map", {})

        # 模拟 vLLM/Transformers 的解析逻辑
        # 如果配置里指定了 AutoModel，系统会尝试去 import 它以获取类对象
        if "AutoModel" in auto_map:
            class_path = auto_map["AutoModel"]
            module_name, class_name = class_path.rsplit(".", 1)

            print(f"[*] 正在解析模型架构: {class_path} ...")

            # 动态加载指定的 Python 文件 (相当于 import malicious.py)
            # 只要这一行执行，恶意代码就会跑起来
            self._dynamic_import(module_name)

    def _dynamic_import(self, module_name):
        # 简单的动态导入模拟
        file_path = os.path.join(self.model_path, f"{module_name}.py")
        if os.path.exists(file_path):
            spec = importlib.util.spec_from_file_location(module_name, file_path)
            module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = module
            try:
                spec.loader.exec_module(module)  # <--- 代码执行发生在这里
            except Exception as e:
                print(f"[!] 模块执行出错 (但这不重要，代码已经跑了): {e}")