mohammedrehan27/pentesting-agent

# PentestingAgent — 网络安全强化学习原型 这是一个教育性质的网络安全强化学习研究原型，旨在训练一个 RL agent 在企业网络图上执行模拟的渗透测试。 该项目将 **Gymnasium 自定义环境**、**Stable-Baselines3 PPO**、**Graph Neural Networks** 和 **Red vs Blue team 模拟**整合在一个可运行的代码库中。 ## 项目概述 Agent 学习在网络图中进行遍历、利用漏洞节点、提升权限并避免被检测——所有这些都在一个安全的模拟环境中进行。项目采用模块化结构，每个模块都可以独立运行，方便对各个组件进行实验。 ## 主要特性 - 基于 Gymnasium 的自定义渗透测试环境 - 使用 NetworkX 进行企业网络建模 - 使用 Stable-Baselines3 的 PPO 强化学习 agent - 针对注重隐蔽性的攻击优化进行 Reward shaping - Graph Neural Network (GCN) 拓扑 embeddings - Red Team vs Blue Team 对抗模拟 - 自动化安全评估报告生成 - 支持多种拓扑结构（小型 / 中型 / 大型企业网络） ## 架构 ``` pentesting-agent/ ├── network.py # Network topology builder (small / medium / large) ├── env.py # Custom Gymnasium environment (PentestEnv) ├── train.py # PPO training script (Stable-Baselines3) ├── visualize.py # Attack path & compromise visualization ├── gnn_agent.py # GNN prototype (PyTorch Geometric) ├── red_blue_sim.py # Red Team vs Blue Team simulation ├── report_generator.py # Automated security report generator ├── requirements.txt # Python dependencies └── README.md # This file ```  ## 强化学习工作流 ``` Network Graph (network.py) | v PentestEnv (env.py) <-- Gymnasium custom environment | observation: [current_node, privilege_level, detection_status] actions: 0=Scan 1=Exploit 2=PrivEsc 3=LateralMove | v PPO Agent (train.py) <-- Stable-Baselines3 | reward signal: +20 successful exploit +40 privilege escalation -10 failed exploit -5 unnecessary scan -50 detection triggered | v pentest_ppo.zip <-- saved trained model ``` Agent 通过试错来探索网络。随着时间的推移，它会学会利用高危漏洞节点，在被攻陷的机器上提升权限，并向高价值目标进行横向移动。 ## GNN 集成概述 `gnn_agent.py` 演示了 **Graph Neural Networks** 如何为 agent 提供更丰富的网络拓扑表示。 **Graph Convolutional Network (GCN)** 不再将每个节点视为扁平的 feature vector，而是聚合来自相邻节点的信息——因此 agent 不仅“知道”当前节点的属性，还了解其周围的环境。 ``` NetworkX Graph | v Node features: [vulnerability_level, privilege_level, compromised] Edge index : all edges (both directions) | v GCNConv (layer 1) -> ReLU -> GCNConv (layer 2) | v Node embeddings [N x 8] <- topology-aware representations ``` 该脚本还运行了一个普通的 MLP baseline（无图结构），并打印出两套 embedding 集合之间的 L2 距离，以展示 GCN 增加了多少结构性信息。 ## Red vs Blue 模拟概述 `red_blue_sim.py` 在以下双方之间运行交替回合的模拟： | 团队 | 操作 | |-----------|------------------------------------------------------| | Red Team | 侦察、漏洞利用、横向移动 | | Blue Team | 补丁修复（降低漏洞）、隔离、恢复节点 | 该模拟在中型拓扑上运行 20 轮，并输出最终的安全评估：攻陷率、获得的总权限数、应用的补丁数、隔离的节点数以及总体结论。 ## 安装说明 ### 1. 创建虚拟环境（推荐） ``` python -m venv venv # Windows venv\Scripts\activate # macOS / Linux source venv/bin/activate ``` ### 2. 安装核心依赖 ``` pip install -r requirements.txt ``` ### 3. 安装 PyTorch Geometric（用于 gnn_agent.py） torch-geometric 需要匹配的 PyTorch 版本。请遵循官方指南： https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html 快速安装（仅限 CPU）： ``` pip install torch-geometric ``` ## 如何运行各个模块 ### 网络拓扑生成器 ``` python network.py ``` 生成 `network_small.png`、`network_medium.png`、`network_large.png`。 ### 训练 PPO agent ``` python train.py ``` 训练 1000 个 timesteps，保存 `pentest_ppo.zip`，然后运行一个评估 episode。 ### 攻击路径可视化 ``` python visualize.py ``` 生成 `attack_path.png` 以及逐步的 `attack_step_N.png` 图像。 ### GNN 原型 ``` python gnn_agent.py ``` 将图转换为 tensors，运行 GCN + MLP，打印 embedding 形状和对比结果。 *（需要 `torch-geometric`；如果未安装，则以 stub 模式运行。）* ### Red vs Blue 模拟 ``` python red_blue_sim.py ``` 向控制台打印完整的事件日志和最终的安全评估。 ### 安全报告生成器 ``` python report_generator.py ``` 运行一个随机动作的 episode，收集指标，并保存 `report.txt`。 ## 示例输出 **PPO 训练** - 环境验证日志 - Reward 跟踪 - 攻击轨迹评估 **GNN Embeddings** - Graph tensor 转换 - 节点 embedding 生成 - GCN vs MLP 对比指标 **Red vs Blue 模拟** - 实时攻防事件日志 - 动态漏洞修补 - 最终安全评估结论 ## 当前局限性 - 出于教育目的，环境设计得故意较小且经过了简化。 - PPO 训练使用扁平的 MLP policy —— GCN embeddings 尚未输入到 RL 循环中。 - 模拟中的 Blue Team 行为是随机生成的，而非学习得来的。 - GNN 原型使用随机权重（未与 RL agent 进行端到端训练）。 - 未使用真实的漏洞数据；所有数值均为合成数据。 ## 未来改进 - [ ] 将 GCN embeddings 连接为 PPO 的 observation 输入（基于图的 policy） - [ ] 添加一个通过学习生成的 Blue Team agent（多智能体 RL） - [ ] 扩展到更大、程序化生成的网络拓扑 - [ ] 整合真实的 CVE 严重性评分作为漏洞级别 - [ ] 添加 Web dashboard 用于实时 episode 可视化 - [ ] 导出训练好的模型，以便部署到网络安全攻防演练平台 ## 依赖项 | 包名 | 用途 | |--------------------|----------------------------------| | gymnasium | 自定义 RL 环境 | | stable-baselines3 | PPO 训练 | | torch | 深度学习骨干网络 | | torch-geometric | Graph neural network | | networkx | 网络图构建 | | matplotlib | 可视化 | | numpy | 数值计算 | | pandas | 数据处理（报告） |

标签：Gymnasium, PE 加载器, Web报告查看器, 凭据扫描, 图神经网络, 强化学习, 特权检测, 网络安全, 逆向工具, 隐私保护