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.
231 lines
12 KiB
231 lines
12 KiB
1 month ago
|
'''
|
||
|
|
实现对大模型调用的封装，隔离具体使用的LLM
|
||
|
|
pip install openai
|
||
|
|
export OPENAI_API_KEY="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"
|
||
|
|
'''
|
||
|
|
import openai
|
||
|
|
import json
|
||
|
|
import threading
|
||
|
|
import re
|
||
|
|
from openai import OpenAI
|
||
|
|
from myutils.MyTime import get_local_timestr
|
||
|
|
from myutils.MyLogger_logger import LogHandler
|
||
|
|
|
||
|
|
class LLMManager:
|
||
|
|
def __init__(self,illm_type=0):
|
||
|
|
self.logger = LogHandler().get_logger("LLMManager")
|
||
|
|
self.api_key = None
|
||
|
|
self.api_url = None
|
||
|
|
self.task_id =0 #一个任务一个id
|
||
|
|
self.llm_sn = 0 # llm执行序列号，--一任务一序列
|
||
|
|
self.llm_sn_lock = threading.Lock() #
|
||
|
|
#temperature设置
|
||
|
|
#DS------代码生成/数学解题：0.0 -- 数据抽取/分析：1.0 -- 通用对话：1.3 -- 翻译：1.3 -- 创意类写作：1.5
|
||
|
|
#腾讯云---
|
||
|
|
self.temperature = 1.0
|
||
|
|
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 ==3: #GPT
|
||
|
|
self.api_key =""
|
||
|
|
self.api_url = ""
|
||
|
|
self.model = ""
|
||
|
|
self.client = OpenAI()
|
||
|
|
|
||
|
|
self.messages = []
|
||
|
|
# 初始化阶段和已知信息
|
||
|
|
current_stage = "信息收集"
|
||
|
|
known_info = {"url": "www.test.com"}
|
||
|
|
results = [] # 存储所有任务结果，用于生成报告
|
||
|
|
|
||
|
|
#******测试使用，设置slef.message的值
|
||
|
|
def test_old_message(self,strMessage):
|
||
|
|
try:
|
||
|
|
self.messages = json.loads(strMessage)
|
||
|
|
except json.JSONDecodeError as e:
|
||
|
|
print(f"JSON解析错误: {str(e)}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"错误: {str(e)}")
|
||
|
|
|
||
|
|
# 构建初始提示 初始化messages
|
||
|
|
def build_initial_prompt(self,target,know_info="无",node=None):
|
||
|
|
# self.messages = [{"role": "system",
|
||
|
|
# "content": "你是一位渗透测试专家,需要动态控制整个渗透测试过程，包括信息收集、漏洞扫描、漏洞利用等阶段，最终生成渗透测试报告。由你规划执行的指令，我会根据你的指令执行并提交结果，你再对结果进行分析，规划并生成下一步指令，直到完成渗透测试，生成测试报告。"
|
||
|
|
# "生成的指令需满足如下约束："
|
||
|
|
# "1.只返回具体的shell指令或Python代码，不要包含注释和说明；"
|
||
|
|
# "2.shell指令以```bash(.*?)```包裹，python代码以```python(.*?)```包裹；"
|
||
|
|
# "3.若提供的是shell指令，需要避免用户再次交互；"
|
||
|
|
# "4.若提供的是python代码，主函数名为dynamic_fun,需要包含错误处理，执行结束返回两个str(status-'success'或'failure'和output-补充输出信息)；"
|
||
|
|
# "5.如果认为渗透测试已完成，请生成生成报告的python代码，并返回'complete'和报告文件路径"}] # 一个messages
|
||
|
|
#return f"现在开始对目标{target}进行渗透测试，已知信息{know_info}，请提供下一步执行的指令。"
|
||
|
|
if not node:
|
||
|
|
return None
|
||
|
|
#根节点初始化message
|
||
|
|
node.messages = [{"role": "system",
|
||
|
|
"content":'''
|
||
|
|
你是一位渗透测试专家,由你动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步的测试指令，推动测试阶段前进，直至完成渗透测试。本地程序会根据你的指令进行执行,然后将执行结果返回给你。
|
||
|
|
整体过程说明：
|
||
|
|
1.以目标系统所在IP为根节点，每个渗透测试点（如端口、服务、漏洞点）作为子节点，形成树型结构（测试树）；
|
||
|
|
2.测试树整体数据由本地程序存储，你只需要关注当前节点的渗透测试推进，节点状态(未完成、已完成)的更新，和是否有子节点新增；
|
||
|
|
3.新增子节点时，同一父节点下的子节点名称不能相同，并提供对该节点的测试指令；
|
||
|
|
3.新增节点示例：{\"action\":\"add_node\", \"parent\": \"80端口\", \"node\": \"http://192.168.1.100/index.php?id=1\", \"status\": \"未完成\"}；
|
||
|
|
4.完成节点测试且发现漏洞示例：{\"action\": \"update_status\", \"node\": \"21端口\", \"status\": \"已完成\"，\"vulnerability\": \"ftp匿名登录\"}；
|
||
|
|
5.发现漏洞后，可根据漏洞类型决定是否新增子节点继续测试，如：{\"action\": \"add_node\", \"parent\": \"21端口\", \"node\": \"ftp配置检查\", \"status\": \"未完成\"}；
|
||
|
|
6.当本地程序确认所有节点状态均为“已完成”时，生成测试报告。
|
||
|
|
生成的渗透测试指令需满足如下约束：
|
||
|
|
1.只返回具体的shell指令或Python代码，不要包含注释和说明；
|
||
|
|
2.shell指令以```bash-[对应节点的路径](.*?)```包裹，python代码以```python-[对应节点的路径](.*?)```包裹；
|
||
|
|
3.若提供的是shell指令，需要避免用户再次交互；
|
||
|
|
4.若提供的是python代码，主函数名为dynamic_fun，需包含错误处理，执行结束后返回一个tuple (status, output)，其中status为'success'或'failure'，output为补充输出信息；
|
||
|
|
示例：
|
||
|
|
```python-[目标系统->192.168.1.100->3306端口]
|
||
|
|
def dynamic_fun():
|
||
|
|
try:
|
||
|
|
result = "扫描完成"
|
||
|
|
return ("success", result)
|
||
|
|
except Exception as e:
|
||
|
|
return ("failure", str(e))
|
||
|
|
```
|
||
|
|
5.如果认为渗透测试已完成，请生成生成报告的python代码，并返回'complete'和报告文件路径。
|
||
|
|
'''}] # 一个messages
|
||
|
|
user_Prompt = f'''
|
||
|
|
当前分支路径：目标系统->{target}
|
||
|
|
当前节点信息：
|
||
|
|
- 节点名称：{target}
|
||
|
|
- 节点状态：未完成
|
||
|
|
- 漏洞类型：未发现
|
||
|
|
上一步结果：已知信息-{know_info}
|
||
|
|
任务：生成下一步渗透测试指令或结束该节点的渗透测试(修改节点状态为：已完成)。
|
||
|
|
'''
|
||
|
|
return user_Prompt
|
||
|
|
|
||
|
|
def init_data(self,task_id=0):
|
||
|
|
#初始化LLM数据
|
||
|
|
self.llm_sn = 0
|
||
|
|
self.task_id = task_id
|
||
|
|
self.messages = []
|
||
|
|
|
||
|
|
# 调用LLM生成指令
|
||
|
|
def get_llm_instruction(self,prompt,th_DBM,node):
|
||
|
|
'''
|
||
|
|
1.由于大模型API不记录用户请求的上下文，一个任务的LLM不能并发！
|
||
|
|
:param prompt:用户本次输入的内容
|
||
|
|
:return: instr_list
|
||
|
|
'''
|
||
|
|
#添加本次输入入该节点的message队列
|
||
|
|
message = {"role":"user","content":prompt}
|
||
|
|
node.messages.append(message)
|
||
|
|
|
||
|
|
#提交LLM
|
||
|
|
post_time = get_local_timestr()
|
||
|
|
response = self.client.chat.completions.create(
|
||
|
|
model=self.model,
|
||
|
|
messages = self.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)
|
||
|
|
else:
|
||
|
|
self.logger.error("处理到未预设的模型！")
|
||
|
|
return None
|
||
|
|
|
||
|
|
#LLM记录存数据库
|
||
|
|
node.llm_sn += 1
|
||
|
|
bres = th_DBM.insert_llm(self.task_id,prompt,reasoning_content,content,post_time,node)
|
||
|
|
if not bres:
|
||
|
|
self.logger.error(f"{node.name}-llm入库失败！")
|
||
|
|
|
||
|
|
#需要对指令进行提取
|
||
|
|
node_list = self.fetch_instruction(content,node)
|
||
|
|
|
||
|
|
#********测试时使用---输出和记录LLM返回指令的message
|
||
|
|
print(f"Messages:{self.messages}")
|
||
|
|
with open("../test", "w", encoding="utf-8") as f: #输出到文件
|
||
|
|
json.dump(self.messages,f,ensure_ascii=False)
|
||
|
|
|
||
|
|
return node_list
|
||
|
|
|
||
|
|
def fetch_instruction(self,response_text,node):
|
||
|
|
'''
|
||
|
|
节点控制指令
|
||
|
|
渗透测试指令
|
||
|
|
提取命令列表，包括：
|
||
|
|
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:
|
||
|
|
#其他的认为是节点操作指令
|
||
|
|
node_cmds.append(part)
|
||
|
|
|
||
|
|
return node_cmds,commands #？存在一个问题：这样分list返回，执行顺序会丢失
|
||
|
|
|
||
|
|
def test_llm(self):
|
||
|
|
with open("../test", "r", encoding="utf-8") as f:
|
||
|
|
messages = json.load(f)
|
||
|
|
text = messages[-1]["content"]
|
||
|
|
list = self.fetch_instruction(text)
|
||
|
|
for itme in list:
|
||
|
|
print("***********")
|
||
|
|
print(itme)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
LM = LLMManager(1)
|
||
|
|
LM.test_llm()
|
||
|
|
|
||
|
|
|