from . import api
from quart import Quart, render_template, redirect, url_for, request,jsonify
from mycode.TargetManager import g_TM
from mycode.DBManager import app_DBM
from mycode.TaskManager import g_TaskM
def is_valid_target(test_target: str) -> bool:
"""
验证 test_target 的逻辑（这里用简单示例代替已有逻辑）
例如：测试目标不能为空且长度大于3
"""
if test_target:
return True
return False
@api.route('/task/start',methods=['POST'])
async def start_task(): #开始任务
data = await request.get_json()
test_target = data.get("testTarget")
cookie_info = data.get("cookieInfo")
llm_type = data.get("curmodel") # //0-腾讯云，1-DS，2-2233.ai,3-GPT 目前只有1-2，2025-4-4
work_type = data.get("workType") #0-人工，1-自动
#新增任务处理
bok,_,_ = g_TM.validate_and_extract(test_target)
if not bok:
# 返回错误信息，状态码 400 表示请求错误
return jsonify({"error": "测试目标验证失败，请检查输入内容！"}), 400
#开始任务
try:
b_success = g_TaskM.create_task(test_target,cookie_info,llm_type,work_type)
#再启动
if not b_success:
return jsonify({"error": "检测任务创建失败，请联系管理员！"}), 500
except:
return jsonify({"error": "该目标已经在测试中，请检查！"}), 400
#跳转到任务管理页面
return redirect(url_for('main.get_html', html='task_manager.html'))
@api.route('/task/getlist',methods=['GET'])
async def get_task_list():
#task_list = app_DBM.get_task_list() #从内存取--2025-4-6
task_list = g_TaskM.get_task_list()
if task_list:
return jsonify(task_list)
else:
return jsonify({"error":"查询任务数据出错！"}),500
@api.route('/task/getinstr',methods=['POST'])
async def get_instr():
data = await request.get_json()
task_id = data.get("task_id")
node_name = data.get("nodeName")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
instrs = app_DBM.get_task_instrs(task_id,node_name)
return jsonify(instrs)
@api.route('/task/getvul',methods=['POST'])
async def get_vul():
data = await request.get_json()
task_id = data.get("task_id")
node_name = data.get("nodeName")
vul_type = data.get("vulType")
vul_level = data.get("vulLevel")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
vuls = app_DBM.get_task_vul(task_id,node_name,vul_type,vul_level)
return jsonify(vuls)
@api.route('/task/gettree',methods=['POST'])
async def get_tree():
data = await request.get_json()
task_id = data.get("task_id")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
tree_dict = g_TaskM.get_node_tree(task_id)
return jsonify({"tree":tree_dict})
@api.route('/task/taskcontrol',methods=['POST'])
async def task_status_control():
'''控制任务状态
1.对于执行时间长的指令，如何处理？强制停止的话，要有个执行中指令的缓存，强制停止该指令返回到待执行，执行完成，该指令到执行完成；
'''
data = await request.get_json()
task_id = data.get("cur_task_id")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
#只做暂停和继续间的切换，以服务器端的状态为准
bsuccess,strerror,new_task_status = g_TaskM.control_taks(task_id)
if bsuccess:
return jsonify({'newstatus':new_task_status})
return jsonify({'error': strerror}), 400
@api.route('/task/taskstep',methods=['POST'])
async def task_one_step():
'''单步推进任务--也就是待处理node 返回bsuccess,error
1.执行单步的前提条件是，工作线程都要在工作；
2.遍历节点把需要处理的节点进入待处理queue，instr和llm只能一个有数据（强制约束）
'''
data = await request.get_json()
task_id = data.get("cur_task_id")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
bsuccess,error = await g_TaskM.task_one_step(task_id)
return jsonify({"bsuccess":bsuccess,"error":error})
@api.route('/task/nodestep',methods=['POST'])
async def node_one_step():
data = await request.get_json()
task_id = data.get("task_id")
node_path = data.get("node_path")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
bsuccess,error = await g_TaskM.node_one_step(task_id,node_path)
return jsonify({"bsuccess":bsuccess,"error":error})
@api.route('/task/taskworktype',methods=['POST'])
async def task_work_type_control():
data = await request.get_json()
task_id = data.get("cur_task_id")
newwork_type = data.get("mode")
if not task_id:
return jsonify({'error': 'Missing task_id or newwork_type'}), 400
bsuccess = g_TaskM.update_task_work_type(task_id,newwork_type)
return jsonify({"bsuccess": bsuccess})
@api.route('/task/nodecontrol',methods=['POST'])
async def node_work_status_control():
data = await request.get_json()
task_id = data.get("task_id")
nodepath = data.get("node_path")
if not task_id or not nodepath:
return jsonify({'error': 'Missing task_id or node_path'}), 400
#修改节点的工作状态
bsuccess,newbwork = g_TaskM.node_bwork_control(task_id,nodepath)
if not bsuccess:
return jsonify({'error': 'node_path not bfind'}), 400
return jsonify({"newbwork":newbwork})
@api.route('/task/nodegetinstr',methods=['POST'])
async def node_get_instr():
data = await request.get_json()
task_id = data.get("task_id")
nodepath = data.get("node_path")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
#返回 { doneInstrs: [...], todoInstrs: [...] }
doneInstrs = app_DBM.get_task_node_done_instr(task_id,nodepath)
todoInstrs = g_TaskM.get_task_node_todo_instr(task_id,nodepath)
return jsonify({"doneInstrs":doneInstrs,"todoInstrs":todoInstrs})
@api.route('/task/nodegetmsg',methods=['POST'])
async def node_get_msg():
data = await request.get_json()
task_id = data.get("task_id")
nodepath = data.get("node_path")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
submitted,pending = g_TaskM.get_task_node_MSG(task_id,nodepath)
return jsonify({"submitted": submitted, "pending": pending})
@api.route('/task/nodeupdatemsg',methods=['POST'])
async def node_update_msg():
data = await request.get_json()
task_id = data.get("task_id")
nodepath = data.get("node_path")
newllm_type = data.get("llmtype")
newcontent = data.get("content")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
bsuccess,error =g_TaskM.update_node_MSG(task_id,nodepath,newllm_type,newcontent)
return jsonify({"bsuccess":bsuccess,"error":error})
@api.route('/task/delnodeinstr',methods=['POST'])
async def node_del_instr():
data = await request.get_json()
task_id = data.get("task_id")
nodepath = data.get("node_path")
instr = data.get("item")
if not task_id:
return jsonify({'error': 'Missing task_id'}), 400
bsuccess,error = g_TaskM.del_node_instr(task_id,nodepath,instr)
return jsonify({"bsuccess": bsuccess, "error": error})