│              │              │           │         │
      └──────────────┴──────────────┴───────────┴─────────┘
                       │
                       ▼
              ┌─────────────────┐
              │   SQLite DB     │
              │  (40+ Columns)  │
              └────────┬────────┘
                       │
    ┌──────────────────┼──────────────────┐
    ▼                  ▼                  ▼

### 数据处理流水线

1. **发现阶段**（可配置间隔，默认30分钟）
   - Shodan扫描（可配置查询）
   - NVD增强（含CVE关联）
   - 内部系统增强（上下文信息）

2. **分析阶段**
   - 基于40+属性的特征工程
   - 集成模型预测（随机森林40% + 梯度提升40% + MLP 20%）
   - 风险标签生成（0=低, 1=中, 2=高）
   - 置信度评分

3. **告警阶段**
   - 检查新增高风险资产
   - 通过SendGrid生成并发送邮件告警
   - 更新通知状态

4. **模型重新训练**（可配置间隔，默认60分钟）
   - 加载所有历史数据
   - 提取特征
   - 使用平衡类别权重训练集成模型
   - 验证性能指标
   - 保存元数据用于审计

---

## 📋 快速开始

### 前提条件
- **Python**：3.9或更高版本
- **API密钥**（必需）：
  - Shodan API密钥（https://www.shodan.io/）
  - SendGrid API密钥（https://sendgrid.com/）- 用于邮件告警
- **可选API密钥**（用于增强丰富化）：
  - CMDB端点及凭据
  - SIEM端点及凭据
  - 补丁管理系统端点及凭据
  - 网络监控端点及凭据

### 安装

#### 步骤1：克隆仓库```bash
git clone https://github.com/EclipseManic/ThreatSentry-AI.git
cd ThreatSentry-AI

#### 步骤3：安装依赖项```bash
pip install -r requirements.txt

使用你的凭据编辑 `.env`：```ini
# Required - Threat Discovery
SHODAN_API_KEY=your_shodan_api_key_here
SHODAN_QUERY=                          # Leave empty to use presets
SHODAN_QUERY_EMPTY_TO_PRESET=True      # Use preset queries when SHODAN_QUERY is empty

# Optional - Email Alerts
SENDGRID_API_KEY=your_sendgrid_key_here
[email protected]    # Must be verified in SendGrid
[email protected],[email protected]

# Optional - Internal Enrichment (Implement in collectors/)
CMDB_API_ENDPOINT=https://cmdb.internal/api
CMDB_API_KEY=your_cmdb_key
SIEM_API_ENDPOINT=https://siem.internal/api
SIEM_API_KEY=your_siem_key

# Configuration
SCAN_INTERVAL_MINUTES=30               # How often to scan for new assets
RETRAIN_INTERVAL_MINUTES=60            # How often to retrain the model
MAX_SHODAN_RESULTS=50                  # Results per Shodan query
LOG_LEVEL=INFO                         # DEBUG, INFO, WARNING, ERROR

# Paths (Optional)
SQLITE_PATH=data/db/threat_sentric_ai.db
MODEL_PATH=data/models/rf_model.pkl
LOG_FILE_PATH=data/logs/threat_sentric_ai.log

#### 步骤6：运行应用程序```bash
python run.py

#### 自定义内部收集器
在 `collectors/` 目录中实现：
1. 从现有收集器复制模板
2. 使用你的 API 调用更新 `_collect()` 方法
3. 返回增强数据
4. 在调度器 (`core/scheduler.py`) 中注册

#### 邮件告警自定义
编辑 `alerts/email_alerts.py`：
- 修改 HTML 模板以实现品牌化邮件
- 添加自定义风险因素描述
- 调整修复建议

---

## 📊 数据模型

### 设备表（40+ 列）```
Core Identifiers:  ip, org, country, first_seen, last_seen
Vulnerability:     cve_count, max_cvss, vulnerabilities (rel)
Security Metrics:  auth_failures_24h, traffic_anomaly_score, patch_lag_days
Risk Assessment:   risk_label, risk_score, confidence_score
Context:          network_segment, service_category, is_critical_service
Compliance:       compliance_requirements, data_sensitivity_level
Historical:       incident_history_count, last_compromise_date, false_positive_count
Alerting:         notified, alert_history

---

## 🔄 工作流程示例

### 场景一：日常威胁狩猎
1. 打开仪表盘 → 查看风险等级为高（红色）的设备
2. 点击设备 → 查看详细的CVE列表
3. 记录组织与基础设施类型
4. 与SIEM进行交叉核对，查找近期可疑活动
5. 根据严重性和补丁延迟优先修复

### 场景二：事件响应
1. 收到关于新的高风险设备的告警
2. 仪表盘显示漏洞详情和风险因素
3. 查看资产管理选项卡 → 判断设备是否已知
4. 向值班SOC团队发送手动告警
5. 修复完成后，当Shodan反映变更时，仪表盘自动更新

### 场景三：风险报告
1. 进入分析选项卡
2. 导出风险分布与趋势图
3. 识别组织内部的风险热点
4. 生成修复路线图
5. 通过定期重新导出跟踪进度

---

## 🔧 故障排除

### 常见问题

#### 问题："Shodan API Key无效"（403 Forbidden）
- **解决方案**：验证 `.env` 文件中的Shodan API密钥。仪表盘将显示错误但继续处理
- **注意**：NVD增强不会运行（防止API故障导致数据丢失）

#### 问题："未找到设备的CVE"
- **解决方案**：检查Banner提取。编辑 `collectors/nvd_collector.py` 中的关键词列表
- **预防**：通过CSV上传手动设置CVE数据进行手动增强

#### 问题："模型置信度太低"
- **解决方案**：使用更多标注数据重新训练。提供关于错误分类设备的反馈
- **预防**：使用反馈系统提高训练数据质量

#### 问题：邮件告警未发送
- **解决方案**：验证SendGrid API密钥及发件人邮箱是否已验证
- **检查**：确保 `ENABLE_EMAIL_ALERTS` 环境变量设置为True

#### 问题：数据库锁定或查询缓慢
- **解决方案**：确保索引已创建（`init_db()` 会执行此操作）
- **检查**：无重复数据库连接（仅刷新UI，而非持续写入）

---

## 🛡️ 安全注意事项

1. **API密钥管理**：使用环境变量，切勿将 `.env` 提交到Git
2. **数据库安全**：SQLite适合单用户；多用户时请迁移到PostgreSQL
3. **网络安全**：在可信网络环境下运行；若暴露API，需实施网络分段
4. **数据隐私**：配置日志轮转以限制磁盘空间。实施数据保留策略
5. **审计日志**：所有模型决策均记录特征值以供审计

---

## 📈 性能优化建议

| 问题 | 解决方案 |
|-------|----------|
| 仪表盘加载缓慢 | 在配置中增加分页大小 |
| 训练期间CPU占用高 | 减少 `model/advanced_model.py` 中的 `n_estimators` |
| 内存占用过高 | 启用日志清理，降低图表分辨率 |
| Shodan扫描缓慢 | 减少 `MAX_SHODAN_RESULTS`，使用更精确的查询 |
| NVD增强缓慢 | 实现API缓存，减少产品关键词提取 |

---

## 🤝 贡献指南

欢迎贡献！请参阅 [CONTRIBUTING.md](https://raw.githubusercontent.com/eclipsemanic/threatsentry-ai/HEAD/CONTRIBUTING.md) 了解指南。

可改进的领域：
- 新增收集器实现（Nessus、Tenable、Qualys集成）
- 提供程序化访问的Web API
- 基于角色的多用户支持
- 高级可视化（Grafana/ELK集成）
- Kubernetes部署支持

---

## 📄 许可证

本项目采用MIT许可证——详见 [LICENSE](https://github.com/eclipsemanic/threatsentry-ai/blob/main/LICENSE) 文件。

---

## 📞 支持与联系

- **问题反馈**：通过GitHub Issues提交Bug报告和功能请求
- **文档**：查看 [docs/](https://github.com/eclipsemanic/threatsentry-ai/blob/main/docs) 目录获取详细技术文档
- **邮箱**：维护者信息见 [CODE_OF_CONDUCT.md](https://github.com/eclipsemanic/threatsentry-ai/blob/main/CODE_OF_CONDUCT.md)

---

## 🙏 致谢

- **Shodan**：提供全面的互联网暴露设备发现
- **NVD/NIST**：提供权威的漏洞数据
- **Scikit-learn**：稳健的机器学习库及文档
- **PyQt5**：优秀的GUI框架
- **社区**：所有提供反馈的贡献者和用户

---

**由 EclipseManic 用 ❤️ 打造 | 守护今日基础设施，迎接明天安全**
pip install -r requirements.txt```

---

## 🔐 Configuration

### 1. Environment Variables (`.env` file)
Create a `.env` file in the root directory with the following structure:

```bash
# --- Shodan API 密钥（必需）---
SHODAN_API_KEY="YOUR_SHODAN_API_KEY"

# --- SendGrid 电子邮件警报（必需）---
SENDGRID_API_KEY="YOUR_SENDGRID_API_KEY"
SENDER_EMAIL="[email protected]"
ALERT_RECIPIENTS="[email protected],[email protected]"

# --- 数据库 ---
SQLITE_PATH="threat_sentric_ai.db"

# --- 模型 ---
MODEL_PATH="threatsentry_model.pkl"

# --- 调度器 ---
SCAN_INTERVAL_MINUTES="30"
RETRAIN_INTERVAL_MINUTES="60"
RETRAIN_ON_SCHEDULE="True"

# --- Shodan 查询行为控制 ---
SHODAN_QUERY=""
SHODAN_QUERY_EMPTY_TO_PRESET="True"

# --- 日志记录 ---
LOG_LEVEL="INFO" # DEBUG, INFO, WARNING, ERROR, CRITICAL

# --- 内部系统凭证（可选 - 根据您的收集器需要添加）---
# CMDB_API_ENDPOINT="..."
# CMDB_API_KEY="..."
# SIEM_API_ENDPOINT="..."
# SIEM_API_KEY="..."
# PATCH_API_ENDPOINT="..."
# PATCH_API_KEY="..."
# NETWORK_MONITOR_ENDPOINT="..."
# NETWORK_MONITOR_KEY="..."```

> The app loads environment variables via `os.getenv()` in `config.py`.  
> Never commit this file to version control.

---

## 🔍 Shodan Query Presets (`config.py`)

Modify the `SHODAN_QUERIES` dictionary to define your custom query presets:

```python
SHODAN_QUERIES = {
    "default": "product:apache",
    "org": 'org:"Your Company Name"',
    "net": 'net:"123.45.67.0/24"',
    "ssl": 'ssl:"yourcompany.com"',
    "hostname": 'hostname:".yourcompany.com"',
    "rdp": 'port:3389 "remote desktop"',
    "mongodb": 'port:27017 "mongodb"',
    "ics_modbus": 'port:502 "modbus"',
    "vuln_example": 'vuln:CVE-2024-12345',
    "http_login": 'http.title:"Login" org:"Your Company"'
}```

If `SHODAN_QUERY` in `.env` is empty **and** `SHODAN_QUERY_EMPTY_TO_PRESET=True`,  
the scheduler will automatically cycle through these presets.

---

## 🧠 Implementing Internal Collectors

> ⚠️ **Important:** The internal collectors (`cmdb_collector.py`, `siem_collector.py`, `patch_collector.py`, `network_monitor_collector.py`) are **placeholders**.  
> Replace the placeholder logic with real integrations to your systems.

You’ll need to:
- Fetch internal data using APIs, databases, or SDKs.
- Map fetched data to the `Device` model attributes.
- Update the database session with this enriched information.

Without these integrations, the model will lack context for accurate predictions.

---

## 🧪 Running the Application

### (Optional) Generate Initial Training Data
```bash
python scripts/generate_realistic_training_data.py --count 1000```
Creates `scripts/my_training_data.json`.  
You can upload this file through the GUI’s **Upload** option to initialize training data.

### Launch the Application
```bash
python run.py```
The GUI will open and the scheduler will start scanning, enriching, and predicting automatically.

---

## ⚡ Quick Demo with Generated Data

1. Complete installation (steps 1–3).  
2. Set up your `.env` file (even with placeholder keys).  
3. Optionally adjust Shodan presets in `config.py`.  
4. Generate data:
   ```bash
python scripts/generate_realistic_training_data.py --count 500   ```
5. Launch the app:
   ```bash
python run.py   ```
6. In the GUI:
   - Select **Upload** → Choose `scripts/my_training_data.json`
   - Click **Refresh** to view populated device data and risk levels.

---

## 🛠️ Scripts and Usage Guide

| Script | Description |
|--------|--------------|
| **run.py** | Main entry point. Starts DB, scheduler, and GUI. |
| **scheduler.py** | Handles periodic scanning, enrichment, prediction, and retraining. |
| **scripts/generate_realistic_training_data.py** | Generates realistic training data for testing or bootstrapping. |
| **scripts/clear_db_enhanced.py** | Interactively clean database or reset notification flags. |
| **scripts/reset_db.py** | Completely resets the database. Use with caution. |

### Example Commands

```bash
# 生成1000条记录
python scripts/generate_realistic_training_data.py --count 1000

# 查看清理选项
python scripts/clear_db_enhanced.py --help

# 删除所有设备和漏洞
python scripts/clear_db_enhanced.py --delete-devices --delete-vulns

# 重置通知标记
python scripts/clear_db_enhanced.py --reset-notified

# 危险的全量重置
python scripts/reset_db.py```

---

## 💬 Feedback & Suggestions

This project was developed by **EclipseManic**.  
While code contributions are currently closed, your **feedback and bug reports** are highly appreciated.

Please open an [Issue](https://github.com/EclipseManic/ThreatSentry-AI/issues) to share your thoughts or report a problem.

---

## 📄 License

This project is licensed under the **MIT License**.  
See the [LICENSE](https://github.com/eclipsemanic/threatsentry-ai/blob/main/LICENSE) file for details.