Browse Source

V0.5.4 修改多批次单步功能前备份

a.页面宽度进行了调整
master
张龙 6 days ago
parent
commit
151fa4746c
  1. 2
      mycode/AttackMap.py
  2. 4
      mycode/TaskManager.py
  3. 34
      mycode/TaskObject.py
  4. 5
      web/API/task.py
  5. 2
      web/main/templates/footer.html
  6. 2
      web/main/templates/header.html
  7. 2
      web/main/templates/index.html
  8. 20
      web/main/templates/task_manager.html

2
mycode/AttackMap.py

@ -197,6 +197,8 @@ class TreeNode:
self._instr_queue = [] #针对当前节点的待执行指令----重要约束:一个节点只能有一个线程在执行指令
self.his_instr = [] #保留执行指令的记录{“instr”:***,"result":***}
#单步相关
self.step_num = 0 #单步执行次数
#用户补充信息
self.cookie = ""
self.ext_info = ""

4
mycode/TaskManager.py

@ -220,10 +220,10 @@ class TaskManager:
return bsuccess,error
#任务单点--只允许web端调用
async def task_one_step(self,task_id):
async def task_one_step(self,task_id,step_num):
task = self.tasks[task_id]
if task:
bsuccess,error = await task.put_one_task()
bsuccess,error = await task.put_one_task(step_num)
return bsuccess,error
else:
return False,"task_id值存在问题!"

34
mycode/TaskObject.py

@ -61,9 +61,11 @@ class TaskObject:
self.llmth_list = [None] * self.llm_max_nums # llm线程list
self.doing_llm_list = [""] * self.llm_max_nums
self.llm_node_queue = queue.Queue() #待提交LLM的节点队列
#自检线程--------
#自检线程-----------
self.check_th = None #自检线程句柄
#-----四队列-----
#单步执行控制-------
self.one_step_num = 0 #单步执行任务次数,默认1
#-----四队列-------
self.run_instr_lock = threading.Lock() # 线程锁
self.runing_instr = {} #执行中指令记录
@ -264,18 +266,30 @@ class TaskObject:
self.is_had_work = True
return bsuccess
def no_work_to_do(self): #任务单步状态控制-- 非工作中
def no_work_to_do(self): #任务单步状态控制-- 非工作中--2025-5-7增加了轮次,需要全面验证执行逻辑的完整和正确性
# 待执行instr-node
instr_node_list = list(self.instr_node_queue.queue) # 待执行指令的node--线程不安全
llm_node_list = list(self.llm_node_queue.queue) # 待提交llm的node--线程不安全
if len(instr_node_list) == 0 and len(llm_node_list) == 0: #没有待办节点了
#遍历在执行
for str_instr in self.doing_instr_list:
if str_instr != "":
return
for str_llm in self.doing_llm_list:
if str_llm != "":
return
#没有在执行任务了
# 没有在执行任务了
self.one_step_num -= 1
if not self.one_step_num: # 执行轮次已结束
pass
else:#发起新的一轮
bok = self.put_one_task()
if not bok:
self.one_step_num = 0
#结束轮次--推送前端
#推送是否有工作任务的状态到前端,
with self.is_had_work_lock:
if self.is_had_work: #如果已经是False那就不需要修改了
@ -483,14 +497,20 @@ class TaskObject:
return False,"当前的任务或节点状态不允许执行单步,请检查!"
#web端提交任务单步--任务单步
async def put_one_task(self):
async def put_one_task(self,step_num):
if self.task_status == 1 and self.work_type == 0:
bsuccess = self.had_work_to_do()
if bsuccess:
nodes = self.attack_tree.traverse_bfs()
b_putwork = False
for node in nodes:
_,_ = await self.put_one_node(node)
return True,"已提交单步任务"
bput,_ = await self.put_one_node(node)
if bput:
b_putwork = True
if b_putwork:
return True,"已提交单步任务"
else:
return False,"该任务已经没有待提交任务"
else:
return False,"当前任务正在执行任务中,不需要提交单步任务!"
else:

5
web/API/task.py

@ -137,9 +137,12 @@ async def task_one_step():
'''
data = await request.get_json()
task_id = data.get("cur_task_id")
step_num = data.get("step_num")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
bsuccess,error = await g_TaskM.task_one_step(task_id)
if not step_num:
step_num = 1 #默认一轮
bsuccess,error = await g_TaskM.task_one_step(task_id,step_num)
return jsonify({"bsuccess":bsuccess,"error":error})
@api.route('/task/nodestep',methods=['POST'])

2
web/main/templates/footer.html

@ -14,7 +14,7 @@
</svg>
<footer class="py-3 my-4 border-top">
<div class="container justify-content-between align-items-center d-flex flex-wrap">
<div class="container-xxl justify-content-between align-items-center d-flex flex-wrap">
<div class="col-md-4 d-flex align-items-center">
<img src="{{ url_for('main.static', filename='images/zf.svg') }}" alt="zfai" width="30" height="24">
<span class="mb-3 mb-md-0 text-muted">&copy; ZFXX All Rights Reserved</span>

Before

Width:  |  Height:  |  Size: 3.8 KiB

After

Width:  |  Height:  |  Size: 3.8 KiB

2
web/main/templates/header.html

@ -1,5 +1,5 @@
<header class="p-3 mb-3 border-bottom">
<div class="container">
<div class="container-xxl">
<div class="d-flex flex-wrap align-items-center justify-content-center justify-content-lg-start">
<a href="/" class="d-flex align-items-center mb-3 mb-md-0 me-md-auto text-dark text-decoration-none">

2
web/main/templates/index.html

@ -11,7 +11,7 @@
<!-- 页面内容块 -->
{% block content %}
<div class="container mt-4">
<div class="container-xxl mt-4">
<!-- 测试目标输入框 -->
<div class="mb-3">
<label for="testTarget" class="form-label">

20
web/main/templates/task_manager.html

@ -151,7 +151,7 @@
{% endblock %}
<!-- 页面内容块 -->
{% block content %}
<div class="container mt-4">
<div class="container-xxl mt-4">
<div class="container-fluid">
<div class="row">
<!-- 左侧:任务列表 -->
@ -164,7 +164,7 @@
<div class="col-10 full-height right-container">
<!-- 上方:基本信息 -->
<div class="row basic-info">
<div class="col-9">
<div class="col-8">
<div class="mb-2">
<label class="fw-bold">测试目标: </label>
<span id="detailTestTarget">192.168.1.110</span>
@ -209,9 +209,21 @@
</div>
</div>
<!-- <div class="col-2" style="display: flex; justify-content: center; align-items: center"> -->
<div class="col-3 d-flex justify-content-center align-items-center">
<div class="col-4 d-flex justify-content-center align-items-center">
<!-- 按钮 (联动测试状态示例: 执行中->暂停, 暂停中->启动,) -->
<button class="btn btn-primary btn-block" id="actionButton">启动</button>
<button class="btn btn-primary btn-block m-2" id="actionButton">启动</button>
<div class="m-2" style="margin-bottom: 5px">
<label class="fw-bold" style="font-size:0.9rem">单步轮次:</label>
<select class="form-select" id="modelSelect" style="font-size:0.9rem">
<option value="1">1</option>
<option value="2">2</option>
<option value="3">3</option>
<option value="4">4</option>
<option value="5">5</option>
<option value="6">6</option>
<option value="7">7</option>
</select>
</div>
<button class="btn btn-primary btn-block m-2" id="one_step">单步</button>
<button class="btn btn-danger btn-block m-2" id="btnTaskOver">结束</button>
</div>

Loading…
Cancel
Save