libervia-backend/libervia/backend/plugins/plugin_xep_0447.py

#!/usr/bin/env python3

# 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 abc import ABC, abstractmethod
from dataclasses import dataclass, field
from functools import partial
import mimetypes
from pathlib import Path
import time
from typing import Any, ClassVar, NamedTuple, Self, cast

from pydantic import BaseModel, Field
import shortuuid
from treq import collect
from twisted.internet import defer
from twisted.web.http_headers import Headers
from twisted.web.iweb import IResponse
from twisted.words.xish import domish

from libervia.backend.core import exceptions
from libervia.backend.core.constants import Const as C
from libervia.backend.core.core_types import MessageData, SatXMPPEntity
from libervia.backend.core.i18n import _
from libervia.backend.core.log import getLogger
from libervia.backend.plugins.plugin_misc_attach import AttachPlugin
from libervia.backend.plugins.plugin_xep_0066 import XEP_0066
from libervia.backend.plugins.plugin_xep_0103 import URLData, XEP_0103
from libervia.backend.plugins.plugin_xep_0334 import XEP_0334
from libervia.backend.plugins.plugin_xep_0358 import JinglePub, XEP_0358
from libervia.backend.plugins.plugin_xep_0363 import XEP_0363
from libervia.backend.plugins.plugin_xep_0367 import XEP_0367
from libervia.backend.plugins.plugin_xep_0428 import XEP_0428
from libervia.backend.plugins.plugin_xep_0446 import FileMetadata, XEP_0446
from libervia.backend.tools import stream

log = getLogger(__name__)


PLUGIN_INFO = {
    C.PI_NAME: "Stateless File Sharing",
    C.PI_IMPORT_NAME: "XEP-0447",
    C.PI_TYPE: "XEP",
    C.PI_MODES: C.PLUG_MODE_BOTH,
    C.PI_PROTOCOLS: ["XEP-0447"],
    C.PI_DEPENDENCIES: [
        "XEP-0103",
        "XEP-0334",
        "XEP-0358",
        "XEP-0446",
        "ATTACH",
        "DOWNLOAD",
    ],
    C.PI_RECOMMENDATIONS: ["XEP-0066", "XEP-0363", "XEP-0367", "XEP-0428"],
    C.PI_MAIN: "XEP_0447",
    C.PI_HANDLER: "no",
    C.PI_DESCRIPTION: _("""Implementation of XEP-0447 (Stateless File Sharing)"""),
}

NS_SFS = "urn:xmpp:sfs:0"


class Source(ABC, BaseModel):

    type: ClassVar[str]
    encrypted: ClassVar[bool] = False

    def __init_subclass__(cls) -> None:
        super().__init_subclass__()
        if not hasattr(cls, "type"):
            raise TypeError(
                f'Can\'t instantiate {cls.__name__} without "type" class attribute.'
            )

    @classmethod
    @abstractmethod
    def from_element(cls, element: domish.Element) -> Self:
        """Parse an element and return corresponding model

        @param element: element to parse
        @raise exceptions.DataError: the element is invalid
        """

    @abstractmethod
    def to_element(self) -> domish.Element:
        """Convert model to an element

        @return: domish.Element representing the model
        """


class FileSharing(BaseModel):
    """
    Model for handling XEP-0447 <file-sharing> element.
    """

    file: FileMetadata
    sources: list[Source]
    disposition: str | None = Field(
        default=None,
        description="Disposition of the file, either 'attachment' or 'inline'.",
    )
    id: str | None = Field(
        default=None, description="Unique identifier for the file-sharing element."
    )
    _sfs: "XEP_0447 | None" = None

    def to_element(self) -> domish.Element:
        """Build the <file-sharing> element from this instance's data.

        @return: <file-sharing> element.
        """
        file_sharing_elt = domish.Element((NS_SFS, "file-sharing"))

        if self.disposition:
            file_sharing_elt["disposition"] = self.disposition

        if self.id:
            file_sharing_elt["id"] = self.id

        file_sharing_elt.addChild(self.file.to_element())

        sources_elt = file_sharing_elt.addElement("sources")
        for source in self.sources:
            sources_elt.addChild(source.to_element())

        return file_sharing_elt

    @classmethod
    def from_element(cls, file_sharing_elt: domish.Element) -> Self:
        """Create a FileSharing instance from a <file-sharing> element or its parent.

        @param file_sharing_elt: The <file-sharing> element or a parent element.
        @return: FileSharing instance.
        @raise exceptions.NotFound: If the <file-sharing> element is not found.
        """
        assert cls._sfs is not None
        if file_sharing_elt.uri != NS_SFS or file_sharing_elt.name != "file-sharing":
            child_file_sharing_elt = next(
                file_sharing_elt.elements(NS_SFS, "file-sharing"), None
            )
            if child_file_sharing_elt is None:
                raise exceptions.NotFound("<file-sharing> element not found")
            else:
                file_sharing_elt = child_file_sharing_elt

        kwargs = {}
        disposition = file_sharing_elt.getAttribute("disposition")
        if disposition:
            kwargs["disposition"] = disposition

        file_id = file_sharing_elt.getAttribute("id")
        if file_id:
            kwargs["id"] = file_id

        kwargs["file"] = FileMetadata.from_element(file_sharing_elt)
        kwargs["sources"] = cls._sfs.parse_sources_elt(file_sharing_elt)

        return cls(**kwargs)


class URLDataSource(URLData, Source):
    type = "url"

    @classmethod
    def from_element(cls, element: domish.Element) -> Self:
        return super().from_element(element)

    def to_element(self) -> domish.Element:
        return super().to_element()


class JinglePubSource(JinglePub, Source):
    type = "jingle"

    @classmethod
    def from_element(cls, element: domish.Element) -> Self:
        return super().from_element(element)

    def to_element(self) -> domish.Element:
        return super().to_element()


class LegacyAttachment(NamedTuple):
    """Single attachment entry for a backward-compatibility message.

    Keys used when building <sources/> elements for legacy clients.
    """

    file_id: str
    url: str


class LegacyAttachments(NamedTuple):
    """Data stored for a pending MUC legacy attach message.

    Kept until the MUC echo arrives with a stanza-id, after which backward-compatibility
    messages are sent using that stanza-id as the <attach-to/> target.
    """

    mess_data: dict[str, Any]
    attachments: list[LegacyAttachment]
    timestamp: float


@dataclass
class SFSClientData:
    """Per-client data for XEP-0447 plugin.

    Stores pending legacy attach entries for MUC messages (waiting for stanza-id echo).
    """

    # origin_id -> LegacyEntry
    pending_legacy: dict[str, LegacyAttachments] = field(default_factory=dict)


class XEP_0447:
    namespace = NS_SFS

    def __init__(self, host):
        self.host = host
        log.info(_("XEP-0447 (Stateless File Sharing) plugin initialization"))
        host.register_namespace("sfs", NS_SFS)
        FileSharing._sfs = self
        self._sources_handlers: dict[tuple[str, str], type[Source]] = {}
        self._u = cast(XEP_0103, host.plugins["XEP-0103"])
        self._jp = cast(XEP_0358, host.plugins["XEP-0358"])
        self._hints = cast(XEP_0334, host.plugins["XEP-0334"])
        self._m = cast(XEP_0446, host.plugins["XEP-0446"])
        self._http_upload = cast(XEP_0363 | None, host.plugins.get("XEP-0363"))
        self._oob = cast(XEP_0066 | None, host.plugins.get("XEP-0066"))
        self._ma = cast(XEP_0367 | None, host.plugins.get("XEP-0367"))
        self._fb = cast(XEP_0428 | None, host.plugins.get("XEP-0428"))
        self._attach = cast(AttachPlugin, host.plugins["ATTACH"])
        self._attach.register(self.can_handle_attachment, self.attach, priority=1000)
        self.register_source(
            self._u.namespace,
            "url-data",
            URLDataSource,
        )
        self.register_source(
            self._jp.namespace,
            "jinglepub",
            JinglePubSource,
        )
        host.plugins["DOWNLOAD"].register_download_handler(
            self._u.namespace, self.download
        )
        host.trigger.add("message_received", self._message_received_trigger)
        host.trigger.add("message_parse", self._message_parse_trigger)
        host.trigger.add("message_sent", self._message_sent_trigger)
        host.trigger.add("disconnected", self._on_connection_closed)

    def register_source(
        self,
        namespace: str,
        element_name: str,
        source: type[Source],
    ) -> None:
        """Register a handler for file source

        @param namespace: namespace of the element supported
        @param element_name: name of the element supported
        @param callback: method to call to parse the element
            get the matching element as argument, must return the parsed data
        @param encrypted: if True, the source is encrypted (the transmitting channel
            should then be end2end encrypted to avoir leaking decrypting data to servers).
        """
        key = (namespace, element_name)
        if key in self._sources_handlers:
            raise exceptions.ConflictError(
                f"There is already a resource handler for namespace {namespace!r} and "
                f"name {element_name!r}"
            )
        self._sources_handlers[key] = source

    async def download(
        self,
        client: SatXMPPEntity,
        attachment: dict[str, Any],
        source: dict[str, Any],
        dest_path: Path | str,
        extra: dict[str, Any] | None = None,
    ) -> tuple[str, defer.Deferred]:
        # TODO: handle url-data headers
        if extra is None:
            extra = {}
        try:
            download_url = source["url"]
        except KeyError:
            raise ValueError(f"{source} has missing URL")

        if extra.get("ignore_tls_errors", False):
            log.warning("TLS certificate check disabled, this is highly insecure")
            treq_client = client.treq_ignore_tls
        else:
            treq_client = client.treq

        try:
            file_size = int(attachment["size"])
        except (KeyError, ValueError):
            head_data = cast(IResponse, await treq_client.head(download_url))
            content_length = cast(Headers, head_data.headers).getRawHeaders(
                "content-length"
            )
            if content_length is None:
                raise ValueError(
                    f"No content-length header in HEAD response for {download_url}"
                )
            file_size = int(content_length[0])

        file_obj = stream.SatFile(
            self.host,
            client,
            dest_path,
            mode="wb",
            size=file_size,
        )

        progress_id = file_obj.uid

        resp = cast(IResponse, await treq_client.get(download_url, unbuffered=True))
        if resp.code == 200:
            d = collect(resp, file_obj.write)
            d.addCallback(lambda __: file_obj.close())
        else:
            d = defer.Deferred()
            self.host.plugins["DOWNLOAD"].errback_download(file_obj, d, resp)
        return progress_id, d

    async def can_handle_attachment(
        self, client: SatXMPPEntity, data: dict[str, Any]
    ) -> bool:
        if self._http_upload is None:
            return False
        try:
            await self._http_upload.get_http_upload_entity(client)
        except exceptions.NotFound:
            return False
        else:
            return True

    def get_sources_elt(
        self, children: list[domish.Element] | None = None
    ) -> domish.Element:
        """Generate <sources> element"""
        sources_elt = domish.Element((NS_SFS, "sources"))
        if children:
            for child in children:
                sources_elt.addChild(child)
        return sources_elt

    def get_file_sharing_elt(
        self,
        sources: list[dict[str, Any]],
        disposition: str | None = None,
        file_id: str | None = None,
        name: str | None = None,
        media_type: str | None = None,
        desc: str | None = None,
        size: int | None = None,
        file_hash: tuple[str, str] | None = None,
        date: float | int | None = None,
        width: int | None = None,
        height: int | None = None,
        length: int | None = None,
        thumbnail: str | None = None,
        **kwargs,
    ) -> domish.Element:
        """Generate the <file-sharing/> element

        @param sources: list of source data dicts
        @param disposition: disposition of the file
        @param file_id: unique identifier for the file-sharing element
        @param name: filename
        @param media_type: media type
        @param desc: description
        @param size: file size
        @param file_hash: tuple of (hash_algo, hash)
        @param date: file date
        @param width: width (for images)
        @param height: height (for images)
        @param length: length (for audio/video)
        @param thumbnail: thumbnail data
        @return: ``<file-sharing/>`` element
        """
        file_sharing_elt = domish.Element((NS_SFS, "file-sharing"))
        if disposition is not None:
            file_sharing_elt["disposition"] = disposition
        if file_id is not None:
            file_sharing_elt["id"] = file_id
        if media_type is None and name:
            media_type = mimetypes.guess_type(name, strict=False)[0]
        file_sharing_elt.addChild(
            self._m.generate_file_metadata(
                name=name,
                media_type=media_type,
                desc=desc,
                size=size,
                file_hash=file_hash,
                date=date,
                width=width,
                height=height,
                length=length,
                thumbnail=thumbnail,
            ).to_element()
        )
        sources_elt = self.get_sources_elt()
        file_sharing_elt.addChild(sources_elt)
        for source_data in sources:
            if "url" in source_data:
                sources_elt.addChild(
                    self._u.generate_url_data(**source_data).to_element()
                )
            else:
                raise NotImplementedError(f"source data not implemented: {source_data}")

        return file_sharing_elt

    def parse_sources_elt(self, sources_elt: domish.Element) -> list[Source]:
        """Parse <sources/> element

        @param sources_elt: <sources/> element, or a direct parent element
        @return: list of found sources data
        @raise: exceptions.NotFound: Can't find <sources/> element
        """
        if sources_elt.name != "sources" or sources_elt.uri != NS_SFS:
            try:
                sources_elt = next(sources_elt.elements(NS_SFS, "sources"))
            except StopIteration:
                raise exceptions.NotFound(
                    f"<sources/> element is missing: {sources_elt.toXml()}"
                )
        sources = []
        for elt in sources_elt.elements():
            if not elt.uri:
                log.warning("ignoring source element {elt.toXml()}")
                continue
            key = (elt.uri, elt.name)
            try:
                source_handler = self._sources_handlers[key]
            except KeyError:
                log.warning(f"unmanaged file sharing element: {elt.toXml}")
                continue
            else:
                source = source_handler.from_element(elt)
                sources.append(source)
        return sources

    def parse_file_sharing_elt(self, file_sharing_elt: domish.Element) -> dict[str, Any]:
        """Parse <file-sharing/> element and return file-sharing data

        @param file_sharing_elt: <file-sharing/> element
        @return: file-sharing data. It a dict whose keys correspond to
            [get_file_sharing_elt] parameters
        """
        if file_sharing_elt.name != "file-sharing" or file_sharing_elt.uri != NS_SFS:
            try:
                file_sharing_elt = next(file_sharing_elt.elements(NS_SFS, "file-sharing"))
            except StopIteration:
                raise exceptions.NotFound
        try:
            data = self._m.parse_file_metadata_elt(file_sharing_elt).model_dump()
        except exceptions.NotFound:
            data = {}
        disposition = file_sharing_elt.getAttribute("disposition")
        if disposition is not None:
            data["disposition"] = disposition
        file_id = file_sharing_elt.getAttribute("id")
        if file_id is not None:
            data["id"] = file_id
        try:
            data["sources"] = self.parse_sources_elt(file_sharing_elt)
        except exceptions.NotFound:
            # XEP-0447 allows <file-sharing/> without <sources/> (e.g. upload in progress)
            data["sources"] = []

        return data

    def _add_file_sharing_attachments(
        self, client: SatXMPPEntity, message_elt: domish.Element, data: dict[str, Any]
    ) -> dict[str, Any]:
        """Check <message> for shared files, and add them as attachments."""
        # check for pending legacy attachments
        origin_id = data["extra"].get("origin_id")
        if origin_id:
            sfs_data = client.get_plugin_data("XEP-0447", SFSClientData)
            if origin_id in sfs_data.pending_legacy:
                stanza_id = data["extra"].get("stanza_id")
                if stanza_id:
                    entry_data = sfs_data.pending_legacy.pop(origin_id)
                    for entry in entry_data.attachments:
                        d = defer.ensureDeferred(
                            self._send_legacy_attach_message(
                                client,
                                entry_data.mess_data,
                                stanza_id,
                                entry.file_id,
                                entry.url,
                            )
                        )
                        d.addErrback(
                            lambda __: log.exception(
                                f"Error sending legacy attach message"
                            )
                        )
                else:
                    log.warning(
                        f"MUC message for origin_id={origin_id!r} has no stanza-id, "
                        "skipping legacy attach messages."
                    )
                    del sfs_data.pending_legacy[origin_id]

            # purge expired entries
            now = time.monotonic()
            expired = [
                k for k, v in sfs_data.pending_legacy.items() if now - v.timestamp > 60
            ]
            if expired:
                log.warning(f"Purging {len(expired)} expired legacy attach entries")
                for k in expired:
                    del sfs_data.pending_legacy[k]

        for file_sharing_elt in message_elt.elements(NS_SFS, "file-sharing"):
            attachment = self.parse_file_sharing_elt(file_sharing_elt)

            if any(
                s.get(C.MESS_KEY_ENCRYPTED, False) for s in attachment["sources"]
            ) and client.encryption.isEncrypted(data):
                # we don't add the encrypted flag if the message itself is not encrypted,
                # because the decryption key is part of the link, so sending it over
                # unencrypted channel is like having no encryption at all.
                attachment[C.MESS_KEY_ENCRYPTED] = True

            attachments = data["extra"].setdefault(C.KEY_ATTACHMENTS, [])
            attachments.append(attachment)

        return data

    async def attach(self, client: SatXMPPEntity, data: dict[str, Any]) -> None:
        """Attach files to a message.

        All files are uploaded and included in a single message with multiple
        <file-sharing/> elements, each identified by a unique id.
        Backward-compatible attach messages are sent for legacy clients.
        """
        attachments = data["extra"][C.KEY_ATTACHMENTS]

        # upload all files
        await self._attach.upload_files(client, data)
        self._hints.add_hint_elements(data["xml"], [self._hints.HINT_STORE])

        # build <file-sharing/> element for each attachment
        for attachment in attachments:
            file_id = attachment.get("name") or shortuuid.uuid()
            try:
                file_hash = (attachment["hash_algo"], attachment["hash"])
            except KeyError:
                file_hash = None
            file_sharing_elt = self.get_file_sharing_elt(
                [{"url": attachment["url"]}],
                file_id=file_id,
                name=attachment.get("name"),
                size=attachment.get("size"),
                desc=attachment.get("desc"),
                media_type=attachment.get("media_type"),
                file_hash=file_hash,
            )
            data["xml"].addChild(file_sharing_elt)

        # store legacy info for deferred sending after main message
        if self._oob is not None and self._ma is not None and self._fb is not None:
            legacy_attachments: list[LegacyAttachment] = []
            for attachment in attachments:
                file_id = attachment.get("name") or shortuuid.uuid()
                legacy_attachments.append(
                    LegacyAttachment(file_id=file_id, url=attachment["url"])
                )
            data["extra"]["_sfs_legacy"] = legacy_attachments

        if (
            not data["extra"]
            and (not data["message"] or data["message"] == {"": ""})
            and not data["subject"]
        ):
            # nothing left to send, we can cancel the message
            raise exceptions.CancelError("Cancelled by XEP_0447 attachment handling")

    async def _send_legacy_attach_message(
        self,
        client: SatXMPPEntity,
        data: dict[str, Any],
        main_msg_id: str,
        file_id: str,
        url: str,
    ) -> None:
        """Send a backward-compatible attach message for legacy clients.

        @param client: Client instance.
        @param data: Original message data (for to, type)
        @param main_msg_id: Uid of the main message containing <file-sharing/>
        @param file_id: Id of the <file-sharing/> element this attaches to
        @param url: Download url for the file
        """
        assert self._oob is not None and self._ma is not None and self._fb is not None
        mess_data: dict[str, Any] = {
            "from": client.jid,
            "to": data["to"],
            "uid": "",
            "message": {"": url},
            "type": data["type"],
            "subject": {},
            "extra": {},
        }
        client.generate_message_xml(cast(MessageData, mess_data))
        message_elt = mess_data["xml"]

        # <attach-to/>
        message_elt.addChild(
            self._ma.get_attach_to_elt(uid=main_msg_id, stanza_id=main_msg_id)
        )

        # <sources/>
        sources_elt = self.get_sources_elt()
        sources_elt["id"] = file_id
        sources_elt.addChild(self._u.generate_url_data(url=url).to_element())
        message_elt.addChild(sources_elt)

        # <fallback/>
        self._fb.add_fallback_elt(message_elt, NS_SFS, url)

        # OOB
        message_elt.addChild(self._oob.generate_oob(url))

        await client.send_message_data(cast(MessageData, mess_data))

    def _message_received_trigger(
        self,
        client: SatXMPPEntity,
        message_elt: domish.Element,
        post_treat: defer.Deferred,
    ) -> bool:
        """Register post-treatment to add file-sharing attachments to message data.

        We use a post_treat callback instead of "message_parse" trigger because we need
        to check if the "encrypted" flag is set to decide if we add the same flag to the
        attachment.
        """
        post_treat.addCallback(
            partial(self._add_file_sharing_attachments, client, message_elt)
        )
        return True

    def _on_connection_closed(self, client: SatXMPPEntity, reason: Any) -> bool:
        """Clean up pending legacy entries when a client disconnects."""
        sfs_data = client.get_plugin_data("XEP-0447", SFSClientData)
        if sfs_data.pending_legacy:
            log.warning(
                f"Clearing {len(sfs_data.pending_legacy)} pending legacy entries for "
                f"{client.jid.userhost()}"
            )
            sfs_data.pending_legacy.clear()
        return True

    async def _message_sent_trigger(
        self, client: SatXMPPEntity, mess_data: dict[str, Any]
    ) -> None:
        """Send backward-compatibility messages after the main message.

        For non-groupchat, send immediately using origin_id.
        For groupchat, defer until we receive the MUC echo with stanza-id.
        """
        legacy_attachments = mess_data["extra"].get("_sfs_legacy")
        if legacy_attachments is None:
            return
        origin_id = mess_data["extra"].get("origin_id") or mess_data["uid"]
        if mess_data["type"] == C.MESS_TYPE_GROUPCHAT:
            # defer: wait for MUC echo to get stanza-id
            sfs_data = client.get_plugin_data("XEP-0447", SFSClientData)
            sfs_data.pending_legacy[origin_id] = LegacyAttachments(
                mess_data=mess_data,
                attachments=legacy_attachments,
                timestamp=time.monotonic(),
            )
            return
        for entry in legacy_attachments:
            await self._send_legacy_attach_message(
                client, mess_data, origin_id, entry.file_id, entry.url
            )

    def _message_parse_trigger(
        self,
        client: SatXMPPEntity,
        message_elt: domish.Element,
        data: dict[str, Any],
    ) -> bool:
        """Block backward-compatibility messages from continuing the workflow.

        If the message contains both <attach-to/> and <sources/>, it is a legacy
        compatibility message generated by our own multi-file sending. We raise
        CancelError to stop the message from being stored in history or sent to
        frontends.
        """
        ns_msg_att = "urn:xmpp:message-attaching:1"
        attach_to_elt = next(message_elt.elements(ns_msg_att, "attach-to"), None)
        sources_elt = next(message_elt.elements(NS_SFS, "sources"), None)

        if attach_to_elt is not None and sources_elt is not None:
            log.debug(
                f"Ignoring backward-compatibility message for "
                f"attach-to id={attach_to_elt['id']!r}"
            )
            raise exceptions.CancelError(
                "Cancelled by XEP_0447 (backward-compatibility message)"
            )
        return True