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/taskover',methods=['POST']) async def over_task(): 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 = g_TaskM.over_task(task_id) return jsonify({"bsuccess": bsuccess, "error": error}) @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})