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()


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__ = [
    "batch_server_display_name",
    "ensure_agent_api_key",
    "generate_agent_api_key_value",
    "install_error_msg",
    "mark_agent_uninstalled",
    "posix_dirname",
    "prefetch_batch_context",
    "result_item_dict",
    "sudo_wrap",
    "update_server_target_path",
]