From 7152b0a81f558d46aacbf4db8d0997844f9bd3f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BD=87=E6=B9=98=E9=9B=A8?= Date: Tue, 30 Sep 2025 07:42:12 +0000 Subject: [PATCH 1/2] =?UTF-8?q?=E9=87=8D=E5=91=BD=E5=90=8D=20README.md=20?= =?UTF-8?q?=E4=B8=BA=20README.zh.md?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md => README.zh.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename README.md => README.zh.md (100%) diff --git a/README.md b/README.zh.md similarity index 100% rename from README.md rename to README.zh.md -- Gitee From 3c840d47c16336c45d09f4593de1ca9dfda97a15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BD=87=E6=B9=98=E9=9B=A8?= Date: Tue, 30 Sep 2025 15:55:51 +0800 Subject: [PATCH 2/2] update --- .DS_Store | Bin 0 -> 6148 bytes LICENSE | 130 +++++++++++++++++++++++++++ README.md | 46 ++++++++++ configs/config.example.yaml | 10 +++ knowledge_base/knowledge_base.yaml | 51 +++++++++++ prompts/prompts.yaml | 26 ++++++ requirements.txt | 3 + scripts/run_pipeline.sh | 8 ++ src/.DS_Store | Bin 0 -> 6148 bytes src/agents/.DS_Store | Bin 0 -> 6148 bytes src/agents/a1_cve_analysis.py | 27 ++++++ src/agents/a2_template_mapping.py | 20 +++++ src/agents/a3_sequence_generation.py | 22 +++++ src/cli.py | 92 +++++++++++++++++++ src/core/.DS_Store | Bin 0 -> 6148 bytes src/core/kb.py | 45 ++++++++++ src/core/logger.py | 9 ++ src/core/nvd_client.py | 20 +++++ src/core/utils.py | 13 +++ tests/.DS_Store | Bin 0 -> 6148 bytes tests/test_kb_load.py | 5 ++ 21 files changed, 527 insertions(+) create mode 100644 .DS_Store create mode 100644 LICENSE create mode 100644 README.md create mode 100644 configs/config.example.yaml create mode 100644 knowledge_base/knowledge_base.yaml create mode 100644 prompts/prompts.yaml create mode 100644 requirements.txt create mode 100755 scripts/run_pipeline.sh create mode 100644 src/.DS_Store create mode 100644 src/agents/.DS_Store create mode 100644 src/agents/a1_cve_analysis.py create mode 100644 src/agents/a2_template_mapping.py create mode 100644 src/agents/a3_sequence_generation.py create mode 100644 src/cli.py create mode 100644 src/core/.DS_Store create mode 100644 src/core/kb.py create mode 100644 src/core/logger.py create mode 100644 src/core/nvd_client.py create mode 100644 src/core/utils.py create mode 100644 tests/.DS_Store create mode 100644 tests/test_kb_load.py diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..453ad30b3200c644bfab4df0e0cf8e0663129388 GIT binary patch literal 6148 zcmeHK&1%~~5T3Q2S|+sk&_a$2xf)^wXjp@8Yx{- z8JDvxQd`l^@GmmJz3bA96y@}d9`sLTXeqs+Yd93ulu<%4+MCdnA_}Rb#5CFjc4=^H|jrwaDXA=LW{%^}OD+ zzqeeT91rAZJYEgtax^>|59IOT(Q4K6-t8ZpUQDmEn_S(Q<_RW+<+Xv$`~#ZZ%|lwO z={6|xOyySy;sQIMD8bf$BWhF3DUJchfMej98L(%WxBJZNxbltx$H4!`0G|&MgVDFx z7}Q4xCjA5ec3{>5o!=6iBQ5$C8-wsbgh>UORAHYO!lc74ZJckhF=*0B*vE&kD+~KV z5qfphmpYt;Z_v$-0mr~&2HNJc&-efN)9?SsN$$%r;23yO42aHIa5lk~?A^MxIlgNp t##am$j%y5VQlRLo7_od6?_p?xU9tl7Ej9*Wfw(^cng%yH2L37ozW}XbZ=e7G literal 0 HcmV?d00001 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d47a3a6 --- /dev/null +++ b/LICENSE @@ -0,0 +1,130 @@ + +木兰宽松许可证, 第2版 + +木兰宽松许可证， 第2版 + +2020年1月 http://license.coscl.org.cn/MulanPSL2 + +您对“软件”的复制、使用、修改及分发受木兰宽松许可证，第2版（“本许可证”）的如下条款的约束： + +0. 定义 + +“软件” 是指由“贡献”构成的许可在“本许可证”下的程序和相关文档的集合。 + +“贡献” 是指由任一“贡献者”许可在“本许可证”下的受版权法保护的作品。 + +“贡献者” 是指将受版权法保护的作品许可在“本许可证”下的自然人或“法人实体”。 + +“法人实体” 是指提交贡献的机构及其“关联实体”。 + +“关联实体” 是指，对“本许可证”下的行为方而言，控制、受控制或与其共同受控制的机构，此处的控制是指有受控方或共同受控方至少50%直接或间接的投票权、资金或其他有价证券。 + +1. 授予版权许可 + +每个“贡献者”根据“本许可证”授予您永久性的、全球性的、免费的、非独占的、不可撤销的版权许可，您可以复制、使用、修改、分发其“贡献”，不论修改与否。 + +2. 授予专利许可 + +每个“贡献者”根据“本许可证”授予您永久性的、全球性的、免费的、非独占的、不可撤销的（根据本条规定撤销除外）专利许可，供您制造、委托制造、使用、许诺销售、销售、进口其“贡献”或以其他方式转移其“贡献”。前述专利许可仅限于“贡献者”现在或将来拥有或控制的其“贡献”本身或其“贡献”与许可“贡献”时的“软件”结合而将必然会侵犯的专利权利要求，不包括对“贡献”的修改或包含“贡献”的其他结合。如果您或您的“关联实体”直接或间接地，就“软件”或其中的“贡献”对任何人发起专利侵权诉讼（包括反诉或交叉诉讼）或其他专利维权行动，指控其侵犯专利权，则“本许可证”授予您对“软件”的专利许可自您提起诉讼或发起维权行动之日终止。 + +3. 无商标许可 + +“本许可证”不提供对“贡献者”的商品名称、商标、服务标志或产品名称的商标许可，但您为满足第4条规定的声明义务而必须使用除外。 + +4. 分发限制 + +您可以在任何媒介中将“软件”以源程序形式或可执行形式重新分发，不论修改与否，但您必须向接收者提供“本许可证”的副本，并保留“软件”中的版权、商标、专利及免责声明。 + +5. 免责声明与责任限制 + +“软件”及其中的“贡献”在提供时不带任何明示或默示的担保。在任何情况下，“贡献者”或版权所有者不对任何人因使用“软件”或其中的“贡献”而引发的任何直接或间接损失承担责任，不论因何种原因导致或者基于何种法律理论，即使其曾被建议有此种损失的可能性。 + +6. 语言 + +“本许可证”以中英文双语表述，中英文版本具有同等法律效力。如果中英文版本存在任何冲突不一致，以中文版为准。 + +条款结束 + +如何将木兰宽松许可证，第2版，应用到您的软件 + +如果您希望将木兰宽松许可证，第2版，应用到您的新软件，为了方便接收者查阅，建议您完成如下三步： + +1， 请您补充如下声明中的空白，包括软件名、软件的首次发表年份以及您作为版权人的名字； + +2， 请您在软件包的一级目录下创建以“LICENSE”为名的文件，将整个许可证文本放入该文件中； + +3， 请将如下声明文本放入每个源文件的头部注释中。 + +Copyright (c) 2025 gongjula +[Software Name] is licensed under Mulan PSL v2. +You can use this software according to the terms and conditions of the Mulan PSL v2. +You may obtain a copy of Mulan PSL v2 at: + http://license.coscl.org.cn/MulanPSL2 +THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, +EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, +MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. +See the Mulan PSL v2 for more details. + +Mulan Permissive Software License，Version 2 + +Mulan Permissive Software License，Version 2 (Mulan PSL v2) + +January 2020 http://license.coscl.org.cn/MulanPSL2 + +Your reproduction, use, modification and distribution of the Software shall be subject to Mulan PSL v2 (this License) with the following terms and conditions: + +0. Definition + +Software means the program and related documents which are licensed under this License and comprise all Contribution(s). + +Contribution means the copyrightable work licensed by a particular Contributor under this License. + +Contributor means the Individual or Legal Entity who licenses its copyrightable work under this License. + +Legal Entity means the entity making a Contribution and all its Affiliates. + +Affiliates means entities that control, are controlled by, or are under common control with the acting entity under this License, ‘control’ means direct or indirect ownership of at least fifty percent (50%) of the voting power, capital or other securities of controlled or commonly controlled entity. + +1. Grant of Copyright License + +Subject to the terms and conditions of this License, each Contributor hereby grants to you a perpetual, worldwide, royalty-free, non-exclusive, irrevocable copyright license to reproduce, use, modify, or distribute its Contribution, with modification or not. + +2. Grant of Patent License + +Subject to the terms and conditions of this License, each Contributor hereby grants to you a perpetual, worldwide, royalty-free, non-exclusive, irrevocable (except for revocation under this Section) patent license to make, have made, use, offer for sale, sell, import or otherwise transfer its Contribution, where such patent license is only limited to the patent claims owned or controlled by such Contributor now or in future which will be necessarily infringed by its Contribution alone, or by combination of the Contribution with the Software to which the Contribution was contributed. The patent license shall not apply to any modification of the Contribution, and any other combination which includes the Contribution. If you or your Affiliates directly or indirectly institute patent litigation (including a cross claim or counterclaim in a litigation) or other patent enforcement activities against any individual or entity by alleging that the Software or any Contribution in it infringes patents, then any patent license granted to you under this License for the Software shall terminate as of the date such litigation or activity is filed or taken. + +3. No Trademark License + +No trademark license is granted to use the trade names, trademarks, service marks, or product names of Contributor, except as required to fulfill notice requirements in section 4. + +4. Distribution Restriction + +You may distribute the Software in any medium with or without modification, whether in source or executable forms, provided that you provide recipients with a copy of this License and retain copyright, patent, trademark and disclaimer statements in the Software. + +5. Disclaimer of Warranty and Limitation of Liability + +THE SOFTWARE AND CONTRIBUTION IN IT ARE PROVIDED WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED. IN NO EVENT SHALL ANY CONTRIBUTOR OR COPYRIGHT HOLDER BE LIABLE TO YOU FOR ANY DAMAGES, INCLUDING, BUT NOT LIMITED TO ANY DIRECT, OR INDIRECT, SPECIAL OR CONSEQUENTIAL DAMAGES ARISING FROM YOUR USE OR INABILITY TO USE THE SOFTWARE OR THE CONTRIBUTION IN IT, NO MATTER HOW IT’S CAUSED OR BASED ON WHICH LEGAL THEORY, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +6. Language + +THIS LICENSE IS WRITTEN IN BOTH CHINESE AND ENGLISH, AND THE CHINESE VERSION AND ENGLISH VERSION SHALL HAVE THE SAME LEGAL EFFECT. IN THE CASE OF DIVERGENCE BETWEEN THE CHINESE AND ENGLISH VERSIONS, THE CHINESE VERSION SHALL PREVAIL. + +END OF THE TERMS AND CONDITIONS + +How to Apply the Mulan Permissive Software License，Version 2 (Mulan PSL v2) to Your Software + +To apply the Mulan PSL v2 to your work, for easy identification by recipients, you are suggested to complete following three steps: + +i. Fill in the blanks in following statement, including insert your software name, the year of the first publication of your software, and your name identified as the copyright owner; +ii. Create a file named "LICENSE" which contains the whole context of this License in the first directory of your software package; +iii. Attach the statement to the appropriate annotated syntax at the beginning of each source file. + +Copyright (c) 2025 gongjula +[Software Name] is licensed under Mulan PSL v2. +You can use this software according to the terms and conditions of the Mulan PSL v2. +You may obtain a copy of Mulan PSL v2 at: + http://license.coscl.org.cn/MulanPSL2 +THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, +EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, +MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. +See the Mulan PSL v2 for more details. diff --git a/README.md b/README.md new file mode 100644 index 0000000..c880fd0 --- /dev/null +++ b/README.md @@ -0,0 +1,46 @@ +# Kernel Vulnerability Reproduction: A Multi-Agent Framework +# 多Agent协同的内核漏洞自动化复现框架 + +## 模块总览 +- **prompts/prompts.yaml**：所有提示词集中管理； +- **knowledge_base/knowledge_base.yaml**：所有 KB（模板 + 关键词映射 + 可选 CWE 映射）集中管理； +- **src/agents/**：五个 Agent 模块代码（I: `a1_cve_analysis.py`，II: `a2_template_mapping.py`，III: `a3_sequence_generation.py`，IV: `a4_syzkaller_orchestration.py`，V: `a5_poc_converter.py`）； +- **src/core/**：通用工具（Logger、NVD 客户端、KB 加载、PolicyGuard）。 + +## 快速开始 +```bash +pip install -r requirements.txt +cp configs/config.example.yaml configs/config.yaml +python -m src.cli analyze --cve CVE-2021-43267 +python -m src.cli map-template --cve CVE-2021-43267 +python -m src.cli gen-skeleton --cve CVE-2021-43267 --out out/skeleton +python -m src.cli orchestrate +python -m src.cli package --case out/skeleton --out out/package +``` + +## 目录 +``` +kernel-vuln-repro-agents-pro/ +├─ configs/ +├─ knowledge_base/ +│ └─ knowledge_base.yaml +├─ prompts/ +│ └─ prompts.yaml +├─ src/ +│ ├─ agents/ +│ │ ├─ cve_analysis.py +│ │ ├─ template_mapping.py +│ │ ├─ sequence_generation.py +│ │ ├─ syzkaller_orchestration.py +│ │ └─ poc_converter.py +│ ├─ core/ +│ │ ├─ logger.py +│ │ ├─ utils.py +│ │ ├─ kb.py +│ │ └─ nvd_client.py +│ └─ cli.py +├─ scripts/ +│ └─ run_pipeline.sh +├─ tests/ +└─ out/ (运行后生成产物) +``` diff --git a/configs/config.example.yaml b/configs/config.example.yaml new file mode 100644 index 0000000..db8e2d2 --- /dev/null +++ b/configs/config.example.yaml @@ -0,0 +1,10 @@ +nvd: + api_base: "https://services.nvd.nist.gov/rest/json/cves/2.0" + timeout: 15 +orchestration: + dry_run: true + max_parallel: 2 + log_dir: "out/logs" +paths: + kb_file: "knowledge_base/knowledge_base.yaml" + prompts_file: "prompts/prompts.yaml" diff --git a/knowledge_base/knowledge_base.yaml b/knowledge_base/knowledge_base.yaml new file mode 100644 index 0000000..f51a36c --- /dev/null +++ b/knowledge_base/knowledge_base.yaml @@ -0,0 +1,51 @@ +templates: + - id: MEM_HEAP_OVERFLOW + name: "堆溢出" + cwe: ["CWE-122"] + risk_level: "high" + context_hints: + - "消息长度/边界校验缺陷" + - "协议解析路径相关" + syz_skeleton: + steps: + - "init: 初始化（占位）" + - "mutate: 非具体化越界输入占位" + - "trigger: 发送占位消息（非真实消息体）" + + - id: UAF_GENERIC + name: "Use-After-Free" + cwe: ["CWE-416"] + risk_level: "high" + context_hints: + - "并发/生命周期管理薄弱" + syz_skeleton: + steps: + - "init" + - "free" + - "reuse" + + - id: INT_OVERFLOW_ALLOC + name: "整数溢出-分配尺寸" + cwe: ["CWE-190"] + risk_level: "medium" + context_hints: + - "长度参与乘加后用于分配" + syz_skeleton: + steps: + - "prepare" + - "alloc" + - "use" + +feature_to_template: + keywords: + - include: ["堆溢出", "heap overflow", "越界写"] + map_to: "MEM_HEAP_OVERFLOW" + - include: ["use-after-free", "UAF", "释放后使用"] + map_to: "UAF_GENERIC" + - include: ["整数溢出", "integer overflow"] + map_to: "INT_OVERFLOW_ALLOC" + +cwe_map: + CWE-122: "MEM_HEAP_OVERFLOW" + CWE-416: "UAF_GENERIC" + CWE-190: "INT_OVERFLOW_ALLOC" diff --git a/prompts/prompts.yaml b/prompts/prompts.yaml new file mode 100644 index 0000000..1957364 --- /dev/null +++ b/prompts/prompts.yaml @@ -0,0 +1,26 @@ +cve_analysis: + system: | + 你是一个 Linux 内核漏洞信息提取专家。只做信息汇总与结构化，不得生成任何攻击性指令或代码。 + 输出：1) 结构化字段；2) ~300 字中文摘要。缺失信息标注 "N/A"。 + user: | + 请提取 ${cve_id} 的关键信息，优先 NVD、发行版公告、官方补丁引用。 + +mapping: + system: | + 你是模板映射助理，仅将“漏洞类型/关键词”映射到抽象模板名称（如 MEM_HEAP_OVERFLOW）， + 可附带非具体化上下文提示，不得包含任何可执行细节。 + +sequence: + system: | + 你是 syzkaller 序列骨架顾问。输出占位骨架步骤（例如 init/mutate/trigger）， + 明确提示需专家审核后才能落地执行。 + +syzkaller: + system: | + 你是 Syzkaller 调度清单助手。输出 dry-run 计划与日志关键字，不含触发样例内容。 + 产出包括资源需求、并发上限、日志路径与异常关键字等。 + +poc: + system: | + 你是 PoC 元数据打包助手。仅整理元信息与说明步骤；不得输出或推导攻击细节。 + 必要时以 [REDACTED] 标注并剔除敏感内容。 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6aba6f8 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +requests>=2.32.0 +PyYAML>=6.0.2 +rich>=13.7.1 diff --git a/scripts/run_pipeline.sh b/scripts/run_pipeline.sh new file mode 100755 index 0000000..be6b008 --- /dev/null +++ b/scripts/run_pipeline.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -e +python -m src.cli analyze --cve ${1:-CVE-2021-43267} +python -m src.cli map-template --cve ${1:-CVE-2021-43267} +python -m src.cli gen-skeleton --cve ${1:-CVE-2021-43267} --out out/skeleton +python -m src.cli orchestrate +python -m src.cli package --case out/skeleton --out out/package +echo "Done." diff --git a/src/.DS_Store b/src/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..6c761300c383016b6bde13eff4fe93fbed50f4f5 GIT binary patch literal 6148 zcmeHK%}T>S5T3QMO%)*rg&r5Y7HpBK;3d@h0!H+pQWFz27_-u(=1>Yb>kIiLK94iI z+hVMOC-G+oX20F}NirYoZUz8Ye;jrJ900IU39S_zJ`nPgF36I$)Df9{Mh;m}WcNWL zf`w>fFbo(5ei{R`ckAH8HAIj?b^ksG_eq@Poz4qes#(kRhTX7N?F;Wl6khJ<?#TH8!4gU zc<@GPBGPNLlh1-H5HN!1%^~)iNtoC3mVmkOUGZ4Ik$=U2-aq#0zpC04YZx#Le1!oz zA2=$ZtFcrlj}9DU3xH^$Sqjdnm!KT3(bZTgL=Os6sfa3-=@x^jbhL9F=V~k!s&rtw z`C$6ZOm`?uz8#;>WjHWbp&1PWhJh*rHMMl;{y+S@|F0&Qm0`dz@TVAH*1#L|u_S%B yE)++1El0gXC84-dp-jO+U&nGnS8)qf3i>!r5M7O>LbRZm9|1{&84Ls8%D^j{DU%cc literal 0 HcmV?d00001 diff --git a/src/agents/.DS_Store b/src/agents/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 GIT binary patch literal 6148 zcmeH~Jr2S!425mzP>H1@V-^m;4Wg<&0T*E43hX&L&p$$qDprKhvt+--jT7}7np#A3 zem<@ulZcFPQ@L2!n>{z**++&mCkOWA81W14cNZlEfg7;MkzE(HCqgga^y>{tEnwC%0;vJ&^%eQ zLs35+`xjp>T0 CVEOutput: + _ = self.nvd.get_cve(cve_id) # 预留：可扩展填充 fields + fields = { + "vuln_id": cve_id, + "software": {"name": "N/A", "version": "N/A"}, + "environment": {"os": "N/A", "deps": []}, + "build": {"system": "kbuild", "steps": ["make defconfig", "make -j8"], "special_flags": []}, + "reproduction": {"poc": {"type": "N/A", "location": "N/A"}, "trigger": {"steps": [], "notes": ""}}, + "verification": {"expected": "kernel warning / KASAN report", "artifacts": ["logs"]}, + "notes": {"patch_diff_refs": []}, + } + summary_zh = f"【占位摘要】{cve_id} 的结构化梳理完成，仅用于观测性规划，不含攻击细节。" + return CVEOutput(fields=fields, summary_zh=summary_zh) diff --git a/src/agents/a2_template_mapping.py b/src/agents/a2_template_mapping.py new file mode 100644 index 0000000..6f5141c --- /dev/null +++ b/src/agents/a2_template_mapping.py @@ -0,0 +1,20 @@ +from __future__ import annotations +from dataclasses import dataclass +from typing import Optional, List +from ..core.kb import KnowledgeBase + +@dataclass +class MappingResult: + template_id: Optional[str] + template_name: Optional[str] + context_hints: List[str] + +class TemplateMappingAgent: + def __init__(self, kb_file: str): + self.kb = KnowledgeBase(kb_file) + + def run(self, feature_text: str) -> MappingResult: + hint = self.kb.match_template_by_text(feature_text or "") + if not hint: + return MappingResult(None, None, []) + return MappingResult(hint.id, hint.name, hint.context_hints) diff --git a/src/agents/a3_sequence_generation.py b/src/agents/a3_sequence_generation.py new file mode 100644 index 0000000..4002edd --- /dev/null +++ b/src/agents/a3_sequence_generation.py @@ -0,0 +1,22 @@ +from __future__ import annotations +from dataclasses import dataclass +from typing import List +from ..core.kb import KnowledgeBase + +@dataclass +class SkeletonSequence: + template_id: str + steps: List[str] + disclaimer: str + +class SequenceGenerationAgent: + def __init__(self, kb_file: str): + self.kb = KnowledgeBase(kb_file) + + def run(self, template_id: str) -> SkeletonSequence: + t = self.kb.get_template(template_id) + if not t: + return SkeletonSequence(template_id, [], "模板不存在，需专家补充。") + steps = [f"{i+1}. {s}" for i, s in enumerate(t.syz_skeleton.get("steps", []))] + disclaimer = "此为占位骨架，不含具体 syscall/参数。需专家审核与隔离环境方可落地。" + return SkeletonSequence(template_id, steps, disclaimer) diff --git a/src/cli.py b/src/cli.py new file mode 100644 index 0000000..93adc8d --- /dev/null +++ b/src/cli.py @@ -0,0 +1,92 @@ +from __future__ import annotations +import argparse, yaml, json +from pathlib import Path +from .core.logger import get_logger +from .core.utils import load_yaml, load_prompts, render_template +from .agents.a1_cve_analysis import CVEAnalysisAgent +from .agents.a2_template_mapping import TemplateMappingAgent +from .agents.a3_sequence_generation import SequenceGenerationAgent +from .agents.a4_syzkaller_orchestration import SyzkallerOrchestrationAgent +from .agents.a5_poc_converter import PoCConverterAgent + +log = get_logger("CLI") + +def load_cfg() -> dict: + p = Path("configs/config.yaml") + if p.exists(): + return yaml.safe_load(p.read_text(encoding="utf-8")) + return yaml.safe_load(Path("configs/config.example.yaml").read_text(encoding="utf-8")) + +def cmd_analyze(args): + cfg = load_cfg() + prompts = load_prompts(cfg["paths"]["prompts_file"]) + # 渲染用户提示（示例，当前未调用 LLM，仅保留痕迹） + _ = render_template(prompts["cve_analysis"]["user"], cve_id=args.cve) + + agent = CVEAnalysisAgent(api_base=cfg["nvd"]["api_base"], timeout=cfg["nvd"]["timeout"]) + out = agent.run(args.cve) + + Path("out").mkdir(exist_ok=True) + Path("out/fields.yaml").write_text(yaml.safe_dump(out.fields, allow_unicode=True, sort_keys=False), encoding="utf-8") + Path("out/summary_zh.txt").write_text(out.summary_zh, encoding="utf-8") + print("已生成：out/fields.yaml, out/summary_zh.txt") + +def cmd_map_template(args): + cfg = load_cfg() + kb_file = cfg["paths"]["kb_file"] + agent = TemplateMappingAgent(kb_file) + feature = f"{args.cve} 堆溢出 协议消息长度校验缺陷（占位）" + res = agent.run(feature) + Path("out").mkdir(exist_ok=True) + Path("out/mapping.json").write_text(json.dumps(res.__dict__, ensure_ascii=False, indent=2), encoding="utf-8") + print("已生成：out/mapping.json") + +def cmd_gen_skeleton(args): + cfg = load_cfg() + kb_file = cfg["paths"]["kb_file"] + mapping = json.loads(Path("out/mapping.json").read_text(encoding="utf-8")) + tid = mapping.get("template_id") or "MEM_HEAP_OVERFLOW" + agent = SequenceGenerationAgent(kb_file) + sk = agent.run(tid) + out_dir = Path(args.out); out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "skeleton.txt").write_text( + "模板: " + tid + "\n" + "\n".join(sk.steps) + "\n\n" + sk.disclaimer, encoding="utf-8" + ) + print(f"已生成：{out_dir}/skeleton.txt") + +def cmd_orchestrate(args): + cfg = load_cfg() + orch = cfg.get("orchestration", {}) + agent = SyzkallerOrchestrationAgent( + dry_run=orch.get("dry_run", True), + max_parallel=int(orch.get("max_parallel", 2)), + log_dir=orch.get("log_dir", "out/logs"), + ) + plan = agent.run() + Path("out").mkdir(exist_ok=True) + Path("out/orchestration.json").write_text(json.dumps(plan.__dict__, ensure_ascii=False, indent=2), encoding="utf-8") + print("已生成：out/orchestration.json") + +def cmd_package(args): + agent = PoCConverterAgent() + res = agent.run(args.case, args.out) + print("已生成包：", res.out_dir) + +def main(): + p = argparse.ArgumentParser(prog="kvra-pro", description="Safe-by-design multi-agent framework (Modular)") + s = p.add_subparsers() + + a = s.add_parser("analyze"); a.add_argument("--cve", required=True); a.set_defaults(func=cmd_analyze) + b = s.add_parser("map-template"); b.add_argument("--cve", required=True); b.set_defaults(func=cmd_map_template) + c = s.add_parser("gen-skeleton"); c.add_argument("--cve", required=True); c.add_argument("--out", required=True); c.set_defaults(func=cmd_gen_skeleton) + d = s.add_parser("orchestrate"); d.set_defaults(func=cmd_orchestrate) + e = s.add_parser("package"); e.add_argument("--case", required=True); e.add_argument("--out", required=True); e.set_defaults(func=cmd_package) + + args = p.parse_args() + if hasattr(args, "func"): + args.func(args) + else: + p.print_help() + +if __name__ == "__main__": + main() diff --git a/src/core/.DS_Store b/src/core/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 GIT binary patch literal 6148 zcmeH~Jr2S!425mzP>H1@V-^m;4Wg<&0T*E43hX&L&p$$qDprKhvt+--jT7}7np#A3 zem<@ulZcFPQ@L2!n>{z**++&mCkOWA81W14cNZlEfg7;MkzE(HCqgga^y>{tEnwC%0;vJ&^%eQ zLs35+`xjp>T0 Optional[TemplateHint]: + tl = (text or "").lower() + for rule in self.keyword_rules: + inc = [k.lower() for k in rule.get("include", [])] + if any(k in tl for k in inc): + return self.templates.get(rule.get("map_to")) + return None + + def get_template(self, template_id: str) -> Optional[TemplateHint]: + return self.templates.get(template_id) diff --git a/src/core/logger.py b/src/core/logger.py new file mode 100644 index 0000000..e4db5e9 --- /dev/null +++ b/src/core/logger.py @@ -0,0 +1,9 @@ +import logging +def get_logger(name: str) -> logging.Logger: + logger = logging.getLogger(name) + if not logger.handlers: + logger.setLevel(logging.INFO) + ch = logging.StreamHandler() + ch.setFormatter(logging.Formatter("[%(levelname)s] %(asctime)s %(name)s: %(message)s")) + logger.addHandler(ch) + return logger diff --git a/src/core/nvd_client.py b/src/core/nvd_client.py new file mode 100644 index 0000000..3210cd5 --- /dev/null +++ b/src/core/nvd_client.py @@ -0,0 +1,20 @@ +from __future__ import annotations +from dataclasses import dataclass +from typing import Dict, Any +import requests + +@dataclass +class NVDConfig: + api_base: str = "https://services.nvd.nist.gov/rest/json/cves/2.0" + timeout: int = 15 + +class NVDClient: + def __init__(self, cfg: NVDConfig = NVDConfig()): + self.cfg = cfg + def get_cve(self, cve_id: str) -> Dict[str, Any]: + try: + r = requests.get(self.cfg.api_base, params={"cveId": cve_id}, timeout=self.cfg.timeout) + r.raise_for_status() + return r.json() + except Exception as e: + return {"offline": True, "cveId": cve_id, "error": str(e)} diff --git a/src/core/utils.py b/src/core/utils.py new file mode 100644 index 0000000..4b68671 --- /dev/null +++ b/src/core/utils.py @@ -0,0 +1,13 @@ +from __future__ import annotations +from pathlib import Path +from string import Template +import yaml + +def load_yaml(path: str) -> dict: + return yaml.safe_load(Path(path).read_text(encoding="utf-8")) + +def load_prompts(path: str) -> dict: + return load_yaml(path) + +def render_template(text: str, **kwargs) -> str: + return Template(text).safe_substitute(**kwargs) diff --git a/tests/.DS_Store b/tests/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..5008ddfcf53c02e82d7eee2e57c38e5672ef89f6 GIT binary patch literal 6148 zcmeH~Jr2S!425mzP>H1@V-^m;4Wg<&0T*E43hX&L&p$$qDprKhvt+--jT7}7np#A3 zem<@ulZcFPQ@L2!n>{z**++&mCkOWA81W14cNZlEfg7;MkzE(HCqgga^y>{tEnwC%0;vJ&^%eQ zLs35+`xjp>T0