Nexus/server/application/server_batch_common.py

"""Shared helpers for server batch SSH operations (Agent install, path detect, etc.)."""

from __future__ import annotations

import logging
import shlex

from server.domain.models import Server
from server.utils.posix_paths import posix_dirname

logger = logging.getLogger("nexus.server_batch")

_INSTALL_ERROR_ZH: dict[str, str] = {
    "Python 3.10+ required but not found and could not be installed automatically.":
        "Python 版本低于 3.10，且无法自动安装",
    "Python 3.10+ required. System install failed — install python3.12 manually and re-run.":
        "Python 3.10+ 不可用且自动安装系统 python3.12 失败，请手动安装后重试",
    "Python 3.10+ required (found only older pyenv/BT Python?). Install system python3.12 or use BT panel pyenv under /www/server/panel/pyenv/bin.":
        "Python 版本低于 3.10；宝塔 pyenv 不可执行时将自动安装系统 python3.12",
    "sudo requires a password. Configure passwordless sudo first:":
        "sudo 需要密码，请先配置免密 sudo",
    "Running as non-root and sudo is not available.":
        "非 root 用户且无 sudo 权限",
    "--url is required":
        "缺少 --url 参数",
    "--id is required (server ID from Nexus dashboard)":
        "缺少 --id 参数（服务器 ID）",
    "--key is required (API key from Nexus settings)":
        "缺少 --key 参数（API Key）",
    "venv creation failed even after installing python3-venv":
        "虚拟环境创建失败（python3-venv 模块不可用）",
    "Cannot download from":
        "无法下载 agent.py",
    "upgrade_ok":
        "升级验证失败",
}


def sudo_wrap(cmd: str, ssh_user: str) -> str:
    """Wrap a command with sudo setup/teardown for non-root SSH users."""
    if ssh_user == "root":
        return cmd
    sudoers_tag = "nexus-agent"
    sudoers_line = (
        f"{ssh_user} ALL=(ALL) NOPASSWD: "
        "/usr/bin/systemctl, "
        "/usr/bin/apt-get, /usr/bin/yum, /usr/bin/dnf, /usr/bin/apk, "
        "/usr/bin/rm, /usr/bin/mkdir, /usr/bin/cp, /usr/bin/tee, "
        "/usr/bin/curl, /usr/bin/fuser, /usr/bin/kill, /usr/bin/pkill, "
        f"{shlex.quote('/opt/nexus-agent/.venv/bin/pip')}*"
    )
    setup = (
        f"echo {shlex.quote(sudoers_line)} "
        f"| sudo tee /etc/sudoers.d/{sudoers_tag} > /dev/null 2>&1; "
    )
    cleanup = f"; sudo rm -f /etc/sudoers.d/{sudoers_tag} 2>/dev/null"
    return setup + cmd + cleanup


def install_error_msg(stdout: str, stderr: str, exit_code: int) -> str:
    """Extract concise error reason from install.sh / upgrade output."""
    for line in reversed(stdout.split("\n")):
        line = line.strip()
        if line.startswith("ERROR:"):
            en = line.replace("ERROR: ", "")
            zh = _INSTALL_ERROR_ZH.get(en)
            if not zh:
                for prefix, translation in _INSTALL_ERROR_ZH.items():
                    if en.startswith(prefix):
                        zh = translation
                        break
            if zh:
                return f"{zh}（{en[:200]}）"
            return en[:300]
    for line in reversed(stdout.split("\n")):
        line = line.strip()
        if line:
            return line[:300]
    if stderr.strip():
        return stderr.strip()[:300]
    return f"exit {exit_code}"


def batch_server_display_name(server: Server | None) -> str:
    if server is None:
        return "未知服务器"
    name = (server.name or "").strip()
    if name:
        return name
    domain = (server.domain or "").strip()
    if domain:
        return domain
    return "未知服务器"


async def prefetch_batch_context(service, server_ids: list[int]) -> tuple[dict[int, Server], dict[int, str]]:
    rows = await service.server_repo.get_by_ids(server_ids)
    server_map = {s.id: s for s in rows}
    labels = {sid: batch_server_display_name(server_map.get(sid)) for sid in server_ids}
    return server_map, labels


def generate_agent_api_key_value() -> str:
    import secrets

    return f"nxs-{secrets.token_urlsafe(32)}"


async def ensure_agent_api_key(server: Server) -> str:
    """Return per-server agent key, auto-generating and persisting when missing."""
    key = (server.agent_api_key or "").strip()
    if key:
        return key

    new_key = generate_agent_api_key_value()
    from server.infrastructure.database.server_repo import ServerRepositoryImpl
    from server.infrastructure.database.session import AsyncSessionLocal

    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        row = await repo.get_by_id(server.id)
        if not row:
            raise ValueError("服务器不存在")
        existing = (row.agent_api_key or "").strip()
        if existing:
            server.agent_api_key = existing
            return existing
        row.agent_api_key = new_key
        await session.commit()
        server.agent_api_key = new_key
    return new_key


async def update_server_target_path(sid: int, target_dir: str) -> None:
    from server.infrastructure.database.session import AsyncSessionLocal
    from server.infrastructure.database.server_repo import ServerRepositoryImpl

    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        server = await repo.get_by_id(sid)
        if not server:
            raise ValueError("服务器不存在")
        server.target_path = target_dir
        await session.commit()


async def mark_agent_uninstalled(sid: int) -> None:
    from server.infrastructure.database.session import AsyncSessionLocal
    from server.infrastructure.database.server_repo import ServerRepositoryImpl

    async with AsyncSessionLocal() as session:
        repo = ServerRepositoryImpl(session)
        server = await repo.get_by_id(sid)
        if not server:
            return
        server.is_online = False
        server.agent_version = None
        server.agent_api_key = None
        await session.commit()


_REMOTE_AGENT_VERSION_CMD = (
    "if [ -f /opt/nexus-agent/agent.py ]; then "
    "grep -m1 '^AGENT_VERSION' /opt/nexus-agent/agent.py 2>/dev/null | "
    "sed -n 's/.*AGENT_VERSION *= *\"\\([^\"]*\\)\".*/\\1/p'; fi"
)


def agent_files_curl_cmds(base_url: str, install_dir: str = "/opt/nexus-agent") -> str:
    """Download agent.py and heartbeat_policy.py from central /agent/."""
    root = base_url.rstrip("/")
    return (
        f"curl -fsSL {shlex.quote(root + '/agent/agent.py')} -o {install_dir}/agent.py "
        f"&& curl -fsSL {shlex.quote(root + '/agent/heartbeat_policy.py')} "
        f"-o {install_dir}/heartbeat_policy.py"
    )


async def _redis_agent_version(server_id: int) -> str:
    from server.infrastructure.redis.client import get_redis

    try:
        redis = get_redis()
        data = await redis.hgetall(f"heartbeat:{server_id}")
        if data:
            return (data.get("agent_version") or "").strip()
    except Exception as e:
        logger.warning("Redis agent_version read failed server_id=%s: %s", server_id, e)
    return ""


async def resolve_remote_agent_version(server: Server, *, ssh_user: str) -> str:
    """Best-effort remote version: Redis heartbeat → MySQL → SSH grep agent.py."""
    from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command

    for candidate in (
        await _redis_agent_version(server.id),
        (server.agent_version or "").strip(),
    ):
        if candidate:
            return candidate
    try:
        cmd = sudo_wrap(_REMOTE_AGENT_VERSION_CMD, ssh_user)
        r = await exec_ssh_command(server, cmd, timeout=15)
        if r["exit_code"] == 0:
            return (r.get("stdout") or "").strip().split("\n")[0].strip()
    except Exception as e:
        logger.warning("SSH agent_version probe failed server_id=%s: %s", server.id, e)
    return ""


async def agent_install_skip_result(
    *,
    server: Server,
    server_id: int,
    server_name: str,
    central_version: str,
) -> dict | None:
    """Return success result dict when install can be skipped; None to proceed with install."""
    from server.utils.agent_version import is_version_at_least

    has_key = bool((server.agent_api_key or "").strip())
    db_version = (server.agent_version or "").strip()
    if not has_key and not db_version:
        return None

    ssh_user = (server.username or "root").strip()
    remote_version = await resolve_remote_agent_version(server, ssh_user=ssh_user)
    if remote_version and is_version_at_least(remote_version, central_version):
        return result_item_dict(
            server_id=server_id,
            server_name=server_name,
            success=True,
            stdout=(
                f"已安装 Agent {remote_version}，不低于主站 {central_version}，已跳过"
            ),
        )
    return None


def result_item_dict(
    *,
    server_id: int,
    server_name: str,
    success: bool,
    stdout: str = "",
    error: str = "",
) -> dict:
    return {
        "server_id": server_id,
        "server_name": server_name,
        "success": success,
        "stdout": stdout,
        "error": error,
    }


__all__ = [
    "agent_files_curl_cmds",
    "agent_install_skip_result",
    "batch_server_display_name",
    "ensure_agent_api_key",
    "generate_agent_api_key_value",
    "install_error_msg",
    "mark_agent_uninstalled",
    "posix_dirname",
    "prefetch_batch_context",
    "resolve_remote_agent_version",
    "result_item_dict",
    "sudo_wrap",
    "update_server_target_path",
]