Nexus/scripts/batch_patch_files_sudoers.py

#!/usr/bin/env python3
"""Batch patch /etc/sudoers.d/nexus-files on sub-servers (add rsync for push elevation).

Usage:
  python scripts/batch_patch_files_sudoers.py --name 温胜夜
  python scripts/batch_patch_files_sudoers.py --non-root-only --limit 50
  python scripts/batch_patch_files_sudoers.py --ids 12,34 --dry-run

Requires .env with DATABASE_URL. SSH creds from servers table.
Run from dev machine or on central host: bash scripts/batch_patch_files_sudoers.sh
"""

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,
    *,
    server_id: int | None,
    server_ids: list[int] | None,
    name_like: str | None,
    non_root_only: bool,
    limit: int,
) -> list[Any]:
    from sqlalchemy import or_, select

    from server.domain.models import Server

    q = select(Server).order_by(Server.id)
    if server_id is not None:
        q = q.where(Server.id == server_id)
    elif server_ids:
        q = q.where(Server.id.in_(server_ids))
    if name_like:
        q = q.where(or_(Server.name.like(f"%{name_like}%"), Server.domain.like(f"%{name_like}%")))
    if non_root_only:
        q = q.where(Server.username.isnot(None)).where(Server.username != "root")

    result = await session.execute(q)
    servers = list(result.scalars().all())
    if limit > 0:
        servers = servers[:limit]
    return servers


async def run_patch(args: argparse.Namespace) -> list[dict[str, Any]]:
    from server.application.services.files_sudoers_service import install_nexus_files_sudoers
    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,
            server_id=args.id,
            server_ids=ids,
            name_like=args.name,
            non_root_only=args.non_root_only,
            limit=args.limit,
        )

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

    sem = asyncio.Semaphore(max(1, args.concurrency))
    results: list[dict[str, Any]] = []

    async def _one(server: Any) -> None:
        async with sem:
            try:
                row = await install_nexus_files_sudoers(server, dry_run=args.dry_run)
            except Exception as exc:
                row = {
                    "server_id": server.id,
                    "server_name": server.name,
                    "ssh_user": server.username,
                    "ok": False,
                    "action": "error",
                    "detail": str(exc)[:500],
                }
            results.append(row)
            status = "OK" if row.get("ok") else "FAIL"
            print(
                f"[{status}] #{row.get('server_id')} {row.get('server_name')} "
                f"({row.get('ssh_user')}) → {row.get('action')}"
                + (f" — {row.get('detail')}" if row.get("detail") else ""),
                flush=True,
            )

    await asyncio.gather(*[_one(s) for s in servers])

    from server.infrastructure.redis.client import close_redis

    await close_redis()

    ok = sum(1 for r in results if r.get("ok"))
    fail = len(results) - ok
    print(f"\nDone: {ok} ok / {fail} fail / {len(results)} total")
    return results


def main() -> int:
    _load_dotenv()
    parser = argparse.ArgumentParser(description="Batch patch nexus-files sudoers (rsync for push)")
    parser.add_argument("--id", type=int, help="Single server id")
    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("--non-root-only", action="store_true", help="Skip username=root")
    parser.add_argument("--limit", type=int, default=0, help="Max servers (0=all)")
    parser.add_argument("--concurrency", type=int, default=5)
    parser.add_argument("--dry-run", action="store_true", help="Report only, no SSH writes")
    args = parser.parse_args()

    if not args.id and not args.ids and not args.name and not args.non_root_only:
        parser.error("Specify --id, --ids, --name, and/or --non-root-only")

    results = asyncio.run(run_patch(args))
    return 0 if results and all(r.get("ok") for r in results) else (1 if results else 2)


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