|
|
|
@ -1,7 +1,6 @@ |
|
|
|
''' |
|
|
|
渗透测试任务管理类 一次任务的闭合性要检查2025-3-10 一次任务后要清理LLM和InstrM的数据 |
|
|
|
''' |
|
|
|
from mycode.TargetManager import TargetManager # 从模块导入类 |
|
|
|
#from LLMManager import LLMManager # 同理修正其他导入 |
|
|
|
from mycode.ControlCenter import ControlCenter #控制中心替代LLM--控制中心要实现一定的基础逻辑和渗透测试树的维护。 |
|
|
|
from mycode.InstructionManager import g_instrM |
|
|
|
@ -15,6 +14,7 @@ from myutils.ConfigManager import myCongif |
|
|
|
from mycode.WebSocketManager import g_WSM |
|
|
|
from mycode.CommandVerify import g_CV |
|
|
|
from mycode.PythonTManager import PythonTManager |
|
|
|
from mycode.DataFilterManager import DataFilterManager |
|
|
|
import asyncio |
|
|
|
import queue |
|
|
|
import time |
|
|
|
@ -26,12 +26,13 @@ import textwrap |
|
|
|
|
|
|
|
class TaskObject: |
|
|
|
|
|
|
|
def __init__(self,test_target,cookie_info,work_type,llm_type,num_threads,local_ip,taskM): |
|
|
|
def __init__(self,test_target,cookie_info,work_type,llm_type,num_threads,local_ip,fake_target,taskM,safe_rank=0): |
|
|
|
#功能类相关 |
|
|
|
self.taskM = taskM |
|
|
|
self.logger = LogHandler().get_logger("TaskObject") |
|
|
|
self.InstrM = g_instrM # 类对象渗透，要约束只读取信息，且不允许有类全局对象--持续检查 |
|
|
|
self.PythonM = PythonTManager(myCongif.get_data("Python_max_procs")) |
|
|
|
self.DataFilter = DataFilterManager(test_target,fake_target) |
|
|
|
self.CCM = ControlCenter() #一个任务一个CCM |
|
|
|
self.LLM = LLMManager(llm_type) # LLM对象调整为一个任务一个对象，这样可以为不同的任务选择不同的LLM |
|
|
|
#全局变量 |
|
|
|
@ -42,10 +43,10 @@ class TaskObject: |
|
|
|
self.cookie = cookie_info |
|
|
|
self.work_type = work_type #工作模式 0-人工，1-自动 |
|
|
|
self.task_id = None |
|
|
|
self.task_status = 0 #0-暂停，1-执行中，2-已完成 |
|
|
|
self.task_status = 0 #0-暂停，1-执行中，2-已完成 3-未启动，2025-4-27为配合批量添加任务增加“未启动”状态 |
|
|
|
self.local_ip = local_ip |
|
|
|
self.attack_tree = None #任务节点树 |
|
|
|
self.safe_rank = 0 #安全级别 0-9 #？暂时还没实现更新逻辑 |
|
|
|
self.safe_rank = safe_rank #安全级别 0-9 #？暂时还没实现更新逻辑 |
|
|
|
self.is_had_work = False |
|
|
|
self.is_had_work_lock = threading.Lock() |
|
|
|
|
|
|
|
@ -88,11 +89,36 @@ class TaskObject: |
|
|
|
"""截取字符串前 max_length 个字符，并添加截断标记（确保总长度不超过限制）""" |
|
|
|
if not s: |
|
|
|
return s |
|
|
|
#一些无效信息的删除 |
|
|
|
s = s.replace("pydev debugger: bytes arguments were passed to a new process creation function. Breakpoints may not work correctly.","") |
|
|
|
|
|
|
|
# 计算保留长度（考虑截断标记占位） |
|
|
|
truncated = s[:max_length - len(ellipsis)] if len(s) > max_length else s |
|
|
|
return truncated + ellipsis if len(s) > max_length else s |
|
|
|
|
|
|
|
def do_instruction(self,instruction): |
|
|
|
instruction = textwrap.dedent(instruction.strip()) |
|
|
|
# 对多shell指令的情况进行处理--也有风险 |
|
|
|
if "python-code" not in instruction: |
|
|
|
if "&&" in instruction: |
|
|
|
instruction = self.mill_instr_preprocess(instruction, "&&") |
|
|
|
elif "||" in instruction: |
|
|
|
instruction = self.mill_instr_preprocess(instruction, "||") |
|
|
|
start_time = get_local_timestr() # 指令执行开始时间 |
|
|
|
|
|
|
|
# 2025-4-27要重新定义bool值作用，初步想法True-作为指令已处理（执行或不执行），False-作为异常，指令还没有处理-可以回Q |
|
|
|
|
|
|
|
if instruction.startswith("python-code"): # python代码--超过子进程数会阻塞等待，但不开始超时记时 |
|
|
|
bsuccess, instr, reslut, source_result, ext_params = self.PythonM.execute_instruction(instruction) |
|
|
|
else: # shell指令 |
|
|
|
bsuccess, instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction) |
|
|
|
|
|
|
|
end_time = get_local_timestr() # 指令执行结束时间 |
|
|
|
# 只取结果的5000长度 |
|
|
|
reslut = self.smart_truncate(reslut, 4000) |
|
|
|
source_result = self.smart_truncate(source_result) |
|
|
|
return start_time,end_time,bsuccess,instr,reslut,source_result,ext_params |
|
|
|
|
|
|
|
def do_worker_th(self,index): |
|
|
|
#线程的dbm需要一个线程一个 |
|
|
|
th_DBM = DBManager() |
|
|
|
@ -113,23 +139,7 @@ class TaskObject: |
|
|
|
break |
|
|
|
self.doing_instr_list[th_index] = instruction |
|
|
|
|
|
|
|
instruction = textwrap.dedent(instruction.strip()) |
|
|
|
#对多shell指令的情况进行处理--也有风险 |
|
|
|
if "python-code" not in instruction: |
|
|
|
if "&&" in instruction: |
|
|
|
instruction = self.mill_instr_preprocess(instruction, "&&") |
|
|
|
elif "||" in instruction: |
|
|
|
instruction = self.mill_instr_preprocess(instruction, "||") |
|
|
|
start_time = get_local_timestr() # 指令执行开始时间 |
|
|
|
|
|
|
|
if instruction.startswith("python-code"):#python代码--超过子进程数会阻塞等待，但不开始超时记时 |
|
|
|
instr, reslut, source_result, ext_params = self.PythonM.execute_instruction(instruction) |
|
|
|
else:#shell指令 |
|
|
|
instr, reslut, source_result, ext_params = self.InstrM.execute_instruction(instruction) |
|
|
|
end_time = get_local_timestr() # 指令执行结束时间 |
|
|
|
#只取结果的5000长度 |
|
|
|
reslut = self.smart_truncate(reslut,4000) |
|
|
|
source_result = self.smart_truncate(source_result) |
|
|
|
start_time, end_time, bsuccess, instr, reslut, source_result, ext_params = self.do_instruction(instruction) |
|
|
|
|
|
|
|
# 入数据库 -- bres True和False 都入数据库2025-3-10---加node_path(2025-3-18)#？ |
|
|
|
if th_DBM.ok: |
|
|
|
@ -139,9 +149,10 @@ class TaskObject: |
|
|
|
ext_params, work_node.path) |
|
|
|
else: |
|
|
|
self.logger.error("数据库连接失败！！") |
|
|
|
#暂存结果 |
|
|
|
oneres = {'执行指令': instr, '结果': reslut} |
|
|
|
results.append(oneres) |
|
|
|
# 暂存结果 |
|
|
|
fake_inst,fake_resul = self.DataFilter.filter_result(instr,reslut) |
|
|
|
oneres = {'执行指令': fake_inst, '结果': fake_resul} |
|
|
|
results.append(oneres) #结果入队列后，指令就不能回退 |
|
|
|
#一条指令执行完成 |
|
|
|
self.doing_instr_list[th_index] = "" |
|
|
|
#指令都执行结束后，入节点待提交队列 |
|
|
|
@ -152,6 +163,7 @@ class TaskObject: |
|
|
|
self.put_node_reslist(work_node, str_res, llm_type) |
|
|
|
# 保存记录 |
|
|
|
g_PKM.WriteData(self.attack_tree,str(self.task_id)) |
|
|
|
|
|
|
|
except queue.Empty: |
|
|
|
if bnode_work: |
|
|
|
bnode_work = False |
|
|
|
@ -187,7 +199,7 @@ class TaskObject: |
|
|
|
- 节点名称：{llm_node.name} |
|
|
|
- 节点状态：未完成 |
|
|
|
- 漏洞类型：{llm_node.vul_type} |
|
|
|
''' |
|
|
|
''' |
|
|
|
while True: |
|
|
|
llm_data = llm_node.get_res() |
|
|
|
if llm_data is None: |
|
|
|
@ -196,10 +208,12 @@ class TaskObject: |
|
|
|
str_res = llm_data["result"] |
|
|
|
#获取提示词 |
|
|
|
prompt = self.get_llm_prompt(llm_type,str_res,user_Prompt) |
|
|
|
fake_prompt = self.DataFilter.filter_prompt(prompt) #目标脱敏 |
|
|
|
|
|
|
|
self.doing_llm_list[th_index] = prompt |
|
|
|
# 提交llm请求返回数据--并对返回数据进行处理，节点指令直接执行，测试指令根据工作模式执行 |
|
|
|
post_time = get_local_timestr() |
|
|
|
bsuccess,node_cmds, commands,reasoning_content, content = self.LLM.get_llm_instruction(prompt,llm_node) # message要更新 |
|
|
|
bsuccess,node_cmds, commands,reasoning_content, content = self.LLM.get_llm_instruction(fake_prompt,llm_node,self.DataFilter) # message要更新 --llm_node只使用messages，都是脱敏后的数据 |
|
|
|
if not bsuccess: |
|
|
|
self.logger.error(f"模型接口调用出错：{content}") |
|
|
|
continue #丢弃 --若需要再次尝试，把llm_data再入队列 |
|
|
|
@ -233,7 +247,7 @@ class TaskObject: |
|
|
|
strdata = "update accack_tree!" |
|
|
|
asyncio.run(g_WSM.send_data(idatatype, strdata)) |
|
|
|
# 先取消当前task,已经通知前端重新获取,这样可以避免后端不必要的数据推送 |
|
|
|
self.taskM.web_cur_task = 0 |
|
|
|
#self.taskM.web_cur_task = 0 |
|
|
|
except queue.Empty: |
|
|
|
if bnode_work: |
|
|
|
bnode_work = False |
|
|
|
@ -279,14 +293,15 @@ class TaskObject: |
|
|
|
while self.brun: |
|
|
|
try: |
|
|
|
cur_time = get_local_timestr() |
|
|
|
print(f"-----------当前时间程序运行情况:{cur_time}") |
|
|
|
print(f"----------{self.task_id}-当前时间程序运行情况:{cur_time}") |
|
|
|
#执行中instr-node |
|
|
|
index = 0 |
|
|
|
for w_th in self.workth_list: |
|
|
|
if not w_th.is_alive():#线程 |
|
|
|
print(f"线程-{index}已处于异常状态，需要重新创建一个工作线程") |
|
|
|
print(f"Work线程-{index}已处于异常状态，需要重新创建一个工作线程") |
|
|
|
else: |
|
|
|
print(f"线程-{index}在执行指令：{self.doing_instr_list[index]}") |
|
|
|
if self.doing_instr_list[index]: |
|
|
|
print(f"Work线程-{index}-在执行指令：{self.doing_instr_list[index]}") |
|
|
|
index += 1 |
|
|
|
|
|
|
|
index = 0 |
|
|
|
@ -294,7 +309,8 @@ class TaskObject: |
|
|
|
if not l_th.is_alive(): |
|
|
|
print(f"LLM线程-{index}已处于异常状态，需要重新创建一个LLM线程") |
|
|
|
else: |
|
|
|
print(f"LLM线程-{index}在执行指令：{self.doing_llm_list[index]}") |
|
|
|
if self.doing_llm_list[index]: |
|
|
|
print(f"LLM线程-{index}-在执行指令：{self.doing_llm_list[index]}") |
|
|
|
index += 1 |
|
|
|
#处理点修复操作 |
|
|
|
icount +=1 |
|
|
|
@ -485,7 +501,7 @@ class TaskObject: |
|
|
|
#更新状态 |
|
|
|
bchange = node.update_work_status(work_status) |
|
|
|
#基于websocket推送到前端 |
|
|
|
if bchange: |
|
|
|
if bchange and work_status != 1: #llm执行完成后会发送单独的指令更新树，所以不发送1更新节点了 |
|
|
|
#判断是否是web端最新获取数据的task |
|
|
|
if self.taskM.web_cur_task == self.task_id: |
|
|
|
idatatype = 1 |
|
|
|
@ -681,8 +697,9 @@ class TaskObject: |
|
|
|
2.这些节点的父节点为当前节点，请正确生成这些节点的节点路径； |
|
|
|
3.只有当还有节点未能生成测试指令或不完整时，才返回未生成指令的节点列表。 |
|
|
|
''' |
|
|
|
bsuccess,node_cmds, commands, reasoning_content, content, post_time = self.LLM.get_llm_instruction(user_Prompt, |
|
|
|
cur_node) # message要更新 |
|
|
|
fake_prompt = self.DataFilter.filter_prompt(user_Prompt) |
|
|
|
bsuccess,node_cmds, commands, reasoning_content, content = self.LLM.get_llm_instruction(fake_prompt, |
|
|
|
cur_node,self.DataFilter) # message要更新 |
|
|
|
if not bsuccess: |
|
|
|
self.logger.error(f"模型接口调用出错：{content}") |
|
|
|
ierror += 1 |
|
|
|
@ -692,6 +709,7 @@ class TaskObject: |
|
|
|
res_str = "" |
|
|
|
# LLM记录存数据库 |
|
|
|
cur_node.llm_sn += 1 |
|
|
|
post_time = get_local_timestr() |
|
|
|
bres = DBM.insert_llm(self.task_id, user_Prompt, reasoning_content, content, post_time, cur_node.llm_sn,cur_node.path) |
|
|
|
if not bres: |
|
|
|
self.logger.error(f"{cur_node.name}-llm入库失败！") |
|
|
|
@ -712,56 +730,99 @@ class TaskObject: |
|
|
|
self.logger.debug("未添加指令的节点，都已完成指令的添加！") |
|
|
|
return new_commands |
|
|
|
|
|
|
|
def start_task(self,task_id,task_status=1,attack_tree=None): |
|
|
|
#-----------------任务的启停-------------------- |
|
|
|
def init_task(self,task_id,attack_tree = None): |
|
|
|
self.task_id = task_id |
|
|
|
''' |
|
|
|
启动该测试任务 |
|
|
|
''' |
|
|
|
#判断目标合法性 |
|
|
|
# bok,target,type = self.TargetM.validate_and_extract(self.target) #是否还需要判断待定? |
|
|
|
# if not bok: |
|
|
|
# return False, "{target}检测目标不合规，请检查！" |
|
|
|
#初始化节点树 |
|
|
|
if attack_tree:#有值的情况是load |
|
|
|
self.attack_tree =attack_tree |
|
|
|
#加载未完成的任务 |
|
|
|
if self.work_type ==1:#自动模式 |
|
|
|
#提交到mq,待线程执行 |
|
|
|
# 初始化节点树 |
|
|
|
if attack_tree: # 有值的情况是load |
|
|
|
self.attack_tree = attack_tree |
|
|
|
# 加载未完成的任务 |
|
|
|
if self.work_type == 1: # 自动模式 |
|
|
|
# 提交到mq,待线程执行 |
|
|
|
self.put_work_node() |
|
|
|
else: #无值的情况是new_create |
|
|
|
root_node = TreeNode(self.target, self.task_id) #根节点 |
|
|
|
self.attack_tree = AttackTree(root_node) #创建测试树，同时更新根节点相关内容 |
|
|
|
self.LLM.build_initial_prompt(root_node) #对根节点初始化system-msg |
|
|
|
#插入一个user消息 |
|
|
|
else: # 无值的情况是new_create |
|
|
|
root_node = TreeNode(self.target, self.task_id) # 根节点 |
|
|
|
self.attack_tree = AttackTree(root_node) # 创建测试树，同时更新根节点相关内容 |
|
|
|
self.LLM.build_initial_prompt(root_node) # 对根节点初始化system-msg |
|
|
|
# 插入一个user消息 |
|
|
|
# 提交第一个llm任务，开始工作 |
|
|
|
know_info = f"本测试主机的IP地址为:{self.local_ip}" |
|
|
|
if self.cookie: |
|
|
|
know_info = know_info + f",本站点的cookie值为{self.cookie}" |
|
|
|
self.put_node_reslist(root_node,know_info,0) #入待提交list |
|
|
|
#初始保存个attack_tree文件 |
|
|
|
g_PKM.WriteData(self.attack_tree,str(self.task_id)) |
|
|
|
#启动工作线程 |
|
|
|
self.task_status = task_status |
|
|
|
self.brun = True #线程正常启动 |
|
|
|
#启动指令工作线程 |
|
|
|
for i in range(self.max_thread_num): |
|
|
|
w_th = threading.Thread(target=self.do_worker_th,args=(i,)) |
|
|
|
w_th.start() |
|
|
|
self.workth_list[i] = w_th |
|
|
|
#启动llm提交线程--llm暂时单线程，多线程处理时attack_tree需要加锁 |
|
|
|
for j in range(self.llm_max_nums): |
|
|
|
l_th = threading.Thread(target=self.th_llm_worker,args=(j,)) |
|
|
|
l_th.start() |
|
|
|
self.llmth_list[j]=l_th |
|
|
|
#启动自检线程 |
|
|
|
self.check_th = threading.Thread(target=self.th_check) |
|
|
|
self.check_th.start() |
|
|
|
self.put_node_reslist(root_node, know_info, 0) # 入待提交list |
|
|
|
# 初始保存个attack_tree文件 |
|
|
|
g_PKM.WriteData(self.attack_tree, str(self.task_id)) |
|
|
|
|
|
|
|
def is_task_stop(self): |
|
|
|
#检查任务是否处于停止状态--防止停止后，线程还没停，又启动工作线程，造成混乱 |
|
|
|
#工作线程 |
|
|
|
for work_th in self.workth_list: |
|
|
|
if work_th: |
|
|
|
if work_th.is_alive(): |
|
|
|
self.logger.debug(f"{self.task_id}有存活工作线程") |
|
|
|
return False |
|
|
|
#llm线程 |
|
|
|
for llm_th in self.llmth_list: |
|
|
|
if llm_th: |
|
|
|
if llm_th.is_alive(): |
|
|
|
self.logger.debug(f"{self.task_id}有存活LLM线程") |
|
|
|
return False |
|
|
|
#自检线程 |
|
|
|
if self.check_th: |
|
|
|
if self.check_th.is_alive(): |
|
|
|
self.logger.debug(f"{self.task_id}有存活自检线程") |
|
|
|
return False |
|
|
|
#工作子进程池 |
|
|
|
if self.PythonM.is_pool_active(): |
|
|
|
self.logger.debug(f"{self.task_id}有存活子进程池") |
|
|
|
return False |
|
|
|
return True |
|
|
|
|
|
|
|
def update_task_work_type(self,new_work_type): |
|
|
|
self.work_type = new_work_type |
|
|
|
#更新数据库 |
|
|
|
app_DBM.update_task_work_type(self.task_id,new_work_type) |
|
|
|
|
|
|
|
def update_task_status(self,new_status): |
|
|
|
self.task_status = new_status |
|
|
|
#更新数据库 |
|
|
|
app_DBM.update_task_status(self.task_id,new_status) |
|
|
|
|
|
|
|
def start_task(self): |
|
|
|
''' |
|
|
|
启动该测试任务， |
|
|
|
:return bool str |
|
|
|
''' |
|
|
|
if self.is_task_stop(): |
|
|
|
#更新状态标识 |
|
|
|
self.update_task_status(1) |
|
|
|
self.brun = True #线程正常启动 |
|
|
|
#启动指令工作线程 |
|
|
|
for i in range(self.max_thread_num): |
|
|
|
w_th = threading.Thread(target=self.do_worker_th,args=(i,),name=f"{self.task_id}-w_th-{i}") |
|
|
|
w_th.start() |
|
|
|
self.workth_list[i] = w_th |
|
|
|
#启动llm提交线程--llm暂时单线程，多线程处理时attack_tree需要加锁 |
|
|
|
for j in range(self.llm_max_nums): |
|
|
|
l_th = threading.Thread(target=self.th_llm_worker,args=(j,),name=f"{self.task_id}-l_th-{i}") |
|
|
|
l_th.start() |
|
|
|
self.llmth_list[j]=l_th |
|
|
|
#启动自检线程 |
|
|
|
self.check_th = threading.Thread(target=self.th_check) |
|
|
|
self.check_th.start() |
|
|
|
#启动python子进程池 --进程池的启动，目前是没全面确认原子进程池的状态，直接None |
|
|
|
self.PythonM.start_pool() |
|
|
|
return True,"启动任务成功！" #就算启动了一部分线程，也要认为是成功 |
|
|
|
else: |
|
|
|
return False,"该任务的工作线程未全面停止，不能重新启动工作，请稍后，若半个小时候还不能启动，请联系技术支持！" |
|
|
|
|
|
|
|
def stop_task(self): #还未处理 |
|
|
|
self.brun = False |
|
|
|
self.InstrM.init_data() |
|
|
|
#结束任务需要收尾处理#？ |
|
|
|
|
|
|
|
if self.brun and self.task_status ==1: #不是正常工作状态就不执行停止工作了 |
|
|
|
#停止子进程池 |
|
|
|
self.PythonM.shutdown_pool() |
|
|
|
#停止线程 |
|
|
|
self.brun = False |
|
|
|
self.update_task_status(0) |
|
|
|
# 结束任务需要收尾处理#？ |
|
|
|
self.InstrM.init_data() #pass |
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
pass |
