diff --git a/TaskManager.py b/TaskManager.py
index d0c8fdd..afd33f5 100644
--- a/TaskManager.py
+++ b/TaskManager.py
@@ -3,7 +3,7 @@
'''
from TargetManager import TargetManager # 从模块导入类
#from LLMManager import LLMManager # 同理修正其他导入
-from mycode import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。
+from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。
from myutils.FileManager import FileManager
from InstructionManager import InstructionManager
from mycode.DBManager import DBManager
@@ -19,7 +19,7 @@ class TaskManager:
def __init__(self):
self.TargetM = TargetManager()
- self.logger = LogHandler.get_logger("TaskManager")
+ self.logger = LogHandler().get_logger("TaskManager")
# 生成功能对象
self.DBM = DBManager() #主进程一个DBM
if not self.DBM.connect():
@@ -35,7 +35,7 @@ class TaskManager:
self.batch_num = 0 #一个批次的指令数量
self.long_instr_num = 0 #耗时指令数量
self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量，不加不减
- self.node_queue = queue.Queue() #线程安全 --待执行指令
+ self.node_queue = [] #
self.lock = threading.Lock() #线程锁
self.node_num = 0 #在处理Node线程的处理
@@ -59,7 +59,7 @@ class TaskManager:
self.logger.error("数据库连接失败！！")
#结果入队列---2025-3-18所有的指令均需返回给LLM便于节点状态的更新，所以bres作用要调整。
- res = {'instr':instr,'reslut':reslut}
+ res = {'执行指令':instr,'结果':reslut}
work_node.add_res(res) #入节点结果队列
def do_worker_th(self):
@@ -67,12 +67,14 @@ class TaskManager:
th_DBM = DBManager()
th_DBM.connect()
while self.brun:
- try:
- with self.lock:
+ work_node = None
+ with self.lock:
+ if self.node_queue:
+ # 获取测试节点
+ work_node = self.node_queue.pop(0)
self.node_num += 1
- #获取测试节点
- work_node = self.node_queue.get(block=False)
+ if work_node:
#开始执行指令
for instruction in work_node.instr_queue:
start_time = get_local_timestr() #指令执行开始时间
@@ -83,7 +85,7 @@ class TaskManager:
#保存记录--测试使用
with self.lock:
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:
pickle.dump(TM.CCM.attack_tree, f)
#针对一个节点的指令执行完成后，提交LLM规划下一步操作
@@ -92,7 +94,7 @@ class TaskManager:
# continue
# for node in node_list:
# self.node_queue.put(node)
- except queue.Empty:
+ else:
time.sleep(10)
# try:
# instruction = self.instr_queue.get(block=False) #quere线程安全，block=false非阻塞get
@@ -122,8 +124,9 @@ class TaskManager:
know_info = "无" #?
#启动--初始化指令
node_list = self.CCM.start_do(target,self.task_id)
- for node in node_list:
- self.node_queue.put(node)
+ with self.lock:
+ for node in node_list:
+ self.node_queue.append(node)
#创建工作线程----2025-3-18调整为一个节点一个线程，
for i in range(self.max_thread_num):
@@ -158,8 +161,11 @@ if __name__ == "__main__":
#测试执行指令
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
- with open("node_list", 'rb') as f:
- TM.node_queue = pickle.load(f)
+ # 遍历node，查看有res的数据
+ 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):
@@ -176,23 +182,27 @@ if __name__ == "__main__":
#遍历node，查看有res的数据
nodes = TM.CCM.attack_tree.traverse_bfs()
for node in nodes:
- if not node.res_quere.empty():
+ if node.res_quere: #list
node_list = TM.CCM.get_llm_instruction(node)
if not node_list:#该节点测试完成
continue
- for do_node in node_list:
- TM.node_queue.put(do_node)
- elif test_type ==3: #新目标测试
+ with TM.lock:
+ for do_node in node_list:
+ TM.node_queue.append(do_node)
+ # 暂存状态--
+ 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)
- for node in node_list:
- TM.node_queue.put(node)
+ with TM.lock:
+ for node in node_list:
+ TM.node_queue.append(node)
#暂存状态--
with open("attack_tree",'wb') as f:
pickle.dump(TM.CCM.attack_tree,f)
- with open("node_list",'wb') as f:
- pickle.dump(TM.node_queue,f)
-
+ elif test_type == 4: #读取messages
+ pass
else:
#完整过程测试---要设定终止条件
pass
diff --git a/mycode/AttackMap.py b/mycode/AttackMap.py
index b928262..ab2bf4d 100644
--- a/mycode/AttackMap.py
+++ b/mycode/AttackMap.py
@@ -57,6 +57,8 @@ class AttackTree:
current_node = self.root #从根节点开始
node_names = node_path.split('->')
for node_name in node_names:
+ if node_name == "目标系统":
+ continue
if node_name == current_node.name:#根节点开始
continue
else:
@@ -114,8 +116,8 @@ class TreeNode:
self.children = [] # 子节点列表
self.parent = None # 父节点引用
self.path = "" #当前节点的路径
- self.instr_queue = queue.Queue() #针对当前节点的执行指令----重要约束：一个节点只能有一个线程在执行指令
- self.res_quere = queue.Queue() #指令执行的结果，一批一批
+ self.instr_queue = [] #queue.Queue() #针对当前节点的执行指令----重要约束：一个节点只能有一个线程在执行指令
+ self.res_quere = [] #queue.Queue() #指令执行的结果，一批一批
self.llm_sn = 0 #针对该节点llm提交次数
self.do_sn = 0 #针对该节点instr执行次数
self.messages = [] #针对当前节点积累的messages -- 针对不同节点提交不同的messages
@@ -127,20 +129,16 @@ class TreeNode:
self.children.append(child_node)
def add_instr(self,instr):
- self.instr_queue.put(instr)
+ self.instr_queue.append(instr)
def get_instr(self):
- if self.instr_queue.empty():
- return None
- return self.instr_queue.get()
+ return self.instr_queue.pop(0) if self.instr_queue else None
def add_res(self,str_res): #结构化结果字串
- self.res_quere.put(str_res)
+ self.res_quere.append(str_res)
def get_res(self):
- if self.res_quere.empty():
- return None
- return self.res_quere.get()
+ return self.res_queue.pop(0) if self.res_queue else None
def __repr__(self):
return f"TreeNode({self.name}, {self.status}, {self.vul_type})"
diff --git a/mycode/ControlCenter.py b/mycode/ControlCenter.py
index 4a05fce..ccce462 100644
--- a/mycode/ControlCenter.py
+++ b/mycode/ControlCenter.py
@@ -11,7 +11,7 @@ from myutils.MyLogger_logger import LogHandler
class ControlCenter:
def __init__(self,DBM):
- self.logger = LogHandler.get_logger("ControlCenter")
+ self.logger = LogHandler().get_logger("ControlCenter")
self.task_id = None
self.target = None
self.attack_tree = None
@@ -55,16 +55,15 @@ class ControlCenter:
#约束1：一个节点只能同时提交一次，未测试的节点不要重复
def get_llm_instruction(self,node):
- #拼接结果字符串
- json_strings = []
- while True:
- try:
- item = node.res_quere.get_nowait()
- json_str = json.dumps(item,ensure_ascii=False)
- json_strings.append(json_str)
- except queue.Empty:
- break
- res_str = ','.join(json_strings)
+ #拼接结果字符串--由于测试需要queue改成了list 2025-3-19
+ # json_strings = []
+ # if node.res_quere:
+ # for item in node.res_quere:
+ # json_str = json.dumps(item, ensure_ascii=False)
+ # json_strings.append(json_str)
+ # res_str = ','.join(json_strings)
+
+ res_str = json.dumps(node.res_quere,ensure_ascii=False)
#构造本次提交的prompt
user_Prompt = f'''
@@ -85,6 +84,8 @@ class ControlCenter:
def tree_manager(self,node_cmds,node):
'''更新渗透测试树'''
+ if not node_cmds or len(node_cmds)==0:
+ return
try:
node_jsons = json.loads(node_cmds)
for node_json in node_jsons:
@@ -125,7 +126,7 @@ class ControlCenter:
find_node = self.attack_tree.find_node_by_nodepath(node_path)
if find_node:
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:
node_list.append(find_node)
@@ -135,7 +136,7 @@ class ControlCenter:
self.logger.error(f"得到的指令格式不符合规范：{command}")
return node_list
-
+ #待修改
def is_user_instr(self,instr):
'''
过滤需要人工确认或手动执行的指令 ---- 待完善
@@ -170,7 +171,7 @@ class ControlCenter:
with self.lock:
self.long_instr_num +=1 #耗时指令数量+1
# 指令入队列
- self.instr_queue.put(instr)
+ self.instr_queue.append(instr)
def stop_do(self):
#清空数据
diff --git a/mycode/LLMManager.py b/mycode/LLMManager.py
index d1a240e..728c61b 100644
--- a/mycode/LLMManager.py
+++ b/mycode/LLMManager.py
@@ -55,53 +55,59 @@ class LLMManager:
# 构建初始提示 初始化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
+ "content":'''
+你是一位渗透测试专家,由你动态控制整个渗透测试过程，根据当前测试状态和返回结果，决定下一步的测试指令，推动测试阶段前进，直至完成渗透测试。本地程序会根据你的指令进行执行,然后将执行结果返回给你。
+总体要求说明：
+1.以目标系统所在IP为根节点，随着信息收集和渗透测试的推进，每个渗透测试点（如端口、服务、漏洞点）作为子节点，形成树型结构（测试树）；
+2.测试树整体数据由本地程序存储，你只需要关注当前节点的渗透测试推进，节点状态(未完成、已完成)的更新，和是否有子节点新增；
+3.你需要返回两类指令：节点指令和测试指令，以空行间隔；
+--示例：
+{\"action\":\"add_node\", \"parent\": \"192.168.1.100\", \"node\": \"21端口\", \"status\": \"未完成\"}
+
+```bash-[目标系统->192.168.1.100->3306端口]
+mysql -u root -p 192.168.1.100```
+若无节点修改或新增，节点指令可以为空，但测试指令必须对应已有节点；
+决策流程：
+1. 如果当前节点是IP地址，且未进行端口扫描，则执行端口扫描。
+2. 如果端口扫描发现开放端口，则为每个开放端口新增一个测试节点。
+3. 如果当前节点是端口，且未进行服务扫描，则执行服务扫描。
+4. 如果服务扫描发现服务版本或漏洞，则新增对应的漏洞测试节点。
+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:
+ result = "扫描完成"
+ return ("success", result)
+ except Exception as e:
+ return ("failure", str(e))
+```
+限制条件：
+1.仅在发现高危漏洞或关键测试路径时新增节点
+ '''}] # 一个messages
user_Prompt = f'''
- 当前分支路径：目标系统->{target}
- 当前节点信息：
- - 节点名称：{target}
- - 节点状态：未完成
- - 漏洞类型：未发现
- 上一步结果：已知信息-{know_info}
- 任务：生成下一步渗透测试指令或结束该节点的渗透测试(修改节点状态为：已完成)。
+当前分支路径：目标系统->{target}
+当前节点信息：
+- 节点名称：{target}
+- 节点状态：未完成
+- 漏洞类型：未发现
+上一步结果：{know_info}
+任务：生成下一步渗透测试指令或结束该节点的渗透测试(修改节点状态为：已完成)。
'''
return user_Prompt
@@ -126,7 +132,7 @@ class LLMManager:
post_time = get_local_timestr()
response = self.client.chat.completions.create(
model=self.model,
- messages = self.messages
+ messages = node.messages
)
#LLM返回结果处理
@@ -156,17 +162,14 @@ class LLMManager:
self.logger.error(f"{node.name}-llm入库失败！")
#需要对指令进行提取
- node_list = self.fetch_instruction(content,node)
+ node_cmds,commands = 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
+ return node_cmds,commands
def fetch_instruction(self,response_text,node):
'''
+ *****该函数很重要，需要一定的容错能力，解析LLM返回内容*****
+ 处理边界：只格式化分析LLM返回内容，指令和节点操作等交其他模块。
节点控制指令
渗透测试指令
提取命令列表，包括：
@@ -211,7 +214,7 @@ class LLMManager:
#其他的认为是节点操作指令
node_cmds.append(part)
- return node_cmds,commands #？存在一个问题：这样分list返回，执行顺序会丢失
+ return node_cmds,commands
def test_llm(self):
with open("../test", "r", encoding="utf-8") as f:
diff --git a/tools/PythoncodeTool.py b/tools/PythoncodeTool.py
index 1aeeb2d..c93618f 100644
--- a/tools/PythoncodeTool.py
+++ b/tools/PythoncodeTool.py
@@ -6,7 +6,7 @@ import builtins
import re
from tools.ToolBase import ToolBase
-class PythonTool(ToolBase):
+class PythoncodeTool(ToolBase):
def is_safe_code(self,code):
# List of high-risk functions to block (can be adjusted based on requirements)
diff --git a/tools/ToolBase.py b/tools/ToolBase.py
index 75557f0..c5b14b6 100644
--- a/tools/ToolBase.py
+++ b/tools/ToolBase.py
@@ -29,6 +29,7 @@ class ToolBase(abc.ABC):
parser = argparse.ArgumentParser(add_help=False) #创建命令行参数解析器对象
parser.add_argument("-o ", "--output", type=str) #添加需要解析的参数规则
parser.add_argument("-oN ", "--Noutput", type=str) #nmap
+ parser.add_argument("-oG ","--NMG",type=str) #nmap
parser.add_argument("-output ", "--nikto", type=str) #nikto
args, _ = parser.parse_known_args(command.split()[1:])
return args
@@ -94,6 +95,10 @@ class ToolBase(abc.ABC):
file_out = self.read_output_file(parsed_arg.nikto)
stderr = ""
stdout = file_out
+ elif parsed_arg.NMG:
+ file_out = self.read_output_file(parsed_arg.NMG)
+ stderr = ""
+ stdout = file_out
else:
stderr = result.stderr
stdout = result.stdout