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

from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
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, IPAddress
from app.models.network import TaskType, TaskStatus
from app.services.network_service import NetworkService
from app.services.enhanced_scan_service import EnhancedScanService
from app.services.scan_service import ScanService

router = APIRouter(
    prefix="/scan",
    tags=["增强扫描"],
    dependencies=[Depends(get_current_user)],
)


@router.post("/comprehensive/{network_id}", response_model=ScanTask, summary="综合扫描网段")
def comprehensive_scan(
    network_id: int,
    background_tasks: BackgroundTasks,
    enable_ping: bool = True,
    enable_arp: bool = True,
    enable_dns: bool = True,
    enable_netbios: bool = True,
    timeout: int = 2,
    db: Session = Depends(get_db)
):
    """
    对指定网段执行综合扫描
    - Ping 扫描检测在线状态
    - ARP 扫描获取MAC地址
    - 反向DNS解析获取主机名
    - NetBIOS 广播获取主机名（Windows/SMB 设备）
    - MAC 厂商识别（IEEE OUI 数据库）
    """
    network = NetworkService.get_by_id(db, network_id)
    if not network:
        raise HTTPException(status_code=404, detail="网段不存在")

    # 创建扫描任务
    task = ScanService.create_scan_task(db, TaskType.PING, network_id=network_id)

    try:
        # 立即执行扫描（后续改为Celery异步）
        ScanService.update_task_status(db, task.id, TaskStatus.RUNNING, progress=0)

        # 执行综合扫描
        results = EnhancedScanService.scan_network(
            network.cidr,
            enable_ping=enable_ping,
            enable_arp=enable_arp,
            enable_dns=enable_dns,
            enable_netbios=enable_netbios,
            timeout=timeout
        )

        # 更新IP状态
        stats = EnhancedScanService.bulk_update_ips_from_scan(db, results)

        # 刷新网段统计字段（used_ips/reserved_ips），保证网段列表利用率能即时反映
        NetworkService.get_stats(db, network_id)

        # 更新任务状态
        ScanService.update_task_status(
            db,
            task.id,
            TaskStatus.COMPLETED,
            progress=100,
            total_count=stats['total_scanned'],
            success_count=stats['online_count']
        )

    except Exception as e:
        ScanService.update_task_status(
            db,
            task.id,
            TaskStatus.FAILED,
            error_message=str(e)
        )
        raise HTTPException(status_code=500, detail=f"扫描失败: {str(e)}")

    # 重新获取任务状态
    task = ScanService.get_task_by_id(db, task.id)
    return task


@router.post("/comprehensive/ip/{ip_address}", summary="单个IP综合扫描")
def comprehensive_scan_ip(
    ip_address: str,
    enable_ping: bool = True,
    enable_arp: bool = True,
    enable_dns: bool = True,
    enable_netbios: bool = True,
    timeout: int = 2,
    db: Session = Depends(get_db)
):
    """
    对单个IP执行综合扫描
    """
    result = EnhancedScanService.comprehensive_scan(
        ip_address,
        enable_ping=enable_ping,
        enable_arp=enable_arp,
        enable_dns=enable_dns,
        enable_netbios=enable_netbios,
        timeout=timeout
    )

    # 如果IP在数据库中，更新信息
    EnhancedScanService.update_ip_from_scan_result(db, result)

    # 刷新所属网段统计字段（used_ips/reserved_ips），保证网段列表利用率能即时反映
    db_ip = db.query(IPAddress).filter(IPAddress.ip_address == ip_address).first()
    if db_ip:
        NetworkService.get_stats(db, db_ip.network_id)

    db.commit()

    return result


@router.get("/mac/vendor/{mac_address}", summary="MAC地址厂商识别")
def get_mac_vendor(mac_address: str):
    """
    根据MAC地址识别厂商
    """
    vendor = EnhancedScanService.get_mac_vendor(mac_address)
    return {
        "mac_address": mac_address,
        "vendor": vendor or "Unknown"
    }


@router.post("/dns/reverse/{ip_address}", summary="反向DNS解析")
def reverse_dns(ip_address: str, timeout: int = 2):
    """
    反向DNS解析获取主机名
    """
    hostname = EnhancedScanService.reverse_dns_lookup(ip_address, timeout)
    return {
        "ip_address": ip_address,
        "hostname": hostname or "Unknown"
    }

@router.post("/netbios/lookup/{ip_address}", summary="NetBIOS 主机名发现")
def netbios_lookup(ip_address: str, timeout: int = 3):
    """
    通过 NetBIOS 广播协议获取主机名（Windows / SMB 设备）
    """
    hostname = EnhancedScanService.netbios_lookup(ip_address, timeout)
    return {
        "ip_address": ip_address,
        "hostname": hostname or "Unknown"
    }

@router.post("/batch/vendor", summary="批量补全厂商信息")
def batch_fill_vendor(
    network_id: Optional[int] = None,
    db: Session = Depends(get_db)
):
    """
    对已有 MAC 地址但缺少厂商信息的 IP 批量补全厂商（使用 IEEE OUI 数据库）
    """
    from app.models.network import IPAddress as IPAddressModel

    query = db.query(IPAddressModel).filter(
        IPAddressModel.mac_address.isnot(None),
        IPAddressModel.vendor.is_(None)
    )
    if network_id:
        query = query.filter(IPAddressModel.network_id == network_id)

    items = query.all()
    updated_count = 0

    for ip in items:
        vendor = EnhancedScanService.get_mac_vendor(ip.mac_address)
        if vendor:
            ip.vendor = vendor
            updated_count += 1

    if updated_count > 0:
        db.commit()

    return {
        "message": f"已补全 {updated_count} 条厂商信息",
        "checked": len(items),
        "updated": updated_count
    }

@router.post("/batch/hostname", summary="批量发现主机名")
def batch_discover_hostname(
    network_id: Optional[int] = None,
    enable_dns: bool = True,
    enable_netbios: bool = True,
    db: Session = Depends(get_db)
):
    """
    对已有 MAC 地址但缺少主机名的 IP 批量执行主机名发现（DNS PTR + NetBIOS）
    """
    from app.models.network import IPAddress as IPAddressModel

    query = db.query(IPAddressModel).filter(
        IPAddressModel.mac_address.isnot(None),
        IPAddressModel.hostname.is_(None)
    )
    if network_id:
        query = query.filter(IPAddressModel.network_id == network_id)

    items = query.all()
    updated_count = 0

    for ip in items:
        hostname = None
        if enable_dns:
            hostname = EnhancedScanService.reverse_dns_lookup(ip.ip_address, timeout=2)
        if not hostname and enable_netbios:
            hostname = EnhancedScanService.netbios_lookup(ip.ip_address, timeout=3)
        if hostname:
            ip.hostname = hostname
            updated_count += 1

    if updated_count > 0:
        db.commit()

    return {
        "message": f"已发现 {updated_count} 条主机名",
        "checked": len(items),
        "updated": updated_count
    }


@router.get("/online/ips", summary="获取所有在线IP列表")
def get_online_ips(
    network_id: Optional[int] = None,
    skip: int = 0,
    limit: int = 100,
    db: Session = Depends(get_db)
):
    """
    获取所有在线的IP地址列表
    """
    from app.models.network import IPAddress as IPAddressModel, IPStatus

    query = db.query(IPAddressModel).filter(IPAddressModel.status == IPStatus.ONLINE)

    if network_id:
        query = query.filter(IPAddressModel.network_id == network_id)

    total = query.count()
    items = query.order_by(IPAddressModel.ip_address).offset(skip).limit(limit).all()

    return {
        "total": total,
        "online_count": total,
        "items": items
    }


@router.get("/statistics/summary", summary="扫描统计摘要")
def get_scan_statistics(db: Session = Depends(get_db)):
    """
    获取扫描统计摘要信息
    """
    from app.models.network import IPAddress as IPAddressModel, IPStatus
    from app.models.network import ScanTask, TaskStatus

    # IP 统计
    total_ips = db.query(IPAddressModel).count()
    online_ips = db.query(IPAddressModel).filter(IPAddressModel.status == IPStatus.ONLINE).count()
    offline_ips = db.query(IPAddressModel).filter(IPAddressModel.status == IPStatus.OFFLINE).count()
    reserved_ips = db.query(IPAddressModel).filter(IPAddressModel.status == IPStatus.RESERVED).count()
    with_mac = db.query(IPAddressModel).filter(IPAddressModel.mac_address.isnot(None)).count()
    with_hostname = db.query(IPAddressModel).filter(IPAddressModel.hostname.isnot(None)).count()

    # 任务统计
    total_tasks = db.query(ScanTask).count()
    completed_tasks = db.query(ScanTask).filter(ScanTask.status == TaskStatus.COMPLETED).count()
    failed_tasks = db.query(ScanTask).filter(ScanTask.status == TaskStatus.FAILED).count()

    return {
        "ip_statistics": {
            "total": total_ips,
            "online": online_ips,
            "offline": offline_ips,
            "reserved": reserved_ips,
            "available": total_ips - online_ips - offline_ips - reserved_ips,
            "with_mac": with_mac,
            "with_hostname": with_hostname,
            "online_rate": round(online_ips / total_ips * 100, 2) if total_ips > 0 else 0
        },
        "task_statistics": {
            "total": total_tasks,
            "completed": completed_tasks,
            "failed": failed_tasks,
            "success_rate": round(completed_tasks / total_tasks * 100, 2) if total_tasks > 0 else 0
        }
    }