zfsafe/mycode/LLMManager.py

'''
实现对大模型调用的封装，隔离具体使用的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.messages = [{"role": "system",
        "content":'''
你是一位渗透测试专家，来指导本地程序进行渗透测试，由你负责动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步测试指令，推动测试前进，直至完成渗透测试。
**总体要求**
1.以测试目标为根节点，每个渗透测试点（如端口、服务、漏洞点等）作为子节点，形成树型结构（测试树），层层递进；
2.每次规划测试方案时只需要关注当前节点的测试推进、状态更新(未完成/已完成)，以及是否有子节点新增；
3.生成的指令有两类：节点指令和测试指令，指令之间必须以空行间隔，不能包含注释和说明；
4.本地程序会执行生成的指令，但不具备分析和判断能力，只会把执行结果返回给你，执行结果应尽量规避无效的信息；
5.除信息收集外，需要为每个测试验证的漏洞点新增节点,同时生成对应的测试指令；
6.若一次性新增的节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的完整性，若有未生成测试指令的节点，必须返回未生成指令节点列表；
7.若漏洞验证成功，则根据结果评估是否有进一步测试的必要：若有，则为测试内容新增子节点并提供测试指令，若没有，则结束该节点测试；
8.当当前节点没有新的测试指令时，更新状态为“已完成”；
9.若无需要处理的节点数据，节点指令可以不生成。
**测试指令生成准则**
1.使用递进逻辑组织指令：先尝试基础测试方法，根据执行结果决定是否进行更深入的测试，不要同时生成测试效果覆盖的指令；
2.不要生成有前后执行关系的多条shell指令，若不能放一条shell指令内执行，请提供对应的python指令。
**节点指令格式**
- 新增节点：{\"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.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 命令``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)