Nexus/scripts/import_servers_from_xls.py

#!/usr/bin/env python3
"""Import servers from Excel (.xls) via Nexus API (name + IP columns).

Usage:
  NEXUS_TEST_BASE=https://api.synaglobal.vip \\
  NEXUS_TEST_ADMIN_USER=admin \\
  NEXUS_TEST_ADMIN_PASSWORD='***' \\
  python scripts/import_servers_from_xls.py /path/to/servers.xls

Options:
  --dry-run     Parse only, do not call API
  --batch       Use /add-by-ips-batch (requires server with name-tab support)
  --chunk N     Batch chunk size (default 200)
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


def _load_test_credentials() -> tuple[str, str, str]:
    sys.path.insert(0, str(ROOT / "tests"))
    import test_api  # noqa: E402

    test_api._load_credentials_from_env_file()
    base = os.environ.get("NEXUS_TEST_BASE", "http://127.0.0.1:8600").rstrip("/")
    user = os.environ.get("NEXUS_TEST_ADMIN_USER", "admin")
    password = os.environ.get("NEXUS_TEST_ADMIN_PASSWORD", "admin")
    return base, user, password


def _api_json(method: str, url: str, token: str | None = None, body: dict | None = None, timeout: int = 300):
    data = None
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    if body is not None:
        data = json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read()
            return resp.status, json.loads(raw) if raw else {}
    except urllib.error.HTTPError as e:
        raw = e.read()
        try:
            payload = json.loads(raw) if raw else {"detail": e.reason}
        except json.JSONDecodeError:
            payload = {"detail": raw.decode(errors="replace")}
        return e.code, payload


def _login(base: str, user: str, password: str) -> str:
    code, data = _api_json("POST", f"{base}/api/auth/login", body={"username": user, "password": password})
    if code != 200 or not data.get("access_token"):
        detail = data.get("detail", data)
        raise SystemExit(f"登录失败 ({code}): {detail}")
    return data["access_token"]


def _read_xls(path: Path) -> list[tuple[str, str]]:
    try:
        import xlrd
    except ImportError as e:
        raise SystemExit("需要 xlrd: pip install xlrd") from e

    wb = xlrd.open_workbook(str(path))
    sh = wb.sheet_by_index(0)
    rows: list[tuple[str, str]] = []
    for r in range(sh.nrows):
        if sh.ncols < 2:
            continue
        name = str(sh.cell_value(r, 0)).strip()
        ip = str(sh.cell_value(r, 1)).strip()
        if not ip:
            continue
        rows.append((name or ip, ip))
    return rows


def _dedupe_rows(rows: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], int]:
    seen: set[str] = set()
    out: list[tuple[str, str]] = []
    removed = 0
    for name, ip in rows:
        key = ip.casefold()
        if key in seen:
            removed += 1
            continue
        seen.add(key)
        out.append((name, ip))
    return out, removed


def _import_single(base: str, token: str, rows: list[tuple[str, str]], port: int, agent_port: int) -> dict:
    stats = {"created": 0, "pending": 0, "skipped": 0, "errors": 0}
    total = len(rows)
    for i, (name, ip) in enumerate(rows, 1):
        code, data = _api_json(
            "POST",
            f"{base}/api/servers/add-by-ip",
            token=token,
            body={
                "domain": ip,
                "name": name,
                "port": port,
                "agent_port": agent_port,
            },
            timeout=300,
        )
        if code == 200 and data.get("success"):
            stats["created"] += 1
            status = "created"
        elif code == 409:
            stats["skipped"] += 1
            status = "skipped"
        elif code == 200 and not data.get("success"):
            stats["pending"] += 1
            status = "pending"
        else:
            stats["errors"] += 1
            status = f"error:{code}"
        print(f"[{i}/{total}] {status} {name} ({ip})", flush=True)
    return stats


def _import_batch(
    base: str,
    token: str,
    rows: list[tuple[str, str]],
    port: int,
    agent_port: int,
    chunk: int,
) -> dict:
    stats = {"created": 0, "pending": 0, "skipped": 0, "errors": 0}
    for start in range(0, len(rows), chunk):
        part = rows[start : start + chunk]
        text = "\n".join(f"{name}\t{ip}" for name, ip in part)
        code, data = _api_json(
            "POST",
            f"{base}/api/servers/add-by-ips-batch",
            token=token,
            body={"text": text, "port": port, "agent_port": agent_port},
            timeout=3600,
        )
        if code != 200:
            stats["errors"] += len(part)
            print(f"batch {start // chunk + 1} failed ({code}): {data.get('detail', data)}", flush=True)
            continue
        stats["created"] += int(data.get("created", 0))
        stats["pending"] += int(data.get("pending", 0))
        stats["skipped"] += int(data.get("skipped", 0))
        print(
            f"batch {start // chunk + 1}: created={data.get('created')} "
            f"pending={data.get('pending')} skipped={data.get('skipped')}",
            flush=True,
        )
    return stats


def main() -> None:
    parser = argparse.ArgumentParser(description="Import servers from Excel via Nexus API")
    parser.add_argument("xls", type=Path, help="Excel file (.xls), column A=name, B=IP")
    parser.add_argument("--dry-run", action="store_true", help="Parse only")
    parser.add_argument("--batch", action="store_true", help="Use batch API instead of single add-by-ip")
    parser.add_argument("--chunk", type=int, default=200, help="Batch chunk size")
    parser.add_argument("--port", type=int, default=22)
    parser.add_argument("--agent-port", type=int, default=8601)
    args = parser.parse_args()

    if not args.xls.is_file():
        raise SystemExit(f"文件不存在: {args.xls}")

    rows = _read_xls(args.xls)
    rows, dup_removed = _dedupe_rows(rows)
    print(f"读取 {args.xls.name}: {len(rows)} 台（去重忽略 {dup_removed} 行重复 IP）")
    if args.dry_run:
        for name, ip in rows[:5]:
            print(f"  {name}\t{ip}")
        if len(rows) > 5:
            print(f"  ... 共 {len(rows)} 行")
        return

    base, user, password = _load_test_credentials()
    print(f"目标 API: {base}  用户: {user}")
    token = _login(base, user, password)
    started = time.time()
    if args.batch:
        stats = _import_batch(base, token, rows, args.port, args.agent_port, args.chunk)
    else:
        stats = _import_single(base, token, rows, args.port, args.agent_port)
    elapsed = int(time.time() - started)
    print(
        f"完成 ({elapsed}s): 成功 {stats['created']}，待连接 {stats['pending']}，"
        f"跳过 {stats['skipped']}，错误 {stats['errors']}"
    )


if __name__ == "__main__":
    main()