pvlking
/
zfsafe
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

V0.1.1 

node_tree_0.2
master
张龙 1 month ago
parent
5c988bb36a
commit
53c3ed4d10
6 changed files with 116 additions and 99 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 48
      TaskManager.py
  2. 18
      mycode/AttackMap.py
  3. 29
      mycode/ControlCenter.py
  4. 91
      mycode/LLMManager.py
  5. 2
      tools/PythoncodeTool.py
  6. 5
      tools/ToolBase.py

48
TaskManager.py
View File

@ -3,7 +3,7 @@
''' '''
from TargetManager import TargetManager # 从模块导入类 from TargetManager import TargetManager # 从模块导入类
#from LLMManager import LLMManager # 同理修正其他导入 #from LLMManager import LLMManager # 同理修正其他导入
from mycode import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。 from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。
from myutils.FileManager import FileManager from myutils.FileManager import FileManager
from InstructionManager import InstructionManager from InstructionManager import InstructionManager
from mycode.DBManager import DBManager from mycode.DBManager import DBManager
@ -19,7 +19,7 @@ class TaskManager:
def __init__(self): def __init__(self):
self.TargetM = TargetManager() self.TargetM = TargetManager()
self.logger = LogHandler.get_logger("TaskManager") self.logger = LogHandler().get_logger("TaskManager")
# 生成功能对象 # 生成功能对象
self.DBM = DBManager() #主进程一个DBM self.DBM = DBManager() #主进程一个DBM
if not self.DBM.connect(): if not self.DBM.connect():
@ -35,7 +35,7 @@ class TaskManager:
self.batch_num = 0 #一个批次的指令数量 self.batch_num = 0 #一个批次的指令数量
self.long_instr_num = 0 #耗时指令数量 self.long_instr_num = 0 #耗时指令数量
self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量，不加不减 self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量，不加不减
self.node_queue = queue.Queue() #线程安全 --待执行指令 self.node_queue = [] #
self.lock = threading.Lock() #线程锁 self.lock = threading.Lock() #线程锁
self.node_num = 0 #在处理Node线程的处理 self.node_num = 0 #在处理Node线程的处理
@ -59,7 +59,7 @@ class TaskManager:
self.logger.error("数据库连接失败！！") self.logger.error("数据库连接失败！！")
#结果入队列---2025-3-18所有的指令均需返回给LLM便于节点状态的更新，所以bres作用要调整。 #结果入队列---2025-3-18所有的指令均需返回给LLM便于节点状态的更新，所以bres作用要调整。
res = {'instr':instr,'reslut':reslut} res = {'执行指令':instr,'结果':reslut}
work_node.add_res(res) #入节点结果队列 work_node.add_res(res) #入节点结果队列
def do_worker_th(self): def do_worker_th(self):
@ -67,12 +67,14 @@ class TaskManager:
th_DBM = DBManager() th_DBM = DBManager()
th_DBM.connect() th_DBM.connect()
while self.brun: while self.brun:
try: work_node = None
with self.lock: with self.lock:
if self.node_queue:
# 获取测试节点
work_node = self.node_queue.pop(0)
self.node_num += 1 self.node_num += 1
#获取测试节点 if work_node:
work_node = self.node_queue.get(block=False)
#开始执行指令 #开始执行指令
for instruction in work_node.instr_queue: for instruction in work_node.instr_queue:
start_time = get_local_timestr() #指令执行开始时间 start_time = get_local_timestr() #指令执行开始时间
@ -83,7 +85,7 @@ class TaskManager:
#保存记录--测试使用 #保存记录--测试使用
with self.lock: with self.lock:
self.node_num -=1 self.node_num -=1
if self.node_num ==0 and self.node_queue.empty(): if self.node_num == 0 and len(self.node_queue) == 0:
with open("attack_tree", 'wb') as f: with open("attack_tree", 'wb') as f:
pickle.dump(TM.CCM.attack_tree, f) pickle.dump(TM.CCM.attack_tree, f)
#针对一个节点的指令执行完成后，提交LLM规划下一步操作 #针对一个节点的指令执行完成后，提交LLM规划下一步操作
@ -92,7 +94,7 @@ class TaskManager:
# continue # continue
# for node in node_list: # for node in node_list:
# self.node_queue.put(node) # self.node_queue.put(node)
except queue.Empty: else:
time.sleep(10) time.sleep(10)
# try: # try:
# instruction = self.instr_queue.get(block=False) #quere线程安全，block=false非阻塞get # instruction = self.instr_queue.get(block=False) #quere线程安全，block=false非阻塞get
@ -122,8 +124,9 @@ class TaskManager:
know_info = "无" #? know_info = "无" #?
#启动--初始化指令 #启动--初始化指令
node_list = self.CCM.start_do(target,self.task_id) node_list = self.CCM.start_do(target,self.task_id)
with self.lock:
for node in node_list: for node in node_list:
self.node_queue.put(node) self.node_queue.append(node)
#创建工作线程----2025-3-18调整为一个节点一个线程， #创建工作线程----2025-3-18调整为一个节点一个线程，
for i in range(self.max_thread_num): for i in range(self.max_thread_num):
@ -158,8 +161,11 @@ if __name__ == "__main__":
#测试执行指令 #测试执行指令
with open("attack_tree", "rb") as f: with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f) TM.CCM.attack_tree = pickle.load(f)
with open("node_list", 'rb') as f: # 遍历node，查看有res的数据
TM.node_queue = pickle.load(f) nodes = TM.CCM.attack_tree.traverse_bfs()
for node in nodes:
if node.instr_queue: # list
TM.node_queue.append(node)
#创建线程执行指令 #创建线程执行指令
for i in range(TM.max_thread_num): for i in range(TM.max_thread_num):
@ -176,23 +182,27 @@ if __name__ == "__main__":
#遍历node，查看有res的数据 #遍历node，查看有res的数据
nodes = TM.CCM.attack_tree.traverse_bfs() nodes = TM.CCM.attack_tree.traverse_bfs()
for node in nodes: for node in nodes:
if not node.res_quere.empty(): if node.res_quere: #list
node_list = TM.CCM.get_llm_instruction(node) node_list = TM.CCM.get_llm_instruction(node)
if not node_list:#该节点测试完成 if not node_list:#该节点测试完成
continue continue
with TM.lock:
for do_node in node_list: for do_node in node_list:
TM.node_queue.put(do_node) TM.node_queue.append(do_node)
elif test_type ==3: #新目标测试 # 暂存状态--
with open("attack_tree", 'wb') as f:
pickle.dump(TM.CCM.attack_tree, f)
elif test_type == 3: #新目标测试
# 启动--初始化指令 # 启动--初始化指令
node_list = TM.CCM.start_do("192.168.204.137", 0) node_list = TM.CCM.start_do("192.168.204.137", 0)
with TM.lock:
for node in node_list: for node in node_list:
TM.node_queue.put(node) TM.node_queue.append(node)
#暂存状态-- #暂存状态--
with open("attack_tree",'wb') as f: with open("attack_tree",'wb') as f:
pickle.dump(TM.CCM.attack_tree,f) pickle.dump(TM.CCM.attack_tree,f)
with open("node_list",'wb') as f: elif test_type == 4: #读取messages
pickle.dump(TM.node_queue,f) pass
else: else:
#完整过程测试---要设定终止条件 #完整过程测试---要设定终止条件
pass pass

18
mycode/AttackMap.py
View File

@ -57,6 +57,8 @@ class AttackTree:
current_node = self.root #从根节点开始 current_node = self.root #从根节点开始
node_names = node_path.split('->') node_names = node_path.split('->')
for node_name in node_names: for node_name in node_names:
if node_name == "目标系统":
continue
if node_name == current_node.name:#根节点开始 if node_name == current_node.name:#根节点开始
continue continue
else: else:
@ -114,8 +116,8 @@ class TreeNode:
self.children = [] # 子节点列表 self.children = [] # 子节点列表
self.parent = None # 父节点引用 self.parent = None # 父节点引用
self.path = "" #当前节点的路径 self.path = "" #当前节点的路径
self.instr_queue = queue.Queue() #针对当前节点的执行指令----重要约束：一个节点只能有一个线程在执行指令 self.instr_queue = [] #queue.Queue() #针对当前节点的执行指令----重要约束：一个节点只能有一个线程在执行指令
self.res_quere = queue.Queue() #指令执行的结果，一批一批 self.res_quere = [] #queue.Queue() #指令执行的结果，一批一批
self.llm_sn = 0 #针对该节点llm提交次数 self.llm_sn = 0 #针对该节点llm提交次数
self.do_sn = 0 #针对该节点instr执行次数 self.do_sn = 0 #针对该节点instr执行次数
self.messages = [] #针对当前节点积累的messages -- 针对不同节点提交不同的messages self.messages = [] #针对当前节点积累的messages -- 针对不同节点提交不同的messages
@ -127,20 +129,16 @@ class TreeNode:
self.children.append(child_node) self.children.append(child_node)
def add_instr(self,instr): def add_instr(self,instr):
self.instr_queue.put(instr) self.instr_queue.append(instr)
def get_instr(self): def get_instr(self):
if self.instr_queue.empty(): return self.instr_queue.pop(0) if self.instr_queue else None
return None
return self.instr_queue.get()
def add_res(self,str_res): #结构化结果字串 def add_res(self,str_res): #结构化结果字串
self.res_quere.put(str_res) self.res_quere.append(str_res)
def get_res(self): def get_res(self):
if self.res_quere.empty(): return self.res_queue.pop(0) if self.res_queue else None
return None
return self.res_quere.get()
def __repr__(self): def __repr__(self):
return f"TreeNode({self.name}, {self.status}, {self.vul_type})" return f"TreeNode({self.name}, {self.status}, {self.vul_type})"

29
mycode/ControlCenter.py
View File

@ -11,7 +11,7 @@ from myutils.MyLogger_logger import LogHandler
class ControlCenter: class ControlCenter:
def __init__(self,DBM): def __init__(self,DBM):
self.logger = LogHandler.get_logger("ControlCenter") self.logger = LogHandler().get_logger("ControlCenter")
self.task_id = None self.task_id = None
self.target = None self.target = None
self.attack_tree = None self.attack_tree = None
@ -55,16 +55,15 @@ class ControlCenter:
#约束1：一个节点只能同时提交一次，未测试的节点不要重复 #约束1：一个节点只能同时提交一次，未测试的节点不要重复
def get_llm_instruction(self,node): def get_llm_instruction(self,node):
#拼接结果字符串 #拼接结果字符串--由于测试需要queue改成了list 2025-3-19
json_strings = [] # json_strings = []
while True: # if node.res_quere:
try: # for item in node.res_quere:
item = node.res_quere.get_nowait() # json_str = json.dumps(item, ensure_ascii=False)
json_str = json.dumps(item,ensure_ascii=False) # json_strings.append(json_str)
json_strings.append(json_str) # res_str = ','.join(json_strings)
except queue.Empty:
break res_str = json.dumps(node.res_quere,ensure_ascii=False)
res_str = ','.join(json_strings)
#构造本次提交的prompt #构造本次提交的prompt
user_Prompt = f''' user_Prompt = f'''
@ -85,6 +84,8 @@ class ControlCenter:
def tree_manager(self,node_cmds,node): def tree_manager(self,node_cmds,node):
'''更新渗透测试树''' '''更新渗透测试树'''
if not node_cmds or len(node_cmds)==0:
return
try: try:
node_jsons = json.loads(node_cmds) node_jsons = json.loads(node_cmds)
for node_json in node_jsons: for node_json in node_jsons:
@ -125,7 +126,7 @@ class ControlCenter:
find_node = self.attack_tree.find_node_by_nodepath(node_path) find_node = self.attack_tree.find_node_by_nodepath(node_path)
if find_node: if find_node:
instruction = re.sub(r'\[.*?\]', "", command,count=1,flags=re.DOTALL) instruction = re.sub(r'\[.*?\]', "", command,count=1,flags=re.DOTALL)
find_node.instr_queue.put(instruction) find_node.instr_queue.append(instruction)
#入输出队列 #入输出队列
if find_node not in node_list: if find_node not in node_list:
node_list.append(find_node) node_list.append(find_node)
@ -135,7 +136,7 @@ class ControlCenter:
self.logger.error(f"得到的指令格式不符合规范：{command}") self.logger.error(f"得到的指令格式不符合规范：{command}")
return node_list return node_list
#待修改
def is_user_instr(self,instr): def is_user_instr(self,instr):
''' '''
过滤需要人工确认或手动执行的指令 ---- 待完善 过滤需要人工确认或手动执行的指令 ---- 待完善
@ -170,7 +171,7 @@ class ControlCenter:
with self.lock: with self.lock:
self.long_instr_num +=1 #耗时指令数量+1 self.long_instr_num +=1 #耗时指令数量+1
# 指令入队列 # 指令入队列
self.instr_queue.put(instr) self.instr_queue.append(instr)
def stop_do(self): def stop_do(self):
#清空数据 #清空数据

91
mycode/LLMManager.py
View File

@ -55,53 +55,59 @@ class LLMManager:
# 构建初始提示 初始化messages # 构建初始提示 初始化messages
def build_initial_prompt(self,target,know_info="无",node=None): 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: if not node:
return None return None
#根节点初始化message #根节点初始化message
node.messages = [{"role": "system", node.messages = [{"role": "system",
"content":''' "content":'''
你是一位渗透测试专家,由你动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步的测试指令，推动测试阶段前进，直至完成渗透测试。本地程序会根据你的指令进行执行,然后将执行结果返回给你。 你是一位渗透测试专家,由你动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步的测试指令，推动测试阶段前进，直至完成渗透测试。本地程序会根据你的指令进行执行,然后将执行结果返回给你。
整体过程说明： 总体要求说明：
1.以目标系统所在IP为根节点，每个渗透测试点（如端口、服务、漏洞点）作为子节点，形成树型结构（测试树）； 1.以目标系统所在IP为根节点，随着信息收集和渗透测试的推进，每个渗透测试点（如端口、服务、漏洞点）作为子节点，形成树型结构（测试树）；
2.测试树整体数据由本地程序存储，你只需要关注当前节点的渗透测试推进，节点状态(未完成、已完成)的更新，和是否有子节点新增； 2.测试树整体数据由本地程序存储，你只需要关注当前节点的渗透测试推进，节点状态(未完成、已完成)的更新，和是否有子节点新增；
3.新增子节点时，同一父节点下的子节点名称不能相同，并提供对该节点的测试指令； 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匿名登录\"}； {\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"node\": \"21端口\", \"status\": \"未完成\"}
5.发现漏洞后，可根据漏洞类型决定是否新增子节点继续测试，如：{\"action\": \"add_node\", \"parent\": \"21端口\", \"node\": \"ftp配置检查\", \"status\": \"未完成\"}；
6.当本地程序确认所有节点状态均为“已完成”时，生成测试报告。 ```bash-[目标系统->192.168.1.100->3306端口]
生成的渗透测试指令需满足如下约束： mysql -u root -p 192.168.1.100```
1.只返回具体的shell指令或Python代码，不要包含注释和说明； 若无节点修改或新增，节点指令可以为空，但测试指令必须对应已有节点；
2.shell指令以```bash-[对应节点的路径](.*?)```包裹，python代码以```python-[对应节点的路径](.*?)```包裹； 决策流程：
3.若提供的是shell指令，需要避免用户再次交互； 1. 如果当前节点是IP地址，且未进行端口扫描，则执行端口扫描。
4.若提供的是python代码，主函数名为dynamic_fun，需包含错误处理，执行结束后返回一个tuple (status, output)，其中status为'success'或'failure'，output为补充输出信息； 2. 如果端口扫描发现开放端口，则为每个开放端口新增一个测试节点。
示例： 3. 如果当前节点是端口，且未进行服务扫描，则执行服务扫描。
```python-[目标系统->192.168.1.100->3306端口] 4. 如果服务扫描发现服务版本或漏洞，则新增对应的漏洞测试节点。
def dynamic_fun(): 5. 如果当前节点是漏洞测试，且漏洞利用成功，则根据利用结果决定是否新增子节点进一步测试。
6. 如果当前节点测试未发现新信息，则更新节点状态为“已完成”，并继续测试其他未完成节点。
生成的节点指令需要满足如下约束：
1.新增节点指令示例：{\"action\":\"add_node\", \"parent\": \"80端口\", \"node\": \"http://192.168.1.100/index.php?id=1\", \"status\": \"未完成\"}；
2.新增子节点时，同一请求返回的子节点名（node）不能相同，且必须同时提供对该节点的测试指令；
3.若认为该节点已完成测试，修改该节点为“已完成”状态，完成节点测试且发现漏洞示例：{\"action\": \"update_status\", \"node\": \"21端口\", \"status\": \"已完成\"，\"vulnerability\": \"ftp匿名登录\"}；
4.发现漏洞后，可根据漏洞类型决定是否新增子节点继续测试，如：{\"action\": \"add_node\", \"parent\": \"21端口\", \"node\": \"ftp配置检查\", \"status\": \"未完成\"}；
生成的渗透测试指令需满足如下约束：
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: try:
result = "扫描完成" result = "扫描完成"
return ("success", result) return ("success", result)
except Exception as e: except Exception as e:
return ("failure", str(e)) return ("failure", str(e))
``` ```
5.如果认为渗透测试已完成，请生成生成报告的python代码，并返回'complete'和报告文件路径。 限制条件：
1.仅在发现高危漏洞或关键测试路径时新增节点
'''}] # 一个messages '''}] # 一个messages
user_Prompt = f''' user_Prompt = f'''
当前分支路径：目标系统->{target} 当前分支路径：目标系统->{target}
当前节点信息： 当前节点信息：
- 节点名称：{target} - 节点名称：{target}
- 节点状态：未完成 - 节点状态：未完成
- 漏洞类型：未发现 - 漏洞类型：未发现
上一步结果：已知信息-{know_info} 上一步结果：{know_info}
任务：生成下一步渗透测试指令或结束该节点的渗透测试(修改节点状态为：已完成)。 任务：生成下一步渗透测试指令或结束该节点的渗透测试(修改节点状态为：已完成)。
''' '''
return user_Prompt return user_Prompt
@ -126,7 +132,7 @@ class LLMManager:
post_time = get_local_timestr() post_time = get_local_timestr()
response = self.client.chat.completions.create( response = self.client.chat.completions.create(
model=self.model, model=self.model,
messages = self.messages messages = node.messages
) )
#LLM返回结果处理 #LLM返回结果处理
@ -156,17 +162,14 @@ class LLMManager:
self.logger.error(f"{node.name}-llm入库失败！") self.logger.error(f"{node.name}-llm入库失败！")
#需要对指令进行提取 #需要对指令进行提取
node_list = self.fetch_instruction(content,node) node_cmds,commands = self.fetch_instruction(content,node)
#********测试时使用---输出和记录LLM返回指令的message return node_cmds,commands
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): def fetch_instruction(self,response_text,node):
''' '''
*****该函数很重要，需要一定的容错能力，解析LLM返回内容*****
处理边界：只格式化分析LLM返回内容，指令和节点操作等交其他模块。
节点控制指令 节点控制指令
渗透测试指令 渗透测试指令
提取命令列表，包括： 提取命令列表，包括：
@ -211,7 +214,7 @@ class LLMManager:
#其他的认为是节点操作指令 #其他的认为是节点操作指令
node_cmds.append(part) node_cmds.append(part)
return node_cmds,commands #？存在一个问题：这样分list返回，执行顺序会丢失 return node_cmds,commands
def test_llm(self): def test_llm(self):
with open("../test", "r", encoding="utf-8") as f: with open("../test", "r", encoding="utf-8") as f:

2
tools/PythoncodeTool.py
View File

@ -6,7 +6,7 @@ import builtins
import re import re
from tools.ToolBase import ToolBase from tools.ToolBase import ToolBase
class PythonTool(ToolBase): class PythoncodeTool(ToolBase):
def is_safe_code(self,code): def is_safe_code(self,code):
# List of high-risk functions to block (can be adjusted based on requirements) # List of high-risk functions to block (can be adjusted based on requirements)

5
tools/ToolBase.py
View File

@ -29,6 +29,7 @@ class ToolBase(abc.ABC):
parser = argparse.ArgumentParser(add_help=False) #创建命令行参数解析器对象‌ parser = argparse.ArgumentParser(add_help=False) #创建命令行参数解析器对象‌
parser.add_argument("-o ", "--output", type=str) #添加需要解析的参数规则‌ parser.add_argument("-o ", "--output", type=str) #添加需要解析的参数规则‌
parser.add_argument("-oN ", "--Noutput", type=str) #nmap parser.add_argument("-oN ", "--Noutput", type=str) #nmap
parser.add_argument("-oG ","--NMG",type=str) #nmap
parser.add_argument("-output ", "--nikto", type=str) #nikto parser.add_argument("-output ", "--nikto", type=str) #nikto
args, _ = parser.parse_known_args(command.split()[1:]) args, _ = parser.parse_known_args(command.split()[1:])
return args return args
@ -94,6 +95,10 @@ class ToolBase(abc.ABC):
file_out = self.read_output_file(parsed_arg.nikto) file_out = self.read_output_file(parsed_arg.nikto)
stderr = "" stderr = ""
stdout = file_out stdout = file_out
elif parsed_arg.NMG:
file_out = self.read_output_file(parsed_arg.NMG)
stderr = ""
stdout = file_out
else: else:
stderr = result.stderr stderr = result.stderr
stdout = result.stdout stdout = result.stdout

Write Preview
Loading…
Cancel
Save