xmpp-discord-bridge/xmpp_discord_bridge/avatar.py

import logging
import hashlib
import os

from slixmpp.exceptions import IqError

class AvatarManager:
    """
    The AvatarManager class is responsible for aquiring the XMPP avatars
    using the means of either XEP-0153 (XEP-0054), XEP-0084 or (in the future)
    XEP-0292, storing and providing a publicly available URL to those avatars.
    """
    def __init__(self, xmpp, config):
        self._xmpp = xmpp
        self._path = config["avatars"]["path"]
        self._public = config["avatars"]["url"]

        self._avatars = {}

        self._logger = logging.getLogger("xmpp.avatar")

    def _save_avatar(self, jid, data, type_):
        if not type_:
            # TODO: Maybe guess?
            self._logger.info("Got VCARD with empty photo type. Assuming image/png")
            type_ = "image/png"

        if len(type_.split("/")) < 2:
            self._logger.debug("Got type_ of '%s'", type_)
            self._logger.debug("COMPAT! Just assuming image/png")
            type_ = "image/png"
        filename = hashlib.sha1(data).hexdigest() + "." + type_.split("/")[1]
        path = os.path.join(self._path, filename)

        if os.path.exists(path):
            self._logger.debug("Avatar for %s already exists, not saving it again",
                               jid)
        else:
            with open(path, "wb") as f:
                f.write(data)

        self._avatars[jid] = filename
        
    async def try_xep_0153(self, jid):
         try:
            iq = await self._xmpp.plugin["xep_0054"].get_vcard(jid=jid,
                                                               ifrom=self._xmpp._bot_jid_full)
            type_ = iq["vcard_temp"]["PHOTO"]["TYPE"]
            data = iq["vcard_temp"]["PHOTO"]["BINVAL"]
            self._save_avatar(jid, data, type_)
            return True
         except IqError:
            self._logger.debug("Avatar retrieval via XEP-0054/XEP-0153 failed. Probably no vCard for XEP-0054 published")
            return False

    async def try_xep_0084(self, jid):
        try:
            iq = await self._xmpp.plugin["xep_0060"].get_items(jid=jid,
                                                               node="urn:xmpp:avatar:data",
                                                               max_items=1,
                                                               ifrom=self._xmpp._bot_jid_full)
            data = iq["pubsub"]["items"]["substanzas"][0]["data"]
            self._save_avatar(jid, data, "image/png")
            return True
        except IqError:
            self._logger.debug("Avatar retrieval via XEP-0084 failed. Probably no avatar published or subscription model not fulfilled.")
            return False
        
    async def aquire_avatar(self, jid):
        """
        Attempt to request the avatar of @jid
        """
        for f in [self.try_xep_0153, self.try_xep_0084]:
            if await f(jid):
                self._logger.debug("Avatar retrieval successful for %s",
                                   jid)
                return
        
        self._logger.debug("Avatar retrieval failed for %s. Giving up.",
                           jid)

    def get_avatar_url(self, jid):
        """
        Returns either the URL to the avatar of @jid or None.
        """
        filename = self._avatars.get(jid, None)
        return self._public + filename if filename else None