ipam/backend/scripts/backfill_vendor.py

#!/usr/bin/env python3
"""
回填 IP 厂商信息（vendor）

对所有 mac_address 有值但 vendor 为空的 IP，调用 EnhancedScanService.get_mac_vendor
回填厂商名。基于 IEEE OUI 官方数据库（mac-vendor-lookup 包），无需重新扫描。

用法：
    # 先在 git pull 之后确保依赖装了
    pip install mac-vendor-lookup

    # 干跑（只看不写）
    python scripts/backfill_vendor.py --dry-run

    # 实际写入
    python scripts/backfill_vendor.py

    # 同时回填主机名（反向 DNS 解析）
    python scripts/backfill_vendor.py --with-hostname

参数：
    --dry-run          只统计能识别的 IP 数量，不写库
    --with-hostname    同时回填 hostname（需要 DNS 解析，会比较慢）
    --timeout N        DNS 解析超时秒数（默认 2）
    --limit N          最多处理 N 条（用于测试）
    --yes              跳过确认提示
"""
import sys
import os
import time
import argparse
from datetime import datetime

# 让脚本能 import app.* 模块
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
BACKEND_DIR = os.path.dirname(SCRIPT_DIR)
sys.path.insert(0, BACKEND_DIR)


def parse_args():
    p = argparse.ArgumentParser(
        description="回填 IP 表的 vendor / hostname 字段",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--dry-run", action="store_true", help="只统计，不写库")
    p.add_argument("--with-hostname", action="store_true", help="同时回填主机名（反向 DNS 解析）")
    p.add_argument("--timeout", type=int, default=2, help="DNS 解析超时秒数（默认 2）")
    p.add_argument("--limit", type=int, default=0, help="最多处理 N 条，0 表示不限")
    p.add_argument("--yes", "-y", action="store_true", help="跳过确认提示")
    return p.parse_args()


def confirm(prompt: str, skip: bool) -> bool:
    if skip:
        return True
    try:
        ans = input(f"{prompt} [y/N] ").strip().lower()
        return ans in ("y", "yes")
    except EOFError:
        return False


def main():
    args = parse_args()

    # 延迟 import，确保 sys.path 设好
    from app.core.database import SessionLocal
    from app.models.network import IPAddress
    from app.services.enhanced_scan_service import EnhancedScanService

    db = SessionLocal()

    try:
        # 找出 mac_address 有值但 vendor 为空的 IP
        query = db.query(IPAddress).filter(
            IPAddress.mac_address.isnot(None),
            IPAddress.vendor.is_(None),
        )
        if args.limit > 0:
            query = query.limit(args.limit)
        candidates = query.all()
        total_candidates = len(candidates)

        print(f"[扫描] 找到 {total_candidates} 个 mac_address 有值但 vendor 为空的 IP")

        if total_candidates == 0:
            print("✅ 没有需要回填的 IP")
            return 0

        # 统计
        filled_vendor = 0
        skipped_vendor = 0
        filled_hostname = 0
        skipped_hostname = 0
        errors = 0
        start = time.time()

        for i, ip in enumerate(candidates, 1):
            mac = ip.mac_address
            vendor = EnhancedScanService.get_mac_vendor(mac)
            if vendor:
                if not args.dry_run:
                    ip.vendor = vendor
                filled_vendor += 1
            else:
                skipped_vendor += 1

            # 反向 DNS（可选）
            hostname = None
            if args.with_hostname:
                try:
                    hostname = EnhancedScanService.reverse_dns_lookup(ip.ip_address, timeout=args.timeout)
                except Exception as e:
                    print(f"  [WARN] {ip.ip_address} DNS 解析失败: {e}")
                    errors += 1
                if hostname:
                    if not args.dry_run:
                        ip.hostname = hostname
                    filled_hostname += 1
                else:
                    skipped_hostname += 1

            # 每 100 条打印进度
            if i % 100 == 0 or i == total_candidates:
                elapsed = time.time() - start
                rate = i / elapsed if elapsed > 0 else 0
                eta = (total_candidates - i) / rate if rate > 0 else 0
                print(
                    f"  进度: {i}/{total_candidates}  "
                    f"已识别 vendor={filled_vendor}, 未识别={skipped_vendor}  "
                    f"elapsed={elapsed:.1f}s, ETA={eta:.1f}s"
                )

        # 提交（如果非 dry-run）
        if not args.dry_run:
            if not confirm(f"确认写入数据库？将更新 {filled_vendor} 个 IP 的 vendor 字段"
                           + (f" + {filled_hostname} 个 IP 的 hostname 字段" if args.with_hostname else ""),
                           args.yes):
                print("已取消")
                db.rollback()
                return 1
            db.commit()
            print()
            print("✅ 已提交数据库")
        else:
            print()
            print(f"[DRY-RUN] 不写库；如需执行去掉 --dry-run")

        # 汇总
        print()
        print("=== 汇总 ===")
        print(f" 候选 IP 总数: {total_candidates}")
        print(f" 识别到厂商: {filled_vendor}")
        print(f" 未识别（OUI 未收录或格式异常）: {skipped_vendor}")
        if args.with_hostname:
            print(f" 解析到主机名: {filled_hostname}")
            print(f" DNS 解析失败或超时: {skipped_hostname}")
        print(f" 异常次数: {errors}")
        return 0

    except KeyboardInterrupt:
        print("\n[中断] 已回滚未提交的修改")
        db.rollback()
        return 130
    except Exception as e:
        print(f"\n[ERROR] {e}")
        import traceback
        traceback.print_exc()
        db.rollback()
        return 1
    finally:
        db.close()


if __name__ == "__main__":
    sys.exit(main())