[WIP] add omemo support

prometheus_xmpp/__main__.py

@@ -22,6 +22,8 @@
 from typing import Optional, Tuple
 
 import slixmpp
+from slixmpp.jid import JID
+from prometheus_xmpp import omemo_plugin
 import yaml
 from aiohttp import web
 from aiohttp_openmetrics import Counter, Gauge
@@ -140,13 +142,16 @@ def __init__(
         password_cb,
         amtool_allowed=None,
         alertmanager_url=None,
+        omemo=False,
+        omemo_keys=None
     ):
         password = password_cb()
 
         slixmpp.ClientXMPP.__init__(self, jid, password)
         self._amtool_allowed = amtool_allowed or []
         self.alertmanager_url = alertmanager_url
         self.auto_authorize = True
+        self.omemo = omemo
         self.add_event_handler("session_start", self.start)
         self.add_event_handler("message", self.message)
         self.add_event_handler("disconnected", self.lost)
@@ -157,6 +162,9 @@ def __init__(
         self.register_plugin("xep_0060")  # PubSub
         self.register_plugin("xep_0199")  # XMPP Ping
         self.register_plugin("xep_0045")  # Multi-User Chat
+        if self.omemo:
+            self.register_plugin("xep_0380")  # Explicit Message Encryption
+            self.register_plugin("xep_0384", pconfig={"storage": omemo_keys}, module=omemo_plugin)  # OMEMO
 
     def failed_auth(self, stanza):
         logging.warning("XMPP Authentication failed: %r", stanza)
@@ -203,6 +211,25 @@ def message(self, msg):
             msg.reply(response).send()
 
 
+    async def send_encrypted_message(self, mto, mtype, mbody, mhtml):
+        msg_stanza = self.make_message(mto=mto, mbody=mbody, mtype=mtype, mhtml=mhtml)
+        msg_stanza.set_to(mto)
+        msg_stanza.set_from(self.boundjid)
+
+        messages, encryption_errors = await self["xep_0384"].encrypt_message(
+            msg_stanza, mto
+        )
+        if len(encryption_errors) > 0:
+            logging.info(
+                f"There were non-critical errors during encryption: {encryption_errors}"
+            )
+
+        for namespace, message in messages.items():
+            message["eme"]["namespace"] = namespace
+            message["eme"]["name"] = self["xep_0380"].mechanisms[namespace]
+            message.send()
+
+
 async def serve_test(request):
     xmpp_app = request.app['xmpp_app']
     try:
@@ -222,11 +249,18 @@ async def serve_test(request):
             request.app['html_template'],
             EXAMPLE_ALERT)
         for to_jid in recipients:
-            xmpp_app.send_message(
-                mto=to_jid,
-                mbody=text,
-                mhtml=html,
-                mtype='chat')
+            if xmpp_app.omemo:
+                await xmpp_app.send_encrypted_message(
+                    mto=JID(to_jid),
+                    mbody=text,
+                    mhtml=html,
+                    mtype="chat")
+            else:
+                xmpp_app.send_message(
+                    mto=to_jid,
+                    mbody=text,
+                    mhtml=html,
+                    mtype='chat')
     except slixmpp.xmlstream.xmlstream.NotConnectedError as e:
         logging.warning("Test alert not posted since we are not online: %s", e)
         return web.Response(body="Did not send message. Not online: %s" % e)
@@ -281,11 +315,18 @@ async def serve_alert(request):
 
             try:
                 for (mto, mtype) in recipients:
-                    xmpp_app.send_message(
-                            mto=mto,
+                    if xmpp_app.omemo:
+                        await xmpp_app.send_encrypted_message(
+                            mto=JID(mto),
                             mbody=text,
                             mhtml=html,
                             mtype=mtype)
+                    else:
+                        xmpp_app.send_message(
+                                mto=mto,
+                                mbody=text,
+                                mhtml=html,
+                                mtype=mtype)
             except slixmpp.xmlstream.xmlstream.NotConnectedError as e:
                 logging.warning("Alert posted but we are not online: %s", e)
                 last_alert_message_succeeded_gauge.set(0)
@@ -413,6 +454,16 @@ def password_cb():
     if config.get('format') not in ('full', 'short', None):
         parser.error("unsupport config format: %s" % config['format'])
 
+    if 'OMEMO' in env:
+        config['omemo'] = env['OMEMO']
+    elif 'omemo' not in config:
+        config['omemo'] = False
+
+    if 'OMEMO_KEYS' in env:
+        config['omemo_keys'] = env['OMEMO_KEYS']
+    elif 'omemo_keys' not in config:
+        parser.error("please specifiy a storage file for omemo keys")
+
     return (
         jid,
         password_cb,
@@ -431,11 +482,14 @@ def main():
 
     amtool_allowed = config.get('amtool_allowed')
     alertmanager_url = config.get('alertmanager_url')
+    omemo = config.get('omemo')
+    omemo_keys = config.get('omemo_keys')
 
     xmpp_app = XmppApp(
         jid, password_cb,
         amtool_allowed,
-        alertmanager_url)
+        alertmanager_url,
+        omemo, omemo_keys)
 
     # Backward compatibility
     text_template = os.environ.get('TEXT_TEMPLATE')

prometheus_xmpp/omemo_plugin.py

@@ -0,0 +1,105 @@
+import json
+import logging
+from typing import Any, Dict, FrozenSet, Optional
+
+from omemo.storage import Just, Maybe, Nothing, Storage
+from omemo.types import DeviceInformation, JSONType
+
+from slixmpp_omemo import TrustLevel, XEP_0384
+
+
+from slixmpp.plugins import register_plugin  # type: ignore[attr-defined]
+
+log = logging.getLogger(__name__)
+
+class StorageImpl(Storage):
+    """
+    storage implementation that stores all data in a single JSON file.
+    Copied from https://github.com/Syndace/slixmpp-omemo/tree/main/examples
+    """
+    def __init__(self, storage_dir) -> None:
+        super().__init__()
+
+        self.JSON_FILE = storage_dir
+        self.__data: Dict[str, JSONType] = {}
+        try:
+            with open(self.JSON_FILE, encoding="utf8") as f:
+                self.__data = json.load(f)
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+
+    async def _load(self, key: str) -> Maybe[JSONType]:
+        if key in self.__data:
+            return Just(self.__data[key])
+
+        return Nothing()
+
+    async def _store(self, key: str, value: JSONType) -> None:
+        self.__data[key] = value
+        with open(self.JSON_FILE, "w", encoding="utf8") as f:
+            json.dump(self.__data, f)
+
+    async def _delete(self, key: str) -> None:
+        self.__data.pop(key, None)
+        with open(self.JSON_FILE, "w", encoding="utf8") as f:
+            json.dump(self.__data, f)
+
+
+class XEP_0384Impl(XEP_0384):  # pylint: disable=invalid-name
+    """
+    implementation of the OMEMO plugin for Slixmpp.
+    Copied from https://github.com/Syndace/slixmpp-omemo/tree/main/examples
+    """
+    def __init__(self, *args: Any, **kwargs: Any) -> None:  # pylint: disable=redefined-outer-name
+        super().__init__(*args, **kwargs)
+
+        self.__storage: Storage
+        #TODO: not sure why pconfig is not available through kwargs ?
+        self.storage_dir: str = args[1]['storage']
+
+    def plugin_init(self) -> None:
+        self.__storage = StorageImpl(self.storage_dir)
+
+        super().plugin_init()
+
+    @property
+    def storage(self) -> Storage:
+        return self.__storage
+
+    def setStorageDir(self, directory: str) -> None:
+        self.storage_dir = directory
+
+    @property
+    def _btbv_enabled(self) -> bool:
+        return True
+
+    async def _devices_blindly_trusted(
+        self,
+        blindly_trusted: FrozenSet[DeviceInformation],
+        identifier: Optional[str]
+    ) -> None:
+        log.info(f"[{identifier}] Devices trusted blindly: {blindly_trusted}")
+
+    async def _prompt_manual_trust(
+        self,
+        manually_trusted: FrozenSet[DeviceInformation],
+        identifier: Optional[str]
+    ) -> None:
+        """
+        In case of manual trust, usually OMEMO use BTBV per default : https://gultsch.de/trust.html
+        """
+        session_mananger = await self.get_session_manager()
+
+        for device in manually_trusted:
+            while True:
+                answer = input(f"[{identifier}] Trust the following device? (yes/no) {device}")
+                if answer in { "yes", "no" }:
+                    await session_mananger.set_trust(
+                        device.bare_jid,
+                        device.identity_key,
+                        TrustLevel.TRUSTED.value if answer == "yes" else TrustLevel.DISTRUSTED.value
+                    )
+                    break
+                print("Please answer yes or no.")
+
+register_plugin(XEP_0384Impl)

pyproject.toml

@@ -24,6 +24,7 @@ classifiers = [
 requires-python = ">= 3.8"
 dependencies = [
     "slixmpp",
+    "slixmpp_omemo",
     "aiohttp",
     "pytz",
     "pyyaml",

xmpp-alerts.yml.example

@@ -14,6 +14,9 @@ amtool_allowed: []
 # muc: yes
 # muc_jid: "[email protected]"
 # muc_bot_nick: "PrometheusAlerts"
+# omemo: yes
+# json file to store the private key of the bot and the publics keys of the clients
+# omemo_keys: "/restricted_directory/prometheus_bot.json"
 
 # HTML message template as jinja2:
 # html_template: |