yuzu-octopus/CTF-LLM

# CTF-LLM 在 Google Colab 免费版上，使用 [Unsloth](https://unsloth.ai) 和 QLoRA 微调开源 LLM（Gemma 4 E4B、Qwen 3.5 9B/4B），使其在网络安全 CTF 竞赛和竞赛编程中表现出色。 ## 项目功能 本项目构建了一个完整的 pipeline，用于创建能够执行以下操作的专业 AI 模型： - **解决 CTF 挑战**：二进制漏洞利用（pwn）、逆向工程、Web 漏洞利用、密码学 - **编写竞赛编程题解**：算法、数据结构、代码优化 - **分析安全漏洞**：识别并解释代码弱点 该 pipeline 从 GitHub 抓取真实的 CTF writeup 和编程题解，将它们与 HuggingFace 上精选的数据集相结合，并使用高效的 QLoRA（4-bit 量化 + LoRA 适配器）对模型进行微调。 ## 快速开始 ``` # 安装 uv (Python package manager) curl -LsSf https://astral.sh/uv/install.sh | sh # 安装依赖 uv sync # 运行完整 pipeline (build data + train) ./finetune.sh gemma4 --all # 或者单独运行各步骤 ./finetune.sh gemma4 --build-data # Just build datasets ./finetune.sh gemma4 --train # Just train (requires data) ./finetune.sh gemma4 --eval # Evaluate trained model ``` ## 工作原理 ### 1. 数据集收集 (`src/build_dataset.py`) 抓取包含 CTF writeup 和安全文档的 GitHub 仓库： | 来源 | 包含内容 | 示例 | |--------|------------------|----------| | `Cajac/picoCTF-Writeups` | 250+ picoCTF 挑战题解 | Markdown + Python/C 漏洞利用代码 | | `DarkCodeOrg/CryptoHack` | 密码学挑战题解 | Python 密码学漏洞利用 | | `Adamkadaban/CTFs` | 多分类 CTF writeup | 涵盖 pwn、rev、web、crypto 的 376 MB 内容 | | pwntools、angr、ROPgadget 文档 | 安全工具文档 | 用法示例和教程 | 此外，还会从 HuggingFace 下载结构化数据集： | 数据集 | 大小 | 内容 | |---------|------|---------| | `nvidia/OpenCodeReasoning` | 73.5 万个样本 | 带有推理过程的竞赛编程题解 | | `AlicanKiraz0/Cybersecurity-Dataset-Fenrir-v2.1` | 9.98 万个示例 | 网络安全问答 | | `Jacqkues/ctf_webserver_v0.1` | 340 个样本 | Web CTF 挑战 | ### 1.5 合成 Rev/PWN 示例 (`src/synthetic_rev_pwn.py`) 包含 23 个手工编写的训练样本，涵盖了难度最高的 CTF 分类： | 分类 | 示例 | 涵盖技术 | |----------|----------|-------------------| | **PWN** (13) | ret2win, ret2libc, format string, heap UAF, stack canary bypass, ROP chain, ret2dlresolve, fastbin attack, heap consolidation, shellcode injection, ASLR bypass, ret2csu, read→shellcode ROP | 漏洞利用代码前的思维链推理 | | **REV** (10) | 二进制分析流程、XOR 密码、反调试、ELF sections、加壳检测、密码识别、keygen、Frida 追踪、LD_PRELOAD hooking、去混淆 | 逐步分析方法，而非仅仅给出答案 | 这些示例使用了思维链推理：模型学会在解决问题之前先进行分析，而不是直接跳到答案。 ### 1.6 训练模式 有关完整的 Fast/Full 模式配置表（包括 LoRA 参数、数据集限制和预计训练时间），请参见 `AGENTS.md`。 ### 2. 数据处理 (`src/process_data.py`) 将原始数据转换为微调所需的聊天格式： ``` Alpaca Format: ChatML Format: { { "instruction": "...", "messages": [ "input": "...", {"role": "system", "content": "..."}, "output": "..." {"role": "user", "content": "..."}, } {"role": "assistant", "content": "..."} ] } ``` 系统提示会根据内容类别（CTF 还是竞赛编程）自动选择。 ### 3. 模型训练 (`src/train.py`) 使用 Unsloth 优化的 QLoRA 训练： - **4-bit 量化**：减少约 75% 的内存使用，可在免费的 T4 GPU 上进行训练 - **LoRA 适配器**：仅训练小型适配器层（约占参数的 1%） - **聊天模板**：应用特定于模型的格式（Gemma 使用 `gemma-4`，Qwen 使用 `chatml`） 训练输出三个版本的模型： - `lora/` - LoRA 适配器（体积小，用于继续训练） - `gguf/` - GGUF 格式（用于 llama.cpp、Ollama、本地推理） - `merged/` - 完整合并的模型（用于 HuggingFace、vLLM） ### 4. 模型评估 (`src/eval.py`) 在包含 200 道题目的 CTF benchmark（每个类别 50 道）上运行训练好的模型，并报告带有 Wilson 95% 置信区间的准确率： ``` # 单个 model uv run src/eval.py --model gemma4 --adapter outputs/gemma4-ctf/lora # 并排比较 models (使用 McNemar's test) uv run src/eval.py --compare \ gemma4:outputs/gemma4-ctf/lora \ qwen35:outputs/qwen35-ctf/lora # 按 category 或 difficulty 筛选 uv run src/eval.py --model gemma4 --adapter outputs/gemma4-ctf/lora --category pwn --difficulty easy # 将结果保存为 JSON uv run src/eval.py --model gemma4 --adapter outputs/gemma4-ctf/lora --output data/eval/results/ ``` 该 benchmark 涵盖 4 个类别（pwn、rev、crypto、web），分为 3 个难度级别，并支持对 flag 提取、多选题和代码生成任务进行自动评分。 **评估器局限性**（N=200 会产生 ±7% 的误差范围）： - `grade_code` 验证语法 + 参考关键字，而非功能正确性 - `grade_mcq` 匹配 `Answer: X` / `(X)` / 后备为最后一个字母 - `grade_flag` 使用正则表达式 `flag\{[^}]+\}` - 结果仅供参考，并不具备决定性 — 扩展至 N=500 可获得 ±4% 的置信区间 ### 5. Colab 集成 (`finetune.sh`) 使用 `colab` CLI 自动化整个工作流程： 1. 创建带有免费 T4 GPU 的 Colab 会话 2. 安装依赖项 3. 上传数据和脚本 4. 运行训练 5. 下载训练好的模型 6. 清理会话 ## 项目结构 ``` finetuning/ ├── finetune.sh # Main entry point - orchestrates everything ├── config.yaml # Central configuration for models and training ├── pyproject.toml # uv project config and dependencies ├── requirements.txt # Local Python deps ├── requirements-colab.txt # Colab-only install set (ordered) ├── AGENTS.md # Instructions for AI agents ├── README.md # This file ├── TRAINING.md # Step-by-step training guide ├── .mimocode/commands/ # Slash-commands for mimocode │ └── ft-diag.md # /ft-diag — Colab failure mode diagnostic │ ├── src/ # Python scripts │ ├── train.py # Unified training script (reads configs/*.yaml) │ ├── build_dataset.py # Scrapes GitHub repos for writeups │ ├── download_datasets.py # Downloads HuggingFace datasets │ ├── process_data.py # Converts data to chat format │ └── eval.py # CTF model evaluator (200-question benchmark) │ ├── configs/ # Model-specific configurations (3 total) │ ├── gemma4.yaml # Gemma 4 E4B settings │ ├── gemma4-12b.yaml # Gemma 4 12B settings │ ├── qwen35.yaml # Qwen 3.5 9B settings │ └── qwen35-4b.yaml # Qwen 3.5 4B settings │ ├── notebooks/ │ └── qwen4b_self_contained.ipynb # ONLY notebook (full pipeline) │ ├── data/ │ ├── eval/ │ │ └── ctf_bench.jsonl # 50 curated CTF challenges (evaluation) │ ├── raw/ # HuggingFace downloads (gitignored) │ ├── processed/ # Converted to chat format (gitignored) │ └── merged/ # Final merged training file (gitignored) ``` ## 配置 ### 主配置 (`config.yaml`) 所有设置的集中管理位置： ``` default_model: gemma4 models: gemma4: name: unsloth/gemma-4-E4B-it r: 32 # LoRA rank lora_alpha: 32 # LoRA scaling max_seq_length: 4096 # Max token length (limited by T4 VRAM) batch_size: 1 # Batch size (limited by T4 VRAM) training: num_train_epochs: 3 learning_rate: 2e-4 gradient_accumulation_steps: 4 ``` ### 模型配置 (`configs/*.yaml`) 覆盖主配置的单个模型配置： ``` # configs/gemma4.yaml model: name: unsloth/gemma-4-E4B-it target_modules: all-linear # Apply LoRA to all linear layers r: 32 lora_alpha: 32 ``` ## 添加新模型 1. 创建包含模型设置的 `configs/newmodel.yaml` 2. 将模型添加到 `config.yaml` 中的 `models:` 下 3. 运行：`./finetune.sh newmodel --all` ## 硬件要求 | 组件 | 要求 | |-----------|-------------| | GPU | NVIDIA T4 (16GB) - 在 Colab 上免费提供 | | VRAM | 使用 QLoRA 时占用约 12GB | | 训练时间（Fast 模式） | 约 500 个样本需 ~30 分钟 | | 训练时间（Full 模式） | 约 2500+ 个样本需 ~50-70 分钟 | | 存储空间 | 数据约 2GB，模型约 4GB | ## 导出格式 | 格式 | 使用场景 | 命令 | |--------|----------|---------| | GGUF | llama.cpp、Ollama、本地推理 | `model.save_pretrained_gguf()` | | SafeTensors | HuggingFace、vLLM、推理 API | `model.save_pretrained_merged()` | | LoRA | 继续训练、微调 | `model.save_pretrained()` | ## 故障排除 **内存不足 (OOM)** - 减少配置中的 `max_seq_length`（尝试 2048） - 确保 `batch_size: 1` - 使用 `use_gradient_checkpointing: unsloth` **训练缓慢** - 确保您使用的是 GPU 运行时 - 检查 `fp16`/`bf16` 设置是否与您的 GPU 匹配 **数据问题** - 验证格式：`head -1 data/merged/train.jsonl | python -m json.tool` - 检查空输出：`grep -c '"output": ""' data/merged/train.jsonl` ## 致谢 - [Unsloth](https://unsloth.ai) - QLoRA 训练框架 - [CTFtime](https://ctftime.org) - CTF writeup 来源 - [picoCTF](https://picoctf.org) - CTF 挑战平台 - [CryptoHack](https://cryptohack.org) - 密码学挑战平台

标签：DLL 劫持, 云资产清单, 人工智能, 凭据扫描, 大语言模型, 微调, 搜索语句（dork）, 用户模式Hook绕过, 逆向工具, 逆向工程