'''
渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据
'''
from mycode.TargetManager import TargetManager # 从模块导入类
#from LLMManager import LLMManager # 同理修正其他导入
from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。
from myutils.FileManager import FileManager
from mycode.InstructionManager import InstructionManager
from mycode.DBManager import DBManager
from myutils.MyTime import get_local_timestr
from myutils.MyLogger_logger import LogHandler
import pickle
import queue
import time
import os
import threading
class TaskManager:
def __init__(self):
self.TargetM = TargetManager()
self.logger = LogHandler().get_logger("TaskManager")
# 生成功能对象
self.DBM = DBManager() #主进程一个DBM
if not self.DBM.connect():
self.logger.error("数据库连接失败！终止工作！")
return
self.CCM = ControlCenter(self.DBM,self)
self.InstrM = InstructionManager(self) # 类对象渗透，要约束只读取信息
# 控制最大并发指令数量
self.max_thread_num = 2
self.task_id = 0 #任务id --
self.workth_list = [] #线程句柄list
# self.long_instr_num = 0 #耗时指令数量
# self.long_time_instr = ['nikto'] #耗时操作不计入批量执行的数量，不加不减
self.node_queue = queue.Queue()
self.lock = threading.Lock() #线程锁
self.node_num = 0 #在处理Node线程的处理
self.brun = True
self.cookie = "" #cookie参数
def res_in_quere(self,bres,instr,reslut,start_time,end_time,th_DBM,source_result,ext_params,work_node):
'''
执行结果入队列
:param bres:
:param instr:
:param reslut:
:return:
'''
#入数据库 -- bres True和False 都入数据库2025-3-10---加node_path(2025-3-18)#？
if th_DBM.ok:
work_node.do_sn += 1 # 指令的执行序列是一个任务共用,要线程锁，错误问题也不大
th_DBM.insetr_result(self.task_id, instr, reslut, work_node.do_sn, start_time, end_time, source_result,
ext_params, work_node.path)
else:
self.logger.error("数据库连接失败！！")
#结果入队列---2025-3-18所有的指令均需返回给LLM便于节点状态的更新，所以bres作用要调整。
res = {'执行指令':instr,'结果':reslut}
str_res = json.dumps(res,ensure_ascii=False) #直接字符串组合也可以-待验证
work_node.llm_type = 1
work_node.add_res(str_res) #入节点结果队列
def do_worker_th(self):
#线程的dbm需要一个线程一个
th_DBM = DBManager()
th_DBM.connect()
while self.brun:
try:
work_node = self.node_queue.get(block=False)
# 测试时使用
with self.lock:
self.node_num += 1
# 开始执行指令
while work_node.instr_queue:
#for instruction in work_node.instr_queue: #这里要重新调整#?
instruction = work_node.instr_queue.pop(0)
start_time = get_local_timestr() # 指令执行开始时间
bres, instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction)
end_time = get_local_timestr() # 指令执行结束时间
self.res_in_quere(bres, instr, reslut, start_time, end_time, th_DBM, source_result,
ext_params, work_node) # 执行结果入队列
# #针对一个节点的指令执行完成后，提交LLM规划下一步操作
#self.CCM.llm_quere.put(work_node)
# 保存记录--测试时使用--后期增加人为停止测试时可以使用
with self.lock:
self.node_num -= 1
if self.node_num == 0 and self.node_queue.empty(): #
self.logger.debug("此批次指令执行完成！")
with open("attack_tree", 'wb') as f:
pickle.dump(self.CCM.attack_tree, f)
except queue.Empty:
time.sleep(20)
def start_task(self,target_name,target_in):
'''
:param target_name: 任务目标名字
:param target_in: 任务目标访问地址
:return:
'''
#判断目标合法性
bok,target,type = self.TargetM.validate_and_extract(target_in)
if bok:
self.target = target
self.type = type #1-IP，2-domain
self.task_id = self.DBM.start_task(target_name,target_in) #数据库新增任务记录
#获取基本信息: 读取数据库或预生成指令，获取基本的已知信息
know_info = "无" #?
#启动--初始化指令
self.CCM.start_do(target,self.task_id)
#创建工作线程----2025-3-18调整为一个节点一个线程，
for i in range(self.max_thread_num):
w_th = threading.Thread(target=self.do_worker_th)
w_th.start()
self.workth_list.append(w_th)
#等待线程结束--执行生成报告
for t in self.workth_list:
t.join()
#生成报告
pass
else:
return False,"{target}检测目标不合规，请检查！"
def stop_task(self):
self.brun = False
self.CCM.stop_do() #清空一些全局变量
self.InstrM.init_data()
#结束任务需要收尾处理#？
if __name__ == "__main__":
import json
TM = TaskManager()
FM = FileManager()
current_path = os.path.dirname(os.path.realpath(__file__))
strMsg = FM.read_file("test",1)
test_type = 2
instr_index = 19
iput_index = -1 # 0是根节点
indexs = []
if test_type == 0: #新目标测试
# 启动--初始化指令
node_list = TM.CCM.start_do("58.216.217.70", 1)
#异步处理，需要等待线程结束了
for th in TM.CCM.llmth_list:
th.join()
elif test_type == 1:
#测试执行指令
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
# 遍历node，查看有instr的ndoe
nodes = TM.CCM.attack_tree.traverse_dfs()
if indexs:
for index in indexs:
node = nodes[index]
if node.instr_queue: # list
TM.node_queue.put(node)
else:
for node in nodes:
if node.instr_queue:
TM.node_queue.put(node)
#创建线程执行指令
for i in range(TM.max_thread_num):
w_th = threading.Thread(target=TM.do_worker_th)
w_th.start()
TM.workth_list.append(w_th)
# 等待线程结束--执行生成报告
for t in TM.workth_list:
t.join()
elif test_type ==2:
#测试LLM返回下一步指令
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
#遍历node，查看有res的数据
iput_max_num = 0
iput_num = 0
nodes = TM.CCM.attack_tree.traverse_dfs()
if indexs:
for index in indexs:
node = nodes[index]
if node.res_quere:
TM.CCM.llm_quere.put(node)
else:
for node in nodes:
if node.res_quere:
TM.CCM.llm_quere.put(node)
#创建llm工作线程
TM.CCM.brun = True
for i in range(TM.CCM.max_thread_num):
l_th = threading.Thread(target=TM.CCM.th_llm_worker())
l_th.start()
TM.CCM.llmth_list.append(l_th)
# 等待线程结束
for t in TM.CCM.llmth_list:
t.join()
elif test_type ==3: #执行指定指令
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
# 遍历node，查看有instr的ndoe
nodes = TM.CCM.attack_tree.traverse_dfs()
instrlist = nodes[instr_index].instr_queue
instrlist = ['''
mkdir -p /tmp/nfs_test && mount -t nfs -o nolock 192.168.204.137:/ /tmp/nfs_test && ls /tmp/nfs_test
''']
for instr in instrlist:
start_time = get_local_timestr() # 指令执行开始时间
bres, instr, reslut, source_result, ext_params = TM.InstrM.execute_instruction(instr)
end_time = get_local_timestr() # 指令执行结束时间
res = {'执行结果': reslut}
str_res = json.dumps(res,ensure_ascii=False) # 直接字符串组合也可以-待验证
print(str_res)
elif test_type == 4: #修改Message
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
#创建一个新的节点
from mycode.AttackMap import TreeNode
testnode = TreeNode("test",0,0)
TM.CCM.LLM.build_initial_prompt(testnode)#新的Message
systems = testnode.messages[0]["content"]
#print(systems)
# 遍历node，查看有instr的ndoe
nodes = TM.CCM.attack_tree.traverse_bfs()
for node in nodes:
node.messages[0]["content"] = systems
with open("attack_tree", 'wb') as f:
pickle.dump(TM.CCM.attack_tree, f)
elif test_type ==5: #显示指令和结果list
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
nodes = TM.CCM.attack_tree.traverse_dfs()
if iput_index == -1:
for node in nodes:
print(f"----{node.path}-{node.status}----\n****instr_quere")
print(f"{','.join(node.instr_queue)}\n****res_quere")
try:
print(f"{','.join(node.res_quere)}")
except:
print(f"{json.dumps(node.res_quere)}")
elif iput_index == -2:#只输出有instruction的数据
index = 0
for node in nodes:
if node.instr_queue:
print(f"----{index}--{node.path}--{node.status}----")
print(f"{','.join(node.instr_queue)}")
index += 1
else:
print(f"********\n{','.join(nodes[iput_index].instr_queue)}\n********")
print(f"&&&&&&&&\n{','.join(nodes[iput_index].res_quere)}\n&&&&&&&&")
elif test_type == 6: #给指定节点添加测试指令
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
nodes = TM.CCM.attack_tree.traverse_dfs()
str_instr = "nmap -sV -p- 192.168.204.137 -T4 -oN nmap_full_scan.txt"
index = 9
nodes[index].instr_queue.append(str_instr)
nodes[index].res_quere = []
with open("attack_tree", 'wb') as f:
pickle.dump(TM.CCM.attack_tree, f)
elif test_type == 7: #给指定节点修改指令的执行结果
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
nodes = TM.CCM.attack_tree.traverse_dfs()
str_instr = "psql -h 192.168.204.137 -U postgres -c '\l'"
start_time = get_local_timestr() # 指令执行开始时间
bres, instr, reslut, source_result, ext_params = TM.InstrM.execute_instruction(str_instr)
end_time = get_local_timestr() # 指令执行结束时间
# 入数据库 -- bres True和False 都入数据库2025-3-10---加node_path(2025-3-18)#？
if TM.DBM.ok:
TM.DBM.insetr_result(0, instr, reslut, 0, start_time, end_time, source_result,
ext_params, "独立命令执行")
index = 9
nodes[index].res_quere.clear()
nodes[index].res_quere.append(reslut)
with open("attack_tree", 'wb') as f:
pickle.dump(TM.CCM.attack_tree, f)
elif test_type ==8: #显示有漏洞信息的数据
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
nodes = TM.CCM.attack_tree.traverse_dfs()
if nodes:
for node in nodes:
if node.vul_type != "未发现":
print(f"{node.path}----{node.vul_type}")
elif test_type == 9: #处理自定义llm回复内容
with open("attack_tree", "rb") as f:
TM.CCM.attack_tree = pickle.load(f)
nodes = TM.CCM.attack_tree.traverse_dfs()
node = nodes[5]
strconent = '''
{'role': 'assistant', 'content': '{"action":"update_status", "node": "25端口", "status": "已完成", "vulnerability": {"name":"SMTP用户枚举漏洞","risk":"中危","info":"VRFY命令可验证有效用户"}}\n\n```bash-[目标系统->192.168.204.137->25端口]\nsmtp-user-enum -M VRFY -U /usr/share/wordlists/metasploit/unix_users.txt -t 192.168.204.137\n```\n\n```bash-[目标系统->192.168.204.137->25端口]\nnc -nv 192.168.204.137 25 << EOF\nEXPN root\nMAIL FROM: attacker@example.com\nRCPT TO: external@example.com\nDATA\nTest open relay\n.\nQUIT\nEOF\n```'}
'''
strjson = json.loads(strconent)
node_cmds,commands = TM.CCM.LLM.fetch_instruction(strjson["content"])
TM.CCM.tree_manager(node_cmds)
# 更新tree
bok, new_commands = TM.CCM.tree_manager(node_cmds, node, commands, TM.DBM)
# 分析指令入对应节点
if bok: # 节点指令若存在错误，测试指令都不处理，需要LLM重新生成
node_list = TM.CCM.instr_in_node(new_commands, node)
#报不保存待定--
with open("attack_tree", 'wb') as f:
pickle.dump(TM.CCM.attack_tree, f)
else:
#完整过程测试---要设定终止条件
pass