libervia-backend/libervia/backend/tools/backup/pie.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 datetime import datetime, timezone
import json
import os
import shutil
import time

from lxml import etree
from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import cast
from sqlalchemy.types import Text

from libervia.backend.memory.sqla_mapping import (
    ParamGen,
    ParamInd,
    PrivateGen,
    PrivateGenBin,
    PrivateInd,
    PrivateIndBin,
    Profile,
    PubsubItem,
    PubsubNode,
)

from .core import (
    BackupContext,
    BackupError,
    NS_PARAMS,
    NS_PRIVATE,
    NS_PRIVATE_BIN,
    NS_ROSTER_VER,
    PARAM_FILE,
    ProfileData,
    ProfileHandler,
    ROSTER_VER_KEY,
    Report,
    Staging,
)
from .xep_0227 import (
    NS_PUBSUB,
    NS_ROSTER,
    append_xml_payload,
    build_pep_items,
    build_roster,
    item_payload,
)


async def _get_profile_id(session, profile_name: str) -> int | None:
    """Get the database id of a profile

    @param session: database session
    @param profile_name: profile to look for
    @return: profile id, or None if the profile is unknown
    """
    return (
        await session.execute(select(Profile.id).where(Profile.name == profile_name))
    ).scalar_one_or_none()


class ParamsHandler(ProfileHandler):
    name = "params"
    uses_db = True

    def known_ns(self) -> tuple[str, ...]:
        return (NS_PARAMS,)

    async def export(
        self,
        ctx: BackupContext,
        staging: Staging,
        profile: ProfileData,
        report: Report,
    ) -> list[etree._Element]:
        section = await self._section(ctx, profile)
        if ctx.param_file.is_file() and PARAM_FILE not in staging:
            staging.add_file(PARAM_FILE, ctx.param_file, 0o644)
        return [section]

    async def import_(
        self,
        ctx: BackupContext,
        session: AsyncSession,
        staging: Staging,
        profile_name: str,
        profile_jid: str,
        sections: dict[str, list[etree._Element]],
        report: Report,
        force: bool = False,
    ) -> None:
        for section in sections.get(NS_PARAMS, []):
            await self._import(session, profile_name, section, report)
        if PARAM_FILE in staging:
            shutil.copyfile(staging.path(PARAM_FILE), ctx.param_file)
            os.chmod(ctx.param_file, 0o600)
            report.done.append("params: parameters file")

    async def _section(self, ctx: BackupContext, profile: ProfileData) -> etree._Element:
        """Build the params section from the database

        @param ctx: backup context
        @param profile: profile to export
        @return: the params section element
        """
        async with ctx.session() as session:
            gen = (
                (
                    await session.execute(
                        select(ParamGen).order_by(ParamGen.category, ParamGen.name)
                    )
                )
                .scalars()
                .all()
            )
            ind = (
                (
                    await session.execute(
                        select(ParamInd)
                        .where(ParamInd.profile_id == profile.id)
                        .order_by(ParamInd.category, ParamInd.name)
                    )
                )
                .scalars()
                .all()
            )
        section = etree.Element(f"{{{NS_PARAMS}}}params", nsmap={None: NS_PARAMS})
        for scope, rows in (("general", gen), ("individual", ind)):
            for row in rows:
                # values are stored as-is in the database, XML only carries strings
                attr = {
                    "category": str(row.category),
                    "name": str(row.name),
                    "scope": scope,
                }
                if row.value is not None:
                    attr["value"] = str(row.value)
                etree.SubElement(section, f"{{{NS_PARAMS}}}param", **attr)
        return section

    async def _import(
        self, session, profile_name: str, section: etree._Element, report: Report
    ) -> None:
        """Import a params section into the database

        @param session: database session
        @param profile_name: profile to import for
        @param section: params section element
        @param report: import report (individual rows of unknown profiles are skipped)
        """
        profile_id = await _get_profile_id(session, profile_name)
        gen_rows = {
            (row.category, row.name): row
            for row in (await session.execute(select(ParamGen))).scalars()
        }
        ind_rows = {}
        if profile_id is not None:
            ind_rows = {
                (row.category, row.name): row
                for row in (
                    await session.execute(
                        select(ParamInd).where(ParamInd.profile_id == profile_id)
                    )
                ).scalars()
            }
        count = 0
        for elt in section.findall(f"{{{NS_PARAMS}}}param"):
            category, name = elt.get("category", ""), elt.get("name", "")
            value = elt.get("value")
            if elt.get("scope", "general") == "general":
                row = gen_rows.get((category, name))
                if row is None:
                    session.add(ParamGen(category=category, name=name, value=value))
                else:
                    row.value = value
                count += 1
            elif profile_id is None:
                report.skipped.append(f"params: unknown profile {profile_name!r}")
            else:
                row = ind_rows.get((category, name))
                if row is None:
                    session.add(
                        ParamInd(
                            category=category,
                            name=name,
                            profile_id=profile_id,
                            value=value,
                        )
                    )
                else:
                    row.value = value
                count += 1
        if count:
            report.done.append(f"params: {count} value(s)")


class PrivateHandler(ProfileHandler):
    name = "private"
    uses_db = True

    def known_ns(self) -> tuple[str, ...]:
        return (NS_PRIVATE, NS_PRIVATE_BIN)

    async def export(
        self,
        ctx: BackupContext,
        staging: Staging,
        profile: ProfileData,
        report: Report,
    ) -> list[etree._Element]:
        sections: list[etree._Element] = []
        for bin_ in (False, True):
            if (section := await self._section(ctx, profile, bin_)) is not None:
                sections.append(section)
        return sections

    async def import_(
        self,
        ctx: BackupContext,
        session: AsyncSession,
        staging: Staging,
        profile_name: str,
        profile_jid: str,
        sections: dict[str, list[etree._Element]],
        report: Report,
        force: bool = False,
    ) -> None:
        for ns, bin_ in ((NS_PRIVATE, False), (NS_PRIVATE_BIN, True)):
            for section in sections.get(ns, []):
                await self._import(
                    session, profile_name, section, bin_=bin_, report=report
                )

    async def _section(
        self, ctx: BackupContext, profile: ProfileData, bin_: bool = False
    ) -> etree._Element | None:
        """Build a private data section (key-value storage, or its binary variant)

        Keys are not necessarily valid XML element names, so each value is stored
        in a <row> element with namespace, key and scope attributes (Libervia
        extension).

        @param ctx: backup context
        @param profile: profile to export
        @param bin_: True for the binary (JSON) values
        @return: the private section element, or None if empty
        """
        ns = NS_PRIVATE_BIN if bin_ else NS_PRIVATE
        gen_model = PrivateGenBin if bin_ else PrivateGen
        ind_model = PrivateIndBin if bin_ else PrivateInd
        async with ctx.session() as session:
            gen = (
                (
                    await session.execute(
                        select(gen_model).order_by(gen_model.namespace, gen_model.key)
                    )
                )
                .scalars()
                .all()
            )
            ind = (
                (
                    await session.execute(
                        select(ind_model)
                        .where(ind_model.profile_id == profile.id)
                        .order_by(ind_model.namespace, ind_model.key)
                    )
                )
                .scalars()
                .all()
            )
        rows = [("general", row) for row in gen] + [("individual", row) for row in ind]
        if not rows:
            return None
        section = etree.Element(f"{{{ns}}}private", nsmap={None: ns})
        for scope, row in rows:
            attr = {
                "namespace": str(row.namespace or ""),
                "key": str(row.key),
                "scope": scope,
            }
            if row.value is None:
                # XML can't distinguish an empty value from a missing one
                attr["nil"] = "true"
            elt = etree.SubElement(section, f"{{{ns}}}row", **attr)
            if row.value is not None:
                if bin_:
                    # JSON values: serialize back with the same serializer
                    # used by the Json column type (byte-identical)
                    elt.text = json.dumps(row.value, ensure_ascii=False)
                else:
                    elt.text = str(row.value)
        return section

    async def _import(
        self,
        session,
        profile_name: str,
        section: etree._Element,
        bin_: bool,
        report: Report,
    ) -> None:
        """Import a private data section into the database

        @param session: database session
        @param profile_name: profile to import for
        @param section: private section element
        @param bin_: True for the binary (JSON) values
        @param report: import report (individual rows of unknown profiles are skipped)
        """
        suffix = "_bin" if bin_ else ""
        gen_model = PrivateGenBin if bin_ else PrivateGen
        ind_model = PrivateIndBin if bin_ else PrivateInd
        profile_id = await _get_profile_id(session, profile_name)
        gen_rows = {
            (row.namespace, row.key): row
            for row in (await session.execute(select(gen_model))).scalars()
        }
        ind_rows = {}
        if profile_id is not None:
            ind_rows = {
                (row.namespace, row.key): row
                for row in (
                    await session.execute(
                        select(ind_model).where(ind_model.profile_id == profile_id)
                    )
                ).scalars()
            }
        count = 0
        for elt in section.findall(f"{{{NS_PRIVATE_BIN if bin_ else NS_PRIVATE}}}row"):
            namespace, key = elt.get("namespace", ""), elt.get("key", "")
            nil = elt.get("nil") == "true"
            if bin_ and not nil:
                try:
                    # the Json column type will serialize it back
                    value = json.loads(elt.text or "")
                except json.JSONDecodeError as e:
                    raise BackupError(
                        f"private{suffix}: invalid JSON for {namespace}/{key!r}: {e}"
                    ) from e
            else:
                value = None if nil else (elt.text or "")
            if elt.get("scope", "general") == "general":
                row = gen_rows.get((namespace, key))
                if row is None:
                    session.add(gen_model(namespace=namespace, key=key, value=value))
                else:
                    row.value = value
                count += 1
            elif profile_id is None:
                report.skipped.append(
                    f"private: unknown profile {profile_name!r}, row ignored"
                )
            else:
                row = ind_rows.get((namespace, key))
                if row is None:
                    session.add(
                        ind_model(
                            namespace=namespace,
                            key=key,
                            value=value,
                            profile_id=profile_id,
                        )
                    )
                else:
                    row.value = value
                count += 1
        if count:
            report.done.append(f"private{suffix}: {count} value(s)")


class RosterHandler(ProfileHandler):
    name = "roster"
    uses_db = True

    def known_ns(self) -> tuple[str, ...]:
        return (NS_ROSTER,)

    async def export(
        self,
        ctx: BackupContext,
        staging: Staging,
        profile: ProfileData,
        report: Report,
    ) -> list[etree._Element]:
        # the roster is exported from its persistent cache in the database, kept
        # up to date on every connection and roster push (roster versioning)
        items: list[etree._Element] = []
        version: str | None = None
        async with ctx.session() as session:
            rows = (
                (
                    await session.execute(
                        select(PrivateInd)
                        .where(
                            PrivateInd.profile_id == profile.id,
                            PrivateInd.namespace == NS_ROSTER_VER,
                        )
                        .order_by(PrivateInd.key)
                    )
                )
                .scalars()
                .all()
            )
        for row in rows:
            if row.key == ROSTER_VER_KEY:
                version = row.value
            else:
                items.append(etree.fromstring(row.value.encode()))
        if not items:
            return []
        return [build_roster(items, version)]

    async def import_(
        self,
        ctx: BackupContext,
        session: AsyncSession,
        staging: Staging,
        profile_name: str,
        profile_jid: str,
        sections: dict[str, list[etree._Element]],
        report: Report,
        force: bool = False,
    ) -> None:
        # importing the roster needs a server connection (the roster has to be pushed to
        # the server), it is not implemented yet
        if sections.get(NS_ROSTER):
            report.skipped.append(
                f"{profile_name}: roster not imported (not implemented yet, "
                "needs a server connection)"
            )


class PepHandler(ProfileHandler):
    name = "pep"
    uses_db = True

    def known_ns(self) -> tuple[str, ...]:
        return (NS_PUBSUB,)

    async def export(
        self,
        ctx: BackupContext,
        staging: Staging,
        profile: ProfileData,
        report: Report,
    ) -> list[etree._Element]:
        section = await self._section(ctx, profile)
        if section is None:
            report.skipped.append("pep: no PEP node in cache")
            return []
        return [section]

    async def import_(
        self,
        ctx: BackupContext,
        session: AsyncSession,
        staging: Staging,
        profile_name: str,
        profile_jid: str,
        sections: dict[str, list[etree._Element]],
        report: Report,
        force: bool = False,
    ) -> None:
        await self._import(
            session, profile_name, profile_jid, sections.get(NS_PUBSUB, []), report
        )

    async def _section(
        self, ctx: BackupContext, profile: ProfileData
    ) -> etree._Element | None:
        """Build the PEP items section from the pubsub cache

        Node configuration is not cached locally, so only items are exported.

        @param ctx: backup context
        @param profile: profile to export
        @return: the PEP section element, or None if the profile has no PEP nodes
        """
        async with ctx.session() as session:
            # the service and data columns carry app-specific types (JID, parsed
            # XML), so they are cast to text to keep the backup byte-identical
            node_rows = (
                await session.execute(
                    select(
                        PubsubNode.id,
                        PubsubNode.name,
                        cast(PubsubNode.service, Text),
                    )
                    .where(PubsubNode.profile_id == profile.id, PubsubNode.pep.is_(True))
                    .order_by(PubsubNode.name)
                )
            ).all()
            if not node_rows:
                return None
            item_rows = (
                await session.execute(
                    select(
                        PubsubItem.node_id,
                        PubsubItem.name,
                        cast(PubsubItem.data, Text),
                    )
                    .where(
                        PubsubItem.node_id.in_([node_id for node_id, _, _ in node_rows])
                    )
                    .order_by(PubsubItem.name)
                )
            ).all()
        by_id = {node_id: (name, service) for node_id, name, service in node_rows}
        nodes: dict[tuple[str, str | None], list[tuple[str, str, str | None]]] = {}
        for node_id, item_name, data in item_rows:
            payload, publisher = self._item_parts(data)
            name, service = by_id[node_id]
            nodes.setdefault((name, service), []).append(
                (item_name or "", payload or "", publisher)
            )
        return build_pep_items(nodes)

    async def _import(
        self,
        session,
        profile_name: str,
        jid: str,
        sections: list[etree._Element],
        report: Report,
    ) -> None:
        """Import the PEP items sections into the pubsub cache

        Rows are written with plain SQL (the columns carry app-specific types)
        to keep the data byte-identical.

        @param session: database session
        @param profile_name: profile to import for
        @param jid: profile JID, used as the node service
        @param sections: PEP items section elements
        @param report: import report (unknown profiles are skipped)
        """
        profile_id = await _get_profile_id(session, profile_name)
        if profile_id is None:
            report.skipped.append(f"pep: unknown profile {profile_name!r}")
            return
        for section in sections:
            for items in section.findall(f"{{{NS_PUBSUB}}}items"):
                node_name = items.get("node")
                if not node_name:
                    continue
                # the node owner is not necessarily the pie user (e.g. service profiles),
                # so the service is stored in the items element (Libervia extension);
                # pie documents from other servers don't have it, and there the node
                # owner is the pie user
                service = items.get("service") or jid
                await session.execute(
                    text(
                        "INSERT INTO pubsub_nodes (profile_id, service, name, pep, "
                        "subscribed, sync_state_updated, max_items) "
                        "VALUES (:profile_id, :service, :name, 1, 0, :now, 0) "
                        "ON CONFLICT (profile_id, service, name, pep) "
                        "DO UPDATE SET pep = 1"
                    ),
                    {
                        "profile_id": profile_id,
                        "service": service,
                        "name": node_name,
                        "now": time.time(),
                    },
                )
                node_id = (
                    await session.execute(
                        text(
                            "SELECT id FROM pubsub_nodes "
                            "WHERE profile_id=:profile_id AND service=:service "
                            "AND name=:name"
                        ),
                        {
                            "profile_id": profile_id,
                            "service": service,
                            "name": node_name,
                        },
                    )
                ).scalar_one()
                now = datetime.now(timezone.utc)
                for item in items.findall(f"{{{NS_PUBSUB}}}item"):
                    data = self._item_data(
                        item.get("id", ""), item.get("publisher"), item_payload(item)
                    )
                    await session.execute(
                        text(
                            "INSERT INTO pubsub_items (node_id, name, data, created, "
                            "updated) VALUES (:node_id, :name, :data, :now, :now) "
                            "ON CONFLICT (node_id, name) DO UPDATE SET "
                            "data = excluded.data, updated = excluded.updated"
                        ),
                        {
                            "node_id": node_id,
                            "name": item.get("id", ""),
                            "data": data,
                            "now": now,
                        },
                    )
                report.done.append(f"pep: node {node_name!r}")

    def _item_parts(self, data: str | None) -> tuple[str | None, str | None]:
        """Extract the payload and publisher of a stored pubsub item

        The data column stores the full <item> element, but the pie document
        carries the payload and the publisher instead.

        @param data: stored item data
        @return: (payload, publisher), payload is an XML string or plain text
        """
        if not data:
            return data, None
        try:
            root = etree.fromstring(data.encode())
        except etree.XMLSyntaxError:
            return data, None
        if etree.QName(root).localname != "item":
            return data, None
        return item_payload(root), root.get("publisher")

    def _item_data(self, item_id: str, publisher: str | None, payload: str) -> str:
        """Build the data to store in the pubsub cache for a pie item

        The payload is wrapped in a pubsub <item> element, as in the cache.

        @param item_id: item id
        @param publisher: item publisher (attribute of the <item> element)
        @param payload: item payload (XML string or plain text)
        @return: data to store
        """
        # the item element is in the pubsub namespace (XEP-0060), as in the cache
        wrapper = etree.Element(f"{{{NS_PUBSUB}}}item", nsmap={None: NS_PUBSUB})
        if item_id:
            wrapper.set("id", item_id)
        if publisher:
            wrapper.set("publisher", publisher)
        append_xml_payload(wrapper, payload)
        return etree.tostring(wrapper, encoding="unicode")