crypto_okx_hub/lib/backup_lib.py

"""中控备份/恢复: .env + hub.db;默认落盘 /root/crypto_okx_hub_backups."""
from __future__ import annotations

import json
import os
import re
import shutil
import tarfile
import tempfile
import threading
import time
from datetime import datetime, timedelta, timezone
from typing import Any, Callable

_BJ = timezone(timedelta(hours=8))
_SAFE_NAME_RE = re.compile(r"^crypto_okx_hub_backup_\d{8}_\d{6}\.tar\.gz$")
_LOCK = threading.Lock()
_LAST_AUTO_DAY = ""


def beijing_now() -> datetime:
    return datetime.now(_BJ)


def backup_keep_days() -> int:
    try:
        return max(1, int(os.getenv("BACKUP_KEEP_DAYS") or 10))
    except (TypeError, ValueError):
        return 10


def backup_dir() -> str:
    root = (os.getenv("BACKUP_DIR") or "/root").strip() or "/root"
    path = os.path.join(root, "crypto_okx_hub_backups")
    os.makedirs(path, exist_ok=True)
    return path


def _safe_backup_name(name: str) -> str | None:
    n = os.path.basename((name or "").strip())
    if not _SAFE_NAME_RE.match(n):
        return None
    return n


def backup_path(name: str) -> str | None:
    safe = _safe_backup_name(name)
    if not safe:
        return None
    path = os.path.abspath(os.path.join(backup_dir(), safe))
    root = os.path.abspath(backup_dir())
    if not path.startswith(root + os.sep) and path != root:
        return None
    return path


def list_backups() -> list[dict[str, Any]]:
    root = backup_dir()
    items: list[dict[str, Any]] = []
    try:
        names = os.listdir(root)
    except OSError:
        return []
    for name in names:
        if not _SAFE_NAME_RE.match(name):
            continue
        path = os.path.join(root, name)
        try:
            st = os.stat(path)
        except OSError:
            continue
        items.append(
            {
                "name": name,
                "size": int(st.st_size),
                "mtime": int(st.st_mtime),
                "mtime_bj": datetime.fromtimestamp(st.st_mtime, _BJ).strftime(
                    "%Y-%m-%d %H:%M:%S"
                ),
            }
        )
    items.sort(key=lambda x: x["mtime"], reverse=True)
    return items


def prune_backups(*, keep_days: int | None = None) -> int:
    days = backup_keep_days() if keep_days is None else max(1, int(keep_days))
    cutoff = time.time() - days * 86400
    removed = 0
    for item in list_backups():
        if float(item["mtime"]) >= cutoff:
            continue
        path = backup_path(item["name"])
        if not path:
            continue
        try:
            os.remove(path)
            removed += 1
        except OSError:
            pass
    return removed


def create_backup(*, root_dir: str, env_path: str, db_path: str, label: str = "manual") -> dict[str, Any]:
    with _LOCK:
        now = beijing_now()
        name = f"crypto_okx_hub_backup_{now.strftime('%Y%m%d_%H%M%S')}.tar.gz"
        out_path = os.path.join(backup_dir(), name)
        manifest = {
            "kind": "crypto_okx_hub",
            "created_bj": now.strftime("%Y-%m-%d %H:%M:%S"),
            "label": (label or "manual").strip() or "manual",
            "includes": [".env", "database"],
        }
        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".tar.gz", prefix="hub_bak_")
        os.close(tmp_fd)
        try:
            with tarfile.open(tmp_path, "w:gz") as tar:
                manifest_path = tmp_path + ".manifest.json"
                with open(manifest_path, "w", encoding="utf-8") as f:
                    json.dump(manifest, f, ensure_ascii=False, indent=2)
                try:
                    tar.add(manifest_path, arcname="manifest.json")
                finally:
                    try:
                        os.remove(manifest_path)
                    except OSError:
                        pass
                if env_path and os.path.isfile(env_path):
                    tar.add(env_path, arcname=".env")
                if db_path and os.path.isfile(db_path):
                    tar.add(db_path, arcname="database/" + os.path.basename(db_path))
            shutil.move(tmp_path, out_path)
        except Exception:
            try:
                os.remove(tmp_path)
            except OSError:
                pass
            raise
        prune_backups()
        st = os.stat(out_path)
        return {
            "ok": True,
            "name": name,
            "path": out_path,
            "size": int(st.st_size),
            "mtime_bj": now.strftime("%Y-%m-%d %H:%M:%S"),
            "keep_days": backup_keep_days(),
        }


def restore_backup(name: str, *, root_dir: str, env_path: str, db_path: str) -> dict[str, Any]:
    path = backup_path(name)
    if not path or not os.path.isfile(path):
        return {"ok": False, "msg": "备份不存在"}
    with _LOCK:
        with tempfile.TemporaryDirectory(prefix="hub_restore_") as tmp:
            with tarfile.open(path, "r:gz") as tar:
                tar.extractall(tmp)
            env_src = os.path.join(tmp, ".env")
            if os.path.isfile(env_src) and env_path:
                shutil.copy2(env_src, env_path)
            db_src_dir = os.path.join(tmp, "database")
            if os.path.isdir(db_src_dir) and db_path:
                for fn in os.listdir(db_src_dir):
                    src = os.path.join(db_src_dir, fn)
                    if os.path.isfile(src):
                        os.makedirs(os.path.dirname(os.path.abspath(db_path)) or ".", exist_ok=True)
                        shutil.copy2(src, db_path)
                        break
        return {
            "ok": True,
            "msg": "已恢复 .env / 数据库,建议重启中控进程使配置完全生效",
            "name": os.path.basename(path),
        }


def maybe_run_auto_backup(
    *,
    root_dir: str,
    env_path: str,
    db_path: str,
    hour: int = 3,
) -> dict[str, Any] | None:
    global _LAST_AUTO_DAY
    now = beijing_now()
    day = now.strftime("%Y-%m-%d")
    if now.hour != int(hour):
        return None
    if _LAST_AUTO_DAY == day:
        return None
    prefix = f"crypto_okx_hub_backup_{now.strftime('%Y%m%d')}_"
    for item in list_backups():
        if str(item["name"]).startswith(prefix):
            _LAST_AUTO_DAY = day
            return None
    result = create_backup(
        root_dir=root_dir,
        env_path=env_path,
        db_path=db_path,
        label="auto",
    )
    _LAST_AUTO_DAY = day
    return result


def start_backup_scheduler(
    *,
    root_dir: str,
    env_path: str,
    db_path_fn: Callable[[], str],
    hour: int = 3,
) -> None:
    def loop() -> None:
        while True:
            try:
                maybe_run_auto_backup(
                    root_dir=root_dir,
                    env_path=env_path,
                    db_path=db_path_fn(),
                    hour=hour,
                )
            except Exception:
                pass
            time.sleep(30)

    t = threading.Thread(target=loop, name="hub-backup-scheduler", daemon=True)
    t.start()