libervia-backend/libervia/backend/tools/backup/__init__.py

#!/usr/bin/env python3

# Libervia: an XMPP client
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from collections.abc import Callable
from datetime import datetime, timezone
import hashlib
import importlib.metadata
import lzma
import os
from pathlib import Path
import stat
import tarfile
import time
from typing import Any

from alembic.runtime import migration as al_migration
from lxml import etree
from pydantic import ValidationError
from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from libervia.backend.core.log import getLogger
from libervia.backend.memory import sqla_config
from libervia.backend.memory.sqla_mapping import ParamInd, Profile
from libervia.backend.tools.common.utils import is_safe_relpath

from .core import (
    BackupContext,
    BackupError,
    CHUNK_SIZE,
    FORMAT,
    FORMAT_VERSION,
    GlobalHandler,
    Handler,
    MANIFEST_FILE,
    PIE_FILE_TPL,
    ProfileData,
    ProfileHandler,
    Report,
    Staging,
    get_db_driver,
    sqlite_db_config,
)
from .handlers import DATA_TYPES, DB_BACKED_TYPES, REGISTRY
from .models import Manifest, ProfileMeta
from .xep_0227 import (
    PieError,
    build_pie,
    parse_pie,
    sections_by_ns,
)

# pie section namespaces handled by this version of Libervia, derived from the handlers
# registry
KNOWN_SECTION_NS = frozenset(ns for handler in REGISTRY for ns in handler.known_ns())

# pie entry name check, derived from the pie file template
_PIE_NAME_PREFIX, _PIE_NAME_SUFFIX = PIE_FILE_TPL.split("{}")

log = getLogger(__name__)

__all__ = [
    "Backup",
    "BackupContext",
    "BackupError",
    "DATA_TYPES",
    "DB_BACKED_TYPES",
    "GlobalHandler",
    "Handler",
    "Manifest",
    "PIE_FILE_TPL",
    "ProfileData",
    "ProfileHandler",
    "ProfileMeta",
    "Report",
    "Staging",
    "get_db_driver",
    "sqlite_db_config",
]


class _ProgressFile:
    """File-like wrapper that reports write progress (for archive compression)"""

    def __init__(
        self, stream, progress_cb: Callable[[int, int], None] | None, total: int
    ):
        self._stream = stream
        self._progress_cb = progress_cb
        self._total = total
        self._pos = 0
        self._reported = 0

    def write(self, data: bytes) -> int:
        count = self._stream.write(data)
        self._pos += count
        if self._progress_cb is not None and (
            self._pos - self._reported >= 256 * 1024 or self._pos >= self._total
        ):
            self._reported = self._pos
            self._progress_cb(min(self._pos, self._total), self._total)
        return count

    def flush(self) -> None:
        self._stream.flush()

    def tell(self) -> int:
        return self._pos

    def close(self) -> None:
        self._stream.close()


class Backup:
    """Backup/restore engine"""

    @staticmethod
    async def open_context(
        db_config: dict,
        param_file: Path,
        local_dir: Path,
        readonly: bool = False,
    ) -> BackupContext:
        """Open a backup context on a database

        @param db_config: database config (see "sqla_config.get_db_config")
        @param param_file: path of the parameters file
        @param local_dir: local data directory (GPG homes, ...)
        @param readonly: open the database in read only mode (export)
        @return: backup context
        """
        path = Path(db_config["path"])
        log.debug(f"opening database {path} for backup (readonly={readonly})")
        url = sqla_config.get_db_url(path, readonly)
        engine = create_async_engine(url)
        return BackupContext(
            engine=engine,
            session_factory=async_sessionmaker(engine, expire_on_commit=False),
            db_url=url,
            param_file=param_file,
            local_dir=local_dir,
            db=get_db_driver(db_config),
        )

    @staticmethod
    async def get_profiles(ctx: BackupContext) -> list[ProfileData]:
        """List profiles with their bare JID

        @param ctx: backup context
        @return: profiles data
        """
        async with ctx.session() as session:
            profs = (
                (await session.execute(select(Profile).order_by(Profile.name)))
                .scalars()
                .all()
            )
            jid_rows = (
                await session.execute(
                    select(ParamInd.profile_id, ParamInd.value).where(
                        ParamInd.category == "Connection",
                        ParamInd.name == "JabberID",
                    )
                )
            ).all()
        jids = {
            profile_id: (value or "").strip().split("/")[0]
            for profile_id, value in jid_rows
        }
        return [
            ProfileData(name=p.name, id=p.id, jid=jids.get(p.id) or None) for p in profs
        ]

    @staticmethod
    def get_app_version() -> str:
        """Libervia version, as reported by the package metadata"""
        try:
            return importlib.metadata.version("libervia-backend")
        except importlib.metadata.PackageNotFoundError:
            return "unknown"

    @staticmethod
    async def get_alembic_rev(ctx: BackupContext) -> str | None:
        """Get the alembic (schema) revision of the database"""

        def _get_heads(conn) -> tuple[str, ...]:
            context = al_migration.MigrationContext.configure(conn)
            return context.get_current_heads()

        async with ctx.engine.connect() as conn:
            try:
                heads = await conn.run_sync(_get_heads)
            except Exception:
                log.error(
                    f"no alembic revision: {ctx.db_url} has no alembic_version table"
                )
                return None
        return heads[0] if heads else None

    @staticmethod
    def get_handlers(types: list[str]) -> list[Handler]:
        """Get the handlers of data types, in export order

        @param types: data types (see DATA_TYPES)
        @return: the matching handlers
        """
        wanted = set(types)
        return [handler for handler in REGISTRY if handler.name in wanted]

    @staticmethod
    async def export_profile(
        ctx: BackupContext,
        staging: Staging,
        profile: ProfileData,
        handlers: list[ProfileHandler],
    ) -> Report:
        """Export a profile data to the staging

        @param ctx: backup context
        @param staging: staging holding the archive entries
        @param profile: profile to export
        @param handlers: profile handlers to run (in export order)
        @return: export report
        """
        report = Report()
        sections: list[etree._Element] = []
        for handler in handlers:
            sections.extend(await handler.export(ctx, staging, profile, report))
        if sections:
            staging.add_bytes(
                PIE_FILE_TPL.format(profile.name),
                build_pie(profile.host_jid, profile.user_name, sections),
                0o600,
            )
        return report

    @staticmethod
    async def import_profile(
        ctx: BackupContext,
        staging: Staging,
        profile_name: str,
        handlers: list[ProfileHandler],
        report: Report,
        force: bool = False,
    ) -> None:
        """Import a profile data into the database

        Unknown pie sections are reported and ignored. The roster section is not imported
        yet.

        @param ctx: backup context
        @param staging: staging holding the archive entries (pie document and parameters
            file)
        @param profile_name: profile to import for
        @param handlers: profile handlers to run (in export order)
        @param report: import report
        @param force: the restore is dangerous, the user confirmed it
        """
        pie_file = PIE_FILE_TPL.format(profile_name)
        if pie_file not in staging:
            report.skipped.append(f"{profile_name}: no data in archive")
            return
        try:
            host_jid, user_name, sections = parse_pie(staging.path(pie_file).read_bytes())
        except PieError as e:
            raise BackupError(f"{profile_name}: {e}") from e
        jid = f"{user_name}@{host_jid}" if host_jid else user_name
        by_ns = sections_by_ns(sections)
        for ns in sorted(set(by_ns) - KNOWN_SECTION_NS):
            report.skipped.append(f"unknown section {ns!r} ignored")
        async with ctx.session() as session:
            for handler in handlers:
                await handler.import_(
                    ctx, session, staging, profile_name, jid, by_ns, report, force=force
                )
            await session.commit()

    @staticmethod
    def build_manifest(
        app_version: str,
        alembic_rev: str | None,
        profiles: dict[str, ProfileMeta],
        staging: Staging,
    ) -> Manifest:
        """Build the archive manifest and stage it

        @param app_version: Libervia version
        @param alembic_rev: database schema revision
        @param profiles: profile name: exported data (jid, types)
        @param staging: staging holding the archive entries (the manifest is added to it)
        @return: the manifest
        """
        manifest = Manifest(
            format=FORMAT,
            format_version=FORMAT_VERSION,
            app="libervia",
            app_version=app_version,
            alembic=alembic_rev,
            date=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            profiles=profiles,
            files={name: {"sha256": staging.get_hash(name)} for name in staging.names()},
        )
        staging.add_bytes(
            MANIFEST_FILE, manifest.model_dump_json(indent=2).encode(), 0o644
        )
        return manifest

    @staticmethod
    def _load_manifest(data: bytes) -> Manifest:
        """Parse and validate the archive manifest

        @raise BackupError: if the manifest is invalid or not from this format
        """
        try:
            manifest = Manifest.model_validate_json(data)
        except ValidationError as e:
            raise BackupError(f"invalid manifest: {e}") from e
        if manifest.format != FORMAT:
            raise BackupError(f"unsupported format: {manifest.format!r}")
        if manifest.format_version != FORMAT_VERSION:
            raise BackupError(f"unsupported format version: {manifest.format_version}")
        return manifest

    @staticmethod
    def _manifest_of(tar: tarfile.TarFile) -> Manifest:
        """Read and validate the manifest of an opened archive

        @param tar: opened backup archive
        @return: archive manifest
        @raise BackupError: if the manifest is missing or invalid
        """
        manifest_file = tar.extractfile(MANIFEST_FILE)
        if manifest_file is None:
            raise BackupError(f"missing manifest: {MANIFEST_FILE!r}")
        return Backup._load_manifest(manifest_file.read())

    @staticmethod
    def _data_members(tar: tarfile.TarFile) -> list[tarfile.TarInfo]:
        """Archive entries holding data (the manifest is not included)"""
        return [m for m in tar.getmembers() if m.isfile() and m.name != MANIFEST_FILE]

    @staticmethod
    def make_archive(
        staging: Staging,
        manifest: Manifest,
        dest_dir: Path,
        progress_cb: Callable[[int, int], None] | None = None,
    ) -> Path:
        """Pack the staged backup entries in a tar.xz archive

        @param staging: staging holding the archive entries (including manifest)
        @param manifest: archive manifest (used for the archive name)
        @param dest_dir: directory where the archive is written
        @param progress_cb: optional progress callback ``callable(written, total)``
        @return: path of the created archive
        """
        dest_dir.mkdir(parents=True, exist_ok=True)
        date = datetime.fromisoformat(manifest.date)
        base_name = f"libervia-backup-{date.strftime('%Y-%m-%d_%H-%M-%S')}"
        path = dest_dir / f"{base_name}.tar.xz"
        counter = 1
        while path.exists():
            path = dest_dir / f"{base_name}-{counter}.tar.xz"
            counter += 1
        mtime = int(time.time())
        # approximate size of the tar stream (for the progress bar): each entry is a
        # 512-byte header plus its data padded to a 512-byte block, and the 20 KB
        # margin covers the end-of-archive marker and the extra headers tar writes
        # for long entry names
        total = sum(
            512 + ((staging.size(name) + 511) // 512) * 512 for name in staging.names()
        )
        total += 20 * 1024
        with lzma.open(path, "wb") as comp:
            wrapper = _ProgressFile(comp, progress_cb, total)
            # the tarfile stubs type fileobj with read/seek, which are not used
            # in write mode
            fileobj: Any = wrapper
            with tarfile.open(fileobj=fileobj, mode="w") as tar:
                for name in staging.names():
                    src = staging.path(name)
                    src_stat = src.stat()
                    info = tarfile.TarInfo(name)
                    info.size = src_stat.st_size
                    info.mtime = mtime
                    # the entry mode is set by the one who staged it (private
                    # entries are owner-only)
                    info.mode = stat.S_IMODE(src_stat.st_mode)
                    with src.open("rb") as f:
                        tar.addfile(info, f)
        wrapper.close()
        if progress_cb is not None:
            progress_cb(total, total)
        os.chmod(path, 0o600)
        return path

    @staticmethod
    def read_archive(
        path: Path,
        staging: Staging,
        progress_cb: Callable[[int, int], None] | None = None,
    ) -> Manifest:
        """Read and check a backup archive, extracting its entries to the staging

        @param path: path to the archive
        @param staging: staging where the entries are extracted
        @param progress_cb: optional progress callback ``callable(done, total)``
        @return: archive manifest
        @raise BackupError: if the archive is invalid or tampered
        """
        try:
            with tarfile.open(path, "r:xz") as tar:
                manifest = Backup._manifest_of(tar)
                members = Backup._data_members(tar)
                total = sum(m.size for m in members)
                done = 0
                for member in members:
                    if not is_safe_relpath(member.name):
                        raise BackupError(f"invalid entry name {member.name!r}")
                    dst = staging.path(member.name)
                    dst.parent.mkdir(parents=True, exist_ok=True)
                    digest = hashlib.sha256()
                    src_f = tar.extractfile(member)
                    if src_f is None:
                        raise BackupError(f"can't extract {member.name!r}")
                    with dst.open("wb") as out:
                        while chunk := src_f.read(CHUNK_SIZE):
                            out.write(chunk)
                            digest.update(chunk)
                            done += len(chunk)
                            if progress_cb:
                                progress_cb(done, total)
                    # we don't want special bits in the permissions (to avoid setuid,
                    # for instance)
                    os.chmod(dst, member.mode & 0o777)
                    expected = manifest.files.get(member.name, {}).get("sha256")
                    if expected is None:
                        # every entry of a valid archive is listed in the
                        # manifest: an unlisted one means the archive was
                        # modified after the export
                        raise BackupError(
                            f"unlisted entry {member.name!r} (archive tampered?)"
                        )
                    if digest.hexdigest() != expected:
                        raise BackupError(
                            f"{member.name}: checksum mismatch (archive tampered?)"
                        )
        except (tarfile.TarError, OSError) as e:
            raise BackupError(f"can't read archive: {e}") from e
        for name in manifest.files:
            if name not in staging:
                raise BackupError(f"{name}: missing from archive")
        return manifest

    @staticmethod
    def verify_archive(
        path: Path,
        progress_cb: Callable[[str, int, int], None] | None = None,
        entry_cb: Callable[[str, dict], None] | None = None,
    ) -> tuple[Manifest, dict[str, dict]]:
        """Stream-check a backup archive (checksums and pie validity)

        The entries are not extracted, only hashed (and parsed, for pie documents).

        @param path: path to the archive
        @param progress_cb: optional (entry name, bytes done, total bytes) callback,
            called on each read chunk
        @param entry_cb: optional (entry name, entry info) callback, called on each
            entry completion
        @return: (manifest, per-entry info: size, sha256, checksum status, pie error if
            any)
        @raise BackupError: if the archive is unreadable or the manifest is missing
        """
        try:
            with tarfile.open(path, "r:xz") as tar:
                manifest = Backup._manifest_of(tar)
                members = Backup._data_members(tar)
                total = sum(m.size for m in members)
                results: dict[str, dict] = {}
                done = 0
                for member in members:
                    is_pie = member.name.startswith(
                        _PIE_NAME_PREFIX
                    ) and member.name.endswith(_PIE_NAME_SUFFIX)
                    digest = hashlib.sha256()
                    data = b""
                    src_f = tar.extractfile(member)
                    if src_f is None:
                        raise BackupError(f"can't extract {member.name!r}")
                    while chunk := src_f.read(CHUNK_SIZE):
                        digest.update(chunk)
                        if is_pie:
                            data += chunk
                        done += len(chunk)
                        if progress_cb:
                            progress_cb(member.name, done, total)
                    pie_error = None
                    if is_pie:
                        try:
                            parse_pie(data)
                        except PieError as e:
                            pie_error = f"invalid pie: {e}"
                    expected = manifest.files.get(member.name, {}).get("sha256")
                    info = {
                        "size": member.size,
                        "sha256": digest.hexdigest(),
                        "checksum": (
                            "ok" if expected == digest.hexdigest() else "mismatch"
                        ),
                        "pie": pie_error,
                    }
                    results[member.name] = info
                    if entry_cb:
                        entry_cb(member.name, info)
                return manifest, results
        except (tarfile.TarError, OSError) as e:
            raise BackupError(f"can't read archive: {e}") from e