Your Name
32feb1b6db
Fixes: - F821: install.py site_url 未定义(NameError)、sync_v2.py os 未导入、script_jobs.py LOG f-string - F401: 移除 22 个未使用导入（models/__init__.py、install.py、auth.py 等） - B904: 20 个 except 块 raise 添加 from e/from None - S110: 20 个有意的 try-except-pass 添加 noqa 注释（SSH清理/DDL幂等/WS断连） - B007: 3 个未使用循环变量重命名（dirs→_dirs、server_id→_server_id） - F541: 3 个无占位符 f-string 修正 - F841: auth.py 未使用 ip_address 变量移除 - S105: auth_service.py Redis key prefix 误报 noqa - B023: script_execution_flush.py 循环变量绑定修复 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
122 lines
3.7 KiB
Python
122 lines
3.7 KiB
Python
"""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
|