PentestAssistant

基于 ChatGLM 等大语言模型，与函数调用技术和多代理架构结合的 Web 渗透检测自动化的助手，解放双手，以对话式的姿态完成渗透任务。

⚠️ 重要提示

• 目前推出 DEMO 版本，项目持续开发迭代中，欢迎小伙伴加入讨论，提供宝贵的意见。
• 支持后端接口调用，前端界面敬请期待

---

目录

• 介绍
• 更新日志
• 快速上手
  • 1. 环境配置
  • 2. 模型下载
  • 4. 一键启动
• 联系我们

介绍

🤖 PentestAssistant 基于 Web 渗透检测工具，构建常用函数调用 API 知识库，通过多代理的方式实现渗透任务建立和维护，基于 RAG 技术实现渗透检测任务的执行，以此完成渗透检测的自动化过程。

✏️ 本项目运行原理如下：

1. 预处理： 构建 Web 渗透检测工具知识库，收集常用方式，并组成可调用调用函数 API
2. 任务建立：根据用户请求和渗透检测知识库，生成渗透检测任务
3. 任务轮询：
   1. 根据任务列表，匹配知识库中可用函数，生成执行序列并运行，反馈结果
   2. 根据执行结果，刷新渗透检测任务状态
4. 任务结束：根据渗透检测任务和工具执行结果，生成任务报告

🥚 支持两种渗透检测使用方式：

1. 单一工具调用：指定渗透检测工具，用户给定自然语言描述任务，系统根据请求自动完成工具调用和执行
2. 多工具联合调用：选择多款渗透检测工具，用户给定自然语言描述任务，系统自动选择工具，执行函数，实现需求

✅ 当前支持工具列表

Tool	Function Number
Nmap	437
w3af	14
zaproxy	14
arachni	14
dirsearch	10
Xsstrike	7
sqlmap	7
tplmap	5
Cmseek	4

更新日志

[24/06/17] 支持UI展示，增加前后端 [24/05/27] 首次发布，支持后端 API 形式访问，支持 Nmap 调用

快速上手

1. 环境配置

首先下载仓库源代码

$ git clone https://github.com/HUSTInfSecLabs/PentestAssistant.git

对于机器人端，我们推荐使用 conda 创建 python 虚拟环境，并使用 python 3.10 版本

$ cd botend
$ conda create -n PentestAssistant python=3.10
$ python --version
Python 3.10.13

在虚拟环境内安装所有依赖

$ conda activate PentestAssistant
$ pip install -r requirements.txt 

如果拉取依赖失败，可以更换 conda 源，推荐清华源或中科大源

对于前后端，使用 node.js 实现，具体安装步骤参考目录内 README.md 文件

2. 模型下载

项目需要下载对话大模型和 Embedding, Reranker 模型两类模型

大模型

项目运行之前，使用本地对话大模型，需要提前下载并存储，常见开源大模型通常可以在 HuggingFace 下载

下载模型后，需要在 .env 文件中修改环境变量，配置模型地址，如 ChatGLM3-6b 模型:

# 使用方式
ChatGLM3_6b_AI__MODEL_TYPE="text-generation"

# 模型地址
ChatGLM3_6b_AI__MODEL_PATH="/home/user/.cache/huggingface/hub/models--THUDM--chatglm3-6b"

# GPU 使用
ChatGLM3_6b_AI__MODEL_DEVICE="0"

目前支持部分大模型，包括：ChatGLM2-6b, ChatGLM3-6b, Llama2-6b, Llama2-13b, Llama3-8b-instruct

Embedding 和 Reranker 模型

与大模型下载、配置方式相同，如 bge-reranker-large 模型：

# 模型地址
RERANKER_EMBEDDING_NAME_OR_PATH="/home/user/.cache/huggingface/hub/models--BAAI--bge-reranker-large"

3. 一键启动

启动前端、后端和机器人端

$ cd frontend
$ pnpm start

$ cd ../backend
$ pnpm start

$ cd ../botend
$ python api_flask.py

联系我们

如果你也对本项目感兴趣，欢迎交流合作至邮箱 [email protected]