Nexus/server/application/services/sync_engine_v2.py

"""Nexus — Sync Engine v2

Rsync push from Nexus server to target servers via SSH.
Supports: local source path / uploaded ZIP extraction.
Parallel execution with Semaphore concurrency control.
"""

import asyncio
import logging
import os
import re
import shutil
import tempfile
import uuid
from datetime import datetime, timezone
from typing import List, Optional

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.database.crypto import decrypt_value

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

MAX_CONCURRENT = 10
RSYNC_TIMEOUT = 300


class SyncEngineV2:
    """Unified sync engine: rsync push from Nexus + preview + ZIP upload"""

    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

    # ── Rsync push from Nexus to target servers ──

    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:
        """Push files from Nexus to target servers via rsync over SSH"""
        if not os.path.isdir(source_path):
            return {"total": 0, "completed": 0, "failed": 0, "results": {},
                    "error": f"源路径不存在: {source_path}"}

        batch_id = uuid.uuid4().hex[:12]
        concurrency = min(concurrency, MAX_CONCURRENT)
        sem = asyncio.Semaphore(concurrency)

        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
        _counter_lock = asyncio.Lock()
        results = {}

        async def _sync_one(server: Server):
            nonlocal completed, failed
            async with sem:
                dest = target_path or server.target_path or "/tmp/sync"
                sync_log = SyncLog(
                    server_id=server.id,
                    source_path=source_path,
                    target_path=dest,
                    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)

                result = await _rsync_push(server, source_path, dest, sync_mode)

                sync_log.status = "success" if result["exit_code"] == 0 else "failed"
                sync_log.finished_at = datetime.now(timezone.utc)
                sync_log.duration_seconds = int(
                    (sync_log.finished_at - sync_log.started_at).total_seconds()
                ) if sync_log.started_at else 0
                sync_log.error_message = result["stderr"][:1000] if result["exit_code"] != 0 else None
                sync_log.diff_summary = result["stdout"][:2000] if result["exit_code"] == 0 else None
                sync_log = await self.sync_log_repo.update(sync_log)

                # Auto-create retry job for failed pushes
                retry_job_id = None
                if sync_log.status == "failed":
                    try:
                        from server.domain.models import PushRetryJob
                        retry_job = PushRetryJob(
                            server_id=server.id,
                            server_name=server.name,
                            operator=operator,
                            source_path=source_path,
                            target_path=dest,
                            status="pending",
                            retry_count=0,
                            max_retries=3,
                            next_retry_at=datetime.now(timezone.utc),
                            last_error=sync_log.error_message[:500] if sync_log.error_message else None,
                        )
                        retry_job = await self.retry_repo.create(retry_job)
                        retry_job_id = retry_job.id
                        logger.info(f"Auto-created retry job #{retry_job.id} for server {server.id}")
                    except Exception as e:
                        logger.warning(f"Failed to create retry job for server {server.id}: {e}")

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

                    # Broadcast per-server progress via WebSocket
                    try:
                        from server.api.websocket import broadcast_sync_progress
                        await broadcast_sync_progress(
                            batch_id=batch_id,
                            server_id=server.id,
                            server_name=server.name,
                            status=sync_log.status,
                            completed=completed,
                            failed=failed,
                            total=total,
                            error_message=sync_log.error_message[:200] if sync_log.error_message else None,
                            duration_seconds=sync_log.duration_seconds,
                            retry_job_id=retry_job_id,
                        )
                    except Exception as e:
                        logger.debug(f"WS broadcast failed for server {server.id}: {e}")

                results[server.id] = (sync_log, retry_job_id)

        tasks = [asyncio.create_task(_sync_one(s)) for s in servers]
        await asyncio.gather(*tasks, return_exceptions=True)

        await self._audit(
            "sync_files", "sync", 0,
            f"文件推送: {source_path} → {total} 台服务器 ({completed} 成功, {failed} 失败)",
            operator,
        )

        # Telegram notification for push completion
        try:
            from server.infrastructure.telegram import send_telegram_sync_complete
            failed_names = []
            for sid, (log, _) in results.items():
                if log.status == "failed":
                    srv = next((s for s in servers if s.id == sid), None)
                    failed_names.append(srv.name if srv else f"#{sid}")
            # Calculate total duration from first start to last finish
            durations = [log.duration_seconds for _, (log, _) in results.items() if log.duration_seconds]
            total_duration = max(durations) if durations else None
            await send_telegram_sync_complete(
                completed=completed, failed=failed, total=total,
                source_path=source_path, operator=operator,
                failed_servers=failed_names if failed_names else None,
                duration_seconds=total_duration,
            )
        except Exception as e:
            logger.warning(f"Telegram sync complete notification failed: {e}")

        # Build results with retry_job_id
        result_dicts = {}
        for sid, (log, retry_jid) in results.items():
            d = _log_to_dict(log)
            if retry_jid:
                d["retry_job_id"] = retry_jid
            result_dicts[sid] = d

        return {
            "total": total,
            "completed": completed,
            "failed": failed,
            "batch_id": batch_id,
            "results": result_dicts,
        }

    # ── Preview (dry-run) ──

    async def preview_sync(
        self,
        server_id: int,
        source_path: str,
        target_path: Optional[str] = None,
        sync_mode: str = "incremental",
        verbose: bool = False,
    ) -> dict:
        """Dry-run rsync --dry-run --stats against one server"""
        server = await self.server_repo.get_by_id(server_id)
        if not server:
            return {"error": "Server not found", "exit_code": -1}

        if not os.path.isdir(source_path):
            return {"error": f"源路径不存在: {source_path}", "exit_code": -1}

        dest = target_path or server.target_path or "/tmp/sync"
        result = await _rsync_push(server, source_path, dest, sync_mode, dry_run=True, verbose=verbose)

        if result["exit_code"] not in (0, 23):
            return {
                "server_id": server_id,
                "server_name": server.name,
                "error": result["stderr"][:500] or f"rsync exited with code {result['exit_code']}",
                "exit_code": result["exit_code"],
            }

        stdout = result.get("stdout", "")
        stats = _parse_rsync_stats(stdout)

        files: list[str] = []
        files_truncated = False
        if verbose:
            all_lines = [
                l for l in stdout.split("\n")
                if l.strip()
                and not l.startswith("Number of")
                and not l.startswith("Total")
                and not l.startswith("Literal")
                and not l.startswith("Matched")
                and not l.startswith("File list")
                and not l.startswith("sent ")
                and not l.startswith("speedup")
                and not l.startswith("sending")
            ]
            files = all_lines[:200]
            files_truncated = len(all_lines) > 200

        return {
            "server_id": server_id,
            "server_name": server.name,
            "source_path": source_path,
            "target_path": dest,
            "sync_mode": sync_mode,
            "dry_run": True,
            "stats": stats,
            "files": files,
            "files_truncated": files_truncated,
            "exit_code": result["exit_code"],
        }

    async def _audit(self, action: str, target_type: str, target_id: int, detail: str, operator: str = "admin"):
        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}")


# ── Rsync execution on Nexus host ──

async def _rsync_push(
    server: Server,
    source_path: str,
    target_path: str,
    sync_mode: str = "incremental",
    dry_run: bool = False,
    verbose: bool = False,
) -> dict:
    """Execute rsync on Nexus host, pushing to target server via SSH.

    Handles both password auth (sshpass) and key auth (temp key file).
    Returns {exit_code, stdout, stderr}.
    """
    import shlex

    ssh_user = server.username or "root"
    ssh_host = server.domain
    ssh_port = server.port or 22
    dest = f"{ssh_user}@{ssh_host}:{target_path}"

    # Build rsync args
    args = ["rsync", "-az"]
    if sync_mode == "full":
        args.append("--delete")
    elif sync_mode == "checksum":
        args.append("--checksum")
    elif sync_mode == "overwrite":
        args.append("--inplace")
    if dry_run:
        args.extend(["--dry-run", "--stats"])
    if verbose:
        args.append("-v")

    # SSH options: port + no strict host key checking (matches asyncssh pool behavior)
    ssh_opts = f"ssh -o StrictHostKeyChecking=no -p {ssh_port}"

    # Auth: prepare temp key file or sshpass
    key_file = None
    env_override = None

    try:
        if server.auth_method == "password" and server.password:
            password = decrypt_value(server.password)
            # sshpass -p requires the password as an argument (visible in /proc on Linux,
            # but acceptable for internal ops tool; sshpass -e from env is slightly safer)
            env_override = {"SSHPASS": password}
            args = ["sshpass", "-e"] + args
            ssh_opts += " -o PubkeyAuthentication=no"
        elif server.ssh_key_configured and server.ssh_key_private:
            key_content = decrypt_value(server.ssh_key_private)
            key_file = tempfile.NamedTemporaryFile(
                mode="w", suffix=".key", prefix=f"nexus_rsync_{server.id}_",
                delete=False,
            )
            key_file.write(key_content)
            key_file.close()
            os.chmod(key_file.name, 0o600)
            ssh_opts += f" -i {shlex.quote(key_file.name)}"
        elif server.ssh_key_path:
            ssh_opts += f" -i {shlex.quote(server.ssh_key_path)}"
        else:
            return {"exit_code": -1, "stdout": "", "stderr": f"服务器 {server.name} 无有效 SSH 凭据"}

        args.extend(["-e", ssh_opts])
        args.append(source_path.rstrip("/") + "/")
        args.append(dest.rstrip("/") + "/")

        # Execute rsync on Nexus host
        proc = await asyncio.create_subprocess_exec(
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env={**os.environ, **(env_override or {})},
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=RSYNC_TIMEOUT + 30
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.communicate()
            return {"exit_code": -1, "stdout": "", "stderr": f"rsync 超时 ({RSYNC_TIMEOUT}s)"}

        return {
            "exit_code": proc.returncode or 0,
            "stdout": stdout.decode("utf-8", errors="replace")[:10000],
            "stderr": stderr.decode("utf-8", errors="replace")[:10000],
        }

    finally:
        if key_file:
            try:
                os.unlink(key_file.name)
            except OSError:
                pass


def _parse_rsync_stats(output: str) -> dict:
    """Extract key numbers from rsync --stats output."""
    stats: dict = {}
    patterns = {
        "files_total":         r"Number of files:\s+([\d,]+)",
        "files_created":       r"Number of created files:\s+([\d,]+)",
        "files_deleted":       r"Number of deleted files:\s+([\d,]+)",
        "files_transferred":   r"Number of regular files transferred:\s+([\d,]+)",
        "total_size_bytes":    r"Total file size:\s+([\d,]+)\s+bytes",
        "transfer_size_bytes": r"Total transferred file size:\s+([\d,]+)\s+bytes",
    }
    for key, pattern in patterns.items():
        m = re.search(pattern, output)
        if m:
            try:
                stats[key] = int(m.group(1).replace(",", ""))
            except ValueError:
                stats[key] = m.group(1)
    return stats


def _log_to_dict(log: SyncLog) -> dict:
    return {
        "id": log.id, "server_id": log.server_id,
        "status": log.status, "sync_mode": log.sync_mode,
        "files_transferred": log.files_transferred,
        "files_skipped": log.files_skipped,
        "bytes_transferred": log.bytes_transferred,
        "diff_summary": log.diff_summary,
        "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,
    }