notes_backup.py

from __future__ import annotations

import hashlib
import json
import os
import secrets
import shutil
import sqlite3
import stat
import tempfile
import zipfile
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Callable


class BackupError(RuntimeError):
    """Privacy-safe verified-backup failure."""


def _is_link_or_reparse_point(path: Path) -> bool:
    try:
        details = os.lstat(path)
    except OSError:
        return False
    reparse_flag = getattr(stat, "FILE_ATTRIBUTE_REPARSE_POINT", 0x400)
    return path.is_symlink() or bool(
        getattr(details, "st_file_attributes", 0) & reparse_flag
    )


@dataclass(frozen=True, slots=True)
class BackupResult:
    path: Path
    size: int
    sha256: str


NowProvider = Callable[[], datetime]


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _safe_attachment_items(raw: object) -> list[dict[str, object]]:
    if not isinstance(raw, str) or not raw:
        return []
    try:
        value = json.loads(raw)
    except json.JSONDecodeError:
        return []
    return [item for item in value if isinstance(item, dict)] if isinstance(value, list) else []


def _normalize_snapshot_attachment_paths(
    connection: sqlite3.Connection,
    *,
    root: Path,
) -> None:
    table_names = {
        str(row[0])
        for row in connection.execute(
            "SELECT name FROM sqlite_master WHERE type='table'"
        )
    }
    if "notes" not in table_names:
        return
    columns = {str(row[1]) for row in connection.execute("PRAGMA table_info(notes)")}
    if "attachments" not in columns:
        return

    for note_id, raw_attachments in connection.execute(
        "SELECT id, attachments FROM notes WHERE attachments <> ''"
    ).fetchall():
        items = _safe_attachment_items(raw_attachments)
        changed = False
        for item in items:
            raw_path = str(item.get("path") or "")
            candidate = Path(raw_path)
            if not raw_path or not candidate.is_absolute():
                continue
            try:
                relative = candidate.resolve().relative_to(root)
            except (OSError, ValueError):
                continue
            if not relative.parts or relative.parts[0].casefold() != "attachments":
                continue
            item["path"] = relative.as_posix()
            changed = True
        if changed:
            connection.execute(
                "UPDATE notes SET attachments=? WHERE id=?",
                (json.dumps(items, ensure_ascii=False), note_id),
            )


def _make_snapshot(root: Path, snapshot_path: Path) -> None:
    db_path = root / "notes.db"
    source: sqlite3.Connection | None = None
    snapshot: sqlite3.Connection | None = None
    try:
        source = sqlite3.connect(
            db_path.resolve().as_uri() + "?mode=ro",
            uri=True,
            timeout=5,
        )
        source.execute("PRAGMA busy_timeout=5000")
        snapshot = sqlite3.connect(snapshot_path)
        source.backup(snapshot)
        _normalize_snapshot_attachment_paths(snapshot, root=root)
        snapshot.commit()
        integrity = snapshot.execute("PRAGMA integrity_check").fetchone()
        if integrity is None or str(integrity[0]).casefold() != "ok":
            raise BackupError("備份資料庫完整性檢查失敗")
    finally:
        if snapshot is not None:
            snapshot.close()
        if source is not None:
            source.close()


def _iter_safe_attachment_files(root: Path) -> Iterable[tuple[Path, Path]]:
    attachment_root = root / "attachments"
    if _is_link_or_reparse_point(attachment_root):
        raise BackupError("attachment root cannot be a symlink, junction, or reparse point")
    if not attachment_root.is_dir():
        return
    resolved_root = attachment_root.resolve()
    for current, directory_names, file_names in os.walk(
        attachment_root,
        topdown=True,
        followlinks=False,
    ):
        current_path = Path(current)
        safe_directories: list[str] = []
        for name in directory_names:
            candidate = current_path / name
            if _is_link_or_reparse_point(candidate):
                continue
            try:
                candidate.resolve(strict=True).relative_to(resolved_root)
            except (OSError, ValueError):
                continue
            safe_directories.append(name)
        directory_names[:] = safe_directories

        for name in file_names:
            candidate = current_path / name
            if _is_link_or_reparse_point(candidate):
                continue
            try:
                resolved = candidate.resolve(strict=True)
                relative = resolved.relative_to(resolved_root)
            except (OSError, ValueError):
                continue
            if resolved.is_file():
                yield resolved, Path("attachments") / relative


def _verify_archive(path: Path) -> None:
    with zipfile.ZipFile(path) as archive:
        names = archive.namelist()
        if (
            names.count("notes.db") != 1
            or len(names) != len(set(names))
            or archive.testzip() is not None
        ):
            raise BackupError("備份壓縮檔驗證失敗")
        for name in names:
            archive_path = Path(name)
            if archive_path.is_absolute() or ".." in archive_path.parts:
                raise BackupError("備份壓縮檔路徑驗證失敗")


def create_verified_backup(
    data_root: str | Path,
    *,
    output: str | Path | None = None,
    now_fn: NowProvider = datetime.now,
) -> BackupResult:
    """Create an atomic, integrity-checked SQLite + attachment backup archive."""

    root = Path(data_root).expanduser().resolve()
    db_path = root / "notes.db"
    if not db_path.is_file():
        raise BackupError("找不到可備份的 Quick Notes 資料庫")

    backup_dir = root / "backups"
    timestamp = now_fn().strftime("%Y%m%d_%H%M%S_%f")
    target = (
        Path(output).expanduser().resolve()
        if output is not None
        else (backup_dir / f"quick_notes_pre_migration_{timestamp}.zip").resolve()
    )
    if target.exists():
        raise BackupError("備份目標已存在")

    temporary_archive = target.parent / f".{target.name}.{secrets.token_hex(8)}.tmp"
    try:
        backup_dir.mkdir(parents=True, exist_ok=True)
        target.parent.mkdir(parents=True, exist_ok=True)
        with tempfile.TemporaryDirectory(prefix=".quick-notes-backup-", dir=backup_dir) as temp_dir:
            snapshot_path = Path(temp_dir) / "notes.db"
            _make_snapshot(root, snapshot_path)
            with zipfile.ZipFile(
                temporary_archive,
                "w",
                compression=zipfile.ZIP_DEFLATED,
                compresslevel=9,
            ) as archive:
                archive.write(snapshot_path, "notes.db")
                for path, archive_name in _iter_safe_attachment_files(root):
                    archive.write(path, archive_name.as_posix())
            _verify_archive(temporary_archive)
            os.replace(temporary_archive, target)
        _verify_archive(target)
        return BackupResult(
            path=target,
            size=target.stat().st_size,
            sha256=_sha256(target),
        )
    except BackupError:
        temporary_archive.unlink(missing_ok=True)
        raise
    except (OSError, sqlite3.Error, RuntimeError, ValueError, zipfile.BadZipFile) as exc:
        temporary_archive.unlink(missing_ok=True)
        raise BackupError("無法建立並驗證 Quick Notes 備份") from exc


def restore_database_from_verified_backup(
    data_root: str | Path,
    backup: BackupResult,
) -> None:
    """Atomically restore only notes.db after verifying the retained archive."""

    root = Path(data_root).expanduser().resolve()
    archive_path = backup.path.expanduser().resolve()
    temporary = root / f".restore-notes-{secrets.token_hex(8)}.tmp"
    descriptor: int | None = None
    try:
        if (
            not archive_path.is_file()
            or archive_path.stat().st_size != backup.size
            or _sha256(archive_path) != backup.sha256
        ):
            raise BackupError("備份檔驗證失敗，無法自動還原")
        _verify_archive(archive_path)
        root.mkdir(parents=True, exist_ok=True)
        descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(descriptor, "wb") as destination:
            descriptor = None
            with zipfile.ZipFile(archive_path) as archive:
                with archive.open("notes.db", "r") as source:
                    shutil.copyfileobj(source, destination, length=1024 * 1024)
            destination.flush()
            os.fsync(destination.fileno())

        restored: sqlite3.Connection | None = None
        try:
            restored = sqlite3.connect(
                temporary.resolve().as_uri() + "?mode=ro",
                uri=True,
                timeout=5,
            )
            integrity = restored.execute("PRAGMA integrity_check").fetchone()
            if integrity is None or str(integrity[0]).casefold() != "ok":
                raise BackupError("備份資料庫完整性檢查失敗，無法自動還原")
        finally:
            if restored is not None:
                restored.close()

        for suffix in ("-wal", "-shm", "-journal"):
            (root / f"notes.db{suffix}").unlink(missing_ok=True)
        os.replace(temporary, root / "notes.db")
    except BackupError:
        if descriptor is not None:
            os.close(descriptor)
        temporary.unlink(missing_ok=True)
        raise
    except (OSError, sqlite3.Error, RuntimeError, ValueError, zipfile.BadZipFile) as exc:
        if descriptor is not None:
            os.close(descriptor)
        temporary.unlink(missing_ok=True)
        raise BackupError("無法從 verified backup 自動還原資料庫") from exc