"""Long-running script jobs — Redis registration and nohup wrapper with completion callback.""" from __future__ import annotations import json import logging import secrets from datetime import datetime, timezone from typing import Any, Optional, Tuple from server.config import settings from server.infrastructure.redis.client import get_redis logger = logging.getLogger("nexus.script_jobs") REDIS_JOB_PREFIX = "script_job:" REDIS_JOB_TTL = 604800 # 7 days def master_callback_url() -> str: """Public Nexus URL for Agent curl callbacks.""" base = (settings.API_BASE_URL or "").strip().rstrip("/") if not base: raise ValueError( "NEXUS_API_BASE_URL 未配置,无法注册长任务回调;请在 .env 或安装向导中设置主站对外 URL" ) lower = base.lower() if lower.startswith("http://"): host = lower[7:].split("/")[0].split(":")[0] if host not in ("localhost", "127.0.0.1", "::1"): raise ValueError( "生产环境 NEXUS_API_BASE_URL 必须使用 HTTPS;" "本地开发可使用 http://localhost" ) return f"{base}/api/agent/script-callback" def build_long_task_command( user_command: str, job_id: str, server_id: int, secret: str, callback_url: str, ) -> str: """Wrap user script in nohup; on exit POST completion to Nexus.""" safe_script = user_command.replace("'", "'\\''") # secret/job_id are url-safe tokens from secrets.token_urlsafe curl_payload = ( f'{{"job_id":"{job_id}","server_id":{server_id},' f'"secret":"{secret}","exit_code":'"'"'$ec'"'"',"message":"done"}}' ) inner = ( f"{safe_script}\n" f"ec=$?\n" f"curl -fsS -X POST '{callback_url}' " f"-H 'Content-Type: application/json' " f"-d '{curl_payload}' " f"|| true\n" f"exit $ec" ) return ( "mkdir -p /var/log/nexus-job && " 'LOG=/var/log/nexus-job/job-$(date +%Y%m%d-%H%M%S).log && ' f"nohup bash -c '{inner}' >> \"$LOG\" 2>&1 & " 'echo $! > "${LOG}.pid" && ' f'echo "started job_id={job_id} pid=$! log=$LOG"' ) def sanitize_long_task_command_for_audit(wrapped: str, secret: str) -> str: """Redact job secret before persisting execution.command.""" return wrapped.replace(secret, "***") async def register_script_job(execution_id: int, server_id: int) -> Tuple[str, str]: """Create a pending job; returns (job_id, secret).""" job_id = secrets.token_urlsafe(16) secret = secrets.token_urlsafe(32) payload = { "execution_id": execution_id, "server_id": server_id, "secret": secret, "created_at": datetime.now(timezone.utc).isoformat(), } redis = get_redis() await redis.set( f"{REDIS_JOB_PREFIX}{job_id}", json.dumps(payload), ex=REDIS_JOB_TTL, ) return job_id, secret async def revoke_script_job(job_id: str) -> None: """Remove pending callback registration (e.g. user cancelled the job).""" if not job_id: return redis = get_redis() await redis.delete(f"{REDIS_JOB_PREFIX}{job_id}") async def consume_script_job( job_id: str, server_id: int, secret: str ) -> Optional[dict[str, Any]]: """Validate callback credentials and remove job (one-time use).""" redis = get_redis() key = f"{REDIS_JOB_PREFIX}{job_id}" raw = await redis.get(key) if not raw: return None try: data = json.loads(raw) except json.JSONDecodeError: await redis.delete(key) return None if data.get("server_id") != server_id: return None if not secrets.compare_digest(data.get("secret", ""), secret): return None await redis.delete(key) return data