''' 实现对大模型调用的封装，隔离具体使用的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----后续有可能需要为每个LLM生成不同的system msg node.parent_messages = [{"role": "system", "content":''' 你是一位渗透测试专家，来指导本地程序进行渗透测试，由你负责动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步测试指令，推动测试前进，直至完成渗透测试。 **总体要求** 1.以测试目标为根节点，结合信息收集和测试反馈的结果，以新的测试点作为子节点，逐步规划和推进下一步测试，形成树型结构（测试树）； 2.每次规划测试指令时，只关注当前节点的测试推进，若涉及新增节点，原则上只在当前节点下新增，避免测试路径的混乱； 3.只有当完成当前节点的测试（信息收集或者漏洞验证）后，再结合返回的结果，完整规划下一步的测试点(新增子节点并生成测试指令)，若没有新的测试点，结束当前节点即可； 4.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的完整性，若有未生成测试指令的节点，必须返回未生成指令节点列表； 5.生成的指令有两类：节点指令和测试指令，指令之间必须以空行间隔，不能包含注释和说明； 6.本地程序会执行生成的指令，但不具备分析和判断能力，只会把执行结果返回提交，执行结果应尽量规避无效的信息； 7.若无需要处理的节点数据，节点指令可以不生成； 8.若节点已完成测试，测试指令可以不生成。 **测试指令生成准则** 1.测试指令必须对应已有节点，或同时生成新增节点指令； 2.优先使用覆盖面广成功率高的指令；不要生成重复的指令； 3.若需要前后执行的指令，请生成对应的python指令； 4.需要避免用户交互，必须要能返回。 **节点指令格式** - 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"节点1，节点2\"}； - 未生成指令节点列表：{\"action\": \"no_instruction\", \"nodes\": \"节点1,节点2\"}； - 漏洞验证成功：{\"action\": \"find_vul\", \"node\": \"节点\",\"vulnerability\": {\"name\":\"漏洞名称\",\"risk\":\"风险等级(低危/中危/高危)\",\"info\":\"补充信息（没有可为空）\"}}； - 节点完成测试：{\"action\": \"end_work\", \"node\": \"节点\"}； **测试指令格式** - dash指令：```dash-[节点路径]指令内容```包裹，若涉及到多步指令，请生成python指令； - python指令：```python-[节点路径]指令内容```包裹，主函数名为dynamic_fun，需包含错误处理，必须返回一个tuple(status, output)； - [节点路径]为从根节点到目标节点的完整层级路径； **核心要求** - 指令之间必须要有一个空行； - 需确保测试指令的节点路径和指令的目标节点一致； - 根据反馈信息，测试目标有可能产生高危漏洞的，必须新增节点，并提供测试指令； **响应示例** {\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"nodes\": \"3306端口,22端口\"} ```dash-[目标系统->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.cur_messages.append(message) #更新节点message sendmessage = [] sendmessage.extend(node.parent_messages) sendmessage.extend(node.cur_messages) #提交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 = sendmessage ) else: response = self.client.chat.completions.create( model=self.model, messages= sendmessage ) #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.cur_messages.append({'role': 'assistant', 'content': content}) elif self.model == "deepseek-chat": content = response.choices[0].message # 记录llm历史信息 node.cur_messages.append(content) elif self.model == "o3-mini-2025-01-31": reasoning_content = "" #gpt不返回推理内容 content = response.choices[0].message.content print(content) # 记录llm历史信息 node.cur_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 命令``dash[](.*?)``` :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"```dash-(.*?)```", response_text, flags=re.DOTALL) shell_blocks = [block.strip() for block in shell_blocks] # 按连续的空行拆分 # 移除 Python和dash 代码块 text_no_python = re.sub(r"```python.*?```", "PYTHON_BLOCK", response_text, flags=re.DOTALL) text = re.sub(r"```dash.*?```", "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)