You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
284 lines
13 KiB
284 lines
13 KiB
'''
|
|
实现对大模型调用的封装，隔离具体使用的LLM
|
|
pip install openai
|
|
export OPENAI_API_KEY="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"
|
|
'''
|
|
import openai
|
|
import json
|
|
import re
|
|
import os
|
|
from openai import OpenAI
|
|
from myutils.ConfigManager import myCongif
|
|
from myutils.MyTime import get_local_timestr
|
|
from myutils.MyLogger_logger import LogHandler
|
|
|
|
class LLMManager:
|
|
def __init__(self,illm_type):
|
|
self.logger = LogHandler().get_logger("LLMManager")
|
|
self.api_key = None
|
|
self.api_url = None
|
|
|
|
#temperature设置
|
|
#DS------代码生成/数学解题：0.0 -- 数据抽取/分析：1.0 -- 通用对话：1.3 -- 翻译：1.3 -- 创意类写作：1.5
|
|
if illm_type == 0: #腾讯云
|
|
self.api_key = "fGBYaQLHykBOQsFwVrQdIFTsYr8YDtDVDQWFU41mFsmvfNPc"
|
|
self.api_url = ""
|
|
elif illm_type == 1: #DS
|
|
self.api_key ="sk-10360148b465424288218f02c87b0e1b"
|
|
self.api_url ="https://api.deepseek.com/v1"
|
|
self.model = "deepseek-reasoner" #model=deepseek-reasoner -- R1 model=deepseek-chat --V3
|
|
# 创建会话对象 -- 一个任务的LLM必须唯一
|
|
self.client = OpenAI(api_key=self.api_key, base_url=self.api_url)
|
|
elif illm_type == 2: #2233.ai
|
|
self.api_key = "sk-J3562ad9aece8fd2855bb495bfa1a852a4e8de8a2a1IOchD"
|
|
self.api_url = "https://api.gptsapi.net/v1"
|
|
self.model = "o3-mini-2025-01-31"
|
|
self.client = OpenAI(api_key=self.api_key,base_url=self.api_url)
|
|
elif illm_type ==3: #GPT
|
|
# 定义代理服务器地址
|
|
proxy_url = "http://192.168.3.102:3128"
|
|
os.environ["HTTP_PROXY"] = proxy_url
|
|
os.environ["HTTPS_PROXY"] = proxy_url
|
|
self.api_key ="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"
|
|
self.api_url = "https://api.openai.com/v1"
|
|
self.model = "o3-mini-2025-01-31"
|
|
|
|
openai.proxy = proxy_url
|
|
openai.api_key = self.api_key
|
|
#self.client = openai
|
|
self.client = OpenAI(api_key=self.api_key,base_url=self.api_url)
|
|
|
|
'''
|
|
**决策原则**
|
|
- 根据节点类型和状态，优先执行基础测试（如端口扫描、服务扫描）。
|
|
- 仅在发现新信息或漏洞时新增子节点。
|
|
- 确保每个新增节点匹配测试指令。
|
|
'''
|
|
# 初始化messages
|
|
def build_initial_prompt(self,node):
|
|
if not node:
|
|
return
|
|
#根节点初始化message
|
|
node.messages = [{"role": "system",
|
|
"content":'''
|
|
你是一位渗透测试专家，基于你的专业知识和经验，负责动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步测试指令，推动测试前进，直至完成渗透测试。本地程序会执行你的指令并返回结果。
|
|
**总体要求**
|
|
1.以目标系统IP为根节点，每个渗透测试点（如端口、服务、漏洞点）作为子节点，形成树型结构（测试树）；
|
|
2.测试树整体数据由本地程序存储，你只需要关注当前节点的测试推进、状态更新(未完成/已完成)及完整新增子节点；
|
|
3.返回两类指令：节点指令和测试指令，以空行间隔，不要包含注释和说明；
|
|
4.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增中高危节点的完整性，若有未生成测试指令的节点，必须返回未生成指令节点列表；
|
|
5.若无需要处理的节点数据，节点指令可以不生成，但测试指令必须对应已有节点；
|
|
**决策流程**
|
|
1. 若当前节点是IP且未进行端口扫描，则对当前节点执行端口扫描；
|
|
2. 若端口扫描发现开放端口，对可能存在中高危以上风险的端口新增节点并提供测试指令；
|
|
3. 若当前节点是端口且未进行服务扫描，则执行服务扫描；
|
|
4. 若服务扫描发现服务版本或漏洞，则新增漏洞测试节点并提供测试指令；
|
|
5. 若漏洞验证成功，则根据结果决定是否需要进一步测试，若需要进一步测试，则必须将进一步测试的内容作为子节点并提供测试指令；
|
|
6. 当当前节点没有新的测试指令时，更新状态为“已完成”。
|
|
**测试指令生成准则**
|
|
1.明确每个测试指令的测试目标，并优先尝试最简单、最直接的办法，不要在同一个请求生成测试效果覆盖的指令；
|
|
2.使用递进逻辑组织指令：先尝试基础测试方法，根据执行结果决定是否进行更深入的测试;
|
|
**节点指令格式**
|
|
- 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\", \"status\": \"未完成\"}；
|
|
- 未生成指令节点列表：{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}；
|
|
- 完成测试未发现漏洞：{\"action\": \"update_status\", \"node\": \"节点\", \"status\": \"已完成\"}；
|
|
- 完成测试且发现漏洞：{\"action\": \"update_status\", \"node\": \"节点\", \"status\": \"已完成\"，\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息（没有可为空）\"}}；
|
|
**测试指令格式**
|
|
- shell指令：```bash-[节点路径]指令内容```包裹，需要避免用户交互,若涉及到多步指令，请生成python代码；
|
|
- python指令：```python-[节点路径]指令内容```包裹，主函数名为dynamic_fun，需包含错误处理，必须返回一个tuple(status, output)；
|
|
- [节点路径]为从根节点到目标节点的完整层级描述。
|
|
**核心要求**
|
|
- 优先保障新增中高危测试节点的完整性。
|
|
- 指令之间必须要有一个空行。
|
|
**响应示例**
|
|
{\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"nodes\": \"3306端口,22端口\", \"status\": \"未完成\"}
|
|
|
|
```bash-[目标系统->192.168.1.100->3306端口]
|
|
mysql -u root -p 192.168.1.100
|
|
```
|
|
'''}] # 一个messages
|
|
|
|
# 调用LLM生成指令
|
|
def get_llm_instruction(self,prompt,node):
|
|
'''
|
|
1.由于大模型API不记录用户请求的上下文，一个任务的LLM不能并发！
|
|
:param prompt:用户本次输入的内容
|
|
:return: instr_list
|
|
'''
|
|
#添加本次输入入该节点的message队列
|
|
message = {"role":"user","content":prompt}
|
|
node.messages.append(message) #更新节点message
|
|
|
|
#提交LLM
|
|
post_time = get_local_timestr()
|
|
if self.model == "o3-mini-2025-01-31":
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
reasoning_effort="high",
|
|
messages = node.messages
|
|
)
|
|
else:
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=node.messages
|
|
)
|
|
|
|
#LLM返回结果处理
|
|
reasoning_content = ""
|
|
content = ""
|
|
#LLM返回处理
|
|
if self.model == "deepseek-reasoner":
|
|
#返回错误码：DS-https://api-docs.deepseek.com/zh-cn/quick_start/error_codes
|
|
reasoning_content = response.choices[0].message.reasoning_content #推理过程
|
|
print(reasoning_content)
|
|
content = response.choices[0].message.content #推理内容
|
|
print(content)
|
|
# 记录llm历史信息
|
|
node.messages.append({'role': 'assistant', 'content': content})
|
|
elif self.model == "deepseek-chat":
|
|
content = response.choices[0].message
|
|
# 记录llm历史信息
|
|
node.messages.append(content)
|
|
elif self.model == "o3-mini-2025-01-31":
|
|
reasoning_content = "" #gpt不返回推理内容
|
|
content = response.choices[0].message.content
|
|
print(content)
|
|
# 记录llm历史信息
|
|
node.messages.append({'role': 'assistant', 'content': content})
|
|
else:
|
|
self.logger.error("处理到未预设的模型！")
|
|
return "","","","",""
|
|
#按格式规定对指令进行提取
|
|
node_cmds,commands = self.fetch_instruction(content)
|
|
return node_cmds,commands,reasoning_content, content, post_time
|
|
|
|
def fetch_instruction(self,response_text):
|
|
'''
|
|
*****该函数很重要，需要一定的容错能力，解析LLM返回内容*****
|
|
处理边界：只格式化分析LLM返回内容，指令和节点操作等交其他模块。
|
|
节点控制指令
|
|
渗透测试指令
|
|
提取命令列表，包括：
|
|
1. Python 代码块 python[](.*?)
|
|
2. Shell 命令``bash[](.*?)```
|
|
:param text: 输入文本
|
|
:return: node_cmds,python_blocks,shell_blocks
|
|
'''
|
|
#针对llm的回复，提取节点操作数据和执行的指令----
|
|
# 正则匹配 Python 代码块
|
|
python_blocks = re.findall(r"```python-(.*?)```", response_text, flags=re.DOTALL)
|
|
# 处理 Python 代码块，去除空行并格式化
|
|
python_blocks = [block.strip() for block in python_blocks]
|
|
|
|
#正则匹配shell指令
|
|
shell_blocks = re.findall(f"```bash-(.*?)```", response_text, flags=re.DOTALL)
|
|
shell_blocks = [block.strip() for block in shell_blocks]
|
|
|
|
# 按连续的空行拆分
|
|
# 移除 Python和bash 代码块
|
|
text_no_python = re.sub(r"```python.*?```", "PYTHON_BLOCK", response_text, flags=re.DOTALL)
|
|
text = re.sub(r"```bash.*?```", "SHELL_BLOCK", text_no_python, flags=re.DOTALL)
|
|
|
|
# 这里用 \n\s*\n 匹配一个或多个空白行
|
|
parts = re.split(r'\n\s*\n', text)
|
|
node_cmds = []
|
|
commands = []
|
|
python_index = 0
|
|
shell_index = 0
|
|
for part in parts:
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
if "PYTHON_BLOCK" in part:
|
|
# 还原 Python 代码块
|
|
commands.append(f"python-code {python_blocks[python_index]}")
|
|
python_index += 1
|
|
elif "SHELL_BLOCK" in part:
|
|
commands.append(shell_blocks[shell_index])
|
|
shell_index +=1
|
|
else:#其他的认为是节点操作指令--指令格式还存在不确定性，需要正则匹配，要求是JSON
|
|
pattern = re.compile(r'\{(?:[^{}]|\{[^{}]*\})*\}')
|
|
# 遍历所有匹配到的 JSON 结构
|
|
|
|
# strlines = part.strip('\n') #按行拆分，避免贪婪模式下，匹配到多行的最后一个}
|
|
# for strline in strlines:
|
|
for match in pattern.findall(part): #正常只能有一个
|
|
try:
|
|
node_cmds.append(json.loads(match)) # 解析 JSON 并添加到列表
|
|
except json.JSONDecodeError as e:#解析不了的不入队列
|
|
self.logger.error(f"LLM-{part}-JSON 解析错误: {e}") #这是需不需要人为介入？
|
|
return node_cmds,commands
|
|
|
|
def test_llm(self):
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "讲个笑话吧。"}
|
|
]
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
reasoning_effort="medium",
|
|
messages=messages
|
|
)
|
|
print(response)
|
|
|
|
if __name__ == "__main__":
|
|
llm = LLMManager(3)
|
|
content = '''
|
|
"```bash-[目标系统->192.168.204.137->22端口]
|
|
ssh -o BatchMode=yes -o ConnectTimeout=5 root@192.168.204.137
|
|
```
|
|
|
|
```bash-[目标系统->192.168.204.137->23端口]
|
|
telnet 192.168.204.137 23 <<<'user anonymous'
|
|
```
|
|
|
|
```bash-[目标系统->192.168.204.137->5432端口]
|
|
PGPASSWORD='' psql -h 192.168.204.137 -U postgres -c '\l'
|
|
```
|
|
|
|
```bash-[目标系统->192.168.204.137->2121端口]
|
|
ftp -nv 192.168.204.137 2121 <<EOF
|
|
user anonymous anonymous
|
|
bye
|
|
EOF
|
|
```
|
|
|
|
```bash-[目标系统->192.168.204.137->139端口]
|
|
smbclient -L //192.168.204.137/ -N -p 139
|
|
```
|
|
|
|
```python-[目标系统->192.168.204.137->513端口]
|
|
import socket
|
|
|
|
def dynamic_fun():
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.settimeout(5)
|
|
s.connect(('192.168.204.137', 513))
|
|
s.send(b'localhost\n')
|
|
response = s.recv(1024)
|
|
s.close()
|
|
return ('open' if response else 'closed', response.decode(errors='ignore'))
|
|
except Exception as e:
|
|
return ('error', str(e))
|
|
```
|
|
|
|
{"action": "no_instruction", "nodes": "514端口"}"
|
|
'''
|
|
node_cmds,commands = llm.fetch_instruction(content)
|
|
print(node_cmds,commands)
|
|
|
|
for node_json in node_cmds:
|
|
action = node_json["action"]
|
|
if action == "update_status":
|
|
print("update")
|
|
elif action == "no_instruction":
|
|
# 返回的未生成指令的数据进行校验：1.要有数据；2.节点不是当前节点就是子节点
|
|
nodes = []
|
|
node_names = node_json["nodes"].split(',')
|
|
for node_name in node_names:
|
|
print(node_name)
|
|
print("111")
|
|
|
|
|
|
|
|
|