Nexus/server/application/services/script_service.py

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

import asyncio
import json
import logging
from typing import Optional, List, Dict

import httpx

from server.domain.models import Script, ScriptExecution, DbCredential, AuditLog
from server.domain.repositories import (
    ScriptRepository, ScriptExecutionRepository,
    DbCredentialRepository, ServerRepository, AuditLogRepository,
)

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


class ScriptService:
    """Business logic for script CRUD and remote command execution via Agent /api/exec"""

    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

    # ── 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"Created script: {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"Updated script: {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"Deleted script ID: {id}")
        return result

    # ── Command Execution (Agent /api/exec) ──

    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",
    ) -> ScriptExecution:
        """Execute a shell command on multiple servers via Agent /api/exec.

        If credential_id is provided, $DB_* variables are substituted in the command
        before sending to each Agent. Variable substitution happens on Nexus side
        so the password never appears in the script text.
        """
        # Resolve DB credential variables (if provided)
        resolved_command = await self._substitute_db_vars(command, credential_id)

        # Create execution record
        execution = ScriptExecution(
            script_id=script_id,
            command=resolved_command,
            server_ids=json.dumps(server_ids),
            status="running",
            operator=operator,
        )
        execution = await self.execution_repo.create(execution)

        # Dispatch to agents concurrently (max 10 at a time)
        results: Dict[str, dict] = {}
        semaphore = asyncio.Semaphore(10)

        async def _exec_on_server(server_id: int):
            server = await self.server_repo.get_by_id(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
            async with semaphore:
                try:
                    result = await self._call_agent_exec(
                        server.domain, server.agent_port, server.agent_api_key,
                        resolved_command, timeout,
                    )
                    results[str(server_id)] = result
                except Exception as e:
                    results[str(server_id)] = {
                        "status": "error",
                        "stdout": "",
                        "stderr": str(e),
                        "exit_code": -1,
                    }

        tasks = [_exec_on_server(sid) for sid in server_ids]
        await asyncio.gather(*tasks)

        # Determine overall status
        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"

        # Update execution record
        execution = await self.execution_repo.update_status(
            execution.id, status, json.dumps(results)
        )

        await self._audit(
            "execute_command", "script_execution", execution.id,
            f"Executed on {len(server_ids)} servers: {status} ({failed_count} failed)",
        )
        return execution

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

    # ── 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"Created credential: {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"Deleted credential ID: {id}")
        return result

    async def test_credential_connection(self, credential: DbCredential) -> dict:
        """Test database connection using provided credentials (before saving)"""
        import shlex
        from server.infrastructure.database.crypto import decrypt_value

        password = credential.encrypted_password  # Not yet encrypted if testing before save
        if credential.id:
            password = decrypt_value(credential.encrypted_password)

        test_cmd = (
            f"mysql -h {shlex.quote(credential.host)} "
            f"-P {credential.port} "
            f"-u {shlex.quote(credential.username)} "
            f"-p{shlex.quote(password)} "
            f"-e 'SELECT 1' 2>/dev/null && echo 'OK' || echo 'FAIL'"
        )
        return {"status": "test_command_ready", "command": test_cmd}

    # ── Private helpers ──

    async def _substitute_db_vars(self, command: str, credential_id: Optional[int]) -> str:
        """Replace $DB_USER/$DB_PASS/$DB_HOST/$DB_PORT/$DB_NAME in command"""
        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)

        replacements = {
            "$DB_USER": credential.username,
            "$DB_PASS": password,
            "$DB_HOST": credential.host,
            "$DB_PORT": str(credential.port),
            "$DB_NAME": credential.database or "",
        }
        for var, value in replacements.items():
            command = command.replace(var, value)

        return command

    async def _call_agent_exec(
        self, host: str, port: int, api_key: str, command: str, timeout: int = 30,
    ) -> dict:
        """Call Agent /api/exec endpoint to run a command on a remote server"""
        url = f"http://{host}:{port}/api/exec"
        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["X-API-Key"] = api_key
        payload = {"command": command, "timeout": timeout}

        async with httpx.AsyncClient(timeout=timeout + 5) as client:
            try:
                resp = await client.post(url, json=payload, headers=headers)
                if resp.status_code == 200:
                    data = resp.json()
                    return {
                        "status": data.get("status", "unknown"),
                        "stdout": data.get("stdout", ""),
                        "stderr": data.get("stderr", ""),
                        "exit_code": data.get("exit_code", -1),
                    }
                else:
                    return {
                        "status": "error",
                        "stdout": "",
                        "stderr": f"Agent returned HTTP {resp.status_code}",
                        "exit_code": -1,
                    }
            except httpx.TimeoutException:
                return {
                    "status": "timeout",
                    "stdout": "",
                    "stderr": f"Agent timeout after {timeout}s",
                    "exit_code": -1,
                }
            except httpx.ConnectError:
                return {
                    "status": "connection_error",
                    "stdout": "",
                    "stderr": f"Cannot connect to Agent at {host}:{port}",
                    "exit_code": -1,
                }

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