NeteaseUrl/auth_system/backup.py

# -*- coding: utf-8 -*-
"""数据库备份、恢复与自动备份线程。

（自 auth_system.py 拆出，行为与原实现一致。）
"""

import math
import os
import shutil
import sqlite3
import sys
import threading
import time
from pathlib import Path
from typing import Any

import auth_system  # 路径常量以包属性为唯一来源（动态读取，便于测试/运维覆盖）

from .db import (
    _DB_LOCK,
    get_conn,
    init_db,
)
from .settings import (
    get_setting,
    set_setting,
)

# ============ 数据库备份 ============


# 备份命名专用锁：create_backup 不再持有全局 _DB_LOCK（不阻塞写请求），
# 但同秒并发的两次备份仍需保证文件名唯一
_BACKUP_NAME_LOCK = threading.Lock()


def create_backup() -> str:
    """创建备份（唯一文件名——同秒冲突自动追加序号，不再互相覆盖）。

    WAL 模式下用独立只读连接走 SQLite 在线备份 API：拿到的是一致性快照，
    且不持有 _DB_LOCK——原实现备份期间阻塞全站写请求，大库时会卡住所有用户。"""
    auth_system.BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    with _BACKUP_NAME_LOCK:
        base = 'backup_' + time.strftime('%Y%m%d_%H%M%S')
        name = base + '.db'
        k = 1
        while (auth_system.BACKUP_DIR / name).exists():
            name = f'{base}-{k}.db'
            k += 1
        # 占位创建目标文件（锁内），保证并发下同秒命名仍唯一
        dest = sqlite3.connect(str(auth_system.BACKUP_DIR / name))
    conn = get_conn()
    try:
        try:
            conn.backup(dest)
        finally:
            dest.close()
        # 任何备份产生后统一裁剪（手动/自动/上传同配额）
        try:
            prune_backups(20, 5)
        except Exception:
            pass
        return name
    except Exception:
        # 备份失败：清掉占位文件，不留半个坏备份
        try:
            (auth_system.BACKUP_DIR / name).unlink(missing_ok=True)
        except OSError:
            pass
        raise
    finally:
        conn.close()


def list_backups() -> list[dict[str, Any]]:
    """备份列表（按文件修改时间倒序——按名字排序会让 upload_ 永远排前，误删新自动备份）"""
    if not auth_system.BACKUP_DIR.exists():
        return []
    out = []
    for p in auth_system.BACKUP_DIR.glob('*.db'):
        st = p.stat()
        out.append({'name': p.name, 'size': st.st_size,
                    'time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(st.st_mtime)),
                    '_mtime': st.st_mtime})
    out.sort(key=lambda x: x['_mtime'], reverse=True)
    for o in out:
        o.pop('_mtime', None)
    return out


def backup_path(name: str) -> Path | None:
    name = Path(name).name  # 防目录穿越
    p = auth_system.BACKUP_DIR / name
    return p if p.exists() and p.suffix == '.db' else None


def _backup_is_valid_db(src: Path) -> bool:
    """深度校验备份文件：SQLite 头 + 能真正打开 + 含 users 表，三者缺一不可"""
    try:
        with open(src, 'rb') as f:
            if f.read(16) != b'SQLite format 3\x00'[:16]:
                return False
        probe = sqlite3.connect(f'file:{src}?mode=ro', uri=True)
        try:
            tables = {r[0] for r in probe.execute("SELECT name FROM sqlite_master WHERE type='table'")}
            return 'users' in tables
        finally:
            probe.close()
    except Exception:
        return False


# 恢复前静止钩子——由应用层注册（取消全部在途下载任务），避免 worker 按旧库账目结算
_before_restore_hook = None


def register_before_restore_hook(fn) -> None:
    global _before_restore_hook
    _before_restore_hook = fn


def restore_backup(name: str) -> None:
    """恢复备份（先验证、留回滚、后替换、失败自动回滚，源备份文件保留不移动）
    恢复前先通过钩子取消全部在途任务（worker 不再引用旧库用户结算）"""
    src = backup_path(name)
    if not src:
        raise ValueError('备份文件不存在')
    if not _backup_is_valid_db(src):
        raise ValueError('不是有效的站点数据库备份（校验失败，未做任何修改）')
    if _before_restore_hook is not None:
        try:
            _before_restore_hook()
        except Exception:
            pass
    safety = auth_system.BACKUP_DIR / ('pre_restore_' + time.strftime('%Y%m%d_%H%M%S') + '.db')
    with _DB_LOCK:
        # 0) 先尽力清 WAL 并合并进主库（被占用导致 checkpoint 忙时明确失败，不假装成功）；
        #    必须在留安全副本之前做——否则副本缺 WAL 里未合并的最近提交，回滚点不完整
        conn = get_conn()
        try:
            row = conn.execute('PRAGMA wal_checkpoint(TRUNCATE)').fetchone()
            busy = bool(row and row[0] == 1)  # 1 = busy
        finally:
            conn.close()
        if busy:
            raise RuntimeError('数据库当前有其他读写在途（checkpoint 忙），请稍后重试恢复')
        # 1) checkpoint 后留安全副本（包含全部已提交数据，可回滚）
        shutil.copy2(str(auth_system.DB_PATH), str(safety))
        # 2) 源备份复制到临时文件（不动源备份本身）
        tmp = auth_system.DB_PATH.with_suffix('.db.restore_tmp')
        shutil.copy2(str(src), str(tmp))
        # 3) 原子替换 + 立即校验
        os.replace(str(tmp), str(auth_system.DB_PATH))
        try:
            init_db()
        except Exception:
            # 替换后初始化失败：回滚安全副本，恢复原库
            shutil.copy2(str(safety), str(auth_system.DB_PATH))
            init_db()
            raise RuntimeError('恢复后数据库初始化失败，已自动回滚到恢复前状态') from None
    # 裁剪移出 _DB_LOCK——glob/unlink 的 I/O 不阻塞全站写
    try:
        prune_backups(20, 5)
    except Exception:
        pass


def delete_backup(name: str) -> bool:
    p = backup_path(name)
    # 允许删除 pre_restore_ 安全副本（配额裁剪已保证至少保留
    # 最近 5 份；原实现一刀切拒删导致恢复频繁时无界累积且管理员手动也删不掉）
    if p:
        p.unlink()
        return True
    return False


# 自动备份线程
_backup_thread_started = False

_BACKUP_ERR_TS = {'ts': 0.0, 'n': 0}   # 备份失败限频记录（失败不能零痕迹）


def prune_backups(keep_normal: int = 20, keep_pre_restore: int = 5) -> None:
    """备份配额裁剪（从自动备份循环抽出为独立函数）——
    非pre_restore 保留最近 N 份、pre_restore_ 安全副本保留最近 M 份。
    在 create_backup/上传备份/恢复备份后统一调用（原实现只在自动备份循环里
    裁剪，手动备份/上传/恢复会让 data/backups 无界累积）。"""
    baks = [b for b in list_backups() if not b['name'].startswith('pre_restore_')]
    for b in baks[keep_normal:]:
        try:
            delete_backup(b['name'])
        except Exception:
            pass
    pre = [b for b in list_backups() if b['name'].startswith('pre_restore_')]
    for b in pre[keep_pre_restore:]:
        try:
            (auth_system.BACKUP_DIR / b['name']).unlink()
        except OSError:
            pass


def _auto_backup_loop() -> None:
    while True:
        try:
            interval = float(get_setting('backup_interval_hours', '0') or 0)
            # 非有限值（NaN/inf）或负数一律视为关闭
            if not math.isfinite(interval):
                interval = 0
                set_setting('backup_interval_hours', '0')
            if interval > 0:
                last = float(get_setting('last_backup_ts', '0') or 0)
                if not math.isfinite(last):
                    last = 0
                if time.time() - last >= interval * 3600:
                    create_backup()
                    set_setting('last_backup_ts', str(time.time()))
                    # 仅保留最近 20 份（按修改时间，不按文件名）
                    # pre_restore_* 安全副本不计入 20 份配额，仅各自保留最近 5 份
                    prune_backups(20, 5)
        except Exception:
            # 定时备份失败绝不能静默（磁盘满/权限问题时 owner 以为有备份、
            # 出事时没有）——限频告警，每 10 分钟最多一条（对齐监控线程口径）
            import traceback as _tb
            _now = time.time()
            _BACKUP_ERR_TS['n'] += 1
            if _now - _BACKUP_ERR_TS['ts'] > 600:
                print(f'[警告] 自动备份失败（近 10 分钟累计 {_BACKUP_ERR_TS["n"]} 次）: '
                      f'{_tb.format_exc()[-400:]}', file=sys.stderr)
                _BACKUP_ERR_TS['ts'] = _now
                _BACKUP_ERR_TS['n'] = 0
        time.sleep(600)


def start_auto_backup() -> None:
    global _backup_thread_started
    if _backup_thread_started:
        return
    t = threading.Thread(target=_auto_backup_loop, daemon=True)
    t.start()
    # 仅在真实启动成功后置位，失败可重试
    _backup_thread_started = True