Aikagra-rgb/cve-severity-classifier

# 🛡️ 基于 AI 的 CVE 严重性分类器与网络威胁情报系统 [](https://python.org) [](https://pytorch.org) [](https://huggingface.co) [](https://streamlit.io) [](LICENSE) ## 🎯 项目功能 给定 CVE 描述和 CWE-ID，系统会： 1. **分类严重性** → `LOW` / `MEDIUM` / `HIGH` / `CRITICAL`，并附带置信度分数 2. **生成威胁 playbook** → 攻击类别、受影响资产以及分步修复方案 3. **映射合规性** → NIST SP 800-61、CIS Critical Controls、ISO 27001 ## 🏗️ 架构 ``` User Input (CVE Description + CWE-ID) │ ├──► DistilBERT + CWE Embedding ──► LayerNorm ──► 512-dim GELU MLP ──► Severity Label │ (fine-tuned on 169K+ NVD CVEs) │ └──► Mistral 7B via Ollama ──► Threat Playbook + Attack Category + Remediation (runs locally, no API key needed) ``` ### 模型升级 (v5) | 组件 | 详情 | |---|---| | 基础模型 | `distilbert-base-uncased` | | CWE Embedding | 针对 19 个顶级 CWE 类别的 32 维学习向量 | | 特征融合 | `[CLS] 768维 ⊕ CWE 32维 = 800维` | | 归一化 | LayerNorm (平衡预训练与随机初始化的规模) | | 分类器头 | 512 维 Linear → GELU → Dropout(0.3) → 4 分类 Softmax | | CWE 前缀 | 在分词时注入 `"Vulnerability Class: CWE-XXX. "` | ## 📂 项目结构 ``` ├── Malware_classification_v5.ipynb # Main training notebook (end-to-end pipeline) ├── app_v2.py # Streamlit web app (inference + playbook UI) ├── run_training.py # Standalone training script ├── ollama_labeler.py # Auto-labeler using local Mistral 7B ├── fix_notebook_id.py # Utility: fixes notebook cell IDs ├── update_notebook_hybrid.py # Utility: applies hybrid model upgrades ├── Cyber-Threat-Intelligence-Custom-Data_new_processed.csv # Custom CTI dataset ├── distilbert_malware_classifier/ # Saved model checkpoint (tokenizer + config) │ ├── config.json │ ├── label_config.json │ ├── tokenizer.json │ └── tokenizer_config.json ├── confusion_matrix_v5.png # Evaluation: confusion matrix ├── training_progress_v5.png # Evaluation: train/val loss curves ├── model_comparison_v5.png # Evaluation: model comparison chart └── .gitignore ``` ## 🚀 快速开始 ### 1. 克隆仓库 ``` git clone https://github.com/Aikagra-rgb/cve-severity-classifier.git cd cve-severity-classifier ``` ### 2. 安装依赖 ``` pip install torch transformers streamlit pandas scikit-learn openpyxl requests ``` ### 3. 安装 Ollama（用于威胁 playbook agent） ``` # 从 https://ollama.com 下载 ollama pull mistral ``` ### 4. 运行 Streamlit 应用 ``` streamlit run app_v2.py ``` ### 5. 从零开始训练（可选） - 在 Jupyter 中打开 `Malware_classification_v5.ipynb` - 从 [NIST NVD](https://nvd.nist.gov/vuln/data-feeds) 下载 `Global_Dataset.xlsx` 并将其放在项目文件夹中 - 按顺序运行所有单元格 ## 📊 数据集与模型权重 | 文件 | 大小 | 来源 | |------|------|--------| | `Global_Dataset.xlsx` | ~16 MB | ✅ 已包含 — NIST NVD CVE 数据库 (16.9万+ 条记录) | | `model.safetensors` | ~267 MB | 本地微调 — 从 [Releases](https://github.com/Aikagra-rgb/cve-severity-classifier/releases) 下载 | 获取模型权重的方式： - 使用 `Malware_classification_v5.ipynb` **自行训练** - 从 [Releases](https://github.com/Aikagra-rgb/cve-severity-classifier/releases) **下载**（如果已单独上传） ## 📈 结果 | 指标 | 分数 | |--------|-------| | 分类准确率 | ~82–88%（在 v5 升级后） | | Macro F1 Score | ~0.80+ | | 类别 | LOW, MEDIUM, HIGH, CRITICAL | | 训练数据集大小 | 169,566 条 CVE 记录 | ## 🛠️ 技术栈 - **Python 3.10+** - **PyTorch 2.x** — 模型训练与推理 - **HuggingFace Transformers** — DistilBERT tokenizer 与基础模型 - **Streamlit** — 交互式 Web UI - **Ollama + Mistral 7B** — 本地生成式威胁 agent - **Pandas / Scikit-learn** — 数据处理与评估 - **CUDA** — GPU 加速（可选） ## 👥 团队 作为一个**基于项目的学习 (PBL)** 学术项目构建。 ## 📄 许可证 MIT 许可证 — 详情请参阅 [LICENSE](LICENSE)。

标签：AI风险缓解, DLL 劫持, GPT, Kubernetes, 人工智能, 凭据扫描, 大语言模型, 威胁情报, 安全运营, 开发者工具, 扫描框架, 漏洞管理, 用户模式Hook绕过, 自动化运维, 逆向工具