opandey1/AI-SOC-Assistant

# 驱动 AI 的 SOC 助手 一个可解释的 AI 分诊管道，可将 NSL-KDD 网络连接分类到特定的攻击家族中，融合有监督和无监督的异常信号，并利用 GenAI 生成可供分析师直接使用的 SOC 事件工单。 [](https://github.com/opandey1/AI-SOC-Assistant/actions/workflows/ci.yml) [](https://www.python.org/) [](LICENSE) [](https://www.unb.ca/cic/datasets/nsl.html) [](https://colab.research.google.com/github/opandey1/AI-SOC-Assistant/blob/main/notebooks/AI_Powered_SOC_Assistant.ipynb)  ## 查看运行情况  ## 优势与独特之处 - **原生多类 SOC 检测：** 区分特定的攻击家族（Normal, DoS, Probe, R2L, U2R），而不是仅输出单一的二元异常标志。 - **双模型分诊：** 使用 Random Forest 预测攻击家族，同时利用 Isolation Forest 为可疑流量模式增加无监督的异常信号。 - **可解释的证据：** SHAP 识别出驱动每个被标记连接的最强特征，并将分析师可读的数值传递到工单中。 - **本地优先的 GenAI：** 支持 Ollama，允许助手在生成工单时无需将原始网络遥测数据发送到外部 LLM API。 - **可操作的输出：** 最终响应是一份结构化的事件工单，包含遏制步骤和可直接复制粘贴的 Splunk SPL 查询。 ## 结果与指标 基于 **`KDDTrain+` 的分层 80/20 留出划分**进行评估（SMOTE 仅应用于训练折）。使用 `python -m src.evaluate` 重新生成以下所有内容。 | 攻击家族 | Precision | Recall | F1-score | Support | | --- | --- | --- | --- | --- | | `normal` | 0.999 | 0.999 | 0.999 | 13,469 | | `dos` | 1.000 | 1.000 | 1.000 | 9,186 | | `probe` | 0.998 | 0.997 | 0.998 | 2,331 | | `r2l` | 0.985 | 0.970 | 0.977 | 199 | | `u2r` | 0.818 | 0.900 | 0.857 | 10 | **总体准确率：99.90% · Macro F1：0.966 · Weighted F1：0.999**  ## 可解释性实战 对于每个被标记的连接，SHAP 会对驱动预测的特征进行排序，并且管道会将它们的**真实的、未缩放的网络数值**（而非 z-scores）转发到工单中。下面的示例是一个典型的半开 SYN 洪水攻击：`flag_S0` 处于活动状态，且 SYN 错误率为 100%。  这些驱动因素在生成的工单中变成了通俗易懂的证据： ``` 3. Why Flagged - Evidence - flag_S0 was active, meaning connections were initiated but never acknowledged - the classic signature of a half-open SYN flood. - serror_rate was 1.0, so every recent connection ended in a SYN error. - dst_host_serror_rate / dst_host_srv_serror_rate were both 1.0 ... ``` 请查看完整的[样本工单](docs/sample_ticket.md)和底层的[SHAP 证据包](docs/shap_example_output.json)。 ## MITRE ATT&CK 映射 每个检测到的攻击家族都映射到一种 MITRE ATT&CK 战术，为分析师提供了一个用于分诊和升级的共享参考框架。 | 攻击家族 | MITRE ATT&CK 战术 | 战术 ID | 代表性技术 | | --- | --- | --- | --- | | `dos` | Impact | [TA0040](https://attack.mitre.org/tactics/TA0040/) | Network Denial of Service (T1498), Endpoint DoS (T1499) | | `probe` | Discovery | [TA0007](https://attack.mitre.org/tactics/TA0007/) | Network Service Scanning (T1046), System Network Discovery (T1016) | | `r2l` | Initial Access | [TA0001](https://attack.mitre.org/tactics/TA0001/) | Valid Accounts (T1078), Exploit Public-Facing Application (T1190) | | `u2r` | Privilege Escalation | [TA0004](https://attack.mitre.org/tactics/TA0004/) | Exploitation for Privilege Escalation (T1068), Abuse Elevation Control (T1548) | | `normal` | — | — | 良性基线流量 | ## 快速开始 ``` git clone https://github.com/opandey1/AI-SOC-Assistant.git cd AI-SOC-Assistant python -m venv .venv source .venv/bin/activate pip install -r requirements.txt ``` 将 `KDDTrain+.txt` 和 `KDDTest+.txt` 放在仓库根目录或 `data/` 下，然后运行一个无 LLM 的冒烟测试： ``` python src/pipeline.py --no-llm ``` 使用本地 Ollama 模型生成工单： ``` ollama pull mistral SOC_LLM_PROVIDER=ollama python src/pipeline.py ``` PowerShell 用户可以这样设置 provider： ``` $env:SOC_LLM_PROVIDER = "ollama" python src/pipeline.py ``` 重新生成指标、混淆矩阵和 SHAP 产物： ``` python -m src.evaluate ``` ### 使用 Docker 运行 一个双容器堆栈在与本地 Ollama 服务器一起运行管道，因此原始遥测数据永远不会离开主机： ``` # 离线 template path（无需 model）： docker compose run --rm soc-assistant python src/pipeline.py --no-llm # 带有本地 model 的完整 GenAI path： docker compose up -d ollama docker compose exec ollama ollama pull mistral docker compose run --rm soc-assistant ``` ### 测试 ``` pip install pytest pytest ``` ## 仓库结构 ``` src/ ingest.py NSL-KDD loading and attack-family mapping preprocess.py one-hot encoding, scaling, and SMOTE balancing train.py Random Forest, Isolation Forest, and fused scoring explain.py SHAP explanation bundle generation agent.py LangGraph SOC analyst agent and threat-intel tools evaluate.py metrics, confusion matrix, and SHAP artifact generation pipeline.py runnable command-line pipeline tests/ pytest unit tests for ingestion, preprocessing, SHAP, tickets notebooks/ AI_Powered_SOC_Assistant.ipynb docs/ architecture, metrics, plots, and sample artifacts .github/workflows/ci.yml black + flake8 + pytest on every push Dockerfile, docker-compose.yml ``` ## 演示产物 - [架构图](docs/soc_architecture.svg) - [动画管道演示](docs/demo.svg) - [混淆矩阵](docs/confusion_matrix.png)和[各类指标](docs/metrics.md) - [SHAP 驱动因素图](docs/shap_drivers.png)和[SHAP 证据包](docs/shap_example_output.json) - [生成的样本工单](docs/sample_ticket.md) - [演进简报](docs/SOC_Assistant_Evolution.pdf) - [Colab notebook](notebooks/AI_Powered_SOC_Assistant.ipynb) ## 演进与成果 ### 1. 特征工程：稳健的分类编码 **初始状态：** 管道对网络分类特征（如 protocol, service 和 flag）使用 `LabelEncoder`。 **挑战：** `LabelEncoder` 会产生虚假的数学序数关系，并且当实时流量引入了训练期间未见的类别时可能会崩溃。 **修复：** 预处理管道现在使用 `OneHotEncoder(handle_unknown="ignore")`。 **成果：** 特征表示在数学上更加合理，并且对新的服务值具有鲁棒性。 ### 2. Agent 编排：迁移到 LangGraph **初始状态：** 助手逻辑依赖于较旧的 LangChain agent 模式。 **挑战：** 现代的工具调用工作流受益于更清晰的状态处理和更可靠的 agent 执行。 **修复：** Agent 层使用 `langgraph.prebuilt.create_react_agent`。 **成果：** 编排层更易于维护，并且更好地与当前的 LangChain/LangGraph 模式保持一致。 ### 3. 环境限制：物理隔离的 LLM 执行 **初始状态：** LLM 的使用与云服务提供商的调用绑定。 **挑战：** SOC 环境通常会限制将原始遥测数据、内部 IP 和安全日志传输到外部 API。 **修复：** `initialize_llm` 支持一个包含 Ollama, OpenAI, Anthropic 和模板模式的 provider 工厂，并且重量级的 LLM 依赖项采用延迟加载，这样离线的 `--no-llm` 路径就不需要它们。 **成果：** 管道可以使用 Ollama 在本地运行，适用于注重隐私的演示和受限环境。 ### 4. 威胁情报集成：容错的 API 工具 **初始状态：** 威胁情报工具只是占位符。 **挑战：** 真实的 API 可能会失败、受到速率限制，或者返回大量淹没 LLM 上下文的冗长嵌套 payload。 **修复：** AbuseIPDB 和 NVD 查询使用超时、显式错误处理和紧凑的响应格式。 **成果：** API 失败变成了工单上下文，而不是 Python 进程崩溃。 ### 5. 可解释 AI：跨版本的 SHAP 兼容性 **初始状态：** SHAP 提取假设为单一的输出形状。 **挑战：** SHAP 在不同版本中改变了多类输出的格式。 **修复：** 解释层同时处理传统的列表输出和较新的 3D 数组输出。 **成果：** 代码在不同的依赖版本之间具有更好的可移植性。 ### 6. 取证准确性：真实世界的指标 **初始状态：** 证据包存在将缩放后的 z-score 值暴露给 LLM 的风险。 **挑战：** SOC 分析师需要真实的数据包、字节、计数和速率值，而不是标准化的模型输入。 **修复：** 模型推理仍然使用缩放后的值，而 SHAP 证据包含未缩放的处理特征值。 **成果：** 生成的工单读起来像分析师的证据，而不是模型的内部数据。 ### 7. GenAI 护栏：消除幻觉和过度生成 **初始状态：** 较小的本地模型可能会偏题生成代码片段或原始的浮点 SHAP 值。 **挑战：** 助手必须生成简洁的事件工单，而不是通用的解释或脚本。 **修复：** 系统 prompt 强制仅输出工单、结构化的章节、精确的 Splunk SPL 代码块，并禁止泄露原始的 SHAP 浮点数。 **成果：** GenAI 层将模型证据转化为可操作的 SOC 工作流输出。

标签：AI安全运营, Apex, DLL 劫持, Python, 入侵检测, 大语言模型, 威胁狩猎, 无后门, 机器学习