from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from typing import List, Dict, Any, Optional from app.core.database import get_db from app.core.security import get_current_user from app.schemas.network import ScanTask from app.models.network import TaskStatus from app.services.network_service import NetworkService from app.services.scan_service import ScanService router = APIRouter( prefix="/async-scan", tags=["异步扫描"], dependencies=[Depends(get_current_user)], ) @router.post("/network/{network_id}", response_model=ScanTask, summary="异步扫描网段") def async_scan_network( network_id: int, enable_ping: bool = True, enable_arp: bool = True, enable_dns: bool = True, timeout: int = 2, db: Session = Depends(get_db) ): """ 提交异步网段扫描任务,立即返回任务ID,后台执行 """ network = NetworkService.get_by_id(db, network_id) if not network: raise HTTPException(status_code=404, detail="网段不存在") # 创建数据库任务记录 db_task = ScanService.create_scan_task(db, 'ping', network_id) # 提交 Celery 任务 from app.tasks.celery_app import scan_network_task celery_task = scan_network_task.delay( network_id=network_id, enable_ping=enable_ping, enable_arp=enable_arp, enable_dns=enable_dns, timeout=timeout ) # 记录 Celery 任务ID db_task.celery_task_id = celery_task.id db.commit() db.refresh(db_task) return db_task @router.post("/ip/{ip_address}", summary="异步扫描单个IP") def async_scan_ip( ip_address: str, enable_ping: bool = True, enable_arp: bool = True, enable_dns: bool = True ): """ 提交异步单个IP扫描任务 """ from app.tasks.celery_app import scan_single_ip_task celery_task = scan_single_ip_task.delay( ip_address=ip_address, enable_ping=enable_ping, enable_arp=enable_arp, enable_dns=enable_dns ) return { "task_id": celery_task.id, "ip_address": ip_address, "status": "pending" } @router.post("/quick", summary="快速批量IP状态更新") def quick_scan_ips(ip_addresses: List[str]): """ 快速批量更新IP状态(只做Ping+ARP,不做DNS) """ from app.tasks.celery_app import quick_status_update celery_task = quick_status_update.delay(ip_addresses=ip_addresses) return { "task_id": celery_task.id, "ips_count": len(ip_addresses), "status": "pending" } @router.get("/task/{task_id}", summary="查询异步任务状态") def get_async_task_status(task_id: str, db: Session = Depends(get_db)): """ 根据 Celery 任务ID 查询任务状态 """ # 先查询数据库任务 from app.models.network import ScanTask as DbScanTask db_task = db.query(DbScanTask).filter(DbScanTask.celery_task_id == task_id).first() # 直接查询 Celery from celery.result import AsyncResult result = AsyncResult(task_id) response = { "task_id": task_id, "celery_status": result.status } if db_task: response["db_status"] = db_task.status.value response["progress"] = db_task.progress response["total_count"] = db_task.total_count response["success_count"] = db_task.success_count response["started_at"] = db_task.started_at response["completed_at"] = db_task.completed_at if result.status == 'SUCCESS': response["result"] = result.result elif result.status == 'FAILURE': response["error"] = str(result.info) return response @router.get("/tasks/running", summary="获取运行中的任务列表") def get_running_tasks(db: Session = Depends(get_db)): """ 获取所有正在运行的扫描任务 """ from app.models.network import ScanTask as DbScanTask running = db.query(DbScanTask).filter(DbScanTask.status == TaskStatus.RUNNING).all() pending = db.query(DbScanTask).filter(DbScanTask.status == TaskStatus.PENDING).all() return { "running_count": len(running), "pending_count": len(pending), "running_tasks": running, "pending_tasks": pending } @router.post("/task/{task_id}/cancel", summary="取消异步任务") def cancel_async_task(task_id: str): """ 取消正在运行的异步任务 """ from celery.result import AsyncResult result = AsyncResult(task_id) if result.status in ['PENDING', 'RUNNING']: result.revoke(terminate=True) return {"status": "cancelled", "task_id": task_id} else: return {"status": "not_running", "current_status": result.status} @router.post("/full-scan", summary="触发全量扫描") def trigger_full_scan( enable_ping: bool = True, enable_arp: bool = True, enable_dns: bool = True ): """ 触发全量网段扫描(所有网段) """ from app.tasks.celery_app import full_network_scan celery_task = full_network_scan.delay( enable_ping=enable_ping, enable_arp=enable_arp, enable_dns=enable_dns ) return { "task_id": celery_task.id, "status": "pending", "message": "全量扫描任务已提交" } @router.get("/celery/inspect", summary="Celery Worker 状态检查") def inspect_celery(): """ 检查 Celery Worker 状态 """ from app.tasks.celery_app import celery_app try: inspector = celery_app.control.inspect() active = inspector.active() scheduled = inspector.scheduled() reserved = inspector.reserved() stats = inspector.stats() return { "status": "healthy" if stats else "no_workers", "active_tasks": active, "scheduled_tasks": scheduled, "reserved_tasks": reserved, "worker_stats": stats } except Exception as e: return { "status": "error", "error": str(e), "message": "Celery Worker 可能未启动" }