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.
124 lines
4.4 KiB
124 lines
4.4 KiB
#自动化测试逻辑规则控制
|
|
#统一控制规则 和 渗透测试树的维护
|
|
import json
|
|
import re
|
|
import queue
|
|
import time
|
|
import threading
|
|
from mycode.AttackMap import AttackTree
|
|
from mycode.AttackMap import TreeNode
|
|
from myutils.MyLogger_logger import LogHandler
|
|
from mycode.DBManager import DBManager
|
|
|
|
class ControlCenter:
|
|
def __init__(self):
|
|
self.logger = LogHandler().get_logger("ControlCenter")
|
|
|
|
def __del__(self):
|
|
self.brun =False
|
|
self.task_id = None
|
|
self.target = None
|
|
self.attack_tree = None
|
|
|
|
def init_cc_data(self):
|
|
#一次任务一次数据
|
|
pass
|
|
|
|
def get_user_init_info(self):
|
|
'''开始任务初，获取用户设定的基础信息，初始信息可以分为两块：
|
|
1.提交llm的补充信息 和 保留在本地的信息（如工具补充参数等，cookie）
|
|
2.用户可以设置全局，和指定节点（端口）
|
|
3.补充测试节点
|
|
'''
|
|
# ?包括是否对目标进行初始化的信息收集
|
|
return {"已知信息":"无"}
|
|
|
|
def verify_node_cmds(self,node_cmds):
|
|
'''
|
|
验证节点指令的合规性，持续维护
|
|
:param node_cmds:
|
|
:param node:
|
|
:return: Flase 存在问题， True 合规
|
|
'''
|
|
strerror = ""
|
|
for node_json in node_cmds:
|
|
if "action" not in node_json:
|
|
self.logger.error(f"缺少action节点：{node_json}")
|
|
strerror = {"节点指令错误":f"{node_json}缺少action节点，不符合格式要求!"}
|
|
break
|
|
action = node_json["action"]
|
|
if action == "add_node":
|
|
if "parent" not in node_json or "status" not in node_json or "nodes" not in node_json:
|
|
strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
|
|
break
|
|
elif action == "update_status":
|
|
if "status" not in node_json or "node" not in node_json:
|
|
strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
|
|
break
|
|
elif action =="no_instruction" or action=="no_create":
|
|
if "nodes" not in node_json:
|
|
strerror = {"节点指令错误": f"{node_json}不符合格式要求,缺少节点！"}
|
|
break
|
|
else:
|
|
strerror = {"节点指令错误": f"{node_json}不可识别的action值！"}
|
|
break
|
|
if not strerror:
|
|
return True,strerror
|
|
return False,strerror
|
|
|
|
def restore_one_llm_work(self,node,llm_type,res_list):
|
|
node.llm_type = llm_type
|
|
node.res_quere = res_list
|
|
|
|
#需要用户确认的信息--待完善
|
|
def need_user_know(self,strinfo,node):
|
|
pass
|
|
|
|
#待修改
|
|
def is_user_instr(self,instr):
|
|
'''
|
|
过滤需要人工确认或手动执行的指令 ---- 待完善
|
|
:param instr:
|
|
:return:
|
|
'''
|
|
#if instr.startswith("curl") or instr.startswith("http") or instr.startswith("wget"):
|
|
if instr.startswith("http") or instr.startswith("wget") or instr.startswith("ssh"):
|
|
return True
|
|
|
|
#指令入队列，待修改
|
|
def instr_in_quere(self,instr_list):
|
|
'''
|
|
对于运行需要较长时间的不强求同一批次返回给LLM
|
|
:param instr_list:
|
|
:return:
|
|
'''
|
|
for instr in instr_list:
|
|
if self.is_user_instr(instr):
|
|
self.user_instr.put(instr)
|
|
print(f"需要人工确认的指令{instr}")
|
|
else:
|
|
matched =False
|
|
for prefix in self.long_time_instr:
|
|
if instr.startswith(prefix):
|
|
matched =True
|
|
if not matched:
|
|
with self.lock:
|
|
self.batch_num += 1 #非耗时指令+1
|
|
print(f"&&&&&&当前batch_num:{self.batch_num}")
|
|
else:
|
|
with self.lock:
|
|
self.long_instr_num +=1 #耗时指令数量+1
|
|
# 指令入队列
|
|
self.instr_queue.append(instr)
|
|
|
|
def stop_do(self):
|
|
#清空数据
|
|
self.task_id = None
|
|
self.target = None
|
|
self.attack_tree = None
|
|
#停止llm处理线程
|
|
self.brun =False
|
|
for th in self.llmth_list:
|
|
th.jion()
|
|
|
|
|
|
|