architway/AI-Threat-Hunting-Query-Generation-Evaluation-System

# AI Strike 威胁狩猎查询生成 AI Strike 可将自然语言的 AWS CloudTrail 威胁狩猎假设转化为可执行的 DuckDB SQL，对该 SQL 在提供的 `nineteenFeaturesDf.csv` 数据集上运行，并根据 `hypotheses_outcomes.json` 对返回的行进行评估。每个生成的查询都包含解释、推理、假设、置信度、执行状态以及 precision/recall/F1 指标。 本项目采用离线优先策略。Mock 模式无需 API key 即可运行完整的本地 pipeline，而真实运行则使用 Cerebras。附加功能层包括 Streamlit 演示、prompt 策略对比、基准测试、执行失败后的一次性 SQL 修复、Docker 文件，以及基于 AWS 官方文档的 AWS 领域上下文 prompt 调优。 本提交策略具有极强的实用性：使用官方 AWS/领域研究、感知 schema 的 prompting、确定性的本地验证、低成本的 Cerebras 托管模型选项，以及离线 mock 路径，从而在不依赖反复调用昂贵前沿模型的情况下获得高评估价值。 ## 快速开始 在 Windows PowerShell 中从仓库根目录运行以下命令。进行 CLI、GUI 和 Docker 相关操作时，指定的 CSV 文件 `nineteenFeaturesDf.csv` 必须存在于仓库根目录中。 ### 1. 创建并激活虚拟环境 ``` python -m venv .venv .\.venv\Scripts\Activate.ps1 ``` 如果 PowerShell 在你的机器上阻止了激活，请在以下命令中直接使用 venv 的 Python： ``` .\.venv\Scripts\python.exe --version ``` ### 2. 安装依赖 ``` python -m pip install --upgrade pip python -m pip install -r requirements.txt ``` ### 3. 运行离线 / Mock CLI 冒烟测试 Mock 模式不进行任何网络调用，并会执行生成器、DuckDB 执行器、评估器、报告及输出产物： ``` .\.venv\Scripts\python.exe launch_check.py .\.venv\Scripts\python.exe main.py --mock-llm --prompt-strategy aws_domain --limit 1 ``` ### 4. 运行真实的 Cerebras CLI 将 `.env.example` 复制到 `.env`，填入 `CEREBRAS_API_KEY`，并保留 Cerebras 设置： ``` Copy-Item .env.example .env notepad .env ``` ``` CEREBRAS_API_KEY=your_cerebras_key_here CEREBRAS_BASE_URL=https://api.cerebras.ai/v1 AI_STRIKE_PROVIDER=cerebras AI_STRIKE_MODEL=gpt-oss-120b ``` 然后在进行更广泛的评估之前，先运行一次真实的单假设检查： ``` .\.venv\Scripts\python.exe launch_check.py --real $env:AI_STRIKE_REQUEST_BUDGET="1"; .\.venv\Scripts\python.exe main.py --hypothesis-id 1 --prompt-strategy aws_domain --model gpt-oss-120b --repair-attempts 0 --verbose ``` 修复/自愈默认设置为 `AI_STRIKE_REPAIR_ATTEMPTS=1`。它仅在第一次生成的 SQL 被 DuckDB 拒绝时才使用第二次模型调用；成功的 SQL 不会触发修复。 ### 5. 运行 Streamlit GUI ``` .\.venv\Scripts\streamlit.exe run app.py ``` Streamlit 启动后会打印一个本地 URL。在浏览器中打开该 URL： ``` http://localhost:8501 ``` GUI 默认以 Mock 模式打开。保持开启 Mock 模式即可进行离线演示。只有在 `.env` 中拥有有效的 `CEREBRAS_API_KEY` 并且在侧边栏中选择了某个 Cerebras 模型后，才能关闭 Mock 模式。在终端中按 `Ctrl+C` 可停止 GUI 服务器，并在修改 `.env` 或代码后重新启动它。 在 GUI 中切换到真实模式之前，请先运行此 CLI 预检： ``` .\.venv\Scripts\python.exe launch_check.py --real ``` ### 6. 使用 Docker Compose 运行 必须单独安装 Docker Desktop 或 Docker Engine。请将 `nineteenFeaturesDf.csv` 保留在仓库根目录中；Compose 会将其以只读方式挂载到容器中。 ``` docker compose build docker compose run --rm ai-strike-cli docker compose run --rm ai-strike-cli python main.py --mock-llm --prompt-strategy aws_domain --limit 1 docker compose up ai-strike-demo ``` Docker GUI 也会在 `http://localhost:8501` 打开。 ## 设置 请使用本地虚拟环境，以便将本作业的依赖保留在此项目文件夹内。 ``` python -m venv .venv .\.venv\Scripts\Activate.ps1 python -m pip install --upgrade pip python -m pip install -r requirements.txt ``` ## Cerebras 配置 Mock 模式不需要凭证。若要进行真实的模型调用，请将 `.env.example` 复制到 `.env` 并设置 Cerebras 的 key 和模型： ``` Copy-Item .env.example .env ``` ``` CEREBRAS_API_KEY=your_cerebras_key_here CEREBRAS_BASE_URL=https://api.cerebras.ai/v1 AI_STRIKE_PROVIDER=cerebras AI_STRIKE_MODEL=gpt-oss-120b AI_STRIKE_PROMPT_STRATEGY=aws_domain AI_STRIKE_REPAIR_ATTEMPTS=1 AI_STRIKE_REQUEST_BUDGET=25 ``` 使用 `AI_STRIKE_MODEL` 指定你要测试的 Cerebras 模型 ID。已保存的对比报告选择 `gpt-oss-120b` 作为主要提交模型，以 `qwen-3-235b-a22b-instruct-2507` 作为最强的备选方案，并以 `llama3.1-8b` 作为更便宜/轻量化的选项。`AI_STRIKE_REPAIR_ATTEMPTS=1` 是默认设置；它仅在 DuckDB 拒绝第一次生成的 SQL 时才会消耗额外的请求。 ## Mock / 离线模式 Mock 模式可在无网络访问的情况下运行生成器、DuckDB 执行器、评估器、产物生成器、报告、基准测试钩子以及 Streamlit 路径： ``` .\.venv\Scripts\python.exe main.py --mock-llm --limit 1 .\.venv\Scripts\python.exe main.py --hypothesis-id 1 --mock-llm --verbose .\.venv\Scripts\python.exe main.py --mock-llm --prompt-strategy aws_domain --limit 2 ``` Mock 生成器是确定性的。它用于冒烟测试和可重复的本地评估，并不代表真实模型的质量。 ## 真实 Cerebras 运行 真实运行会消耗服务商的请求和 token。请从小规模开始，在单个假设上验证 JSON/SQL 行为，只有在启动检查通过后才进行扩展。 ``` .\.venv\Scripts\python.exe launch_check.py --real $env:AI_STRIKE_REQUEST_BUDGET="1"; .\.venv\Scripts\python.exe main.py --hypothesis-id 1 --prompt-strategy aws_domain --model gpt-oss-120b --repair-attempts 0 --verbose .\.venv\Scripts\python.exe main.py --hypothesis-id 4 --model gpt-oss-120b .\.venv\Scripts\python.exe main.py --prompt-strategy aws_domain ``` 也可使用离线启动检查： ``` .\.venv\Scripts\python.exe launch_check.py ``` ## Streamlit GUI Streamlit 演示是一个轻量级的 UI，运行在与 CLI 相同的 pipeline 之上。它支持假设选择、Mock 或真实模式、prompt 策略选择、修复尝试控制、SQL/解释显示、指标、查询建议以及结果预览。 ``` .\.venv\Scripts\streamlit.exe run app.py ``` GUI 默认为 Mock 模式，因此无需 API key 即可打开。 真实模式故障排除： - 如果 `.env` 发生更改，请使用 `Ctrl+C` 停止所有正在运行的 Streamlit 终端，然后重新启动： `.\.venv\Scripts\streamlit.exe run app.py` - 如果你看到 `[WinError 10061]` / `ConnectionRefusedError`，说明本地机器无法从运行中的 Streamlit 进程连接到 `https://api.cerebras.ai/v1`。常见原因包括： 无网络连接、防火墙/代理策略、DNS/网络问题，或 Streamlit 进程已过期。 这是一个服务商连接问题，而不是 SQL/prompt/自定义指令失败。 - 某些环境会导出无效的本地代理变量（例如： `HTTP_PROXY=http://127.0.0.1:9`, `HTTPS_PROXY=http://127.0.0.1:9`, `ALL_PROXY=http://127.0.0.1:9`）。本项目中的 Cerebras/OpenAI 客户端现在默认忽略代理环境变量（`trust_env=False`）以避免出现该故障。 - 要诊断 shell 中的代理变量： `Get-ChildItem Env: | Where-Object { $_.Name -match 'proxy' }` - 在重试真实 GUI 模式之前，请运行 `.\.venv\Scripts\python.exe launch_check.py --real`。 修复/自愈默认尝试一次。它仅在 DuckDB 拒绝第一次生成的 SQL 时触发，因此首次通过的查询成功不会消耗额外的模型请求。你依然可以在每次运行时覆盖它： ``` .\.venv\Scripts\python.exe main.py --hypothesis-id 4 --repair-attempts 1 .\.venv\Scripts\python.exe main.py --hypothesis-id 4 --repair-attempts 0 $env:AI_STRIKE_REPAIR_ATTEMPTS = "1"; .\.venv\Scripts\streamlit.exe run app.py ``` ## 实验 使用相同的评估器对比 prompt 策略和模型标签： ``` .\.venv\Scripts\python.exe experiments.py --mock-llm --strategies base,structured,multi_step,aws_domain --limit 2 .\.venv\Scripts\python.exe experiments.py --models gpt-oss-120b,qwen-3-235b-a22b-instruct-2507 --strategies aws_domain --limit 3 --request-budget 20 ``` 输出： - `comparison_results.json` - `comparison_summary.csv` - `MODEL_COMPARISON_REPORT.md` 真实实验会在开始前估算请求数量，并拒绝超过 `AI_STRIKE_REQUEST_BUDGET` 或 `--request-budget` 的限制。 ## 基准测试 对生成、DuckDB 执行、评估及总运行时间进行基准测试： ``` .\.venv\Scripts\python.exe benchmark.py --mock-llm --limit 2 .\.venv\Scripts\python.exe benchmark.py --mock-llm --prompt-strategy aws_domain --limit 2 .\.venv\Scripts\python.exe benchmark.py --hypothesis-id 4 --prompt-strategy aws_domain ``` 输出： - `benchmark_results.json` - `BENCHMARK_REPORT.md` ## 输出文件 每次 CLI 运行都会在 `outputs/` 下写入一个带有时间戳的文件夹： ``` outputs/ run_YYYYMMDD_HHMMSS/ generated_queries.json evaluation_results.json summary.csv errors.json ``` 根目录下的 `evaluation_results.json` 也会被写入，因为作业将其指定为交付物。在已提交的状态下，它反映了最新的离线冒烟测试运行结果。`EVALUATION_REPORT.md` 保留了该当前产物，并总结了保存在 `benchmark_results.json` 和 `MODEL_COMPARISON_REPORT.md` 中的最终全假设 Cerebras 对比结果。 主要提交交付物： - `main.py`, `pipeline.py`, `query_generator.py`, `evaluator.py`, `executor.py` - `README.md`, `APPROACH.md`, `EVALUATION_REPORT.md` - `evaluation_results.json` - `requirements.txt` - 附加产物：`app.py`, `Dockerfile`, `docker-compose.yml`, `benchmark.py`, `experiments.py`, `BENCHMARK_REPORT.md`, `MODEL_COMPARISON_REPORT.md` ## 提交清单 所需交付物： - [x] 核心查询生成逻辑：`query_generator.py`, `prompts.py`, `llm_client.py`, 和 `pipeline.py` - [x] 评估框架实现：`evaluator.py` - [x] 运行完整评估的入口点：`main.py` - [x] 依赖项：`requirements.txt` - [x] 关键组件的单元测试 - [x] 包含设置说明、架构概述、设计决策、权衡取舍和扩展指南的 README - [x] 包含 prompting 策略、迭代过程、挑战、解决方案、局限性和未来工作的 APPROACH.md - [x] `evaluation_results.json` 包含的评估输出产物 - [x] `EVALUATION_REPORT.md` 包含总体指标、按假设划分的明细以及修改前后的说明 - [x] 可解释的生成输出：解释、推理、假设、威胁说明和置信度 - [x] 确保预期结果仅由评估器在 SQL 执行后使用的防护机制 可选和附加项： - [x] 交互式演示：`app.py` 中的 Streamlit UI - [x] 容器化：`Dockerfile` 和 `docker-compose.yml` - [x] 查询优化建议：`query_advisor.py` 中确定性的评估后建议 - [x] prompt 策略和多步 prompting：`base`, `structured`, `multi_step`, 和 `aws_domain` - [x] 每个生成的查询中都包含带有解释的置信度评分 - [x] 在 DuckDB 失败后，通过默认的 `AI_STRIKE_REPAIR_ATTEMPTS=1`、`--repair-attempts` 或 Streamlit 进行自动化 SQL 修复/自愈 - [x] 扩展评估：使用 `experiments.py` 进行 prompt 策略对比 - [x] 通过 `experiments.py --models` 支持 A/B 模型测试 - [x] 使用 `benchmark.py` 和 `BENCHMARK_REPORT.md` 进行性能基准测试 - [x] 在 `MODEL_COMPARISON_REPORT.md` 中进行 Cerebras 模型对比 - [x] 用于测试、演示、Docker 冒烟检查和可重复本地评估的离线 Mock 模式 ## 架构 ``` hypotheses.json + CSV schema + optional AWS domain context | v prompts.py -> query_generator.py -> Cerebras or MockLLMClient | v DuckDB SQL against cloudtrail view over nineteenFeaturesDf.csv | v executor.py -> evaluator.py -> query_advisor.py | v outputs/, evaluation_results.json, reports, Streamlit, experiments, benchmarks ``` 预期结果仅在 SQL 执行后由评估器加载。它们不会包含在 prompt、修复 prompt、领域上下文或重用的聊天历史中。 ## Docker 可选选项 Docker 是可选的，并且独立于 `.venv` 工作流。必须在外部安装 Docker Desktop 或 Docker Engine；正常的本地运行不需要它。 首先将 `nineteenFeaturesDf.csv` 放在项目根目录中。该大型 CSV 会通过 `.dockerignore` 被排除在镜像之外，并在运行时由 Compose 以只读方式挂载。 Mock CLI 冒烟测试： ``` docker compose build docker compose run --rm ai-strike-cli docker compose run --rm ai-strike-cli python main.py --mock-llm --prompt-strategy aws_domain --limit 2 ``` Streamlit 演示： ``` docker compose up ai-strike-demo ``` 演示服务在 `http://localhost:8501` 上监听。 在 PowerShell 中使用宿主机上已加载的环境变量运行真实的 Cerebras 单假设测试： ``` $env:CEREBRAS_API_KEY = "your_cerebras_key_here" docker compose run --rm -e CEREBRAS_API_KEY -e AI_STRIKE_PROVIDER=cerebras -e AI_STRIKE_MODEL=gpt-oss-120b ai-strike-cli python main.py --hypothesis-id 1 --prompt-strategy aws_domain ``` ## 扩展至其他数据集 要使本项目适应另一个安全数据集： 1. 将 `AI_STRIKE_DATA_PATH` 指向新的 CSV，或更新执行器以适配其他数据源。 2. 更新 `hypotheses.json` 和 `hypotheses_outcomes.json`。 3. 调整 `aws_domain_prompt_context.md`，或针对新的 schema 和平台添加新的领域上下文文件。 4. 如果数据集使用的是与 `eventID`、`row_id` 或聚合 `count` 不同的稳定标识符，请更新评估器的身份验证规则。 5. 仅从安全的领域知识中添加 prompt 示例或策略说明，切勿使用预期结果行中的内容。 ## 测试 ``` .\.venv\Scripts\python.exe -m pytest .\.venv\Scripts\python.exe main.py --mock-llm --limit 1 ```

标签：AI代码生成, AWS CloudTrail, DLL 劫持, DuckDB, Kubernetes, Streamlit, 反取证, 大语言模型, 安全评估, 访问控制, 请求拦截, 逆向工具