@ -4,6 +4,7 @@ from mycode.DBManager import app_DBM
from myutils . PickleManager import g_PKM
from myutils . MyLogger_logger import LogHandler
from mycode . TargetManager import TargetManager # 从模块导入类
from mycode . PythonTManager import PythonTManager
import time
import threading
@ -16,20 +17,30 @@ class TaskManager:
data = app_DBM . get_system_info ( )
self . local_ip = data [ 0 ]
self . version = data [ 1 ]
self . tasks_lock = threading . Lock ( )
self . brun = True
self . tm_th = None #任务控制线程
#----临时任务相关-------
self . web_cur_task = 0 #web端当前显示的 -- web端可以操作的都为临时任务
self . tasks_lock = threading . Lock ( )
self . tasks = { } # 执行中的任务，test_id为key
self . max_run_task = myCongif . get_data ( " max_run_task " )
self . cur_task_run_num = 0
#-----巡检任务相关-------
self . polling_llmtype = myCongif . get_data ( " polling_llm_type " )
self . polling_tasks_lock = threading . Lock ( )
self . polling_tasks = { } #独立一个队列
self . max_polling_task = myCongif . get_data ( " max_polling_task " )
self . cur_polling_task = 0
self . cur_polling_task_num = 0
self . work_type = 1 #巡检工作模式就是自动了
#python 进程池 提到程序主线程层面 --不随任务进行启停
self . PythonM = PythonTManager ( myCongif . get_data ( " Python_max_procs " ) )
# 启动python子进程池 --进程池的启动，目前是没全面确认原子进程池的状态，直接None
self . PythonM . start_pool ( ) #开启进程池
def __del__ ( self ) :
# 停止子进程池
self . PythonM . shutdown_pool ( )
#判断目标是不是在当前执行任务中，---没加锁，最多跟删除有冲突，问题应该不大
def is_target_in_tasks ( self , task_target , itype = 1 ) :
@ -46,7 +57,7 @@ class TaskManager:
#程序启动后，加载未完成的测试任务
def load_tasks ( self ) :
''' 程序启动时，加载未执行完成的,未点击结束的任务 -- task_status<>2
''' 程序启动时，加载未执行完成的,未点击结束的任务 -- task_status<>4
#若不是异常停止，正常情况下，任务都应该是没有待办MQ的
'''
datas = app_DBM . get_run_tasks ( ) #
@ -62,27 +73,43 @@ class TaskManager:
target_id = data [ 8 ]
# 创建任务对象
task = TaskObject ( task_target , cookie_info , work_type , llm_type , self . num_threads , self . local_ip , fake_target , target_id , self , safe_rank )
#这里要做区分.临时任务和巡检任务
#?
task = TaskObject ( task_target , cookie_info , work_type , llm_type , self . num_threads , self . local_ip , fake_target , self , target_id , safe_rank )
#读取attact_tree---load的任务应该都要有attact_tree
attack_tree = g_PKM . ReadData ( str ( task_id ) )
if attack_tree :
#恢复数据
#恢复数据--遍历节点恢复状态 --临时和polling适用
task . init_task ( task_id , attack_tree )
#开始任务 ---根据task_status来判断是否需要启动工作线程
if task_status == 1 :
if task_status == 1 : #执行中的任务
bcan_start = False
if target_id == 0 : #临时任务
if self . cur_task_run_num < self . max_run_task : #load 是程序刚起，只有主线程，不加锁
bsuc , strout = task . start_task ( )
if bsuc :
bcan_start = True
else : #巡检任务
if self . cur_polling_task_num < self . max_polling_task :
bcan_start = True
if bcan_start :
isuc , strout = task . start_task ( )
if isuc == 1 : #启动工作线程成功
if target_id == 0 :
self . cur_task_run_num + = 1
else :
task . update_task_status ( 0 )
self . cur_polling_task_num + = 1
elif isuc == - 1 : # 存在未结束的线程--load这边应该不会发生
self . logger . debug ( f " load-task遇到未结束线程的任务-- { task_id } " )
task . update_task_status ( 2 )
else : # 0 --没有待办事项
self . logger . debug ( f " load-task遇到无待办实现的任务-- { task_id } " )
task . update_task_status ( 3 )
else :
self . logger . error ( " 重载未结束任务，不应该超过最大运行数量的task_status为启动状态 " )
task . update_task_status ( 0 ) #尝试硬恢复
task . update_task_status ( 0 ) #等其他task结束后会启动
# 内存保留task对象
if target_id == 0 :
self . tasks [ task_id ] = task
else :
self . polling_tasks [ task_id ] = task
else :
self . logger . error ( f " { task_id } 任务的节点树数据缺失，需要检查! " )
@ -102,7 +129,7 @@ class TaskManager:
fail_list . append ( target )
continue
#判断是否已在执行列表
if self . is_target_in_tasks ( target ) :
if self . is_target_in_tasks ( target , 1 ) :
fail_list . append ( target )
continue
#raise ValueError(f"Task {test_target} already exists")
@ -116,7 +143,7 @@ class TaskManager:
#保留task对象
self . tasks [ task_id ] = task
#尝试启动task
self . start_task_TM ( task_id )
_ , _ = self . start_task_TM ( task_id )
else :
fail_list . append ( target )
result = " , " . join ( fail_list )
@ -134,22 +161,46 @@ class TaskManager:
else :
fake_terget = self . TargetM . get_fake_target ( 2 )
#创建任务实例
polling_task = TaskObject ( target , " " , self . work_type , self . polling_llmtype , self . num_threads , self . local_ip , fake_target , self )
polling_task = TaskObject ( target , " " , self . work_type , self . polling_llmtype , self . num_threads , self . local_ip , fake_target , self , target [ 0 ] )
# 获取task_id -- test_target,cookie_info,work_type,llm_type 入数据库
task_id = app_DBM . start_task ( target , " " , self . work_type , self . polling_llmtype , fake_target , target [ 0 ] ) #增加一个target_id
if task_id > 0 :
polling_task . init_task ( task_id ) #初始化任务信息，并提交到待处理节点队列
#保留task对象--巡检任务
self . polling_tasks [ task_id ] = polling_task
#尝试启动task
_ , _ = self . start_polling_task_TM ( polling_task )
def is_can_start_task ( self , itype ) :
if itype == 1 :
if self . cur_task_run_num < self . max_run_task :
return True
else :
if self . cur_polling_task_num < self . max_polling_task :
return True
return False
def th_control_takes ( self ) :
while self . brun :
#遍历临时任务--启停任务
pass
#遍历巡检任务-启停
pass
#休眠
time . sleep ( 60 * 5 )
def start_next_task ( self , task_id , target_id ) :
if target_id == 0 : #临时任务
with self . tasks_lock :
self . cur_task_run_num - = 1
for task in self . tasks . values ( ) :
if task . task_status == 0 : #只有是0状态会自动启动，2暂停状态需要人工启动
isuc , error = task . start_task ( )
if isuc == 1 :
self . cur_task_run_num + = 1
break
else : #巡检任务--是直接结束置4
with self . polling_tasks_lock :
del self . polling_tasks [ task_id ] # 删除缓存
self . cur_polling_task_num - = 1
for pl_task in self . polling_tasks :
if pl_task . task_status == 0 :
isuc , error = pl_task . start_task ( )
if isuc == 1 :
self . cur_polling_task_num + = 1
break
#开启task任务
def start_task_TM ( self , task_id ) :
@ -158,8 +209,8 @@ class TaskManager:
with self . tasks_lock :
if self . cur_task_run_num < self . max_run_task :
#启动工作线程和进程
b suc, error = task . start_task ( )
if bsuc :
i suc, error = task . start_task ( )
if isuc == 1 :
self . cur_task_run_num + = 1
return True , " 启动成功 "
else :
@ -168,7 +219,20 @@ class TaskManager:
return False , f " 已到达最大的启动数-- { self . max_run_task } "
return False , " 该任务不存在，程序逻辑存在问题！ "
#停止task任务--暂停
def start_polling_task_TM ( self , pl_task ) :
if pl_task :
with self . polling_tasks_lock :
if self . cur_polling_task_num < self . max_polling_task :
isuc , error = pl_task . start_task ( )
if isuc == 1 :
self . cur_polling_task_num + = 1
return True , " 启动成功 "
else :
return False , error
else :
return False , f " 已到达最大的启动数-- { self . max_run_task } "
#停止task任务--暂停---要启动下一个task---只针对临时任务
def stop_task_TM ( self , task_id ) :
task = self . tasks [ task_id ]
if task :
@ -176,25 +240,33 @@ class TaskManager:
with self . tasks_lock :
task . stop_task ( ) #停止线程应该就没什么失败需要处理的
self . cur_task_run_num - = 1
#启动下一个任务-只针对临时任务
for tmptask in self . tasks . values ( ) :
if tmptask . task_status == 0 : # 只有是0状态会自动启动，2暂停状态需要人工启动
isuc , error = tmptask . start_task ( )
if isuc == 1 :
self . cur_task_run_num + = 1
break
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 ）成功，都执行结束
bsuccess = app_DBM . over_task ( task_id ) # 不管是不是修改（置4 ）成功，都执行结束
del self . tasks [ task_id ] # 删除缓存
return True , " 结束任务成功 "
else :
return False , " 该任务不存在，程序逻辑存在问题！ "
#删除任务
#删除任务--历史任务处理
def del_task ( self , task_id ) :
#删除数据记录
app_DBM . del_task ( task_id )
@ -206,7 +278,7 @@ class TaskManager:
def control_taks ( self , task_id ) :
task = self . tasks [ task_id ]
if task :
if task . task_status == 0 : # 0-暂停，1-执行中，2-已完成
if task . task_status in ( 0 , 2 ) : # 0-未启动，1-执行中，2-暂停中，3-已完成，4-已结束
bsuc , error = self . start_task_TM ( task_id ) #任务是否存在的状态有点重复
elif task . task_status == 1 :
bsuc , error = self . stop_task_TM ( task_id )
@ -243,11 +315,11 @@ class TaskManager:
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 :
if task . task_status in ( 0 , 2 ) :
task . update_task_work_type ( new_work_type )
return True
return False
