ForcePilot/backend/package/yuxi/channels/adapters/qqbot/credential_backup.py

from __future__ import annotations

import base64
import hashlib
import json
import logging
import os
import tempfile
import time
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

DEFAULT_BACKUP_DIR = os.path.join(tempfile.gettempdir(), "yuxi_qqbot_credentials")


def _get_fernet():
    try:
        from cryptography.fernet import Fernet
    except ImportError:
        return None

    key = os.environ.get("QQBOT_CREDENTIAL_ENCRYPTION_KEY")
    if not key:
        return None
    digest = hashlib.sha256(key.encode()).digest()
    return Fernet(base64.urlsafe_b64encode(digest))


def _encrypt_field(fernet, value: str) -> str:
    if not fernet or not value:
        return value
    return fernet.encrypt(value.encode()).decode()


def _decrypt_field(fernet, value: str) -> str:
    if not fernet or not value:
        return value
    try:
        return fernet.decrypt(value.encode()).decode()
    except Exception:
        return value


@dataclass
class CredentialSnapshot:
    app_id: str = ""
    app_secret: str = ""
    access_token: str = ""
    expires_at: float = 0
    token_obtained_at: float = 0
    session_id: str = ""
    sandbox: bool = False
    metadata: dict = field(default_factory=dict)

    def is_valid(self) -> bool:
        return bool(self.app_id and self.app_secret)

    def token_expired(self) -> bool:
        if not self.access_token or not self.expires_at:
            return True
        return time.monotonic() > self.expires_at - 300


class CredentialBackup:
    def __init__(self, app_id: str, backup_dir: str | None = None):
        self._app_id = app_id
        self._backup_dir = backup_dir or DEFAULT_BACKUP_DIR
        self._backup_path = os.path.join(self._backup_dir, f"{app_id}.json")

    def save(self, snapshot: CredentialSnapshot) -> bool:
        try:
            os.makedirs(self._backup_dir, exist_ok=True)

            fernet = _get_fernet()
            data = {
                "app_id": snapshot.app_id,
                "app_secret": _encrypt_field(fernet, snapshot.app_secret),
                "access_token": _encrypt_field(fernet, snapshot.access_token),
                "expires_at": snapshot.expires_at,
                "token_obtained_at": snapshot.token_obtained_at,
                "session_id": snapshot.session_id,
                "sandbox": snapshot.sandbox,
                "metadata": snapshot.metadata,
                "saved_at": time.time(),
            }

            tmp_path = self._backup_path + ".tmp"
            with open(tmp_path, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False)
            os.replace(tmp_path, self._backup_path)
            os.chmod(self._backup_path, 0o600)

            logger.info(
                "CredentialBackup: saved snapshot for app_id=%s%s",
                self._app_id[:6] + "..." if len(self._app_id) > 6 else self._app_id,
                " (encrypted)" if fernet else "",
            )
            return True
        except OSError:
            logger.exception("CredentialBackup: failed to save snapshot")
            return False

    def restore(self) -> CredentialSnapshot | None:
        try:
            if not os.path.exists(self._backup_path):
                return None

            with open(self._backup_path, encoding="utf-8") as f:
                data = json.load(f)

            fernet = _get_fernet()
            snapshot = CredentialSnapshot(
                app_id=data.get("app_id", ""),
                app_secret=_decrypt_field(fernet, data.get("app_secret", "")),
                access_token=_decrypt_field(fernet, data.get("access_token", "")),
                expires_at=data.get("expires_at", 0),
                token_obtained_at=data.get("token_obtained_at", 0),
                session_id=data.get("session_id", ""),
                sandbox=data.get("sandbox", False),
                metadata=data.get("metadata", {}),
            )

            if not snapshot.is_valid():
                logger.warning("CredentialBackup: restored snapshot is invalid for app_id=%s", self._app_id[:6] + "...")
                return None

            logger.info("CredentialBackup: restored snapshot for app_id=%s", self._app_id[:6] + "...")
            return snapshot
        except (OSError, json.JSONDecodeError, KeyError):
            logger.exception("CredentialBackup: failed to restore snapshot")
            return None

    def clear(self) -> bool:
        try:
            if os.path.exists(self._backup_path):
                os.remove(self._backup_path)
            tmp_path = self._backup_path + ".tmp"
            if os.path.exists(tmp_path):
                os.remove(tmp_path)
            logger.info("CredentialBackup: cleared backup for app_id=%s", self._app_id[:6] + "...")
            return True
        except OSError:
            logger.exception("CredentialBackup: failed to clear backup")
            return False

    @staticmethod
    def cleanup_expired(backup_dir: str | None = None, max_age_s: float = 86400 * 7) -> int:
        directory = backup_dir or DEFAULT_BACKUP_DIR
        removed = 0

        if not os.path.exists(directory):
            return 0

        now = time.time()
        try:
            for filename in os.listdir(directory):
                if not filename.endswith(".json"):
                    continue
                filepath = os.path.join(directory, filename)
                try:
                    stat = os.stat(filepath)
                    if now - stat.st_mtime > max_age_s:
                        os.remove(filepath)
                        removed += 1
                        logger.debug("CredentialBackup: removed expired backup %s", filename)
                except OSError:
                    pass
        except OSError:
            logger.exception("CredentialBackup: cleanup failed")

        return removed