"""Background server batch jobs (category, health check, Agent ops, path detect).""" from __future__ import annotations import asyncio import logging import shlex from typing import Any, Optional from server.application.server_batch_common import ( batch_server_display_name, ensure_agent_api_key, install_error_msg, mark_agent_uninstalled, posix_dirname, prefetch_batch_context, result_item_dict, sudo_wrap, update_server_target_path, ) from server.config import settings from server.domain.models import AuditLog, Server from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl from server.infrastructure.database.server_repo import ServerRepositoryImpl from server.infrastructure.database.session import AsyncSessionLocal from server.infrastructure.redis import server_batch_store as sbs from server.infrastructure.redis.client import get_redis from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command logger = logging.getLogger("nexus.server_batch_service") _BG_TASKS: set[asyncio.Task] = set() CONCURRENCY = 5 OP_LABELS = { "category": "批量修改分类", "health-check": "健康检查", "detect-path": "自动检测目标路径", "install-agent": "批量安装 Agent", "upgrade-agent": "批量升级 Agent", "uninstall-agent": "批量卸载 Agent", } AUDIT_ACTIONS = { "category": "batch_update_category", "health-check": "batch_health_check", "detect-path": "batch_detect_path", "install-agent": "batch_install_agent", "upgrade-agent": "batch_upgrade_agent", "uninstall-agent": "batch_uninstall_agent", } HEARTBEAT_PREFIX = "heartbeat:" async def _notify_progress(live: dict[str, Any]) -> None: try: from server.api.websocket import broadcast_script_progress stats = sbs.progress_stats(live) await broadcast_script_progress({ "execution_id": int(live["id"]), "task_kind": "server_batch", "op": live.get("op") or "", "label": live.get("label") or "", "status": live.get("status") or "running", "long_task": False, "operator": live.get("operator") or "", **stats, }) except Exception as e: logger.warning("server batch progress notify failed: %s", e) async def _notify_complete(live: dict[str, Any]) -> None: try: from server.api.websocket import broadcast_script_complete stats = sbs.progress_stats(live) await broadcast_script_complete({ "execution_id": int(live["id"]), "task_kind": "server_batch", "op": live.get("op") or "", "label": live.get("label") or "", "status": live.get("status") or "completed", "long_task": False, "operator": live.get("operator") or "", **stats, }) except Exception as e: logger.warning("server batch complete notify failed: %s", e) async def start_batch_job( *, op: str, server_ids: list[int], operator: str, params: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: if op not in OP_LABELS: raise ValueError(f"Unknown batch op: {op}") if not server_ids: raise ValueError("server_ids 不能为空") job_id = await sbs.next_job_id() label = OP_LABELS[op] if op == "category" and params and params.get("category"): label = f"{label} → {params['category']}" live = sbs.init_live( job_id=job_id, op=op, label=label, server_ids=server_ids, operator=operator, params=params, ) await sbs.save_live(live) await _persist_job_start(live) await _notify_progress(live) task = asyncio.create_task(_run_background(job_id)) _BG_TASKS.add(task) task.add_done_callback(_BG_TASKS.discard) stats = sbs.progress_stats(live) return { "job_id": job_id, "status": "running", "label": label, "op": op, **stats, } async def get_batch_job(job_id: int) -> Optional[dict[str, Any]]: live = await sbs.get_live(job_id) if live: return sbs.live_to_api(live) return await _get_batch_job_from_mysql(job_id) async def list_batch_jobs( *, limit: int = 50, offset: int = 0, op: Optional[str] = None, ) -> dict[str, Any]: from server.infrastructure.database.server_batch_job_repo import ( ServerBatchJobRepositoryImpl, job_row_to_api, ) async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) rows, total = await repo.list_recent(limit=limit, offset=offset, op=op) items = [job_row_to_api(row) for row in rows] return {"items": items, "total": total, "limit": limit, "offset": offset} async def _persist_job_start(live: dict[str, Any]) -> None: from server.infrastructure.database.server_batch_job_repo import ServerBatchJobRepositoryImpl try: async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) await repo.upsert_running( job_id=int(live["id"]), op=str(live.get("op") or ""), label=str(live.get("label") or ""), server_ids=list(live.get("server_ids") or []), operator=str(live.get("operator") or ""), params=live.get("params") if isinstance(live.get("params"), dict) else {}, ) except Exception as e: logger.warning("Server batch MySQL create failed job_id=%s: %s", live.get("id"), e) async def _persist_job_finalize(live: dict[str, Any]) -> None: from server.infrastructure.database.server_batch_job_repo import ServerBatchJobRepositoryImpl try: api = sbs.live_to_api(live) async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) await repo.finalize( int(live["id"]), status=str(live.get("status") or "completed"), results=list(api.get("results") or []), ) except Exception as e: logger.warning("Server batch MySQL finalize failed job_id=%s: %s", live.get("id"), e) async def _get_batch_job_from_mysql(job_id: int) -> Optional[dict[str, Any]]: from server.infrastructure.database.server_batch_job_repo import ( ServerBatchJobRepositoryImpl, job_row_to_api, ) try: async with AsyncSessionLocal() as session: repo = ServerBatchJobRepositoryImpl(session) row = await repo.get_by_id(job_id) if not row: return None return job_row_to_api(row) except Exception as e: logger.warning("Server batch MySQL read failed job_id=%s: %s", job_id, e) return None async def _run_background(job_id: int) -> None: live = await sbs.get_live(job_id) if not live: return op = live.get("op") or "" try: if op == "category": await _run_category(live) elif op == "health-check": await _run_health_check(live) elif op == "detect-path": await _run_detect_path(live) elif op == "install-agent": await _run_install_agent(live) elif op == "upgrade-agent": await _run_upgrade_agent(live) elif op == "uninstall-agent": await _run_uninstall_agent(live) else: live["status"] = "failed" await sbs.save_live(live) return live = await sbs.get_live(job_id) if live and live.get("status") == "running": sbs.finalize_status(live) await sbs.save_live(live) except Exception: logger.exception("Server batch job failed job_id=%s op=%s", job_id, op) live = await sbs.get_live(job_id) if live: live["status"] = "failed" await sbs.save_live(live) live = await sbs.get_live(job_id) if live: await _persist_job_finalize(live) await _write_audit(live) await _notify_complete(live) async def _write_audit(live: dict[str, Any]) -> None: import json op = live.get("op") or "" action = AUDIT_ACTIONS.get(op, "batch_server_op") stats = sbs.progress_stats(live) api = sbs.live_to_api(live) failures = [ {"server_name": (e.get("server_name") or "").strip() or "未知服务器", "error": (e.get("error") or "")[:300]} for e in (api.get("results") or []) if not e.get("success") ] summary_map = { "category": f"批量修改分类 更新{stats['success']}台", "health-check": f"批量健康检查: {stats['success']}/{stats['total']} 在线", "detect-path": f"批量检测路径: {stats['success']}/{stats['total']} 成功", "install-agent": f"批量安装Agent: {stats['success']}/{stats['total']} 成功", "upgrade-agent": f"批量升级Agent: {stats['success']}/{stats['total']} 成功", "uninstall-agent": f"批量卸载Agent: {stats['success']}/{stats['total']} 成功", } summary = summary_map.get(op, f"{live.get('label')} {stats['success']}/{stats['total']}") detail = json.dumps({ "job_id": int(live["id"]), "summary": summary, "success": stats["success"], "failed": stats["failed"], "total": stats["total"], "failures": failures[:30], }, ensure_ascii=False) try: async with AsyncSessionLocal() as session: audit_repo = AuditLogRepositoryImpl(session) await audit_repo.create(AuditLog( admin_username=live.get("operator") or "admin", action=action, target_type="server_batch_job", target_id=int(live["id"]), detail=detail, ip_address="", )) await session.commit() except Exception as e: logger.warning("Server batch audit failed job_id=%s: %s", live.get("id"), e) async def _finish_one(live: dict[str, Any], item: dict[str, Any]) -> None: sbs.record_result(live, item) await sbs.save_live(live) await _notify_progress(live) async def _run_with_semaphore(live: dict[str, Any], handler) -> None: from server.application.services.server_service import ServerService from server.infrastructure.database.server_repo import ServerRepositoryImpl from server.infrastructure.database.sync_log_repo import SyncLogRepositoryImpl server_ids = live.get("server_ids") or [] sem = asyncio.Semaphore(CONCURRENCY) async with AsyncSessionLocal() as session: service = ServerService( server_repo=ServerRepositoryImpl(session), sync_log_repo=SyncLogRepositoryImpl(session), ) server_map, labels = await prefetch_batch_context(service, server_ids) async def _one(sid: int) -> None: async with sem: item = await handler(sid, server_map, labels, live) await _finish_one(live, item) await asyncio.gather(*[_one(sid) for sid in server_ids]) async def _run_category(live: dict[str, Any]) -> None: params = live.get("params") or {} category = (params.get("category") or "").strip() or None server_ids = live.get("server_ids") or [] async with AsyncSessionLocal() as session: repo = ServerRepositoryImpl(session) servers = await repo.get_by_ids(server_ids) found_ids = {s.id for s in servers} labels = {s.id: batch_server_display_name(s) for s in servers} for sid in server_ids: labels.setdefault(sid, "未知服务器") for sid in server_ids: if sid not in found_ids: await _finish_one(live, result_item_dict( server_id=sid, server_name=labels[sid], success=False, error="服务器不存在", )) if found_ids: await repo.update_category_batch(list(found_ids), category) await session.commit() cat_label = category or "(清除)" for sid in found_ids: await _finish_one(live, result_item_dict( server_id=sid, server_name=labels[sid], success=True, stdout=f"分类 → {cat_label}", )) async def _run_health_check(live: dict[str, Any]) -> None: from server.infrastructure.ssh.remote_probe import ssh_health_probe async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: probe = await ssh_health_probe(server, timeout=15) online = probe.get("status") == "online" detail = probe.get("error") or probe.get("channel") or probe.get("status") or "" return result_item_dict( server_id=sid, server_name=label, success=online, stdout=str(detail)[:500] if online else "", error="" if online else str(probe.get("error") or "离线")[:300], ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler) async def _run_detect_path(live: dict[str, Any]) -> None: async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: ssh_user = (server.username or "root").strip() cmd = "find /www/wwwroot -iname 'workerman.bat' -type f -print -quit 2>/dev/null" cmd = sudo_wrap(cmd, ssh_user) r = await exec_ssh_command(server, cmd, timeout=30) if r["exit_code"] != 0 and not r.get("stdout", "").strip(): return result_item_dict( server_id=sid, server_name=label, success=False, error=f"搜索失败 (exit {r['exit_code']})", ) found = r["stdout"].strip() if not found: return result_item_dict(server_id=sid, server_name=label, success=False, error="未找到 workerman.bat") first_match = found.split("\n")[0].strip() target_dir = posix_dirname(first_match) await update_server_target_path(sid, target_dir) return result_item_dict( server_id=sid, server_name=label, success=True, stdout=f"target_path → {target_dir}", ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler) async def _run_install_agent(live: dict[str, Any]) -> None: base_url = (settings.API_BASE_URL or "").strip().rstrip("/") if not base_url: for sid in live.get("server_ids") or []: await _finish_one(live, result_item_dict( server_id=sid, server_name=str(sid), success=False, error="NEXUS_API_BASE_URL 未配置", )) return async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: api_key = await ensure_agent_api_key(server) ssh_user = (server.username or "root").strip() install_cmd = ( f"TMP=$(mktemp) && curl -fsSL {shlex.quote(base_url)}/agent/install.sh -o \"$TMP\" && bash \"$TMP\" -- " f"--url {shlex.quote(base_url)} --key {shlex.quote(api_key)} " f"--id {sid} --port {shlex.quote(str(server.agent_port or 8601))}" f"; RC=$?; rm -f \"$TMP\"; exit $RC" ) install_cmd = sudo_wrap(install_cmd, ssh_user) r = await exec_ssh_command(server, install_cmd, timeout=180) ok = r["exit_code"] == 0 if ok and "Status: FAILED" in r.get("stdout", ""): ok = False return result_item_dict( server_id=sid, server_name=label, success=ok, stdout=r["stdout"][:2000] if ok else r["stdout"][:1000], error="" if ok else install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]), ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler) async def _run_upgrade_agent(live: dict[str, Any]) -> None: base_url = (settings.API_BASE_URL or "").strip().rstrip("/") if not base_url: for sid in live.get("server_ids") or []: await _finish_one(live, result_item_dict( server_id=sid, server_name=str(sid), success=False, error="NEXUS_API_BASE_URL 未配置", )) return async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: agent_url = f"{base_url}/agent/agent.py" install_dir = "/opt/nexus-agent" venv_python = f"{install_dir}/.venv/bin/python" ssh_user = (server.username or "root").strip() batch_systemctl_prefix = "sudo " if ssh_user != "root" else "" batch_kill_port = "fuser -k 8601/tcp 2>/dev/null || true; " upgrade_cmd = sudo_wrap( f"(command -v rsync >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y -qq rsync) || (yum install -y -q rsync) || (dnf install -y -q rsync)) " f"&& {venv_python} -c 'import sys; v=sys.version_info; sys.exit(0 if v.major==3 and v.minor>=10 else 1)' " f"&& {venv_python} -m pip install -q fastapi==0.115.6 uvicorn==0.34.0 httpx==0.28.1 psutil python-multipart==0.0.19 " f"&& cp {install_dir}/agent.py {install_dir}/agent.py.bak " f"&& curl -fsSL {shlex.quote(agent_url)} -o {install_dir}/agent.py " f"&& {batch_kill_port}{batch_systemctl_prefix}systemctl restart nexus-agent " f"&& sleep 2 " f"&& {batch_systemctl_prefix}systemctl is-active nexus-agent " f"&& echo 'upgrade_ok'", ssh_user, ) r = await exec_ssh_command(server, upgrade_cmd, timeout=120) ok = r["exit_code"] == 0 and "upgrade_ok" in r["stdout"] if not ok: try: rollback = sudo_wrap( f"cp {install_dir}/agent.py.bak {install_dir}/agent.py && {batch_systemctl_prefix}systemctl restart nexus-agent", ssh_user, ) await exec_ssh_command(server, rollback, timeout=30) except Exception as rb_err: logger.warning("Rollback failed for server %s: %s", sid, rb_err) return result_item_dict( server_id=sid, server_name=label, success=ok, stdout=r["stdout"][:500] if ok else "", error=install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]) if not ok else "", ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:200]) await _run_with_semaphore(live, handler) async def _run_uninstall_agent(live: dict[str, Any]) -> None: base_url = (settings.API_BASE_URL or "").strip().rstrip("/") if not base_url: for sid in live.get("server_ids") or []: await _finish_one(live, result_item_dict( server_id=sid, server_name=str(sid), success=False, error="NEXUS_API_BASE_URL 未配置", )) return async def handler(sid: int, server_map: dict[int, Server], labels: dict[int, str], _live) -> dict: server = server_map.get(sid) label = labels[sid] if not server: return result_item_dict(server_id=sid, server_name=label, success=False, error="服务器不存在") try: ssh_user = (server.username or "root").strip() uninstall_cmd = ( f"TMP=$(mktemp) && curl -fsSL {shlex.quote(base_url)}/agent/uninstall.sh -o \"$TMP\" " f"&& bash \"$TMP\" ; RC=$?; rm -f \"$TMP\"; exit $RC" ) uninstall_cmd = sudo_wrap(uninstall_cmd, ssh_user) r = await exec_ssh_command(server, uninstall_cmd, timeout=60) ok = r["exit_code"] == 0 if ok: try: redis = get_redis() await redis.delete(f"{HEARTBEAT_PREFIX}{sid}") except Exception as redis_err: logger.warning("Failed to clear Redis heartbeat for server %s: %s", sid, redis_err) await mark_agent_uninstalled(sid) return result_item_dict( server_id=sid, server_name=label, success=ok, stdout=r["stdout"][:500] if ok else "", error="" if ok else install_error_msg(r.get("stdout", ""), r.get("stderr", ""), r["exit_code"]), ) except Exception as e: return result_item_dict(server_id=sid, server_name=label, success=False, error=str(e)[:300]) await _run_with_semaphore(live, handler)