Add existing code

setup.py

@@ -0,0 +1,28 @@
+from setuptools import setup, find_packages
+
+setup(
+    name = "xmpp_discord_bridge",
+    version = "0.1.0",
+    url = "https://git.polynom.me/PapaTutuWawa/xmpp-discord-bridge",
+    author = "Alexander \"PapaTutuWawa\"",
+    author_email = "papatutuwawa <at> polynom.me",
+    license = "GPLc3",
+    packages = find_packages(),
+    install_requires = [
+        "requests>=2.26.0",
+        "slixmpp>=1.7.1",
+        "discord.py>=1.7.3",
+        "toml>=0.10.2"
+    ],
+    extra_require = {
+        "dev": [
+            "black"
+        ]
+    },
+    zip_safe = True,
+    entry_points = {
+        "console_scripts": [
+            "discord-xmpp-bridge = discord_xmpp_bridge.main:main"
+        ]
+    }
+)

xmpp_discord_bridge/main.py

@@ -0,0 +1,578 @@
+import logging
+import os
+import sys
+import asyncio
+import signal
+import base64
+import hashlib
+import hmac
+import urllib.parse
+from optparse import OptionParser
+
+import toml
+import slixmpp
+from slixmpp import Message
+from slixmpp.componentxmpp import ComponentXMPP
+from slixmpp.exceptions import XMPPError, IqError
+from slixmpp.xmlstream import ElementBase, register_stanza_plugin
+from slixmpp.jid import JID
+import discord
+import requests
+
+class OOBData(ElementBase):
+    name = "x"
+    namespace = "jabber:x:oob"
+    plugin_attrib = "oob"
+    interfaces = {"url"}
+    sub_interfaces = interfaces
+
+class AvatarManager:
+    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_):
+        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)
+        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):
+        # First try vCard via 0054/0153
+        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(self, jid):
+        return self._avatars.get(jid, None)
+
+class DiscordClient(discord.Client):
+    def __init__(self, xmpp, config):
+        intents = discord.Intents.default()
+        intents.members = True
+        intents.presences = True
+        intents.messages = True
+        intents.reactions = True
+
+        discord.Client.__init__(self, intents=intents)
+
+        self._xmpp = xmpp
+        self._config = config
+        self._logger = logging.getLogger("discord.client")
+
+    async def on_ready(self):
+        await self._xmpp.on_discord_ready()
+
+    async def on_message(self, message):
+        await self._xmpp.on_discord_message(message)
+
+    async def on_member_update(self, before, after):
+        await self._xmpp.on_discord_member_update(before, after)
+
+    async def on_member_join(self, member):
+        await self._xmpp.on_discord_member_join(member)
+
+    async def on_member_leave(self, member):
+        await self._xmpp.on_discord_member_leave(member)
+
+    async def on_guild_channel_update(self, before, after):
+        await self._xmpp.on_discord_channel_update(before, after)
+
+    async def on_reaction(self, payload):
+        message = await (self.get_guild(payload.guild_id)
+                         .get_channel(payload.channel_id)
+                         .fetch_message(payload.message_id))
+        
+        await self._xmpp.on_discord_reaction(payload.guild_id,
+                                             payload.channel_id,
+                                             payload.emoji.name,
+                                             message,
+                                             payload.user_id,
+                                             payload.event_type)
+       
+    async def on_raw_reaction_add(self, payload):
+        await self.on_reaction(payload)
+
+    async def on_raw_reaction_remove(self, payload):
+        await self.on_reaction(payload)
+
+def discord_status_to_xmpp_show(status):
+    return {
+        discord.Status.online: "available",
+        discord.Status.idle: "away",
+        discord.Status.dnd: "dnd",
+        discord.Status.do_not_disturb: "dnd",
+        discord.Status.invisible: "xa", # TODO: Kinda
+        discord.Status.offline: "unavailable"
+    }.get(status, "available")
+        
+class BridgeComponent(ComponentXMPP):
+    def __init__(self, jid, secret, server, port, token, config):
+        ComponentXMPP.__init__(self, jid, secret, server, port)
+
+        self._config = config
+        self._logger = logging.getLogger("xmpp.bridge")
+        self._domain = jid
+        self._bot_nick = "Bot"
+        self._bot_jid_bare = JID("bot@" + jid)
+        self._bot_jid_full = JID("bot@" + jid + "/" + self._bot_nick)
+
+        self._token = token
+        self._discord = None
+
+        self._avatars = AvatarManager(self, config)
+        
+        # State tracking
+        self._virtual_muc_users = {} # MUC -> [Resources]
+        self._virtual_muc_nicks = {} # MUC -> User ID -> Nick
+        self._real_muc_users = {} # MUC -> [Resources]
+        self._guild_map = {} # Guild ID -> Channel ID -> MUC
+        self._muc_map = {} # MUC -> (Guild ID, Channel ID)
+        self._mucs = [] # List of known MUCs
+        self._webhooks = {} # MUC -> Webhook URL
+
+        register_stanza_plugin(Message, OOBData)
+        
+        self.add_event_handler("session_start", self.on_session_start)
+        self.add_event_handler("groupchat_message", self.on_groupchat_message)
+        self.add_event_handler("groupchat_presence", self.on_groupchat_presence)
+
+    async def send_oob_data(self, url, muc, member):
+        """
+        Send a message using XEP-0066 OOB data
+        """
+        proxy = self._config["general"].get("proxy_discord_urls_to", None)
+        if proxy:
+            hmac_str = urllib.parse.quote(base64.b64encode(hmac.digest(self._config["general"]["hmac_secret"].encode(),
+                                                                   url.encode(),
+                                                                   "sha256")), safe="")
+            proxy = proxy.replace("<hmac>", hmac_str).replace("<url>", urllib.parse.quote(url, safe=""))
+            url = proxy
+
+        self._logger.debug("OOB URL: %s", url)
+        msg = self.make_message(muc,
+                                mbody=url,
+                                mtype="groupchat",
+                                mfrom=self.spoof_member_jid(member.id))
+        msg["oob"]["url"] = url
+        msg.send()
+        
+    def spoof_member_jid(self, id_):
+        """
+        Return a full JID that we use for the puppets
+        """
+        return JID(str(id_) + "@" + self._domain + "/discord")
+
+    async def on_sigint(self):
+        await self._discord.close()
+
+        # Remove all virtual users
+        # NOTE: We cannot use leave_muc as this would also remove the MUC
+        #       from slixmpp's internal tracking, which would probably break
+        #       later leaves
+        for muc in self._mucs:
+            for uid in self._virtual_muc_nicks[muc]:
+                nick = self._virtual_muc_nicks[muc][uid]
+                self.send_presence(pshow='unavailable',
+                                   pto="%s/%s" % (muc, nick),
+                                   pfrom=self.spoof_member_jid(uid))
+
+            # Remove the Bot user
+            self.send_presence(pshow='unavailable',
+                               pto="%s/%s" % (muc, "Bot"),
+                               pfrom=self._bot_jid_full)
+
+        # Disconnect and close
+        await self.disconnect()
+    
+    async def on_discord_ready(self):
+        asyncio.get_event_loop().add_signal_handler(signal.SIGINT,
+                                                      lambda: asyncio.create_task(self.on_sigint()))
+
+        for ch in self._config["discord"]["channels"]:
+            muc = ch["muc"]
+            channel = ch["channel"]
+            guild = ch["guild"]
+            dchannel = self._discord.get_channel(channel)
+
+            # Initialise state tracking
+            self._muc_map[muc] = (guild, channel)
+            self._virtual_muc_users[muc] = []
+            self._virtual_muc_nicks[muc] = {}
+            for member in dchannel.members:
+                if member.status == discord.Status.offline:
+                    continue
+                
+                self._virtual_muc_users[muc].append(member.display_name)
+                self._virtual_muc_nicks[muc][member.id] = member.display_name
+            self._real_muc_users[muc] = []
+            self._mucs.append(muc)
+            if not guild in self._guild_map:
+                self._guild_map[guild] = {
+                    channel: muc
+                }
+            else:
+                self._guild_map[guild][channel] = muc
+
+            self._logger.debug("Joining %s", muc)
+            self.plugin["xep_0045"].join_muc(muc,
+                                             nick=self._bot_nick,
+                                             pfrom=self._bot_jid_full)
+
+            # Set the subject
+            subject = dchannel.topic or ""
+            self.plugin["xep_0045"].set_subject(muc,
+                                                subject,
+                                                mfrom=self._bot_jid_full)
+
+            # TODO: Is this working?
+            # Mirror the guild's icon
+            icon = await dchannel.guild.icon_url_as(static_format="png",
+                                                    format="png",
+                                                    size=128).read()
+            vcard = self.plugin["xep_0054"].make_vcard()
+            vcard["PHOTO"]["TYPE"] = "image/png"
+            vcard["PHOTO"]["BINVAL"] = base64.b64encode(icon)
+            self.send_raw("""
+<iq type="set" from="{}" to="{}">
+    <vCard xmlns="vcard-temp">
+{}
+    </vCard>
+</iq>
+            """.format(self._bot_jid_full,
+                       muc,
+                       str(vcard)))
+            
+            # Aquire a webhook
+            webhook_url = ""
+            for webhook in await dchannel.webhooks():
+                if webhook.name == "discord-xmpp-bridge":
+                    webhook_url = webhook.url
+                    break
+            if not webhook_url:
+                webhook = dchannel.create_webhook(name="discord-xmpp-bridge",
+                                                  reason="Bridging Discord and XMPP")
+                webhook_url = webhook.url
+            self._webhooks[muc] = webhook_url
+
+            # Make sure our virtual users can join
+            affiliation_delta = [
+                (self.spoof_member_jid(member.id).bare, "member") for member in dchannel.members
+            ]
+            await self.plugin["xep_0045"].send_affiliation_list(muc,
+                                                                affiliation_delta,
+                                                                ifrom=self._bot_jid_full)
+            
+            for member in dchannel.members:
+                self.virtual_user_join_muc(muc, member, update_state_tracking=False)
+
+            self._logger.info("%s is ready", muc)
+        self._logger.info("Bridge is ready")
+
+    async def on_groupchat_message(self, message):
+        muc = message["from"].bare
+        if not message["body"]:
+            return
+        if not message["from"].resource in self._real_muc_users[muc]:
+            return
+        # Prevent the message being reflected back into Discord
+        if not message["to"] == self._bot_jid_full:
+            return
+
+        webhook = {
+            "content": message["body"],
+            "username": message["from"].resource
+        }
+
+        if self._config["general"]["relay_xmpp_avatars"] and self._avatars.get_avatar(message["from"]):
+            webhook["avatar_url"] = self._avatar.get_avatar(message["from"])
+
+        # Look for mentions and replace them
+        guild, channel = self._muc_map[muc]
+        for member in self._discord.get_guild(guild).get_channel(channel).members:
+            self._logger.debug("Checking %s", member.display_name)
+            if "@" + member.display_name in webhook["content"]:
+                self._logger.debug("Found mention for %s. Replaceing.",
+                                   member.display_name)
+                webhook["content"] = webhook["content"].replace("@" + member.display_name,
+                                                                member.mention)
+            
+        if message["oob"]["url"] and message["body"] == message["oob"]["url"]:
+            webhook["embed"] = [{
+                "type": "rich",
+                "url": message["oob"]["url"]
+            }]
+            
+        requests.post(self._webhooks[muc],
+                      data=webhook)
+
+    def virtual_user_update_presence(self, muc, uid, pshow):
+        """
+        Change the status of a virtual MUC member.
+        NOTE: This assumes that the user is in the MUC
+        """
+        self.send_presence(pshow=pshow,
+                           pto="%s/%s" % (muc, self._virtual_muc_nicks[muc][uid]),
+                           pfrom=self.spoof_member_jid(uid))
+        
+    def virtual_user_join_muc(self, muc, member, update_state_tracking=False):
+        """
+        Makes the a puppet of member (@discord.Member) join the
+        MUC. Does nothing if the member is offline.
+
+        If update_state_tracking is True, then _virtual_muc_... gets updates.
+        """
+        if member.status == discord.Status.offline and not self._config["general"]["dont_ignore_offline"]:
+            return
+
+        if update_state_tracking:
+            self._virtual_muc_users[muc].append(member.display_name)
+            self._virtual_muc_nicks[muc][member.id] = member.display_name
+
+        # Prevent offline users from getting an unavailable presence by
+        # accident
+        pshow = discord_status_to_xmpp_show(member.status)
+        if member.status == discord.Status.offline:
+            pshow = "xa"
+
+        self.plugin["xep_0045"].join_muc(muc,
+                                         nick=member.display_name,
+                                         pfrom=self.spoof_member_jid(member.id),
+                                         pshow=pshow)
+        
+    async def on_discord_member_join(self, member):
+        guild = member.guild.id
+        if not guild in self._guild_map:
+            return
+
+        self._logger.debug("%s joined a known guild. Updating channels.",
+                           member.display_name)
+        for channel in self._guild_map[guild]:
+            muc = self._guild_map[guild][channel]
+            self.virtual_user_join_muc(muc, member, update_state_tracking=True)
+
+    async def on_discord_member_leave(self, member):
+        guild = member.guild.id
+        if not guild in self._guild_map:
+            return
+
+        self._logger.debug("%s left a known guild. Updating channels.",
+                           member.display_name)
+        for channel in self._guild_map[guild]:
+            muc = self._guild_map[member.guild.id][channel]
+            self._virtual_muc_users[muc].remove(member.display_name)
+            del self._virtual_muc_nicks[muc][member.id]
+
+            self.virtual_user_update_presence(muc, member.id, "unavailable")
+
+            
+    async def on_discord_member_update(self, before, after):
+        guild = after.guild.id
+        if not guild in self._guild_map:
+            return
+
+        # TODO: Handle nick changes
+        if before.status != after.status:
+            # Handle a status change
+            for channel in self._guild_map[guild]:
+                muc = self._guild_map[guild][channel]
+                if after.status == discord.Status.offline:
+                    if self._config["general"]["dont_ignore_offline"]:
+                        self.virtual_user_update_presence(muc,
+                                                          after.id,
+                                                          "xa")
+                    else:
+                        self._logger.debug("%s went offline. Removing from state tracking.",
+                                           after.display_name)
+                        self.virtual_user_update_presence(muc,
+                                                          after.id,
+                                                          discord_status_to_xmpp_show(after.status))
+
+                        # Remove from all state tracking
+                        self._virtual_muc_users[muc].remove(after.display_name)
+                        del self._virtual_muc_nicks[muc][after.id]
+                elif before.status == discord.Status.offline and after.status != discord.Status.offline:
+                    self.virtual_user_join_muc(muc, after, update_state_tracking=True)
+                else:
+                    self.virtual_user_update_presence(muc,
+                                                      after.id,
+                                                      discord_status_to_xmpp_show(after.status))
+
+
+    async def on_discord_channel_update(self, before, after):
+        if after.type != discord.ChannelType.text:
+            return
+
+        guild = after.guild.id
+        channel = after.id
+        if not guild in self._guild_map:
+            return
+        if not channel in self._guild_map[guild]:
+            return
+
+        # NOTE: We can only really handle description changes
+        muc = self._guild_map[guild][channel]
+        if before.topic != after.topic:
+            self._logger.debug("Channel %s changed the topic. Relaying.",
+                               after.name)
+            self.plugin["xep_0045"].set_subject(muc,
+                                                after.topic or "",
+                                                mfrom=self._bot_jid_full)
+
+    async def on_discord_reaction(self, guild, channel, emoji_str, msg, uid, kind):
+        """
+        Handle a Discord reaction.
+
+        reaction: discord.Reaction
+        user: discord.Member
+        kind: Either "add" or "remove"
+        """
+        if not self._config["general"]["reactions_compat"]:
+            self._logger.debug("Got a reaction but reactions_compat is turned off. Ignoring.")
+            return
+        if not guild in self._guild_map:
+            return
+        if not channel in self._guild_map[guild]:
+            return
+
+        muc = self._guild_map[guild][channel]
+
+        # TODO: Handle attachments
+        content = "> " + msg.clean_content.replace("\n", "\n> ") + "\n"
+        content += "+" if kind == "REACTION_ADD" else "-"
+        content += " " + emoji_str
+
+        self.send_message(mto=muc,
+                          mbody=content,
+                          mtype="groupchat",
+                          mfrom=self.spoof_member_jid(uid))
+
+    async def on_discord_message(self, msg):
+        guild, channel = msg.guild.id, msg.channel.id
+        if not (guild in self._guild_map and channel in self._guild_map[guild]):
+            return
+
+        muc = self._guild_map[guild][channel]
+        if msg.author.bot and msg.author.display_name in self._real_muc_users[muc]:
+            return
+
+        # TODO: Handle embeds
+        for attachment in msg.attachments:
+            await self.send_oob_data(attachment.url, muc, msg.author)
+
+        if not msg.clean_content:
+            self._logger.debug("Message empty. Not relaying.")
+            return
+
+        if self._config["general"]["muc_mention_compat"]:
+            mentions = [mention.display_name for mention in msg.mentions]
+            content = ", ".join(mentions) + ": " + msg.clean_content
+        else:
+            content = msg.clean_content
+        
+        self.send_message(mto=muc,
+                          mbody=content,
+                          mtype="groupchat",
+                          mfrom=self.spoof_member_jid(msg.author.id))
+
+    async def on_groupchat_presence(self, presence):
+        muc = presence["from"].bare
+        resource = presence["from"].resource
+
+        if not muc in self._mucs:
+            self._logger.warn("Received presence in unknown MUC %s", muc)
+            return
+        if resource in self._virtual_muc_users[muc]:
+            return
+        if not presence["to"] == self._bot_jid_full:
+            return
+        if resource == self._bot_nick:
+            return
+        
+        if presence["type"] == "unavailable":
+            try:
+                self._real_muc_users.remove(resource)
+            except:
+                self._logger.debug("Trying to remove %s from %s, but user is not in list. Skipping...",
+                                   muc,
+                                   resource)
+        else:
+            await self._avatars.aquire_avatar(presence["from"])
+            self._real_muc_users[muc].append(resource)
+        
+    async def on_session_start(self, event):
+        self._discord = DiscordClient(self, config)
+        asyncio.ensure_future(self._discord.start(self._token))
+
+def main():
+    if os.path.exists("./config.toml"):
+        config = toml.load("./config.toml")
+    elif os.path.exists("/etc/discord-xmpp-bridge/config.toml"):
+        config = toml.load("/etc/discord-xmpp-bridge/config.toml")
+    else:
+        raise Exception("config.toml not found")
+
+    parser = OptionParser()
+    parser.add_option(
+        "-d", "--debug", dest="debug", help="Enable debug logging", action="store_true"
+    )
+
+    (option, args) = parser.parse_args()
+    verbosity = logging.DEBUG if options.debug else logging.INFO
+    
+    general = config["general"]
+    xmpp = BridgeComponent(general["jid"],
+                           general["secret"],
+                           general["server"],
+                           general["port"],
+                           general["discord_token"],
+                           config)
+    for xep in [ "0030", "0199", "0045", "0084", "0153", "0054", "0060" ]:
+        xmpp.register_plugin(f"xep_{xep}")
+
+    logging.basicConfig(stream=sys.stdout, level=verbosity)
+    
+    xmpp.connect()
+    xmpp.process(forever=False)
+        
+if __name__ == "__main__":
+    main()