slidge 0.1.0rc1__py3-none-any.whl → 0.1.2__py3-none-any.whl

slidge/core/pubsub.py

@@ -1,116 +1,185 @@
-import asyncio
 import hashlib
 import io
 import logging
 from copy import copy
 from pathlib import Path
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Type, Union
 
 from PIL import Image, UnidentifiedImageError
-from slixmpp import JID, ComponentXMPP, CoroutineCallback, Iq, Presence, StanzaPath
+from PIL.Image import Image as PILImage
+from slixmpp import (
+    JID,
+    CoroutineCallback,
+    Iq,
+    Presence,
+    StanzaPath,
+    register_stanza_plugin,
+)
 from slixmpp.exceptions import XMPPError
 from slixmpp.plugins.base import BasePlugin, register_plugin
 from slixmpp.plugins.xep_0060.stanza import Event, EventItem, EventItems, Item
 from slixmpp.plugins.xep_0084 import Data as AvatarData
 from slixmpp.plugins.xep_0084 import MetaData as AvatarMetadata
 from slixmpp.plugins.xep_0172 import UserNick
+from slixmpp.plugins.xep_0292.stanza import VCard4
 from slixmpp.types import JidStr, OptJidStr
 
-from ..util.db import user_store
-from ..util.types import AvatarType, PepItemType
-from ..util.xep_0292.stanza import VCard4
-from . import config
-from .cache import avatar_cache
+from ..contact.contact import LegacyContact
+from ..contact.roster import ContactIsUser
+from ..util.db import GatewayUser, user_store
+from ..util.sql import db
+from ..util.types import AvatarType, LegacyFileIdType, PepItemType
+from .cache import CachedAvatar, avatar_cache
+from .mixins.lock import NamedLockMixin
+
+if TYPE_CHECKING:
+    from slidge import BaseGateway
 
 VCARD4_NAMESPACE = "urn:xmpp:vcard4"
 
 
 class PepItem:
-    def __init__(self, authorized_jid: Optional[JidStr] = None):
-        self.authorized_jid = authorized_jid
+    @staticmethod
+    def from_db(jid: JID, user: Optional[GatewayUser] = None) -> Optional["PepItem"]:
+        raise NotImplementedError
+
+    def to_db(self, jid: JID, user: Optional[GatewayUser] = None):
+        raise NotImplementedError
 
 
 class PepAvatar(PepItem):
-    def __init__(self, authorized_jid: Optional[JidStr] = None):
-        super().__init__(authorized_jid)
+    def __init__(self, jid: JID):
         self.metadata: Optional[AvatarMetadata] = None
-        self.data: Optional[AvatarData] = None
         self.id: Optional[str] = None
+        self.jid = jid
+        self._avatar_data_path: Optional[Path] = None
+
+    @property
+    def data(self) -> Optional[AvatarData]:
+        if self._avatar_data_path is None:
+            return None
+        data = AvatarData()
+        data.set_value(self._avatar_data_path.read_bytes())
+        return data
+
+    @staticmethod
+    def _sha(b: bytes):
+        return hashlib.sha1(b).hexdigest()
+
+    def _get_weak_unique_id(self, avatar: AvatarType):
+        if isinstance(avatar, str):
+            return avatar  # url
+        elif isinstance(avatar, bytes):
+            return self._sha(avatar)
+        elif isinstance(avatar, Path):
+            return self._sha(avatar.read_bytes())
 
-    async def set_avatar(self, avatar: AvatarType):
+    @staticmethod
+    async def _get_image(avatar: AvatarType) -> PILImage:
         if isinstance(avatar, str):
-            return await self.set_avatar_from_url(avatar)
+            # async with aiohttp.ClientSession() as session:
+            async with avatar_cache.http.get(avatar) as response:
+                return Image.open(io.BytesIO(await response.read()))
         elif isinstance(avatar, bytes):
-            img = Image.open(io.BytesIO(avatar))
+            return Image.open(io.BytesIO(avatar))
         elif isinstance(avatar, Path):
-            img = Image.open(avatar)
+            return Image.open(avatar)
         else:
             raise TypeError("Avatar must be bytes, a Path or a str (URL)", avatar)
 
-        metadata = AvatarMetadata()
+    async def set_avatar(
+        self, avatar: AvatarType, unique_id: Optional[LegacyFileIdType] = None
+    ):
+        if unique_id is None:
+            if isinstance(avatar, str):
+                await self._set_avatar_from_url_alone(avatar)
+                return
+            unique_id = self._get_weak_unique_id(avatar)
 
-        resampled = False
-        if (size := config.AVATAR_SIZE) and any(x > size for x in img.size):
-            img.thumbnail((size, size))
-            log.debug("Resampled image to %s", img.size)
-            resampled = True
-
-        if not resampled and img.format == "PNG" and isinstance(avatar, bytes):
-            avatar_bytes = avatar
-        elif not resampled and img.format == "PNG" and isinstance(avatar, Path):
-            with avatar.open("rb") as f:
-                avatar_bytes = f.read()
+        await self._set_avatar_from_unique_id(avatar, unique_id)
+
+    async def _set_avatar_from_unique_id(
+        self, avatar: AvatarType, unique_id: LegacyFileIdType
+    ):
+        cached_avatar = avatar_cache.get(unique_id)
+        if cached_avatar:
+            # this shouldn't be necessary but here to re-use avatars downloaded
+            # before the change introducing the JID to unique ID mapping
+            avatar_cache.store_jid(self.jid, unique_id)
         else:
-            with io.BytesIO() as f:
-                img.save(f, format="PNG")
-                avatar_bytes = f.getvalue()
+            img = await self._get_image(avatar)
+            cached_avatar = await avatar_cache.convert_and_store(
+                img, unique_id, self.jid
+            )
 
-        hash_ = hashlib.sha1(avatar_bytes).hexdigest()
-        self.id = hash_
-        metadata.add_info(
-            id=hash_,
-            itype="image/png",
-            ibytes=len(avatar_bytes),
-            height=str(img.height),
-            width=str(img.width),
-        )
-        self.metadata = metadata
+        self._set_avatar_from_cache(cached_avatar)
 
-        data = AvatarData()
-        data.set_value(avatar_bytes)
-        self.data = data
+    async def _set_avatar_from_url_alone(self, url: str):
+        cached_avatar = await avatar_cache.get_avatar_from_url_alone(url, self.jid)
+        self._set_avatar_from_cache(cached_avatar)
 
-    async def set_avatar_from_url(self, url: str):
-        avatar = await avatar_cache.get_avatar(url)
+    def _set_avatar_from_cache(self, cached_avatar: CachedAvatar):
         metadata = AvatarMetadata()
-        self.id = avatar.hash
+        self.id = cached_avatar.hash
         metadata.add_info(
-            id=avatar.hash,
+            id=cached_avatar.hash,
             itype="image/png",
-            ibytes=len(avatar.data),
-            height=str(avatar.height),
-            width=str(avatar.width),
+            ibytes=cached_avatar.path.stat().st_size,
+            height=str(cached_avatar.height),
+            width=str(cached_avatar.width),
         )
         self.metadata = metadata
+        self._avatar_data_path = cached_avatar.path
 
-        data = AvatarData()
-        data.set_value(avatar.data)
-        self.data = data
+    @staticmethod
+    def from_db(jid: JID, user: Optional[GatewayUser] = None) -> Optional["PepAvatar"]:
+        cached_id = db.avatar_get(jid)
+        if cached_id is None:
+            return None
+        item = PepAvatar(jid)
+        cached_avatar = avatar_cache.get(cached_id)
+        if cached_avatar is None:
+            raise XMPPError("internal-server-error")
+        item._set_avatar_from_cache(cached_avatar)
+        return item
+
+    def to_db(self, jid: JID, user=None):
+        cached_id = avatar_cache.get_cached_id_for(jid)
+        if cached_id is None:
+            log.warning("Could not store avatar for %s", jid)
+            return
+        db.avatar_store(jid, cached_id)
+
+    @staticmethod
+    def remove_from_db(jid: JID):
+        db.avatar_delete(jid)
 
 
 class PepNick(PepItem):
-    def __init__(
-        self, authorized_jid: Optional[JidStr] = None, nick: Optional[str] = None
-    ):
-        super().__init__(authorized_jid)
+    def __init__(self, nick: Optional[str] = None):
         nickname = UserNick()
         if nick is not None:
             nickname["nick"] = nick
         self.nick = nickname
+        self.__nick_str = nick
 
+    @staticmethod
+    def from_db(jid: JID, user: Optional[GatewayUser] = None) -> Optional["PepNick"]:
+        if user is None:
+            raise XMPPError("not-allowed")
+        nick = db.nick_get(jid, user)
+        if nick is None:
+            return None
+        return PepNick(nick)
 
-class PubSubComponent(BasePlugin):
-    xmpp: ComponentXMPP
+    def to_db(self, jid: JID, user: Optional[GatewayUser] = None):
+        if user is None:
+            raise XMPPError("not-allowed")
+        db.nick_store(jid, str(self.__nick_str), user)
+
+
+class PubSubComponent(NamedLockMixin, BasePlugin):
+    xmpp: "BaseGateway"
 
     name = "pubsub"
     description = "Pubsub component"
@@ -126,8 +195,7 @@ class PubSubComponent(BasePlugin):
 
     def __init__(self, *a, **kw):
         super(PubSubComponent, self).__init__(*a, **kw)
-        self._avatars = dict[JID, PepAvatar]()
-        self._nicks = dict[JID, PepNick]()
+        register_stanza_plugin(EventItem, UserNick)
 
     def plugin_init(self):
         self.xmpp.register_handler(
@@ -151,7 +219,7 @@ class PubSubComponent(BasePlugin):
                 self._get_vcard,  # type:ignore
             )
         )
-        self.xmpp.add_event_handler("got_online", self._on_got_online)
+        self.xmpp.add_event_handler("presence_available", self._on_presence_available)
 
         disco = self.xmpp.plugin["xep_0030"]
         disco.add_identity("pubsub", "pep", self.component_name)
@@ -160,23 +228,55 @@ class PubSubComponent(BasePlugin):
         disco.add_feature("http://jabber.org/protocol/pubsub#retrieve-items")
         disco.add_feature("http://jabber.org/protocol/pubsub#persistent-items")
 
-    async def _on_got_online(self, p: Presence):
+    async def _on_presence_available(self, p: Presence):
+        if p.get_plugin("muc_join", check=True) is not None:
+            log.debug("Ignoring MUC presence here")
+            return
+
         from_ = p.get_from()
         ver_string = p["caps"]["ver"]
         info = None
+
+        to = p.get_to()
+
+        # we don't want to push anything for contacts that are not in the user's roster
+        if to != self.xmpp.boundjid.bare:
+            session = self.xmpp.get_session_from_stanza(p)
+
+            if session is None:
+                return
+
+            await session.contacts.ready
+            try:
+                contact = await session.contacts.by_jid(to)
+            except XMPPError as e:
+                log.debug(
+                    "Could not determine if %s was added to the roster: %s", to, e
+                )
+                return
+            except Exception as e:
+                log.warning("Could not determine if %s was added to the roster.", to)
+                log.exception(e)
+                return
+            if not contact.is_friend:
+                return
+
         if ver_string:
-            await asyncio.sleep(5)
             info = await self.xmpp.plugin["xep_0115"].get_caps(from_)
         if info is None:
-            info = await self.xmpp.plugin["xep_0030"].get_info(from_)
+            async with self.lock(from_):
+                iq = await self.xmpp.plugin["xep_0030"].get_info(from_)
+            info = iq["disco_info"]
         features = info["features"]
         if AvatarMetadata.namespace + "+notify" in features:
             try:
-                pep_avatar = self._get_authorized_avatar(p)
+                pep_avatar = await self._get_authorized_avatar(p)
             except XMPPError:
                 pass
             else:
-                self._broadcast(
+                if pep_avatar.metadata is None:
+                    raise XMPPError("internal-server-error", "Avatar but no metadata?")
+                await self._broadcast(
                     data=pep_avatar.metadata,
                     from_=p.get_to(),
                     to=from_,
@@ -184,26 +284,26 @@ class PubSubComponent(BasePlugin):
                 )
         if UserNick.namespace + "+notify" in features:
             try:
-                pep_nick = self._get_authorized_nick(p)
+                pep_nick = await self._get_authorized_nick(p)
             except XMPPError:
                 pass
             else:
-                self._broadcast(data=pep_nick.nick, from_=p.get_to(), to=from_)
+                await self._broadcast(data=pep_nick.nick, from_=p.get_to(), to=from_)
 
         if VCARD4_NAMESPACE + "+notify" in features:
-            self.broadcast_vcard_event(p.get_to(), to=from_)
+            await self.broadcast_vcard_event(p.get_to(), to=from_)
 
-    def broadcast_vcard_event(self, from_, to):
+    async def broadcast_vcard_event(self, from_, to):
         item = Item()
         item.namespace = VCARD4_NAMESPACE
         item["id"] = "current"
-        vcard: VCard4 = self.xmpp["xep_0292_provider"].get_vcard(from_, to)
+        vcard: VCard4 = await self.xmpp["xep_0292_provider"].get_vcard(from_, to)
         # The vcard content should NOT be in this event according to the spec:
         # https://xmpp.org/extensions/xep-0292.html#sect-idm45669698174224
         # but movim expects it to be here, and I guess
 
         log.debug("Broadcast vcard4 event: %s", vcard)
-        self._broadcast(
+        await self._broadcast(
             data=vcard,
             from_=JID(from_).bare,
             to=to,
@@ -211,28 +311,29 @@ class PubSubComponent(BasePlugin):
             node=VCARD4_NAMESPACE,
         )
 
-    @staticmethod
-    def _get_authorized_item(
-        store: dict[JID, PepItemType], stanza: Union[Iq, Presence]
+    async def _get_authorized_item(
+        self, cls: Type[PepItemType], stanza: Union[Iq, Presence]
     ) -> PepItemType:
-        item = store.get(stanza.get_to())
+        sto = stanza.get_to()
+        user = user_store.get_by_jid(stanza.get_from())
+        item = cls.from_db(sto, user)
         if item is None:
             raise XMPPError("item-not-found")
 
-        if item.authorized_jid is not None:
-            if stanza.get_from().bare != item.authorized_jid:
-                raise XMPPError("item-not-found")
+        if sto != self.xmpp.boundjid.bare:
+            session = self.xmpp.get_session_from_stanza(stanza)
+            await session.contacts.by_jid(sto)
 
-        return item
+        return item  # type:ignore
 
-    def _get_authorized_avatar(self, stanza: Union[Iq, Presence]):
-        return self._get_authorized_item(self._avatars, stanza)
+    async def _get_authorized_avatar(self, stanza: Union[Iq, Presence]) -> PepAvatar:
+        return await self._get_authorized_item(PepAvatar, stanza)
 
-    def _get_authorized_nick(self, stanza: Union[Iq, Presence]):
-        return self._get_authorized_item(self._nicks, stanza)
+    async def _get_authorized_nick(self, stanza: Union[Iq, Presence]) -> PepNick:
+        return await self._get_authorized_item(PepNick, stanza)
 
     async def _get_avatar_data(self, iq: Iq):
-        pep_avatar = self._get_authorized_avatar(iq)
+        pep_avatar = await self._get_authorized_avatar(iq)
 
         requested_items = iq["pubsub"]["items"]
         if len(requested_items) == 0:
@@ -246,7 +347,7 @@ class PubSubComponent(BasePlugin):
                 raise XMPPError("item-not-found")
 
     async def _get_avatar_metadata(self, iq: Iq):
-        pep_avatar = self._get_authorized_avatar(iq)
+        pep_avatar = await self._get_authorized_avatar(iq)
 
         requested_items = iq["pubsub"]["items"]
         if len(requested_items) == 0:
@@ -263,7 +364,7 @@ class PubSubComponent(BasePlugin):
         # this is not the proper way that clients should retrieve VCards, but
         # gajim does it this way.
         # https://xmpp.org/extensions/xep-0292.html#sect-idm45669698174224
-        vcard: VCard4 = self.xmpp["xep_0292_provider"].get_vcard(
+        vcard: VCard4 = await self.xmpp["xep_0292_provider"].get_vcard(
             iq.get_to().bare, iq.get_from().bare
         )
         log.debug("VCARD: %s -- %s -- %s", iq.get_to().bare, iq.get_from().bare, vcard)
@@ -271,25 +372,43 @@ class PubSubComponent(BasePlugin):
             raise XMPPError("item-not-found")
         self._reply_with_payload(iq, vcard, "current", VCARD4_NAMESPACE)
 
+    @staticmethod
+    def get_avatar(jid: JID):
+        return PepAvatar.from_db(jid)
+
     @staticmethod
     def _reply_with_payload(
         iq: Iq,
-        payload: Union[AvatarMetadata, AvatarData, VCard4],
-        id_: str,
+        payload: Optional[Union[AvatarMetadata, AvatarData, VCard4]],
+        id_: Optional[str],
         namespace: Optional[str] = None,
     ):
         result = iq.reply()
         item = Item()
         if payload:
             item.set_payload(payload.xml)
-        item["id"] = id_
-        result["pubsub"]["items"]["node"] = (
-            namespace if namespace else payload.namespace
-        )
+            item["id"] = id_
+            result["pubsub"]["items"]["node"] = (
+                namespace if namespace else payload.namespace
+            )
         result["pubsub"]["items"].append(item)
         result.send()
 
-    def _broadcast(self, data, from_: JidStr, to: OptJidStr = None, **kwargs):
+    async def _broadcast(self, data, from_: JidStr, to: OptJidStr = None, **kwargs):
+        from_ = JID(from_)
+        if from_ != self.xmpp.boundjid.bare and to is not None:
+            to = JID(to)
+            session = self.xmpp.get_session_from_jid(to)
+            if session is None:
+                return
+            await session.ready
+            try:
+                entity = await session.get_contact_or_group_or_participant(from_)
+            except ContactIsUser:
+                return
+            if isinstance(entity, LegacyContact) and not entity.is_friend:
+                return
+
         item = EventItem()
         if data:
             item.set_payload(data.xml)
@@ -321,45 +440,85 @@ class PubSubComponent(BasePlugin):
         self,
         jid: JidStr,
         avatar: Optional[AvatarType] = None,
-        restrict_to: OptJidStr = None,
+        broadcast_to: OptJidStr = None,
+        unique_id=None,
+        broadcast=True,
     ):
         jid = JID(jid)
         if avatar is None:
-            try:
-                del self._avatars[jid]
-            except KeyError:
-                pass
-            self._broadcast(AvatarMetadata(), jid, restrict_to)
+            PepAvatar.remove_from_db(jid)
+            await self._broadcast(AvatarMetadata(), jid, broadcast_to)
+            avatar_cache.delete_jid(jid)
         else:
-            pep_avatar = PepAvatar()
+            pep_avatar = PepAvatar(jid)
             try:
-                await pep_avatar.set_avatar(avatar)
-            except UnidentifiedImageError as e:
+                await pep_avatar.set_avatar(avatar, unique_id)
+            except (UnidentifiedImageError, FileNotFoundError) as e:
                 log.warning("Failed to set avatar for %s: %r", self, e)
                 return
-
-            pep_avatar.authorized_jid = restrict_to
-            self._avatars[jid] = pep_avatar
+            pep_avatar.to_db(jid)
             if pep_avatar.metadata is None:
                 raise RuntimeError
-            self._broadcast(
+            if not broadcast:
+                return
+            await self._broadcast(
                 pep_avatar.metadata,
                 jid,
-                restrict_to,
+                broadcast_to,
                 id=pep_avatar.metadata["info"]["id"],
             )
 
+    async def set_avatar_from_cache(
+        self, jid: JID, send_empty: bool, broadcast_to: OptJidStr = None, broadcast=True
+    ):
+        uid = avatar_cache.get_cached_id_for(jid)
+        if uid is None:
+            if not send_empty:
+                return
+            self.xmpp.loop.create_task(
+                self.set_avatar(jid, None, broadcast_to, uid, broadcast)
+            )
+            return
+        cached_avatar = avatar_cache.get(str(uid))
+        if cached_avatar is None:
+            # should not happen but well…
+            log.warning(
+                "Something is wrong with the avatar, %s won't have an "
+                "avatar because avatar not found in cache",
+                jid,
+            )
+            return
+        self.xmpp.loop.create_task(
+            self.set_avatar(jid, cached_avatar.path, broadcast_to, uid, broadcast)
+        )
+
     def set_nick(
         self,
+        user: GatewayUser,
         jid: JidStr,
         nick: Optional[str] = None,
-        restrict_to: OptJidStr = None,
     ):
         jid = JID(jid)
-        nickname = PepNick(restrict_to, nick)
-        self._nicks[jid] = nickname
-        log.debug("NICK: %s", nickname.nick)
-        self._broadcast(nickname.nick, jid, restrict_to)
+        nickname = PepNick(nick)
+        nickname.to_db(jid, user)
+        log.debug("New nickname: %s", nickname.nick)
+        self.xmpp.loop.create_task(self._broadcast(nickname.nick, jid, user.bare_jid))
+
+    async def broadcast_all(self, from_: JID, to: JID):
+        """
+        Force push avatar and nick for a stored JID.
+        """
+        a = PepAvatar.from_db(from_)
+        if a:
+            if a.metadata:
+                await self._broadcast(
+                    a.metadata, from_, to, id=a.metadata["info"]["id"]
+                )
+            else:
+                log.warning("No metadata associated to this cached avatar?!")
+        n = PepNick.from_db(from_, user_store.get_by_jid(to))
+        if n:
+            await self._broadcast(n.nick, from_, to)
 
 
 log = logging.getLogger(__name__)