AstrBot/astrbot/dashboard/services/backup_service.py

from __future__ import annotations

import asyncio
import json
import math
import os
import re
import shutil
import time
import traceback
import uuid
import zipfile
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

import jwt

from astrbot.core import logger
from astrbot.core.backup.exporter import AstrBotExporter
from astrbot.core.backup.importer import AstrBotImporter
from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
from astrbot.core.db import BaseDatabase
from astrbot.core.utils.astrbot_path import (
    get_astrbot_backups_path,
    get_astrbot_data_path,
)

CHUNK_SIZE = 1024 * 1024
UPLOAD_EXPIRE_SECONDS = 3600


class BackupServiceError(Exception):
    pass


@dataclass
class BackupDownload:
    path: str
    filename: str


def secure_filename(filename: str) -> str:
    filename = filename.replace("\\", "/")
    filename = os.path.basename(filename)
    filename = filename.replace("..", "_")
    filename = re.sub(r"[^\w\-.]", "_", filename)
    filename = filename.strip(".")
    if not filename or filename.replace("_", "") == "":
        filename = "backup"
    return filename


def generate_unique_filename(original_filename: str) -> str:
    name, ext = os.path.splitext(original_filename)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return f"{name}_{timestamp}{ext}"


class BackupService:
    def __init__(
        self,
        db: BaseDatabase,
        core_lifecycle: AstrBotCoreLifecycle,
    ) -> None:
        self.db = db
        self.core_lifecycle = core_lifecycle
        self.config = core_lifecycle.astrbot_config
        self.backup_dir = get_astrbot_backups_path()
        self.data_dir = get_astrbot_data_path()
        self.chunks_dir = os.path.join(self.backup_dir, ".chunks")
        self.backup_tasks: dict[str, dict] = {}
        self.backup_progress: dict[str, dict] = {}
        self.upload_sessions: dict[str, dict] = {}
        self._cleanup_task: asyncio.Task | None = None

    @staticmethod
    def _payload(data: object) -> dict[str, Any]:
        return data if isinstance(data, dict) else {}

    @staticmethod
    async def _save_upload(file: Any, target_path: str) -> None:
        if hasattr(file, "save"):
            result = file.save(target_path)
            if hasattr(result, "__await__"):
                await result
            return

        if hasattr(file, "read"):
            data = file.read()
            if hasattr(data, "__await__"):
                data = await data
            Path(target_path).write_bytes(data)
            return

        raise BackupServiceError("无效的上传文件")

    @staticmethod
    def _validate_backup_filename(filename: str | None, *, missing: str) -> str:
        if not filename:
            raise BackupServiceError(missing)
        if ".." in filename or "/" in filename or "\\" in filename:
            raise BackupServiceError("无效的文件名")
        return filename

    def _init_task(self, task_id: str, task_type: str, status: str = "pending") -> None:
        self.backup_tasks[task_id] = {
            "type": task_type,
            "status": status,
            "result": None,
            "error": None,
        }
        self.backup_progress[task_id] = {
            "status": status,
            "stage": "waiting",
            "current": 0,
            "total": 100,
            "message": "",
        }

    def _set_task_result(
        self,
        task_id: str,
        status: str,
        result: dict | None = None,
        error: str | None = None,
    ) -> None:
        if task_id in self.backup_tasks:
            self.backup_tasks[task_id]["status"] = status
            self.backup_tasks[task_id]["result"] = result
            self.backup_tasks[task_id]["error"] = error
        if task_id in self.backup_progress:
            self.backup_progress[task_id]["status"] = status

    def _update_progress(
        self,
        task_id: str,
        *,
        status: str | None = None,
        stage: str | None = None,
        current: int | None = None,
        total: int | None = None,
        message: str | None = None,
    ) -> None:
        if task_id not in self.backup_progress:
            return
        progress = self.backup_progress[task_id]
        if status is not None:
            progress["status"] = status
        if stage is not None:
            progress["stage"] = stage
        if current is not None:
            progress["current"] = current
        if total is not None:
            progress["total"] = total
        if message is not None:
            progress["message"] = message

    def _make_progress_callback(self, task_id: str):
        async def _callback(
            stage: str,
            current: int,
            total: int,
            message: str = "",
        ) -> None:
            self._update_progress(
                task_id,
                status="processing",
                stage=stage,
                current=current,
                total=total,
                message=message,
            )

        return _callback

    def ensure_cleanup_task_started(self) -> None:
        if self._cleanup_task is None or self._cleanup_task.done():
            try:
                self._cleanup_task = asyncio.create_task(
                    self._cleanup_expired_uploads()
                )
            except RuntimeError:
                pass

    async def _cleanup_expired_uploads(self) -> None:
        while True:
            try:
                await asyncio.sleep(300)
                current_time = time.time()
                expired_sessions = []

                for upload_id, session in self.upload_sessions.items():
                    last_activity = session.get("last_activity", session["created_at"])
                    if current_time - last_activity > UPLOAD_EXPIRE_SECONDS:
                        expired_sessions.append(upload_id)

                for upload_id in expired_sessions:
                    await self.cleanup_upload_session(upload_id)
                    logger.info(f"清理过期的上传会话: {upload_id}")

            except asyncio.CancelledError:
                break
            except Exception as exc:
                logger.error(f"清理过期上传会话失败: {exc}")

    async def cleanup_upload_session(self, upload_id: str) -> None:
        if upload_id in self.upload_sessions:
            session = self.upload_sessions[upload_id]
            chunk_dir = session.get("chunk_dir")
            if chunk_dir and os.path.exists(chunk_dir):
                try:
                    shutil.rmtree(chunk_dir)
                except Exception as exc:
                    logger.warning(f"清理分片目录失败: {exc}")
            del self.upload_sessions[upload_id]

    def get_backup_manifest(self, zip_path: str) -> dict | None:
        try:
            with zipfile.ZipFile(zip_path, "r") as zf:
                if "manifest.json" in zf.namelist():
                    manifest_data = zf.read("manifest.json")
                    return json.loads(manifest_data.decode("utf-8"))
                return None
        except Exception as exc:
            logger.debug(f"读取备份 manifest 失败: {exc}")
        return None

    def list_backups(self, *, page: int, page_size: int) -> dict:
        self.ensure_cleanup_task_started()
        Path(self.backup_dir).mkdir(parents=True, exist_ok=True)

        backup_files = []
        for filename in os.listdir(self.backup_dir):
            if not filename.endswith(".zip") or filename.startswith("."):
                continue

            file_path = os.path.join(self.backup_dir, filename)
            if not os.path.isfile(file_path):
                continue

            manifest = self.get_backup_manifest(file_path)
            if manifest is None:
                logger.debug(f"跳过无效备份文件: {filename}")
                continue

            stat = os.stat(file_path)
            backup_files.append(
                {
                    "filename": filename,
                    "size": stat.st_size,
                    "created_at": stat.st_mtime,
                    "type": manifest.get("origin", "exported"),
                    "astrbot_version": manifest.get("astrbot_version", "未知"),
                    "exported_at": manifest.get("exported_at"),
                }
            )

        backup_files.sort(key=lambda x: x["created_at"], reverse=True)
        start = (page - 1) * page_size
        end = start + page_size

        return {
            "items": backup_files[start:end],
            "total": len(backup_files),
            "page": page,
            "page_size": page_size,
        }

    def export_backup(self) -> dict:
        task_id = str(uuid.uuid4())
        self._init_task(task_id, "export", "pending")
        asyncio.create_task(self.background_export_task(task_id))
        return {
            "task_id": task_id,
            "message": "export task created, processing in background",
        }

    async def background_export_task(self, task_id: str) -> None:
        try:
            self._update_progress(task_id, status="processing", message="正在初始化...")
            kb_manager = getattr(self.core_lifecycle, "kb_manager", None)
            exporter = AstrBotExporter(
                main_db=self.db,
                kb_manager=kb_manager,
                config_path=os.path.join(self.data_dir, "cmd_config.json"),
            )
            zip_path = await exporter.export_all(
                output_dir=self.backup_dir,
                progress_callback=self._make_progress_callback(task_id),
            )
            self._set_task_result(
                task_id,
                "completed",
                result={
                    "filename": os.path.basename(zip_path),
                    "path": zip_path,
                    "size": os.path.getsize(zip_path),
                },
            )
        except Exception as exc:
            logger.error(f"后台导出任务 {task_id} 失败: {exc}")
            logger.error(traceback.format_exc())
            self._set_task_result(task_id, "failed", error=str(exc))

    async def upload_backup(self, file: Any | None) -> dict:
        if not file:
            raise BackupServiceError("缺少备份文件")
        if not file.filename or not file.filename.endswith(".zip"):
            raise BackupServiceError("请上传 ZIP 格式的备份文件")

        safe_filename = secure_filename(file.filename)
        unique_filename = generate_unique_filename(safe_filename)

        Path(self.backup_dir).mkdir(parents=True, exist_ok=True)
        zip_path = os.path.join(self.backup_dir, unique_filename)
        await self._save_upload(file, zip_path)

        logger.info(
            f"上传的备份文件已保存: {unique_filename} (原始名称: {file.filename})"
        )
        return {
            "filename": unique_filename,
            "original_filename": file.filename,
            "size": os.path.getsize(zip_path),
        }

    def upload_init(self, data: object) -> dict:
        payload = self._payload(data)
        filename = payload.get("filename")
        total_size = payload.get("total_size", 0)

        if not filename:
            raise BackupServiceError("缺少 filename 参数")
        if not filename.endswith(".zip"):
            raise BackupServiceError("请上传 ZIP 格式的备份文件")
        if total_size <= 0:
            raise BackupServiceError("无效的文件大小")

        total_chunks = math.ceil(total_size / CHUNK_SIZE)
        upload_id = str(uuid.uuid4())
        chunk_dir = os.path.join(self.chunks_dir, upload_id)
        Path(chunk_dir).mkdir(parents=True, exist_ok=True)

        safe_filename = secure_filename(filename)
        unique_filename = generate_unique_filename(safe_filename)
        current_time = time.time()
        self.upload_sessions[upload_id] = {
            "filename": unique_filename,
            "original_filename": filename,
            "total_size": total_size,
            "total_chunks": total_chunks,
            "received_chunks": set(),
            "created_at": current_time,
            "last_activity": current_time,
            "chunk_dir": chunk_dir,
        }

        logger.info(
            f"初始化分片上传: upload_id={upload_id}, "
            f"filename={unique_filename}, total_chunks={total_chunks}"
        )

        return {
            "upload_id": upload_id,
            "chunk_size": CHUNK_SIZE,
            "total_chunks": total_chunks,
            "filename": unique_filename,
        }

    async def upload_chunk(
        self,
        *,
        upload_id: str | None,
        chunk_index_str: str | None,
        chunk_file: Any | None,
    ) -> dict:
        if not upload_id or chunk_index_str is None:
            raise BackupServiceError("缺少必要参数")

        try:
            chunk_index = int(chunk_index_str)
        except ValueError as exc:
            raise BackupServiceError("无效的分片索引") from exc

        if not chunk_file:
            raise BackupServiceError("缺少分片数据")
        if upload_id not in self.upload_sessions:
            raise BackupServiceError("上传会话不存在或已过期")

        session = self.upload_sessions[upload_id]
        if chunk_index < 0 or chunk_index >= session["total_chunks"]:
            raise BackupServiceError("分片索引超出范围")

        chunk_path = os.path.join(session["chunk_dir"], f"{chunk_index}.part")
        await self._save_upload(chunk_file, chunk_path)
        session["received_chunks"].add(chunk_index)
        session["last_activity"] = time.time()

        received_count = len(session["received_chunks"])
        total_chunks = session["total_chunks"]
        logger.debug(
            f"接收分片: upload_id={upload_id}, chunk={chunk_index + 1}/{total_chunks}"
        )

        return {
            "received": received_count,
            "total": total_chunks,
            "chunk_index": chunk_index,
        }

    def mark_backup_as_uploaded(self, zip_path: str) -> None:
        try:
            manifest = {"origin": "uploaded", "uploaded_at": datetime.now().isoformat()}
            with zipfile.ZipFile(zip_path, "r") as zf:
                if "manifest.json" in zf.namelist():
                    manifest_data = zf.read("manifest.json")
                    manifest = json.loads(manifest_data.decode("utf-8"))
                    manifest["origin"] = "uploaded"
                    manifest["uploaded_at"] = datetime.now().isoformat()

            with zipfile.ZipFile(zip_path, "a") as zf:
                new_manifest = json.dumps(manifest, ensure_ascii=False, indent=2)
                zf.writestr("manifest.json", new_manifest)

            logger.debug(f"已标记备份为上传来源: {zip_path}")
        except Exception as exc:
            logger.warning(f"标记备份来源失败: {exc}")

    async def upload_complete(self, data: object) -> dict:
        payload = self._payload(data)
        upload_id = payload.get("upload_id")

        if not upload_id:
            raise BackupServiceError("缺少 upload_id 参数")
        if upload_id not in self.upload_sessions:
            raise BackupServiceError("上传会话不存在或已过期")

        session = self.upload_sessions[upload_id]
        received = session["received_chunks"]
        total = session["total_chunks"]

        if len(received) != total:
            missing = set(range(total)) - received
            raise BackupServiceError(f"分片不完整，缺少: {sorted(missing)[:10]}...")

        chunk_dir = session["chunk_dir"]
        filename = session["filename"]

        Path(self.backup_dir).mkdir(parents=True, exist_ok=True)
        output_path = os.path.join(self.backup_dir, filename)

        try:
            with open(output_path, "wb") as outfile:
                for i in range(total):
                    chunk_path = os.path.join(chunk_dir, f"{i}.part")
                    with open(chunk_path, "rb") as chunk_file:
                        while True:
                            data_block = chunk_file.read(8192)
                            if not data_block:
                                break
                            outfile.write(data_block)

            file_size = os.path.getsize(output_path)
            self.mark_backup_as_uploaded(output_path)
            logger.info(f"分片上传完成: {filename}, size={file_size}, chunks={total}")
            await self.cleanup_upload_session(upload_id)

            return {
                "filename": filename,
                "original_filename": session["original_filename"],
                "size": file_size,
            }
        except Exception:
            if os.path.exists(output_path):
                os.remove(output_path)
            raise

    async def upload_abort(self, data: object) -> tuple[dict | None, str | None]:
        payload = self._payload(data)
        upload_id = payload.get("upload_id")
        if not upload_id:
            raise BackupServiceError("缺少 upload_id 参数")

        if upload_id in self.upload_sessions:
            await self.cleanup_upload_session(upload_id)
            logger.info(f"取消分片上传: {upload_id}")

        return None, "上传已取消"

    def check_backup(self, data: object) -> dict:
        payload = self._payload(data)
        filename = self._validate_backup_filename(
            payload.get("filename"),
            missing="缺少 filename 参数",
        )
        zip_path = os.path.join(self.backup_dir, filename)
        if not os.path.exists(zip_path):
            raise BackupServiceError(f"备份文件不存在: {filename}")

        kb_manager = getattr(self.core_lifecycle, "kb_manager", None)
        importer = AstrBotImporter(
            main_db=self.db,
            kb_manager=kb_manager,
            config_path=os.path.join(self.data_dir, "cmd_config.json"),
        )
        return importer.pre_check(zip_path).to_dict()

    def import_backup(self, data: object) -> dict:
        payload = self._payload(data)
        filename = self._validate_backup_filename(
            payload.get("filename"),
            missing="缺少 filename 参数",
        )
        confirmed = payload.get("confirmed", False)
        if not confirmed:
            raise BackupServiceError(
                "请先确认导入。导入将会清空并覆盖现有数据，此操作不可撤销。"
            )

        zip_path = os.path.join(self.backup_dir, filename)
        if not os.path.exists(zip_path):
            raise BackupServiceError(f"备份文件不存在: {filename}")

        task_id = str(uuid.uuid4())
        self._init_task(task_id, "import", "pending")
        asyncio.create_task(self.background_import_task(task_id, zip_path))

        return {
            "task_id": task_id,
            "message": "import task created, processing in background",
        }

    async def background_import_task(self, task_id: str, zip_path: str) -> None:
        try:
            self._update_progress(task_id, status="processing", message="正在初始化...")
            kb_manager = getattr(self.core_lifecycle, "kb_manager", None)
            importer = AstrBotImporter(
                main_db=self.db,
                kb_manager=kb_manager,
                config_path=os.path.join(self.data_dir, "cmd_config.json"),
            )
            result = await importer.import_all(
                zip_path=zip_path,
                mode="replace",
                progress_callback=self._make_progress_callback(task_id),
            )

            if result.success:
                self._set_task_result(task_id, "completed", result=result.to_dict())
            else:
                self._set_task_result(
                    task_id,
                    "failed",
                    error="; ".join(result.errors),
                )
        except Exception as exc:
            logger.error(f"后台导入任务 {task_id} 失败: {exc}")
            logger.error(traceback.format_exc())
            self._set_task_result(task_id, "failed", error=str(exc))

    def get_progress(self, task_id: str | None) -> dict:
        if not task_id:
            raise BackupServiceError("缺少参数 task_id")
        if task_id not in self.backup_tasks:
            raise BackupServiceError("找不到该任务")

        task_info = self.backup_tasks[task_id]
        status = task_info["status"]
        response_data = {
            "task_id": task_id,
            "type": task_info["type"],
            "status": status,
        }

        if status == "processing" and task_id in self.backup_progress:
            response_data["progress"] = self.backup_progress[task_id]
        if status == "completed":
            response_data["result"] = task_info["result"]
        if status == "failed":
            response_data["error"] = task_info["error"]

        return response_data

    def prepare_download(
        self,
        *,
        filename: str | None,
        token: str | None,
        jwt_secret: str | None,
    ) -> BackupDownload:
        if not filename:
            raise BackupServiceError("缺少参数 filename")
        if not token:
            raise BackupServiceError("缺少参数 token")
        if not jwt_secret:
            raise BackupServiceError("服务器配置错误")

        try:
            jwt.decode(
                token,
                jwt_secret,
                algorithms=["HS256"],
                options={
                    "require": ["exp"],
                    "verify_signature": True,
                    "verify_exp": True,
                },
            )
        except jwt.ExpiredSignatureError as exc:
            raise BackupServiceError("Token 已过期，请刷新页面后重试") from exc
        except jwt.InvalidTokenError as exc:
            raise BackupServiceError("Token 无效") from exc

        filename = self._validate_backup_filename(filename, missing="缺少参数 filename")
        file_path = os.path.join(self.backup_dir, filename)
        if not os.path.exists(file_path):
            raise BackupServiceError("备份文件不存在")
        return BackupDownload(path=file_path, filename=filename)

    def delete_backup(self, data: object) -> tuple[dict | None, str | None]:
        payload = self._payload(data)
        filename = self._validate_backup_filename(
            payload.get("filename"),
            missing="缺少参数 filename",
        )
        file_path = os.path.join(self.backup_dir, filename)
        if not os.path.exists(file_path):
            raise BackupServiceError("备份文件不存在")

        os.remove(file_path)
        return None, "删除备份成功"

    def rename_backup(self, data: object) -> dict:
        payload = self._payload(data)
        filename = self._validate_backup_filename(
            payload.get("filename"),
            missing="缺少参数 filename",
        )
        new_name = payload.get("new_name")
        if not new_name:
            raise BackupServiceError("缺少参数 new_name")

        new_name = secure_filename(new_name)
        if new_name.endswith(".zip"):
            new_name = new_name[:-4]
        if not new_name or new_name.replace("_", "") == "":
            raise BackupServiceError("新文件名无效")

        new_filename = f"{new_name}.zip"
        old_path = os.path.join(self.backup_dir, filename)
        if not os.path.exists(old_path):
            raise BackupServiceError("备份文件不存在")

        new_path = os.path.join(self.backup_dir, new_filename)
        if os.path.exists(new_path):
            raise BackupServiceError(f"文件名 '{new_filename}' 已存在")

        os.rename(old_path, new_path)
        logger.info(f"备份文件重命名: {filename} -> {new_filename}")
        return {
            "old_filename": filename,
            "new_filename": new_filename,
        }