ipam/backend/app/api/v1/async_scan.py

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 可能未启动"
        }