mirror/app/services/applier.py

from __future__ import annotations

import logging
from datetime import timedelta

from app.config import Config
from app.extensions import db
from app.models import ApplyLog, Profile, Setting, SystemState
from app.models import utcnow
from app.services import apt, backup, dns, docker_svc, github, pip_npm
from app.services.host import run_on_host

logger = logging.getLogger(__name__)

APPLIERS = {
    "dns": dns.apply_dns,
    "apt": apt.apply_apt,
    "docker": docker_svc.apply_docker,
    "github": github.apply_github,
    "pip": pip_npm.apply_pip,
    "npm": pip_npm.apply_npm,
}

APPLY_MODE_TEST = "test"
APPLY_MODE_FINAL = "final"


def ensure_initial_backup() -> SystemState:
    state = SystemState.query.first()
    if not state:
        state = SystemState()
        db.session.add(state)
        db.session.commit()

    if not state.initial_backup_done:
        path = backup.create_initial_backup()
        state.initial_backup_path = str(path)
        state.initial_backup_done = True
        db.session.commit()
        logger.info("Initial backup created at %s", path)

    return state


def _rollback_minutes_for_mode(mode: str, rollback_minutes: int | None) -> int | None:
    if mode == APPLY_MODE_FINAL:
        return None
    if rollback_minutes is not None:
        return rollback_minutes
    default = Setting.get("test_rollback_minutes", str(Config.TEST_ROLLBACK_MINUTES))
    return int(default or Config.TEST_ROLLBACK_MINUTES)


def apply_profile(
    profile: Profile,
    mode: str = APPLY_MODE_TEST,
    rollback_minutes: int | None = None,
) -> ApplyLog:
    state = ensure_initial_backup()
    rollback_minutes = _rollback_minutes_for_mode(mode, rollback_minutes)

    previous_backup = backup.create_backup(f"before_apply_{profile.id}")
    apply_log = ApplyLog(
        profile_id=profile.id,
        status="running",
        backup_path=str(previous_backup),
        previous_backup_path=state.initial_backup_path,
        rollback_at=utcnow() + timedelta(minutes=rollback_minutes) if rollback_minutes else None,
    )
    db.session.add(apply_log)
    db.session.flush()

    details: dict = {
        "steps": [],
        "profile": profile.name,
        "mode": mode,
        "rollback_minutes": rollback_minutes,
    }
    mirrors = profile.mirrors_by_category()
    order = ["dns", "apt", "docker", "github", "pip", "npm"]

    for category in order:
        mirror = mirrors.get(category)
        if not mirror:
            continue
        applier = APPLIERS.get(category)
        if not applier:
            continue
        try:
            result = applier(mirror)
            details["steps"].append({"category": category, "mirror": mirror.name, **result})
            if not result.get("success", True):
                break
        except Exception as exc:
            logger.exception("Apply failed for %s", category)
            details["steps"].append(
                {"category": category, "mirror": mirror.name, "success": False, "error": str(exc)}
            )
            break

    failed = [s for s in details["steps"] if not s.get("success", True)]
    if failed and apply_log.backup_path:
        rollback_result = backup.restore_backup(apply_log.backup_path)
        details["auto_rollback"] = rollback_result

    apply_log.set_details(details)
    apply_log.finished_at = utcnow()

    if failed:
        apply_log.status = "failed"
        apply_log.rollback_at = None
    elif mode == APPLY_MODE_FINAL:
        apply_log.confirmed = True
        apply_log.status = "confirmed"
        apply_log.rollback_at = None
        state.current_profile_id = profile.id
        state.pending_apply_log_id = None
    else:
        apply_log.status = "pending_confirm"
        state.current_profile_id = profile.id
        state.pending_apply_log_id = apply_log.id

    db.session.commit()

    return apply_log


def confirm_apply(apply_log_id: int) -> ApplyLog:
    apply_log = ApplyLog.query.get_or_404(apply_log_id)
    apply_log.confirmed = True
    apply_log.status = "confirmed"
    apply_log.rollback_at = None

    details = apply_log.get_details()
    details["confirmed_from"] = APPLY_MODE_TEST
    apply_log.set_details(details)

    state = SystemState.query.first()
    if state:
        state.pending_apply_log_id = None
        if apply_log.profile_id:
            state.current_profile_id = apply_log.profile_id
    db.session.commit()
    return apply_log


def rollback_apply(apply_log_id: int) -> dict:
    apply_log = ApplyLog.query.get_or_404(apply_log_id)
    if not apply_log.backup_path:
        return {"success": False, "error": "مسیر backup یافت نشد"}

    result = backup.restore_backup(apply_log.backup_path)
    apply_log.status = "rolled_back"
    apply_log.finished_at = utcnow()
    apply_log.rollback_at = None

    state = SystemState.query.first()
    if state:
        state.pending_apply_log_id = None
    db.session.commit()
    return result


def emergency_restore_last_apply() -> dict:
    """One-shot restore of the most recent pre-apply backup."""
    path = backup.get_latest_apply_backup()
    if not path:
        return {"success": False, "error": "backup قبل از apply یافت نشد"}
    result = backup.restore_backup(path)
    if result.get("success"):
        result["backup_path"] = str(path)

        state = SystemState.query.first()
        pending_id = state.pending_apply_log_id if state else None
        if pending_id:
            log = ApplyLog.query.get(pending_id)
            if log and log.status == "pending_confirm":
                log.status = "rolled_back"
                log.finished_at = utcnow()
                log.rollback_at = None
                state.pending_apply_log_id = None
                db.session.commit()
    return result


def restore_initial() -> dict:
    state = SystemState.query.first()
    if not state or not state.initial_backup_path:
        return {"success": False, "error": "backup اولیه یافت نشد"}
    return backup.restore_backup(state.initial_backup_path)


def get_system_status() -> dict:
    codename_result = run_on_host(["lsb_release", "-cs"])
    version_result = run_on_host(["lsb_release", "-d"])

    return {
        "ubuntu_codename": codename_result.stdout.strip() if codename_result.returncode == 0 else "unknown",
        "ubuntu_description": version_result.stdout.strip() if version_result.returncode == 0 else "unknown",
        "dns": dns.read_current_dns(),
        "apt": apt.read_current_apt(),
        "docker": docker_svc.read_current_docker(),
        "github": github.read_current_github(),
        "host_root": Config.HOST_ROOT or "(local mode)",
    }