Nexus/server/application/services/sync_engine_v2.py

"""Nexus — Sync Engine v2 (Step 4 Upgrade)

S1: Existing rsync push retained as Sync file mode
S2: New Sync command mode — SSH exec batch execution
S3: New Sync config mode — system parameters batch push
S4: SFTP file transfer channel
S5: Parallel execution + result aggregation with Semaphore
"""

import asyncio
import json
import logging
from datetime import datetime, timezone
from typing import List, Optional

from sqlalchemy.ext.asyncio import AsyncSession

from server.domain.models import Server, SyncLog, AuditLog
from server.infrastructure.database.server_repo import ServerRepositoryImpl
from server.infrastructure.database.sync_log_repo import SyncLogRepositoryImpl
from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl
from server.infrastructure.database.push_schedule_repo import PushRetryJobRepositoryImpl
from server.infrastructure.redis.client import get_redis
from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command, ssh_pool

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

# S5: Semaphore for concurrency control
MAX_CONCURRENT = 10


class SyncEngineV2:
    """Unified sync engine: file + command + config modes, with parallel execution (S5)"""

    def __init__(
        self,
        server_repo: ServerRepositoryImpl,
        sync_log_repo: SyncLogRepositoryImpl,
        audit_repo: AuditLogRepositoryImpl,
        retry_repo: PushRetryJobRepositoryImpl,
    ):
        self.server_repo = server_repo
        self.sync_log_repo = sync_log_repo
        self.audit_repo = audit_repo
        self.retry_repo = retry_repo

    # ── S1: Sync File Mode (rsync-based) ──

    async def sync_files(
        self,
        server_ids: List[int],
        source_path: str,
        target_path: Optional[str] = None,
        sync_mode: str = "incremental",
        operator: str = "admin",
        batch_size: int = 50,
        concurrency: int = 10,
        trigger_type: str = "manual",
    ) -> dict:
        """S1: File sync using rsync over SSH"""
        concurrency = min(concurrency, MAX_CONCURRENT)
        sem = asyncio.Semaphore(concurrency)

        results = {}
        # Build server list
        servers = []
        for sid in server_ids:
            server = await self.server_repo.get_by_id(sid)
            if server:
                servers.append(server)

        total = len(servers)
        completed = 0
        failed = 0

        async def _sync_one(server: Server):
            nonlocal completed, failed
            async with sem:
                # Create sync log
                sync_log = SyncLog(
                    server_id=server.id,
                    source_path=source_path,
                    target_path=target_path or server.target_path or "/tmp/sync",
                    trigger_type=trigger_type,
                    operator=operator,
                    status="running",
                    sync_mode=sync_mode,
                    started_at=datetime.now(timezone.utc),
                )
                sync_log = await self.sync_log_repo.create(sync_log)

                # Execute rsync
                import shlex
                rsync_cmd = f"rsync -az --delete {shlex.quote(source_path)}/ {shlex.quote(target_path or server.target_path or '/tmp/sync')}/"
                result = await exec_ssh_command(server, rsync_cmd, timeout=300)

                # Update sync log
                sync_log.status = "success" if result["exit_code"] == 0 else "failed"
                sync_log.finished_at = datetime.now(timezone.utc)
                sync_log.duration_seconds = (sync_log.finished_at - sync_log.started_at).seconds if sync_log.started_at else 0
                sync_log.error_message = result["stderr"][:1000] if result["exit_code"] != 0 else None
                sync_log = await self.sync_log_repo.update(sync_log)

                if sync_log.status == "success":
                    completed += 1
                else:
                    failed += 1

                results[server.id] = sync_log
                return sync_log

        # S5: Parallel execution with concurrency control
        tasks = [asyncio.create_task(_sync_one(s)) for s in servers]
        await asyncio.gather(*tasks, return_exceptions=True)

        # Audit
        await self._audit(
            "sync_files", "sync", 0,
            f"File sync: {source_path} → {total} servers ({completed} ok, {failed} failed)",
            operator,
        )

        return {
            "total": total,
            "completed": completed,
            "failed": failed,
            "results": {sid: _log_to_dict(log) for sid, log in results.items()},
        }

    # ── S2: Sync Command Mode (SSH exec batch) ──

    async def sync_commands(
        self,
        server_ids: List[int],
        commands: List[str],
        operator: str = "admin",
        timeout: int = 60,
        concurrency: int = 10,
    ) -> dict:
        """S2: Execute commands on multiple servers via SSH, collect results"""
        concurrency = min(concurrency, MAX_CONCURRENT)
        sem = asyncio.Semaphore(concurrency)

        results = {}
        errors = []

        async def _exec_one(server_id: int):
            server = await self.server_repo.get_by_id(server_id)
            if not server:
                errors.append({"server_id": server_id, "error": "Server not found"})
                return

            async with sem:
                server_results = []
                for cmd in commands:
                    result = await exec_ssh_command(server, cmd, timeout=timeout)
                    server_results.append({
                        "command": cmd[:200],
                        "exit_code": result["exit_code"],
                        "stdout": result["stdout"][:5000],
                        "stderr": result["stderr"][:2000],
                        "status": result["status"],
                    })

                # Create sync log
                sync_log = SyncLog(
                    server_id=server.id,
                    source_path=f"#cmd:{len(commands)} commands",
                    target_path=server.name,
                    trigger_type="manual",
                    operator=operator,
                    status="success" if all(r["exit_code"] == 0 for r in server_results) else "failed",
                    sync_mode="command",
                    started_at=datetime.now(timezone.utc),
                    finished_at=datetime.now(timezone.utc),
                    diff_summary=json.dumps(server_results),
                )
                sync_log = await self.sync_log_repo.create(sync_log)
                results[server_id] = {**{"server_name": server.name}, **server_results, "log": _log_to_dict(sync_log)}

        tasks = [asyncio.create_task(_exec_one(sid)) for sid in server_ids]
        await asyncio.gather(*tasks, return_exceptions=True)

        await self._audit(
            "sync_commands", "sync", 0,
            f"Command sync: {len(commands)} commands → {len(server_ids)} servers",
            operator,
        )

        return {"total_servers": len(server_ids), "results": results, "errors": errors}

    # ── S3: Sync Config Mode ──

    BACKUP_CONFIG_PATH = "/etc/sysctl.d/99-nexus.conf"

    async def sync_config(
        self,
        server_ids: List[int],
        config_updates: dict,  # {key: value} system parameters to push
        operator: str = "admin",
        concurrency: int = 10,
        rollback: bool = False,
    ) -> dict:
        """S3: Push configuration changes to multiple servers

        Config updates are applied as shell commands:
        - `sysctl -w` for immediate kernel param changes
        - Deduplicated writes to /etc/sysctl.d/99-nexus.conf (sed removes old entries first)

        Rollback mode: restores from backup created during last config push.
        """
        import re, shlex

        if rollback:
            return await self._rollback_config(server_ids, operator, concurrency)

        # Build deduplicated config commands
        commands = []
        config_path = self.BACKUP_CONFIG_PATH
        backup_path = f"{config_path}.bak.{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}"

        for key, value in config_updates.items():
            # Sanitize: only allow alphanumeric/dot/underscore/dash keys
            if not re.match(r'^[a-zA-Z0-9._-]+$', str(key)):
                logger.warning(f"Skipping invalid config key: {key!r}")
                continue
            safe_value = shlex.quote(str(value))

            # Deduplicate: remove existing lines for this key before appending
            commands.append(f"sed -i '/^{re.escape(key)}\\s*=/d' {config_path} 2>/dev/null || true")
            # Append new value
            commands.append(f"echo {key}={safe_value} >> {config_path}")
            # Apply immediately
            commands.append(f"sysctl -w {key}={safe_value} 2>/dev/null || true")

        # Pre-push: backup existing config file
        backup_commands = [
            f"cp {config_path} {backup_path} 2>/dev/null || true",
            f"touch {config_path}",
        ]

        # Combine: backup first, then deduplicated config commands
        all_commands = backup_commands + commands

        result = await self.sync_commands(
            server_ids=server_ids,
            commands=all_commands,
            operator=operator,
            timeout=30,
            concurrency=concurrency,
        )

        # Store backup path in result for rollback reference
        result["backup_path"] = backup_path
        result["config_path"] = config_path
        return result

    async def _rollback_config(
        self,
        server_ids: List[int],
        operator: str = "admin",
        concurrency: int = 10,
    ) -> dict:
        """Rollback config: find the most recent backup and restore it"""
        config_path = self.BACKUP_CONFIG_PATH

        async def _rollback_one(server_id: int):
            server = await self.server_repo.get_by_id(server_id)
            if not server:
                return {"server_id": server_id, "status": "error", "message": "Server not found"}

            # Find the latest backup file
            find_cmd = f"ls -t {config_path}.bak.* 2>/dev/null | head -1"
            result = await exec_ssh_command(server, find_cmd, timeout=10)

            if result["exit_code"] != 0 or not result["stdout"].strip():
                return {"server_id": server_id, "status": "error", "message": "No backup found for rollback"}

            backup_file = result["stdout"].strip().split("\n")[0]

            # Restore from backup
            restore_cmd = f"cp {backup_file} {config_path} && sysctl --system 2>/dev/null || true"
            restore_result = await exec_ssh_command(server, restore_cmd, timeout=30)

            status = "success" if restore_result["exit_code"] == 0 else "failed"
            return {
                "server_id": server_id,
                "server_name": server.name,
                "status": status,
                "backup_file": backup_file,
                "error": restore_result["stderr"][:500] if status == "failed" else None,
            }

        sem = asyncio.Semaphore(min(concurrency, MAX_CONCURRENT))
        results = []

        async def _limited_rollback(sid):
            async with sem:
                return await _rollback_one(sid)

        tasks = [asyncio.create_task(_limited_rollback(sid)) for sid in server_ids]
        task_results = await asyncio.gather(*tasks, return_exceptions=True)
        results = [r for r in task_results if isinstance(r, dict)]

        await self._audit(
            "rollback_config", "sync", 0,
            f"Config rollback on {len(server_ids)} servers: {sum(1 for r in results if r.get('status')=='success')} ok, {sum(1 for r in results if r.get('status')=='failed')} failed",
            operator,
        )

        return {"total": len(server_ids), "results": results}

    # ── S4: SFTP File Transfer ──

    async def sftp_transfer(
        self,
        server_id: int,
        operation: str,  # "put" or "get"
        local_path: str,
        remote_path: str,
        operator: str = "admin",
    ) -> dict:
        """S4: SFTP file transfer to/from a server using asyncssh"""
        server = await self.server_repo.get_by_id(server_id)
        if not server:
            return {"status": "error", "message": "Server not found"}

        try:
            conn = await ssh_pool.acquire(server)
            try:
                async with conn.start_sftp_client() as sftp:
                    if operation == "put":
                        await sftp.put(local_path, remote_path)
                        message = f"Uploaded {local_path} → {server.name}:{remote_path}"
                    elif operation == "get":
                        await sftp.get(remote_path, local_path)
                        message = f"Downloaded {server.name}:{remote_path} → {local_path}"
                    else:
                        return {"status": "error", "message": f"Unknown SFTP operation: {operation}"}

                    # Audit
                    await self._audit(
                        f"sftp_{operation}", "sync", server_id,
                        message, operator,
                    )

                    return {"status": "success", "message": message}

            finally:
                await ssh_pool.release(server.id)

        except Exception as e:
            logger.error(f"SFTP transfer failed: {e}")
            return {"status": "error", "message": str(e)}

    async def _audit(self, action: str, target_type: str, target_id: int, detail: str, operator: str = "admin"):
        """Create audit log entry"""
        try:
            audit_log = AuditLog(
                admin_username=operator,
                action=action,
                target_type=target_type,
                target_id=target_id,
                detail=detail,
            )
            await self.audit_repo.create(audit_log)
        except Exception as e:
            logger.error(f"Audit log failed: {e}")


def _log_to_dict(log: SyncLog) -> dict:
    return {
        "id": log.id, "server_id": log.server_id,
        "status": log.status, "sync_mode": log.sync_mode,
        "error_message": log.error_message,
        "started_at": str(log.started_at) if log.started_at else None,
        "finished_at": str(log.finished_at) if log.finished_at else None,
        "duration_seconds": log.duration_seconds,
    }