Nexus/scripts/batch_detect_target_path.py

#!/usr/bin/env python3
"""Batch detect target_path via SSH (workerman.bat under /www/wwwroot).

Targets servers with「未设路径」(empty or ``/www/wwwroot`` placeholder) by default.

Usage:
  python scripts/batch_detect_target_path.py --unset-only
  python scripts/batch_detect_target_path.py --ids 245,246
  python scripts/batch_detect_target_path.py --name 温胜夜 --dry-run
"""

from __future__ import annotations

import argparse
import asyncio
import os
import sys
from pathlib import Path
from typing import Any

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


def _load_dotenv() -> None:
    env_path = ROOT / ".env"
    if not env_path.is_file():
        return
    for line in env_path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        key, value = key.strip(), value.strip().strip("\"'")
        if key.startswith("NEXUS_") and key[6:] not in os.environ:
            os.environ[key[6:]] = value
        elif key and key not in os.environ:
            os.environ[key] = value


async def _fetch_servers(
    session: Any,
    *,
    ids: list[int] | None,
    name_like: str | None,
    unset_only: bool,
    limit: int,
) -> list[Any]:
    from sqlalchemy import or_, select

    from server.domain.models import Server
    from server.infrastructure.database.server_repo import _target_path_unset_filter

    q = select(Server).order_by(Server.id)
    if ids:
        q = q.where(Server.id.in_(ids))
    if name_like:
        q = q.where(or_(Server.name.like(f"%{name_like}%"), Server.domain.like(f"%{name_like}%")))
    if unset_only:
        q = q.where(_target_path_unset_filter(True))
    result = await session.execute(q)
    servers = list(result.scalars().all())
    if limit > 0:
        servers = servers[:limit]
    return servers


async def run_detect(args: argparse.Namespace) -> list[dict[str, Any]]:
    from server.application.services.server_batch_service import _run_detect_path
    from server.config import settings
    from server.infrastructure.database.session import AsyncSessionLocal, init_db
    from server.infrastructure.redis.client import close_redis, init_redis

    await init_db()
    await init_redis(settings.REDIS_URL)

    ids: list[int] | None = None
    if args.ids:
        ids = [int(x.strip()) for x in args.ids.split(",") if x.strip()]

    async with AsyncSessionLocal() as session:
        servers = await _fetch_servers(
            session,
            ids=ids,
            name_like=args.name,
            unset_only=args.unset_only,
            limit=args.limit,
        )

    if not servers:
        print("No servers matched.", file=sys.stderr)
        return []

    if args.dry_run:
        for s in servers:
            print(f"[DRY] #{s.id} {s.name} target_path={s.target_path!r}")
        print(f"\nWould detect: {len(servers)} servers")
        await close_redis()
        return [{"server_id": s.id, "dry_run": True} for s in servers]

    live = {
        "server_ids": [s.id for s in servers],
        "results": [],
        "done": 0,
        "total": len(servers),
    }
    await _run_detect_path(live)
    await close_redis()

    results = live.get("results") or []
    ok = sum(1 for r in results if r.get("success"))
    fail = len(results) - ok
    print(f"\nDetect: {ok} ok / {fail} fail / {len(results)} total")
    for r in results:
        status = "OK" if r.get("success") else "FAIL"
        print(
            f"[{status}] #{r.get('server_id')} {r.get('server_name')}: "
            f"{r.get('stdout') or r.get('error') or ''}"
        )
    return results


def main() -> int:
    _load_dotenv()
    parser = argparse.ArgumentParser(description="Batch detect server target_path")
    parser.add_argument("--ids", type=str, help="Comma-separated server ids")
    parser.add_argument("--name", type=str, help="Filter name/domain LIKE")
    parser.add_argument("--unset-only", action="store_true", help="Only「未设路径」servers")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    if not args.ids and not args.name and not args.unset_only:
        parser.error("Specify --ids, --name, or --unset-only")
    results = asyncio.run(run_detect(args))
    if args.dry_run:
        return 0 if results else 2
    return 0 if results and all(r.get("success") for r in results) else 1


if __name__ == "__main__":
    raise SystemExit(main())