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.
230 lines
9.7 KiB
230 lines
9.7 KiB
'''
|
|
渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据
|
|
'''
|
|
from TargetManager import TargetManager # 从模块导入类
|
|
#from LLMManager import LLMManager # 同理修正其他导入
|
|
from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。
|
|
from myutils.FileManager import FileManager
|
|
from 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
|
|
|
|
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}
|
|
work_node.llm_type = 1
|
|
work_node.add_res(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(): #
|
|
with open("attack_tree", 'wb') as f:
|
|
pickle.dump(TM.CCM.attack_tree, f)
|
|
|
|
except queue.Empty:
|
|
self.logger.debug("暂无需要执行指令的节点！")
|
|
time.sleep(20)
|
|
|
|
def start_task(self,target_name,target_in):
|
|
#判断目标合法性
|
|
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 = 5
|
|
iput_index = 6 # 0是根节点
|
|
if test_type == 0: #新目标测试
|
|
# 启动--初始化指令
|
|
node_list = TM.CCM.start_do("192.168.204.137", 0)
|
|
#异步处理，需要等待线程结束了
|
|
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_bfs()
|
|
for node in nodes:
|
|
if node.instr_queue: # list
|
|
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 = 1
|
|
iput_num = 0
|
|
nodes = TM.CCM.attack_tree.traverse_bfs()
|
|
if iput_index != -1:#index 不为-1就是指定节点返回，人为保障不越界
|
|
node = nodes[iput_index]
|
|
if node.res_quere:
|
|
TM.CCM.llm_quere.put(node)
|
|
else:
|
|
for node in nodes:
|
|
if node.res_quere: #有结果需要提交LLM
|
|
TM.CCM.llm_quere.put(node)
|
|
iput_num += 1
|
|
if iput_max_num > 0: #0是有多少提交多少
|
|
if iput_num == iput_max_num:
|
|
break
|
|
|
|
#创建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: #执行指定指令
|
|
instrlist=[
|
|
"msfconsole -q -x \"use auxiliary/scanner/smb/smb_version; set RHOSTS 192.168.204.137; run; exit\""]
|
|
|
|
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() # 指令执行结束时间
|
|
# 入数据库 -- 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, "独立命令执行")
|
|
else:
|
|
TM.logger.error("数据库连接失败！！")
|
|
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)
|
|
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_bfs()
|
|
print(f"********\n{','.join(nodes[iput_index].instr_queue)}\n********")
|
|
print(f"&&&&&&&&\n{','.join(nodes[iput_index].res_quere)}\n&&&&&&&&")
|
|
else:
|
|
#完整过程测试---要设定终止条件
|
|
pass
|
|
|