zfsafe/web/API/task.py

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})