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.
282 lines
12 KiB
282 lines
12 KiB
from myutils.ConfigManager import myCongif
|
|
from mycode.TaskObject import TaskObject
|
|
from mycode.DBManager import app_DBM
|
|
from myutils.PickleManager import g_PKM
|
|
from myutils.MyLogger_logger import LogHandler
|
|
from mycode.TargetManager import TargetManager # 从模块导入类
|
|
import time
|
|
import threading
|
|
|
|
class TaskManager:
|
|
def __init__(self):
|
|
self.logger = LogHandler().get_logger("TaskManager")
|
|
self.TargetM = TargetManager()
|
|
self.tasks = {} # 执行中的任务，test_id为key
|
|
self.num_threads = myCongif.get_data("Task_max_threads")
|
|
self.max_run_task = myCongif.get_data("max_run_task")
|
|
self.cur_task_run_num = 0
|
|
#获取系统信息 -- 用户可修改的都放在DB中，不修改的放config
|
|
data = app_DBM.get_system_info()
|
|
self.local_ip = data[0]
|
|
self.version = data[1]
|
|
self.tasks_lock = threading.Lock()
|
|
self.web_cur_task = 0 #web端当前显示的
|
|
|
|
#判断目标是不是在当前执行任务中，---没加锁，最多跟删除有冲突，问题应该不大
|
|
def is_target_in_tasks(self,task_target):
|
|
for task in self.tasks.values():
|
|
if task_target == task.target:
|
|
return True
|
|
return False
|
|
|
|
#程序启动后，加载未完成的测试任务
|
|
def load_tasks(self):
|
|
'''程序启动时，加载未执行完成的,未点击结束的任务 -- task_status<>2'''
|
|
datas = app_DBM.get_run_tasks()
|
|
for data in datas:
|
|
task_id = data[0]
|
|
task_target = data[1]
|
|
task_status = data[2]
|
|
work_type = data[3]
|
|
cookie_info = data[4]
|
|
llm_type = data[5]
|
|
safe_rank = data[6]
|
|
fake_target = data[7]
|
|
# 创建任务对象
|
|
task = TaskObject(task_target, cookie_info, work_type, llm_type, self.num_threads, self.local_ip,fake_target,self,safe_rank)
|
|
#读取attact_tree---load的任务应该都要有attact_tree
|
|
attack_tree = g_PKM.ReadData(str(task_id))
|
|
if attack_tree:
|
|
#恢复数据
|
|
task.init_task(task_id,attack_tree)
|
|
#开始任务 ---根据task_status来判断是否需要启动工作线程
|
|
if task_status == 1:
|
|
if self.cur_task_run_num < self.max_run_task: #load 是程序刚起，只有主线程，不加锁
|
|
bsuc,strout = task.start_task()
|
|
if bsuc:
|
|
self.cur_task_run_num +=1
|
|
else:
|
|
task.update_task_status(0)
|
|
else:
|
|
self.logger.error("重载未结束任务，不应该超过最大运行数量的task_status为启动状态")
|
|
task.update_task_status(0)#尝试硬恢复
|
|
# 内存保留task对象
|
|
self.tasks[task_id] = task
|
|
else:
|
|
self.logger.error(f"{task_id}任务的节点树数据缺失，需要检查!")
|
|
|
|
#新建测试任务--2025-4-28调整为可以批量添加--cookie_info信息取消
|
|
def create_task(self, test_targets,llm_type,work_type):
|
|
"""创建新任务--create和load复用？--
|
|
1.创建和初始化task_object;
|
|
2.创建task_id
|
|
return 失败的list
|
|
"""
|
|
fail_list = []
|
|
target_list = test_targets.split(",")
|
|
for target in target_list:
|
|
#这里判断目标的合法性
|
|
bok,target,type,fake_target = self.TargetM.validate_and_extract(target) #是否还需要判断待定?
|
|
if not bok:#目标不合法
|
|
fail_list.append(target)
|
|
continue
|
|
#判断是否已在执行列表
|
|
if self.is_target_in_tasks(target):
|
|
fail_list.append(target)
|
|
continue
|
|
#raise ValueError(f"Task {test_target} already exists")
|
|
#创建任务对象--cookie参数取消
|
|
task = TaskObject(target,"",work_type,llm_type,self.num_threads,self.local_ip,fake_target,self)
|
|
#获取task_id -- test_target,cookie_info,work_type,llm_type 入数据库
|
|
task_id = app_DBM.start_task(target,"",work_type,llm_type,fake_target)
|
|
if task_id >0:
|
|
#2025-4-28调整批量添加任务，默认不启动
|
|
task.init_task(task_id)
|
|
#保留task对象
|
|
self.tasks[task_id] = task
|
|
else:
|
|
fail_list.append(target)
|
|
result = ",".join(fail_list)
|
|
return result
|
|
|
|
#开启task任务--正常只应该有web触发调用
|
|
def start_task_TM(self,task_id):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
with self.tasks_lock:
|
|
if self.cur_task_run_num < self.max_run_task:
|
|
#启动工作线程和进程
|
|
bsuc,error = task.start_task()
|
|
if bsuc:
|
|
self.cur_task_run_num += 1
|
|
return True,"启动成功"
|
|
else:
|
|
return False,error
|
|
else:
|
|
return False,f"已到达最大的启动数--{self.max_run_task}"
|
|
return False,"该任务不存在，程序逻辑存在问题！"
|
|
|
|
#停止task任务
|
|
def stop_task_TM(self,task_id):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
if task.brun and task.task_status ==1: #是运行状态
|
|
with self.tasks_lock:
|
|
task.stop_task() #停止线程应该就没什么失败需要处理的
|
|
self.cur_task_run_num -= 1
|
|
return True,"停止任务成功"
|
|
else:
|
|
return True,"该任务已处于停止状态"
|
|
return False,"该任务不存在，程序逻辑存在问题！"
|
|
|
|
#结束任务
|
|
def over_task(self,task_id):
|
|
#先尝试停止
|
|
bsuccess,_ = self.stop_task_TM(task_id)
|
|
time.sleep(1)
|
|
if bsuccess:
|
|
#再结束
|
|
bsuccess = app_DBM.over_task(task_id) # 不管是不是修改（置2）成功，都执行结束
|
|
del self.tasks[task_id] # 删除缓存
|
|
return True,"结束任务成功"
|
|
else:
|
|
return False,"该任务不存在，程序逻辑存在问题！"
|
|
|
|
#删除任务
|
|
def del_task(self,task_id):
|
|
#删除数据记录
|
|
app_DBM.del_task(task_id)
|
|
#删除节点树
|
|
g_PKM.DelData(str(task_id))
|
|
return True,"删除任务成功！"
|
|
|
|
#控制task启停----线程不停 --2025-4-28 配合批量任务，需要停止工作线程了
|
|
def control_taks(self,task_id):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
if task.task_status == 0: # 0-暂停，1-执行中，2-已完成
|
|
bsuc,error = self.start_task_TM(task_id) #任务是否存在的状态有点重复
|
|
elif task.task_status == 1:
|
|
bsuc,error = self.stop_task_TM(task_id)
|
|
else:
|
|
return False,"当前任务状态不允许修改，请联系管理员！",task.task_status
|
|
else:
|
|
return False,"没有找到对应的任务",None
|
|
return bsuc,error,task.task_status
|
|
|
|
# 获取任务list
|
|
def get_task_list(self):
|
|
tasks = []
|
|
for task in self.tasks.values():
|
|
one_json = {"taskID": task.task_id, "testTarget": task.target, "taskStatus": task.task_status,
|
|
"safeRank": task.safe_rank,"workType": task.work_type}
|
|
tasks.append(one_json)
|
|
rejson = {"tasks": tasks}
|
|
return rejson
|
|
|
|
#获取节点树
|
|
def get_node_tree(self,task_id):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
self.web_cur_task = task_id
|
|
tree_dict = task.attack_tree.get_node_dict()
|
|
return tree_dict
|
|
return None
|
|
|
|
#获取历史节点树数据
|
|
def get_his_node_tree(self,task_id):
|
|
attack_tree = g_PKM.ReadData(str(task_id))
|
|
if attack_tree:
|
|
tree_dict = attack_tree.get_node_dict()
|
|
return tree_dict
|
|
return None
|
|
|
|
#修改任务的工作模式，只有在暂停状态才能修改
|
|
def update_task_work_type(self,task_id,new_work_type):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
if task.task_status == 0:
|
|
task.update_task_work_type(new_work_type)
|
|
return True
|
|
return False
|
|
|
|
#------------节点操作相关-------还未二次走查-------------
|
|
#控制节点的工作状态
|
|
def node_bwork_control(self,task_id,node_path):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
bsuccess,new_bwork = task.attack_tree.update_node_bwork(node_path)
|
|
if bsuccess:
|
|
pass #是否要更新IO数据？----待验证是否有只修改部分数据的方案
|
|
return bsuccess,new_bwork
|
|
return False,False
|
|
|
|
#节点单步--只允許web端调用
|
|
async def node_one_step(self,task_id,node_path):
|
|
task = self.tasks[task_id]
|
|
node = task.attack_tree.find_node_by_nodepath(node_path)
|
|
#web端触发的任务，需要判断当前的执行状态
|
|
bsuccess,error = await task.put_one_node(node)
|
|
return bsuccess,error
|
|
|
|
#任务单点--只允许web端调用
|
|
async def task_one_step(self,task_id):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
bsuccess,error = await task.put_one_task()
|
|
return bsuccess,error
|
|
else:
|
|
return False,"task_id值存在问题！"
|
|
|
|
#获取节点待执行任务
|
|
def get_task_node_todo_instr(self,task_id,nodepath):
|
|
todoinstr = []
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
node = task.attack_tree.find_node_by_nodepath(nodepath)
|
|
if node:
|
|
todoinstr = node.get_instr_user()
|
|
return todoinstr
|
|
|
|
#获取节点的MSG信息
|
|
def get_task_node_MSG(self,task_id,nodepath):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
node = task.attack_tree.find_node_by_nodepath(nodepath)
|
|
if node:
|
|
tmpMsg = node.get_res_user()
|
|
if tmpMsg:
|
|
return node.cur_messages,tmpMsg[0] #待提交消息正常应该只有一条
|
|
else:
|
|
return node.cur_messages,{}
|
|
return [],{}
|
|
|
|
def update_node_MSG(self,task_id,nodepath,newtype,newcontent):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
node = task.attack_tree.find_node_by_nodepath(nodepath)
|
|
if node:
|
|
work_status = node.get_work_status()
|
|
if work_status == 0 or work_status == 3:
|
|
if work_status == 0:
|
|
if not node.parent_messages: #如果messages为空--且不会是根节点
|
|
node.copy_messages(node.parent.parent_messages,node.parent.cur_messages)
|
|
bsuccess,error = node.updatemsg(newtype,newcontent,0) #取的第一条，也就修改第一条
|
|
return bsuccess,error
|
|
else:
|
|
return False,"当前节点的工作状态不允许修改MSG！"
|
|
return False,"找不到对应节点！"
|
|
|
|
def del_node_instr(self,task_id,nodepath,instr):
|
|
task = self.tasks[task_id]
|
|
if task:
|
|
node = task.attack_tree.find_node_by_nodepath(nodepath)
|
|
if node:
|
|
return node.del_instr(instr)
|
|
return False,"找不到对应节点！"
|
|
|
|
def get_his_tasks(self,target_name,safe_rank,llm_type,start_time,end_time):
|
|
tasks = app_DBM.get_his_tasks(target_name,safe_rank,llm_type,start_time,end_time)
|
|
return tasks
|
|
|
|
g_TaskM = TaskManager() #单一实例
|