slidge 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

slidge/core/mixins/attachment.py

@@ -1,3 +1,4 @@
+import base64
 import functools
 import logging
 import os
@@ -7,22 +8,23 @@ import stat
 import tempfile
 import warnings
 from datetime import datetime
-from mimetypes import guess_type
+from itertools import chain
+from mimetypes import guess_extension, guess_type
 from pathlib import Path
-from typing import IO, Collection, Optional, Sequence, Union
+from typing import IO, AsyncIterator, Collection, Optional, Sequence, Union
 from urllib.parse import quote as urlquote
 from uuid import uuid4
 from xml.etree import ElementTree as ET
 
-import blurhash
-from PIL import Image
+import thumbhash
+from PIL import Image, ImageOps
 from slixmpp import JID, Message
 from slixmpp.exceptions import IqError
 from slixmpp.plugins.xep_0363 import FileUploadError
-from slixmpp.plugins.xep_0385.stanza import Sims
 from slixmpp.plugins.xep_0447.stanza import StatelessFileSharing
 
-from ...util.sql import db
+from ...db.avatar import avatar_cache
+from ...slixfix.xep_0264.stanza import Thumbnail
 from ...util.types import (
     LegacyAttachment,
     LegacyMessageType,
@@ -31,13 +33,13 @@ from ...util.types import (
 )
 from ...util.util import fix_suffix
 from .. import config
-from ..cache import avatar_cache
-from .message_maker import MessageMaker
+from .message_text import TextMessageMixin
 
 
-class AttachmentMixin(MessageMaker):
-    def send_text(self, *_, **k) -> Optional[Message]:
-        raise NotImplementedError
+class AttachmentMixin(TextMessageMixin):
+    def __init__(self, *a, **kw):
+        super().__init__(*a, **kw)
+        self.__store = self.xmpp.store.attachments
 
     async def __upload(
         self,
@@ -138,6 +140,7 @@ class AttachmentMixin(MessageMaker):
     async def __get_url(
         self,
         file_path: Optional[Path] = None,
+        async_data_stream: Optional[AsyncIterator[bytes]] = None,
         data_stream: Optional[IO[bytes]] = None,
         data: Optional[bytes] = None,
         file_url: Optional[str] = None,
@@ -146,13 +149,13 @@ class AttachmentMixin(MessageMaker):
         legacy_file_id: Optional[Union[str, int]] = None,
     ) -> tuple[bool, Optional[Path], str]:
         if legacy_file_id:
-            cache = db.attachment_get_url(legacy_file_id)
+            cache = self.__store.get_url(str(legacy_file_id))
             if cache is not None:
                 async with self.session.http.head(cache) as r:
                     if r.status < 400:
                         return False, None, cache
                     else:
-                        db.attachment_remove(legacy_file_id)
+                        self.__store.remove(str(legacy_file_id))
 
         if file_url and config.USE_ATTACHMENT_ORIGINAL_URLS:
             return False, None, file_url
@@ -165,7 +168,12 @@ class AttachmentMixin(MessageMaker):
             )
 
         if file_path is None:
-            file_name = str(uuid4()) if file_name is None else file_name
+            if file_name is None:
+                file_name = str(uuid4())
+                if content_type is not None:
+                    ext = guess_extension(content_type, strict=False)  # type:ignore
+                    if ext is not None:
+                        file_name += ext
             temp_dir = Path(tempfile.mkdtemp())
             file_path = temp_dir / file_name
             if file_url:
@@ -173,14 +181,23 @@ class AttachmentMixin(MessageMaker):
                     with file_path.open("wb") as f:
                         f.write(await r.read())
 
-            else:
-                if data_stream is not None:
-                    data = data_stream.read()
+            elif data_stream is not None:
+                data = data_stream.read()
                 if data is None:
                     raise RuntimeError
 
                 with file_path.open("wb") as f:
                     f.write(data)
+            elif async_data_stream is not None:
+                # TODO: patch slixmpp to allow this as data source for
+                #       upload_file() so we don't even have to write anything
+                #       to disk.
+                with file_path.open("wb") as f:
+                    async for chunk in async_data_stream:
+                        f.write(chunk)
+            elif data is not None:
+                with file_path.open("wb") as f:
+                    f.write(data)
 
             is_temp = not bool(config.NO_UPLOAD_PATH)
         else:
@@ -198,7 +215,7 @@ class AttachmentMixin(MessageMaker):
             local_path = file_path
             new_url = await self.__upload(file_path, file_name, content_type)
         if legacy_file_id:
-            db.attachment_store_url(legacy_file_id, new_url)
+            self.__store.set_url(self.session.user_pk, str(legacy_file_id), new_url)
 
         return is_temp, local_path, new_url
 
@@ -210,37 +227,44 @@ class AttachmentMixin(MessageMaker):
         content_type: Optional[str] = None,
         caption: Optional[str] = None,
         file_name: Optional[str] = None,
-    ):
-        cache = db.attachment_get_sims(uploaded_url)
+    ) -> Thumbnail | None:
+        cache = self.__store.get_sims(uploaded_url)
         if cache:
-            msg.append(Sims(xml=ET.fromstring(cache)))
-            return
+            ref = self.xmpp["xep_0372"].stanza.Reference(xml=ET.fromstring(cache))
+            msg.append(ref)
+            if ref["sims"]["file"].get_plugin("thumbnail", check=True):
+                return ref["sims"]["file"]["thumbnail"]
+            else:
+                return None
 
         if not path:
-            return
+            return None
 
-        sims = self.xmpp["xep_0385"].get_sims(
+        ref = self.xmpp["xep_0385"].get_sims(
             path, [uploaded_url], content_type, caption
         )
         if file_name:
-            sims["sims"]["file"]["name"] = file_name
+            ref["sims"]["file"]["name"] = file_name
+        thumbnail = None
         if content_type is not None and content_type.startswith("image"):
             try:
                 h, x, y = await self.xmpp.loop.run_in_executor(
-                    avatar_cache._thread_pool, get_blurhash, path
+                    avatar_cache._thread_pool, get_thumbhash, path
                 )
             except Exception as e:
-                log.debug("Could not generate a blurhash", exc_info=e)
+                log.debug("Could not generate a thumbhash", exc_info=e)
             else:
-                thumbnail = sims["sims"]["file"]["thumbnail"]
+                thumbnail = ref["sims"]["file"]["thumbnail"]
                 thumbnail["width"] = x
                 thumbnail["height"] = y
-                thumbnail["media-type"] = "image/blurhash"
-                thumbnail["uri"] = "data:image/blurhash," + urlquote(h)
+                thumbnail["media-type"] = "image/thumbhash"
+                thumbnail["uri"] = "data:image/thumbhash;base64," + urlquote(h)
 
-        db.attachment_store_sims(uploaded_url, str(sims))
+        self.__store.set_sims(uploaded_url, str(ref))
 
-        msg.append(sims)
+        msg.append(ref)
+
+        return thumbnail
 
     def __set_sfs(
         self,
@@ -250,8 +274,9 @@ class AttachmentMixin(MessageMaker):
         content_type: Optional[str] = None,
         caption: Optional[str] = None,
         file_name: Optional[str] = None,
+        thumbnail: Optional[Thumbnail] = None,
     ):
-        cache = db.attachment_get_sfs(uploaded_url)
+        cache = self.__store.get_sfs(uploaded_url)
         if cache:
             msg.append(StatelessFileSharing(xml=ET.fromstring(cache)))
             return
@@ -262,7 +287,9 @@ class AttachmentMixin(MessageMaker):
         sfs = self.xmpp["xep_0447"].get_sfs(path, [uploaded_url], content_type, caption)
         if file_name:
             sfs["file"]["name"] = file_name
-        db.attachment_store_sfs(uploaded_url, str(sfs))
+        if thumbnail is not None:
+            sfs["file"].append(thumbnail)
+        self.__store.set_sfs(uploaded_url, str(sfs))
 
         msg.append(sfs)
 
@@ -274,6 +301,7 @@ class AttachmentMixin(MessageMaker):
         caption: Optional[str] = None,
         carbon=False,
         when: Optional[datetime] = None,
+        correction=False,
         **kwargs,
     ) -> list[Message]:
         msg["oob"]["url"] = uploaded_url
@@ -281,11 +309,19 @@ class AttachmentMixin(MessageMaker):
         if caption:
             m1 = self._send(msg, carbon=carbon, **kwargs)
             m2 = self.send_text(
-                caption, legacy_msg_id=legacy_msg_id, when=when, carbon=carbon, **kwargs
+                caption,
+                legacy_msg_id=legacy_msg_id,
+                when=when,
+                carbon=carbon,
+                correction=correction,
+                **kwargs,
             )
             return [m1, m2] if m2 else [m1]
         else:
-            self._set_msg_id(msg, legacy_msg_id)
+            if correction:
+                msg["replace"]["id"] = self._replace_id(legacy_msg_id)
+            else:
+                self._set_msg_id(msg, legacy_msg_id)
             return [self._send(msg, carbon=carbon, **kwargs)]
 
     async def send_file(
@@ -293,6 +329,7 @@ class AttachmentMixin(MessageMaker):
         file_path: Optional[Union[Path, str]] = None,
         legacy_msg_id: Optional[LegacyMessageType] = None,
         *,
+        async_data_stream: Optional[AsyncIterator[bytes]] = None,
         data_stream: Optional[IO[bytes]] = None,
         data: Optional[bytes] = None,
         file_url: Optional[str] = None,
@@ -309,6 +346,7 @@ class AttachmentMixin(MessageMaker):
         Send a single file from this :term:`XMPP Entity`.
 
         :param file_path: Path to the attachment
+        :param async_data_stream: Alternatively (and ideally) an AsyncIterator yielding bytes
         :param data_stream: Alternatively, a stream of bytes (such as a File object)
         :param data: Alternatively, a bytes object
         :param file_url: Alternatively, a URL
@@ -326,6 +364,16 @@ class AttachmentMixin(MessageMaker):
         carbon = kwargs.pop("carbon", False)
         mto = kwargs.pop("mto", None)
         store_multi = kwargs.pop("store_multi", True)
+        correction = kwargs.get("correction", False)
+        if correction and (original_xmpp_id := self._legacy_to_xmpp(legacy_msg_id)):
+            xmpp_ids = self.xmpp.store.multi.get_xmpp_ids(
+                self.session.user_pk, original_xmpp_id
+            )
+
+            for xmpp_id in xmpp_ids:
+                if xmpp_id == original_xmpp_id:
+                    continue
+                self.retract(xmpp_id, thread)
         msg = self._make_message(
             when=when,
             reply_to=reply_to,
@@ -339,6 +387,7 @@ class AttachmentMixin(MessageMaker):
 
         is_temp, local_path, new_url = await self.__get_url(
             Path(file_path) if file_path else None,
+            async_data_stream,
             data_stream,
             data,
             file_url,
@@ -355,10 +404,12 @@ class AttachmentMixin(MessageMaker):
             self._set_msg_id(msg, legacy_msg_id)
             return None, [self._send(msg, **kwargs)]
 
-        await self.__set_sims(
+        thumbnail = await self.__set_sims(
             msg, new_url, local_path, content_type, caption, file_name
         )
-        self.__set_sfs(msg, new_url, local_path, content_type, caption, file_name)
+        self.__set_sfs(
+            msg, new_url, local_path, content_type, caption, file_name, thumbnail
+        )
         if is_temp and isinstance(local_path, Path):
             local_path.unlink()
             local_path.parent.rmdir()
@@ -472,35 +523,23 @@ class AttachmentMixin(MessageMaker):
                 ids.append(stanza_id["id"])
             else:
                 ids.append(msg.get_id())
-        db.attachment_store_legacy_to_multi_xmpp_msg_ids(legacy_msg_id, ids)
-
-
-def get_blurhash(path: Path, n=9) -> tuple[str, int, int]:
-    img = Image.open(path)
-    width, height = img.size
-    n = min(width, height, n)
-    if width == height:
-        x = y = n
-    elif width > height:
-        x = n
-        y = round(n * height / width)
-    else:
-        x = round(n * width / height)
-        y = n
-    # There are 2 blurhash-python packages:
-    # https://github.com/woltapp/blurhash-python
-    # https://github.com/halcy/blurhash-python
-    # With this hack we're compatible with both, which is useful for packaging
-    # without using pyproject.toml, as most distro do
-    try:
-        hash_ = blurhash.encode(img, x, y)
-    except TypeError:
-        # We are using halcy's blurhash which expects
-        # the 1st argument to be a 3-dimensional array
-        import numpy  # type:ignore
-
-        hash_ = blurhash.encode(numpy.array(img.convert("RGB")), x, y)
-    return hash_, width, height
+        self.xmpp.store.multi.set_xmpp_ids(
+            self.session.user_pk, str(legacy_msg_id), ids
+        )
+
+
+def get_thumbhash(path: Path) -> tuple[str, int, int]:
+    with path.open("rb") as fp:
+        img = Image.open(fp)
+        width, height = img.size
+        img = img.convert("RGBA")
+        if width > 100 or height > 100:
+            img.thumbnail((100, 100))
+    img = ImageOps.exif_transpose(img)
+    rgba_2d = list(img.getdata())
+    rgba = list(chain(*rgba_2d))
+    ints = thumbhash.rgba_to_thumb_hash(img.width, img.height, rgba)
+    return base64.b64encode(bytes(ints)).decode(), width, height
 
 
 log = logging.getLogger(__name__)

slidge/core/mixins/avatar.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional
 
 from slixmpp import JID
 
+from ...db.avatar import CachedAvatar, avatar_cache
 from ...util.types import (
     URL,
     AnyBaseSession,
@@ -12,7 +13,6 @@ from ...util.types import (
     AvatarType,
     LegacyFileIdType,
 )
-from ..cache import avatar_cache
 
 if TYPE_CHECKING:
     from ..pubsub import PepAvatar
@@ -28,13 +28,14 @@ class AvatarMixin:
 
     jid: JID = NotImplemented
     session: AnyBaseSession = NotImplemented
-    _avatar_pubsub_broadcast: bool = NotImplemented
     _avatar_bare_jid: bool = NotImplemented
 
     def __init__(self) -> None:
         super().__init__()
         self._set_avatar_task: Optional[Task] = None
+        self.__broadcast_task: Optional[Task] = None
         self.__avatar_unique_id: Optional[AvatarIdType] = None
+        self._avatar_pk: Optional[int] = None
 
     @property
     def __avatar_jid(self):
@@ -72,6 +73,10 @@ class AvatarMixin:
             name=f"Set avatar of {self} from property",
         )
 
+    @property
+    def avatar_pk(self) -> int | None:
+        return self._avatar_pk
+
     @staticmethod
     def __get_uid(a: Optional[AvatarType]) -> Optional[AvatarIdType]:
         if isinstance(a, str):
@@ -84,17 +89,39 @@ class AvatarMixin:
             return None
         raise TypeError("Bad avatar", a)
 
-    async def __set_avatar(self, a: Optional[AvatarType], uid: Optional[AvatarIdType]):
+    async def __set_avatar(
+        self, a: Optional[AvatarType], uid: Optional[AvatarIdType], delete: bool
+    ):
         self.__avatar_unique_id = uid
-        await self.session.xmpp.pubsub.set_avatar(
-            jid=self.__avatar_jid,
-            avatar=a,
-            unique_id=None if isinstance(uid, URL) else uid,
-            broadcast_to=self.session.user.jid.bare,
-            broadcast=self._avatar_pubsub_broadcast,
-        )
+
+        if a is None:
+            cached_avatar = None
+            self._avatar_pk = None
+        else:
+            try:
+                cached_avatar = await avatar_cache.convert_or_get(a)
+            except Exception as e:
+                self.session.log.error("Failed to set avatar %s", a, exc_info=e)
+                self._avatar_pk = None
+                self.__avatar_unique_id = uid
+                return
+            self._avatar_pk = cached_avatar.pk
+
+        if self.__should_pubsub_broadcast():
+            await self.session.xmpp.pubsub.broadcast_avatar(
+                self.__avatar_jid, self.session.user_jid, cached_avatar
+            )
+
+        if delete and isinstance(a, Path):
+            a.unlink()
+
         self._post_avatar_update()
 
+    def __should_pubsub_broadcast(self):
+        return getattr(self, "is_friend", False) and getattr(
+            self, "added_to_roster", False
+        )
+
     async def _no_change(self, a: Optional[AvatarType], uid: Optional[AvatarIdType]):
         if a is None:
             return self.__avatar_unique_id is None
@@ -103,23 +130,27 @@ class AvatarMixin:
         if isinstance(uid, URL):
             if self.__avatar_unique_id != uid:
                 return False
-            return not await avatar_cache.url_has_changed(uid)
+            return not await avatar_cache.url_modified(uid)
         return self.__avatar_unique_id == uid
 
     async def set_avatar(
         self,
         a: Optional[AvatarType],
         avatar_unique_id: Optional[LegacyFileIdType] = None,
+        delete: bool = False,
         blocking=False,
         cancel=True,
     ) -> None:
         """
         Set an avatar for this entity
 
-        :param a:
-        :param avatar_unique_id:
-        :param blocking:
-        :param cancel:
+        :param a: The avatar, in one of the types slidge supports
+        :param avatar_unique_id: A globally unique ID for the avatar on the
+            legacy network
+        :param delete: If the avatar is provided as a Path, whether to delete
+            it once used or not.
+        :param blocking: Internal use by slidge for tests, do not use!
+        :param cancel: Internal use by slidge, do not use!
         """
         if avatar_unique_id is None and a is not None:
             avatar_unique_id = self.__get_uid(a)
@@ -128,7 +159,7 @@ class AvatarMixin:
         if cancel and self._set_avatar_task:
             self._set_avatar_task.cancel()
         awaitable = create_task(
-            self.__set_avatar(a, avatar_unique_id),
+            self.__set_avatar(a, avatar_unique_id, delete),
             name=f"Set pubsub avatar of {self}",
         )
         if not self._set_avatar_task or self._set_avatar_task.done():
@@ -136,32 +167,70 @@ class AvatarMixin:
         if blocking:
             await awaitable
 
+    def get_cached_avatar(self) -> Optional["CachedAvatar"]:
+        if self._avatar_pk is None:
+            return None
+        return avatar_cache.get_by_pk(self._avatar_pk)
+
     def get_avatar(self) -> Optional["PepAvatar"]:
-        if not self.__avatar_unique_id:
+        cached_avatar = self.get_cached_avatar()
+        if cached_avatar is None:
             return None
-        return self.session.xmpp.pubsub.get_avatar(self.__avatar_jid)
+        from ..pubsub import PepAvatar
+
+        item = PepAvatar()
+        item.set_avatar_from_cache(cached_avatar)
+        return item
 
     def _post_avatar_update(self) -> None:
         return
 
+    def __get_cached_avatar_id(self):
+        i = self._get_cached_avatar_id()
+        if i is None:
+            return None
+        return self.session.xmpp.AVATAR_ID_TYPE(i)
+
+    def _get_cached_avatar_id(self) -> Optional[str]:
+        raise NotImplementedError
+
     async def avatar_wrap_update_info(self):
-        cached_id = avatar_cache.get_cached_id_for(self.__avatar_jid)
+        cached_id = self.__get_cached_avatar_id()
         self.__avatar_unique_id = cached_id
         try:
             await self.update_info()  # type:ignore
         except NotImplementedError:
             return
         new_id = self.avatar
-        if isinstance(new_id, URL) and not await avatar_cache.url_has_changed(new_id):
+        if isinstance(new_id, URL) and not await avatar_cache.url_modified(new_id):
             return
         elif new_id != cached_id:
             # at this point it means that update_info set the avatar, and we don't
             # need to do anything else
             return
 
-        await self.session.xmpp.pubsub.set_avatar_from_cache(
-            self.__avatar_jid,
-            new_id is None and cached_id is not None,
-            self.session.user.jid.bare,
-            self._avatar_pubsub_broadcast,
+        if self.__should_pubsub_broadcast():
+            if new_id is None and cached_id is None:
+                return
+            if self._avatar_pk is not None:
+                cached_avatar = avatar_cache.get_by_pk(self._avatar_pk)
+            else:
+                cached_avatar = None
+            self.__broadcast_task = self.session.xmpp.loop.create_task(
+                self.session.xmpp.pubsub.broadcast_avatar(
+                    self.__avatar_jid, self.session.user_jid, cached_avatar
+                )
+            )
+
+    def _set_avatar_from_store(self, stored):
+        if stored.avatar_id is None:
+            return
+        if stored.avatar is None:
+            # seems to happen after avatar cleanup for some reason?
+            self.__avatar_unique_id = None
+            return
+        self.__avatar_unique_id = (
+            stored.avatar.legacy_id
+            if stored.avatar.legacy_id is not None
+            else URL(stored.avatar.url)
         )

slidge/core/mixins/base.py

@@ -6,9 +6,8 @@ from slixmpp import JID
 from ...util.types import MessageOrPresenceTypeVar
 
 if TYPE_CHECKING:
-    from slidge.core.gateway import BaseGateway
-    from slidge.core.session import BaseSession
-    from slidge.util.db import GatewayUser
+    from ..gateway import BaseGateway
+    from ..session import BaseSession
 
 
 class MetaBase(ABCMeta):
@@ -18,11 +17,18 @@ class MetaBase(ABCMeta):
 class Base:
     session: "BaseSession" = NotImplemented
     xmpp: "BaseGateway" = NotImplemented
-    user: "GatewayUser" = NotImplemented
 
     jid: JID = NotImplemented
     name: str = NotImplemented
 
+    @property
+    def user_jid(self):
+        return self.session.user_jid
+
+    @property
+    def user_pk(self):
+        return self.session.user_pk
+
 
 class BaseSender(Base):
     def _send(

slidge/core/mixins/db.py

@@ -0,0 +1,18 @@
+from contextlib import contextmanager
+
+
+class UpdateInfoMixin:
+    """
+    This mixin just adds a context manager that prevents commiting to the DB
+    on every attribute change.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._updating_info = False
+
+    @contextmanager
+    def updating_info(self):
+        self._updating_info = True
+        yield
+        self._updating_info = False

slidge/core/mixins/disco.py

@@ -13,10 +13,6 @@ class BaseDiscoMixin(Base):
     DISCO_NAME: str = NotImplemented
     DISCO_LANG = None
 
-    def __init__(self):
-        super().__init__()
-        self.__caps_cache: Optional[str] = None
-
     def _get_disco_name(self):
         if self.DISCO_NAME is NotImplemented:
             return self.xmpp.COMPONENT_NAME
@@ -44,17 +40,11 @@ class BaseDiscoMixin(Base):
         return info
 
     async def get_caps_ver(self, jid: OptJid = None, node: Optional[str] = None):
-        if self.__caps_cache:
-            return self.__caps_cache
         info = await self.get_disco_info(jid, node)
         caps = self.xmpp.plugin["xep_0115"]
         ver = caps.generate_verstring(info, caps.hash)
-        self.__caps_cache = ver
         return ver
 
-    def reset_caps_cache(self):
-        self.__caps_cache = None
-
 
 class ChatterDiscoMixin(BaseDiscoMixin):
     AVATAR = True

slidge/core/mixins/lock.py

@@ -15,14 +15,16 @@ class NamedLockMixin:
         locks = self.__locks
         if not locks.get(id_):
             locks[id_] = asyncio.Lock()
-        async with locks[id_]:
-            log.trace("acquired %s", id_)  # type:ignore
-            yield
-        log.trace("releasing %s", id_)  # type:ignore
-        waiters = locks[id_]._waiters  # type:ignore
-        if not waiters:
-            del locks[id_]
-            log.trace("erasing %s", id_)  # type:ignore
+        try:
+            async with locks[id_]:
+                log.trace("acquired %s", id_)  # type:ignore
+                yield
+        finally:
+            log.trace("releasing %s", id_)  # type:ignore
+            waiters = locks[id_]._waiters  # type:ignore
+            if not waiters:
+                del locks[id_]
+                log.trace("erasing %s", id_)  # type:ignore
 
     def get_lock(self, id_: Hashable):
         return self.__locks.get(id_)