Nexus/server/application/services/script_jobs.py

"""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