data_portability.py

"""SQLite/设置便携备份；可选 Windows DPAPI 本地账户加密。"""
from __future__ import annotations

import ctypes
from ctypes import wintypes
import datetime
import json
import os
import shutil
import tempfile
import zipfile

import history_store
import paths

MAGIC = b"NEKO-DPAPI-v1\0"


class _Blob(ctypes.Structure):
    _fields_ = [("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_byte))]


def _blob(data: bytes):
    buf = ctypes.create_string_buffer(data)
    return _Blob(len(data), ctypes.cast(buf, ctypes.POINTER(ctypes.c_byte))), buf


def protect_bytes(data: bytes) -> bytes:
    if os.name != "nt":
        raise OSError("DPAPI 仅支持 Windows")
    src, keep = _blob(data); out = _Blob()
    crypt = ctypes.windll.crypt32.CryptProtectData
    crypt.argtypes = [ctypes.POINTER(_Blob), wintypes.LPCWSTR, ctypes.c_void_p,
                      ctypes.c_void_p, ctypes.c_void_p, wintypes.DWORD, ctypes.POINTER(_Blob)]
    crypt.restype = wintypes.BOOL
    if not crypt(ctypes.byref(src), "NekoCardReader backup", None, None, None, 0,
                 ctypes.byref(out)):
        raise ctypes.WinError()
    try:
        return ctypes.string_at(out.pbData, out.cbData)
    finally:
        ctypes.windll.kernel32.LocalFree(out.pbData)


def unprotect_bytes(data: bytes) -> bytes:
    if os.name != "nt":
        raise OSError("DPAPI 仅支持 Windows")
    src, keep = _blob(data); out = _Blob()
    crypt = ctypes.windll.crypt32.CryptUnprotectData
    crypt.argtypes = [ctypes.POINTER(_Blob), ctypes.c_void_p, ctypes.c_void_p,
                      ctypes.c_void_p, ctypes.c_void_p, wintypes.DWORD, ctypes.POINTER(_Blob)]
    crypt.restype = wintypes.BOOL
    if not crypt(ctypes.byref(src), None, None, None, None, 0, ctypes.byref(out)):
        raise ctypes.WinError()
    try:
        return ctypes.string_at(out.pbData, out.cbData)
    finally:
        ctypes.windll.kernel32.LocalFree(out.pbData)


def create_bundle(destination: str, encrypted=False, include_settings=True) -> str:
    destination = os.path.abspath(destination)
    os.makedirs(os.path.dirname(destination), exist_ok=True)
    with tempfile.TemporaryDirectory() as td:
        db = os.path.join(td, "history_v3.sqlite3")
        export = os.path.join(td, "history_export.json")
        history_store.create_backup(db); history_store.export_json(export)
        metadata = {"format": 1, "created_at": datetime.datetime.now().isoformat(),
                    "encrypted": bool(encrypted), "stats": history_store.database_stats()}
        with open(os.path.join(td, "metadata.json"), "w", encoding="utf-8") as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        settings_path = paths.data_path("settings.json")
        if include_settings and os.path.exists(settings_path):
            shutil.copy2(settings_path, os.path.join(td, "settings.json"))
        raw_zip = os.path.join(td, "bundle.zip")
        with zipfile.ZipFile(raw_zip, "w", zipfile.ZIP_DEFLATED) as z:
            for name in ("history_v3.sqlite3", "history_export.json", "metadata.json", "settings.json"):
                p = os.path.join(td, name)
                if os.path.exists(p):
                    z.write(p, name)
        with open(raw_zip, "rb") as f:
            data = f.read()
        if encrypted:
            data = MAGIC + protect_bytes(data)
        with open(destination, "wb") as f:
            f.write(data)
    return destination


def restore_bundle(source: str, restore_settings=True) -> dict:
    source = os.path.abspath(source)
    with open(source, "rb") as f:
        data = f.read()
    encrypted = data.startswith(MAGIC)
    if encrypted:
        data = unprotect_bytes(data[len(MAGIC):])
    with tempfile.TemporaryDirectory() as td:
        raw_zip = os.path.join(td, "bundle.zip")
        with open(raw_zip, "wb") as f:
            f.write(data)
        with zipfile.ZipFile(raw_zip) as z:
            allowed = {"history_v3.sqlite3", "history_export.json", "metadata.json", "settings.json"}
            if any(name not in allowed or "/" in name or "\\" in name for name in z.namelist()):
                raise ValueError("备份包含不安全路径")
            z.extractall(td)
        db = os.path.join(td, "history_v3.sqlite3")
        if not os.path.exists(db):
            raise ValueError("备份中缺少数据库")
        backup_dir = os.path.join(os.path.dirname(os.path.abspath(history_store._DB)), "backups")
        os.makedirs(backup_dir, exist_ok=True)
        emergency = os.path.join(backup_dir, "pre_restore_" +
                                 datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + ".sqlite3")
        history_store.create_backup(emergency)
        history_store.restore_backup(db)
        settings_src = os.path.join(td, "settings.json")
        if restore_settings and os.path.exists(settings_src):
            shutil.copy2(settings_src, paths.data_path("settings.json"))
        metadata_path = os.path.join(td, "metadata.json")
        if os.path.exists(metadata_path):
            with open(metadata_path, encoding="utf-8") as f:
                metadata = json.load(f)
        else:
            metadata = {}
    return {"encrypted": encrypted, "emergency_backup": emergency,
            "metadata": metadata, "stats": history_store.database_stats()}


def bundle_info(source: str) -> dict:
    with open(source, "rb") as f:
        data = f.read()
    encrypted = data.startswith(MAGIC)
    if encrypted:
        data = unprotect_bytes(data[len(MAGIC):])
    with tempfile.TemporaryDirectory() as td:
        p = os.path.join(td, "bundle.zip")
        with open(p, "wb") as f:
            f.write(data)
        with zipfile.ZipFile(p) as z:
            if "metadata.json" not in z.namelist():
                return {"encrypted": encrypted, "metadata": {}}
            return {"encrypted": encrypted,
                    "metadata": json.loads(z.read("metadata.json").decode("utf-8"))}