Nexus/scripts/verify_push_elevation.py

#!/usr/bin/env python3
"""Verify rsync push elevation for selected servers (diagnose + dry-run + optional probe push).

Usage:
  python scripts/verify_push_elevation.py --ids 313,314
  python scripts/verify_push_elevation.py --name 温胜夜 --push
"""

from __future__ import annotations

import argparse
import asyncio
import os
import secrets
import sys
import tempfile
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,
    non_root_only: bool = False,
) -> list[Any]:
    from sqlalchemy import or_, select

    from server.domain.models import Server

    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 non_root_only:
        q = q.where(Server.username.isnot(None)).where(Server.username != "root")
    result = await session.execute(q)
    return list(result.scalars().all())


async def _diagnose_write(server: Any, dest: str) -> dict[str, Any]:
    import shlex

    from server.infrastructure.ssh.remote_shell import exec_ssh_command_with_fallback
    from server.utils.posix_paths import remote_join

    test_file = remote_join(dest, ".nexus_verify_probe")
    cmd = f"touch {shlex.quote(test_file)} && rm -f {shlex.quote(test_file)}"
    r = await exec_ssh_command_with_fallback(server, cmd, timeout=15)
    return {
        "path_writable": r.get("exit_code") == 0,
        "elevated": bool(r.get("elevated")),
        "stderr": (r.get("stderr") or "")[:200],
    }


async def run_verify(args: argparse.Namespace) -> list[dict[str, Any]]:
    from server.application.services.files_sudoers_service import probe_rsync_sudo
    from server.application.services.sync_engine_v2 import _rsync_push
    from server.config import settings
    from server.infrastructure.database.session import AsyncSessionLocal, init_db
    from server.infrastructure.redis.client import close_redis, init_redis
    from server.utils.posix_paths import UPLOAD_STAGING_PREFIX, normalize_remote_abs_path, resolve_push_target_path

    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, non_root_only=args.non_root_only
        )

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

    staging = f"{UPLOAD_STAGING_PREFIX}{secrets.token_hex(6)}"
    os.makedirs(staging, exist_ok=True)
    probe_name = ".nexus_push_probe.txt"
    probe_path = os.path.join(staging, probe_name)
    with open(probe_path, "w", encoding="utf-8") as fh:
        fh.write("nexus push elevation verify\n")

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

    async def _verify_one(server: Any) -> None:
        async with sem:
            dest = resolve_push_target_path(server.target_path)
            row: dict[str, Any] = {
                "server_id": server.id,
                "server_name": server.name,
                "ssh_user": server.username,
                "target_path": dest,
            }

            row["rsync_sudo"] = await probe_rsync_sudo(server)
            write = await _diagnose_write(server, dest)
            row["write_check"] = write

            preview = await _rsync_push(
                server, staging, dest, sync_mode="incremental", dry_run=True
            )
            row["dry_run_exit"] = preview.get("exit_code")
            row["dry_run_ok"] = preview.get("exit_code") in (0, 23)
            if not row["dry_run_ok"]:
                row["dry_run_stderr"] = (preview.get("stderr") or "")[:300]

            if args.push and row["dry_run_ok"]:
                push = await _rsync_push(server, staging, dest, sync_mode="incremental")
                row["push_exit"] = push.get("exit_code")
                row["push_ok"] = push.get("exit_code") == 0
                row["push_elevated"] = bool(push.get("elevated"))
                if not row["push_ok"]:
                    row["push_stderr"] = (push.get("stderr") or "")[:300]

            ok = (
                row.get("rsync_sudo")
                and write.get("path_writable")
                and row.get("dry_run_ok")
                and (not args.push or row.get("push_ok"))
            )
            row["ok"] = ok
            results.append(row)

            status = "OK" if ok else "FAIL"
            print(
                f"[{status}] #{server.id} {server.name} user={server.username} dest={dest}\n"
                f"  rsync_sudo={row['rsync_sudo']} write={write.get('path_writable')} "
                f"elevated={write.get('elevated')} dry_run={row.get('dry_run_exit')}"
                + (
                    f" push={row.get('push_exit')} push_elevated={row.get('push_elevated')}"
                    if args.push
                    else ""
                ),
                flush=True,
            )
            if row.get("dry_run_stderr"):
                print(f"  dry_run_err: {row['dry_run_stderr']}", flush=True)
            if row.get("push_stderr"):
                print(f"  push_err: {row['push_stderr']}", flush=True)

    try:
        await asyncio.gather(*[_verify_one(s) for s in servers])
    finally:
        import shutil

        shutil.rmtree(staging, ignore_errors=True)
        await close_redis()

    ok_n = sum(1 for r in results if r.get("ok"))
    print(f"\nVerify: {ok_n}/{len(results)} passed")
    return results


def main() -> int:
    _load_dotenv()
    parser = argparse.ArgumentParser(description="Verify push elevation on servers")
    parser.add_argument("--ids", type=str, help="Comma-separated server ids")
    parser.add_argument("--name", type=str, help="Filter name/domain")
    parser.add_argument("--non-root-only", action="store_true", help="All servers except username=root")
    parser.add_argument("--concurrency", type=int, default=8, help="Parallel servers (default 8)")
    parser.add_argument("--push", action="store_true", help="Run real probe file push after dry-run")
    args = parser.parse_args()
    if not args.ids and not args.name and not args.non_root_only:
        parser.error("Specify --ids, --name, or --non-root-only")
    results = asyncio.run(run_verify(args))
    return 0 if results and all(r.get("ok") for r in results) else 1


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