Nexus/server/application/services/script_service.py

"""Nexus — Script Service (Business logic for script library + remote execution)
Application layer — orchestrates Repository + SSH exec + DB credential substitution.
"""

import asyncio
import json
import logging
import re
import shlex
from datetime import datetime, timezone
from typing import Optional, List, Dict, Any

from server.config import settings
from server.domain.models import Server, Script, ScriptExecution, DbCredential, AuditLog
from server.infrastructure.redis import script_execution_store as ses
from server.domain.repositories import (
    ScriptRepository, ScriptExecutionRepository,
    DbCredentialRepository, ServerRepository, AuditLogRepository,
)

logger = logging.getLogger("nexus.script_service")

_NEXUS_JOB_LOG_RE = re.compile(r"log=(/var/log/nexus-job/[^\s\"']+)")
_NEXUS_JOB_PID_RE = re.compile(r"pid=(\d+)")


class ScriptService:
    """Business logic for script CRUD and remote command execution via SSH (port 22)"""

    def __init__(
        self,
        script_repo: ScriptRepository,
        execution_repo: ScriptExecutionRepository,
        credential_repo: DbCredentialRepository,
        server_repo: ServerRepository,
        audit_repo: AuditLogRepository,
    ):
        self.script_repo = script_repo
        self.execution_repo = execution_repo
        self.credential_repo = credential_repo
        self.server_repo = server_repo
        self.audit_repo = audit_repo

    async def _agent_api_key_for_server(self, server_id: int) -> str:
        server = await self.server_repo.get_by_id(server_id)
        if not server or not server.agent_api_key:
            raise ValueError(
                f"服务器 {server_id} 未设置 Agent API Key，无法注册长任务回调，请先生成"
            )
        return server.agent_api_key

    # ── Script CRUD ──

    async def list_scripts(self, category: Optional[str] = None) -> List[Script]:
        if category:
            return await self.script_repo.get_by_category(category)
        return await self.script_repo.get_all()

    async def get_script(self, id: int) -> Optional[Script]:
        return await self.script_repo.get_by_id(id)

    async def create_script(self, script: Script) -> Script:
        created = await self.script_repo.create(script)
        await self._audit("create_script", "script", created.id, f"名称={script.name}")
        return created

    async def update_script(self, script: Script) -> Script:
        updated = await self.script_repo.update(script)
        await self._audit("update_script", "script", updated.id, f"名称={script.name}")
        return updated

    async def delete_script(self, id: int) -> bool:
        result = await self.script_repo.delete(id)
        if result:
            await self._audit("delete_script", "script", id, f"脚本ID={id}")
        return result

    # ── Command Execution (SSH) ──

    async def execute_command(
        self,
        command: str,
        server_ids: List[int],
        script_id: Optional[int] = None,
        credential_id: Optional[int] = None,
        timeout: int = 30,
        operator: str = "admin",
        long_task: bool = False,
    ) -> ScriptExecution:
        """Execute a shell command on multiple servers via SSH.

        If credential_id is provided, $DB_* variables are substituted on Nexus
        before the command is sent over SSH.
        """
        # Resolve DB credential variables (if provided)
        resolved_command = await self._substitute_db_vars(command, credential_id)

        from server.application.services.script_jobs import (
            build_long_task_command,
            master_callback_url,
            register_script_job,
        )

        callback_url = None
        if long_task:
            callback_url = master_callback_url()

        # Create execution record — use sanitized command (password redacted)
        sanitized_command = await self._substitute_db_vars(command, credential_id, redact=True)
        if long_task:
            sanitized_command = f"[long_task] {sanitized_command}"
        execution = ScriptExecution(
            script_id=script_id,
            command=sanitized_command,
            server_ids=json.dumps(server_ids),
            credential_id=credential_id,
            status="running",
            operator=operator,
        )
        execution = await self.execution_repo.create(execution)
        await ses.init_live_from_execution(execution, long_task=long_task)
        await self._audit(
            "execute_started", "script_execution", execution.id,
            f"开始执行 {len(server_ids)} 台"
            + (" [long_task]" if long_task else ""),
            operator=operator,
        )

        batch_size = max(1, settings.SCRIPT_EXEC_BATCH_SIZE)
        results: Dict[str, dict] = {}
        batches = [
            server_ids[i : i + batch_size]
            for i in range(0, len(server_ids), batch_size)
        ]
        for batch_idx, batch in enumerate(batches, start=1):
            if len(batches) > 1:
                logger.info(
                    "Script exec batch %s/%s (%s servers, execution_id=%s)",
                    batch_idx,
                    len(batches),
                    len(batch),
                    execution.id,
                )
            per_server_commands: Optional[Dict[int, str]] = None
            job_ids_by_server: Dict[int, str] = {}
            if long_task:
                per_server_commands = {}
                for sid in batch:
                    job_id, secret = await register_script_job(execution.id, sid)
                    job_ids_by_server[sid] = job_id
                    agent_key = await self._agent_api_key_for_server(sid)
                    per_server_commands[sid] = build_long_task_command(
                        resolved_command, job_id, sid, secret, callback_url, agent_key
                    )
            batch_results = await self._dispatch_to_servers(
                batch,
                resolved_command,
                timeout,
                execution_id=execution.id,
                per_server_commands=per_server_commands,
                long_task=long_task,
                job_ids_by_server=job_ids_by_server or None,
            )
            results.update(batch_results)
            await ses.apply_update(
                self.execution_repo,
                execution.id,
                status="running",
                results=results,
                event_action="batch_progress",
                event_detail=f"批次 {batch_idx}/{len(batches)} 完成",
                operator=operator,
                terminal=False,
                long_task=long_task,
            )

        # Determine overall status
        if long_task:
            start_failed = sum(
                1 for r in results.values() if r.get("exit_code", -1) != 0
            )
            status = (
                "failed"
                if start_failed == len(server_ids)
                else "running"
            )
            failed_count = start_failed
        else:
            failed_count = sum(1 for r in results.values() if r.get("exit_code", -1) != 0)
            status = (
                "completed"
                if failed_count == 0
                else "failed"
                if failed_count == len(server_ids)
                else "partial"
            )

        terminal = status in ses.TERMINAL_STATUSES
        await ses.apply_update(
            self.execution_repo,
            execution.id,
            status=status,
            results=results,
            event_action="execution_finished",
            event_detail=f"执行结束: {status} ({failed_count} 台失败)",
            operator=operator,
            terminal=terminal,
            long_task=long_task,
        )

        await self._audit(
            "execute_command", "script_execution", execution.id,
            f"执行于 {len(server_ids)} 台: {status} ({failed_count} 台失败)",
            operator=operator,
        )
        return await self.execution_repo.get_by_id(execution.id)

    async def get_execution(self, id: int) -> Optional[ScriptExecution]:
        return await self.execution_repo.get_by_id(id)

    async def list_executions(
        self,
        limit: int = 50,
        offset: int = 0,
        status: Optional[str] = None,
    ) -> dict:
        """List execution history (MySQL list + Redis overlay when live)."""
        rows, total = await self.execution_repo.list_recent(limit, offset, status)
        items = []
        for ex in rows:
            view = await ses.get_execution_view(self.execution_repo, ex.id)
            if view:
                items.append(view)
            else:
                results, events = ses.unpack_results_blob(ex.results)
                try:
                    sids = json.loads(ex.server_ids or "[]")
                except json.JSONDecodeError:
                    sids = []
                items.append({
                    "id": ex.id,
                    "script_id": ex.script_id,
                    "command": ex.command,
                    "server_ids": ex.server_ids,
                    "status": ex.status,
                    "results": results,
                    "events": events,
                    "operator": ex.operator,
                    "started_at": str(ex.started_at) if ex.started_at else None,
                    "completed_at": str(ex.completed_at) if ex.completed_at else None,
                    "progress": ses.execution_progress_summary(results, sids),
                    "source": "mysql",
                })
        return {"items": items, "total": total, "limit": limit, "offset": offset}

    async def get_execution_detail(
        self, id: int, fetch_logs: bool = False, log_tail_lines: int = 100,
    ) -> Optional[dict]:
        """Return execution dict (Redis first); optionally tail nexus-job logs."""
        view = await ses.get_execution_view(self.execution_repo, id)
        if not view:
            return None
        results: Dict[str, Any] = dict(view.get("results") or {})

        if fetch_logs:
            await self._attach_remote_logs(results, log_tail_lines)
            live = await ses.load_live(id)
            if live:
                live["results"] = results
                await ses.save_live(live)

        return {
            "id": view["id"],
            "script_id": view.get("script_id"),
            "command": view.get("command"),
            "server_ids": view.get("server_ids"),
            "status": view.get("status"),
            "results": results,
            "events": view.get("events") or [],
            "operator": view.get("operator"),
            "started_at": view.get("started_at"),
            "completed_at": view.get("completed_at"),
            "long_task": view.get("long_task", False),
            "progress": view.get("progress"),
            "source": view.get("source"),
        }

    @staticmethod
    def _parse_execution_meta(command: str) -> tuple[str, bool]:
        """Strip [long_task] prefix from stored command."""
        if command.startswith("[long_task] "):
            return command[len("[long_task] "):], True
        return command, False

    @staticmethod
    def _extract_nexus_job_pid(stdout: str) -> Optional[int]:
        m = _NEXUS_JOB_PID_RE.search(stdout or "")
        if not m:
            return None
        try:
            return int(m.group(1))
        except ValueError:
            return None

    @staticmethod
    def _build_kill_command(pid: Optional[int], log_path: Optional[str]) -> str:
        parts = []
        if pid:
            parts.append(f"kill -TERM {pid} 2>/dev/null || kill -KILL {pid} 2>/dev/null")
        if log_path:
            # log_path is validated to /var/log/nexus-job/[^\s"']+ (no spaces/quotes).
            # Use single-quote the path directly — do NOT wrap shlex.quote() inside
            # double-quotes, which would embed the single-quote chars literally and
            # cause the cat to target a non-existent filename like "'path'.pid".
            q = shlex.quote(log_path)  # produces '/var/log/nexus-job/...'
            parts.append(
                f'P=$(cat {q}.pid 2>/dev/null); '
                f'[ -n "$P" ] && (kill -TERM "$P" 2>/dev/null || kill -KILL "$P" 2>/dev/null)'
            )
        if not parts:
            return "echo 'cannot stop: pid unknown'; exit 1"
        return " ; ".join(parts) + '; echo "stopped"'

    async def stop_execution(self, execution_id: int, operator: str = "admin") -> Optional[ScriptExecution]:
        """Stop pending long-task processes on child hosts."""
        from server.application.services.script_jobs import revoke_script_job

        view = await ses.get_execution_view(self.execution_repo, execution_id)
        if not view:
            return None
        results: Dict[str, Any] = dict(view.get("results") or {})

        stopped = 0
        for sid_str, entry in list(results.items()):
            if not isinstance(entry, dict) or entry.get("phase") != "pending":
                continue
            try:
                server_id = int(sid_str)
            except ValueError:
                continue
            server = await self.server_repo.get_by_id(server_id)
            if not server or not server.is_online:
                entry["phase"] = "cancelled"
                entry["status"] = "cancelled"
                entry["exit_code"] = -1
                entry["message"] = "服务器离线，标记为已停止"
                stopped += 1
                continue

            job_id = entry.get("job_id")
            if job_id:
                await revoke_script_job(job_id)

            pid = entry.get("pid") or self._extract_nexus_job_pid(entry.get("stdout", ""))
            log_path = self._extract_nexus_job_log(entry.get("stdout", ""))
            try:
                kill_cmd = self._build_kill_command(
                    int(pid) if pid else None,
                    log_path,
                )
                out = await self._call_ssh_exec(server, kill_cmd, timeout=15)
                entry["stop_stdout"] = out.get("stdout", "") or out.get("status", "")
                entry["stop_stderr"] = out.get("stderr", "")
            except Exception as e:
                entry["stop_stderr"] = str(e)

            entry["phase"] = "cancelled"
            entry["status"] = "cancelled"
            entry["exit_code"] = -1
            entry["message"] = "用户停止"
            entry["completed_at"] = datetime.now(timezone.utc).isoformat()
            stopped += 1

        try:
            server_ids = json.loads(view["server_ids"]) if isinstance(view["server_ids"], str) else view.get("server_ids", [])
        except json.JSONDecodeError:
            server_ids = []

        status = _aggregate_from_results(server_ids, results)
        terminal = status in ses.TERMINAL_STATUSES

        await ses.apply_update(
            self.execution_repo,
            execution_id,
            status=status,
            results=results,
            event_action="stopped",
            event_detail=f"用户停止，共停止 {stopped} 台 pending 任务",
            operator=operator,
            terminal=terminal,
        )
        await self._audit(
            "stop_execution", "script_execution", execution_id,
            f"用户停止，共停止 {stopped} 台 pending 任务",
            operator=operator,
        )
        return await self.execution_repo.get_by_id(execution_id)

    async def mark_execution_stuck(
        self,
        execution_id: int,
        operator: str = "admin",
        detail: str = "",
    ) -> Optional[ScriptExecution]:
        """Record suspected stuck state — persisted to Redis + MySQL + audit."""
        view = await ses.get_execution_view(self.execution_repo, execution_id)
        if not view:
            return None
        results = dict(view.get("results") or {})
        msg = detail.strip() or "用户标记执行可能卡住"
        await ses.apply_update(
            self.execution_repo,
            execution_id,
            status=view.get("status") or "running",
            results=results,
            event_action="stuck_marked",
            event_detail=msg,
            operator=operator,
            terminal=False,
        )
        await self._audit(
            "mark_stuck", "script_execution", execution_id, msg, operator=operator,
        )
        return await self.execution_repo.get_by_id(execution_id)

    async def retry_execution(
        self,
        execution_id: int,
        server_ids: Optional[List[int]] = None,
        credential_id: Optional[int] = None,
        timeout: int = 120,
        long_task: Optional[bool] = None,
        operator: str = "admin",
    ) -> ScriptExecution:
        """Re-run command for failed / pending / cancelled servers."""
        execution = await self.execution_repo.get_by_id(execution_id)
        if not execution:
            raise ValueError("Execution not found")

        command, was_long = self._parse_execution_meta(execution.command or "")
        if not command.strip():
            raise ValueError("原执行命令为空，无法重试")

        try:
            all_ids = json.loads(execution.server_ids or "[]")
        except json.JSONDecodeError:
            all_ids = []

        if server_ids is None:
            view = await ses.get_execution_view(self.execution_repo, execution_id)
            results = dict((view or {}).get("results") or {})
            server_ids = []
            for sid in all_ids:
                r = results.get(str(sid), {})
                phase = r.get("phase", "done")
                if phase == "pending" or r.get("exit_code", -1) != 0:
                    server_ids.append(sid)
            if not server_ids:
                server_ids = list(all_ids)

        if not server_ids:
            raise ValueError("没有可重试的服务器")

        retry_results, _ = ses.unpack_results_blob(execution.results)
        await ses.apply_update(
            self.execution_repo,
            execution_id,
            status=execution.status,
            results=retry_results,
            event_action="retry_requested",
            event_detail=f"重试 {len(server_ids)} 台: {server_ids}",
            operator=operator,
            terminal=False,
        )
        await self._audit(
            "retry_execution", "script_execution", execution_id,
            f"重试服务器 {server_ids}",
            operator=operator,
        )

        use_long = was_long if long_task is None else long_task
        cred_id = credential_id if credential_id is not None else getattr(
            execution, "credential_id", None
        )
        return await self.execute_command(
            command=command,
            server_ids=server_ids,
            script_id=execution.script_id,
            credential_id=cred_id,
            timeout=timeout,
            operator=operator,
            long_task=use_long,
        )

    @staticmethod
    def _extract_nexus_job_log(stdout: str) -> Optional[str]:
        m = _NEXUS_JOB_LOG_RE.search(stdout or "")
        if not m:
            return None
        path = m.group(1)
        if not path.startswith("/var/log/nexus-job/"):
            return None
        return path

    async def _attach_remote_logs(self, results: Dict[str, Any], tail_lines: int) -> None:
        """In-place: add remote_log for servers with a nexus-job log path."""
        tail_lines = max(10, min(tail_lines, 500))
        server_ids: set[int] = set()
        for sid_str, entry in results.items():
            if not isinstance(entry, dict):
                continue
            if not self._extract_nexus_job_log(entry.get("stdout", "")):
                continue
            try:
                server_ids.add(int(sid_str))
            except ValueError:
                continue

        server_map: Dict[int, Server] = {}
        for sid in server_ids:
            server = await self.server_repo.get_by_id(sid)
            if server:
                server_map[sid] = server

        async def _tail_one(sid_str: str, entry: dict) -> None:
            log_path = self._extract_nexus_job_log(entry.get("stdout", ""))
            if not log_path:
                return
            try:
                server_id = int(sid_str)
            except ValueError:
                return
            server = server_map.get(server_id)
            if not server or not server.is_online:
                entry["remote_log"] = "(服务器离线，无法拉取日志)"
                return
            cmd = f"tail -n {tail_lines} {shlex.quote(log_path)} 2>&1"
            try:
                out = await self._call_ssh_exec(server, cmd, timeout=30)
                parts = []
                if out.get("stdout"):
                    parts.append(out["stdout"])
                if out.get("stderr"):
                    parts.append(out["stderr"])
                entry["remote_log"] = "\n".join(parts).strip() or "(日志为空)"
                entry["remote_log_fetched_at"] = datetime.now(timezone.utc).isoformat()
            except Exception as e:
                entry["remote_log"] = f"(拉取失败: {e})"

        await asyncio.gather(
            *[_tail_one(sid, entry) for sid, entry in results.items() if isinstance(entry, dict)]
        )

    async def _dispatch_to_servers(
        self,
        server_ids: List[int],
        resolved_command: str,
        timeout: int,
        execution_id: Optional[int] = None,
        per_server_commands: Optional[Dict[int, str]] = None,
        long_task: bool = False,
        job_ids_by_server: Optional[Dict[int, str]] = None,
    ) -> Dict[str, dict]:
        """Run command on server_ids with concurrency limit (one batch wave)."""
        results: Dict[str, dict] = {}
        semaphore = asyncio.Semaphore(settings.SCRIPT_EXEC_CONCURRENCY)

        server_map: Dict[int, Server] = {}
        for sid in server_ids:
            server = await self.server_repo.get_by_id(sid)
            if server:
                server_map[sid] = server

        async def _exec_on_server(server_id: int):
            server = server_map.get(server_id)
            if not server or not server.is_online:
                results[str(server_id)] = {
                    "status": "skipped",
                    "stdout": "",
                    "stderr": f"Server offline or not found (ID={server_id})",
                    "exit_code": -1,
                }
                return
            cmd = (
                per_server_commands.get(server_id, resolved_command)
                if per_server_commands
                else resolved_command
            )
            async with semaphore:
                try:
                    exec_timeout = min(timeout, 30) if long_task else timeout
                    result = await self._call_ssh_exec(server, cmd, timeout=exec_timeout)
                    if long_task and result.get("exit_code", -1) == 0:
                        pid = result.get("pid") or self._extract_nexus_job_pid(
                            result.get("stdout", ""),
                        )
                        if pid:
                            result["pid"] = int(pid)
                        result = {
                            **result,
                            "phase": "pending",
                            "background": True,
                            "job_id": (job_ids_by_server or {}).get(server_id, ""),
                            "status": "started",
                            "started_at": datetime.now(timezone.utc).isoformat(),
                            "channel": "ssh",
                        }
                    results[str(server_id)] = result
                except Exception as e:
                    results[str(server_id)] = {
                        "status": "error",
                        "stdout": "",
                        "stderr": str(e),
                        "exit_code": -1,
                    }

        await asyncio.gather(*[_exec_on_server(sid) for sid in server_ids])
        return results

    # ── DB Credential Management ──

    async def list_credentials(self) -> List[DbCredential]:
        return await self.credential_repo.get_all()

    async def create_credential(self, credential: DbCredential) -> DbCredential:
        from server.infrastructure.database.crypto import encrypt_value
        credential.encrypted_password = encrypt_value(credential.encrypted_password)
        created = await self.credential_repo.create(credential)
        await self._audit("create_credential", "db_credential", created.id, f"名称={credential.name}")
        return created

    async def delete_credential(self, id: int) -> bool:
        result = await self.credential_repo.delete(id)
        if result:
            await self._audit("delete_credential", "db_credential", id, f"凭据ID={id}")
        return result

    # ── Private helpers ──

    async def _substitute_db_vars(self, command: str, credential_id: Optional[int], redact: bool = False) -> str:
        """Replace $DB_USER/$DB_PASS/$DB_HOST/$DB_PORT/$DB_NAME in command.

        If redact=True, replaces $DB_PASS with '***' instead of the actual password.
        The redacted version is stored in ScriptExecution.command (DB record).
        The actual substitution (redact=False) is sent over SSH for execution.
        """
        if not credential_id:
            return command

        credential = await self.credential_repo.get_by_id(credential_id)
        if not credential:
            return command

        from server.infrastructure.database.crypto import decrypt_value
        password = decrypt_value(credential.encrypted_password)

        # Use unique sentinels to avoid nested substitution if a credential value
        # itself contains another $DB_* placeholder (e.g. password = "$DB_USER").
        sentinels = {
            "$DB_USER": "\x00DB_USER\x00",
            "$DB_PASS": "\x00DB_PASS\x00",
            "$DB_HOST": "\x00DB_HOST\x00",
            "$DB_PORT": "\x00DB_PORT\x00",
            "$DB_NAME": "\x00DB_NAME\x00",
        }
        values = {
            "\x00DB_USER\x00": credential.username or "",
            "\x00DB_PASS\x00": "***" if redact else password,
            "\x00DB_HOST\x00": credential.host or "",
            "\x00DB_PORT\x00": str(credential.port or ""),
            "\x00DB_NAME\x00": credential.database or "",
        }
        for placeholder, sentinel in sentinels.items():
            command = command.replace(placeholder, sentinel)
        for sentinel, value in values.items():
            command = command.replace(sentinel, value)

        return command

    async def _call_ssh_exec(
        self,
        server: Server,
        command: str,
        timeout: int = 30,
    ) -> dict:
        from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command

        data = await exec_ssh_command(server, command, timeout=timeout)
        out = {
            "status": data.get("status", "unknown"),
            "stdout": data.get("stdout", ""),
            "stderr": data.get("stderr", ""),
            "exit_code": data.get("exit_code", -1),
            "channel": "ssh",
        }
        pid = self._extract_nexus_job_pid(out.get("stdout", ""))
        if pid:
            out["pid"] = int(pid)
        return out

    async def _audit(
        self,
        action: str,
        target_type: str,
        target_id: int,
        detail: str,
        operator: str = "system",
    ):
        """Write audit log entry"""
        log = AuditLog(
            action=action,
            target_type=target_type,
            target_id=target_id,
            detail=detail,
            admin_username=operator,
        )
        await self.audit_repo.create(log)


def _aggregate_from_results(server_ids: List[int], results: Dict[str, Any]) -> str:
    pending = failed = 0
    for sid in server_ids:
        r = results.get(str(sid), {})
        if r.get("phase") == "pending":
            pending += 1
            continue
        if r.get("exit_code", -1) != 0:
            failed += 1
    if pending:
        return "running"
    if failed == 0:
        return "completed"
    if failed == len(server_ids):
        return "failed"
    return "partial"