XMPP Probe

Pipfile

@@ -10,6 +10,7 @@ aioscheduler = "*"
 pydantic = "*"
 pyyaml = "*"
 jsonpath-ng = "*"
+slixmpp = "*"
 
 [dev-packages]
 

README.md

@@ -116,7 +116,7 @@ For example, to test a `/status` endpoint that responds with `{"status":"ok"}` i
 
 #### probes.script
 
-This is broadly a chained set of the requests above. Keys carry
+This is broadly a chained set of the requests above. Keys created by `set_keys` carry
 over between requests, and unlike the simple  probe above,
 both URLs and request bodies can be interpolated using them.
 
@@ -125,10 +125,10 @@ an `init_keys` mapping, containing each key's initial value.
 
 Finally, a `steps` array has the same keys as the `probes.http` probe's keys above, except:
 
-* Generic probe keys are not permitted here.
+* Generic probe keys, like `title`, or `every`, are not permitted here.
 * `url` and `body` keys are interpolated.
 
-Interpolation operates by the Python str.format_map, in fact, with the current keys passed in directly as a map.
+Interpolation operates by the Python `str.format_map`, in fact, with the current keys passed in directly as a map.
 
 However, if a string to be interpolated precisely matches the name of a key, it's simply replaced. Therefore if there is a key called `test` currently set to `true`, and the body looks like:
 

main.py

@@ -43,6 +43,9 @@ def probe_group(group: ProbeGroup):
 async def main():
     import sentry_sdk
     import sys
+    import logging
+
+    logging.basicConfig(level=logging.DEBUG)
 
     config = load(sys.argv[1])
 

probes/xmpp/__init__.py

@@ -0,0 +1,4 @@
+from .xmpp import probe
+from .config import check_config
+
+__all__ = ['probe']

probes/xmpp/config.py

@@ -0,0 +1,27 @@
+from config.base import ProbeGroup, ProbeConfig
+from typing import Optional
+
+
+class Config(ProbeConfig):
+    jid: str
+    password: str
+    starttls: bool = False
+    address: Optional[str] = None
+    port: Optional[int] = None
+    remote: Optional[str] = None
+
+
+def check_config(config: dict, group: ProbeGroup):
+    conf = Config(**config)
+    conf.tags.update(group.tags)
+    if not conf.title:
+        xmpp = 'xmpps'
+        if conf.starttls:
+            xmpp = 'xmpp'
+            conf.port = conf.port or 5222
+        else:
+            conf.port = conf.port or 5223
+            if not conf.address:
+                raise ValueError('Must supply address if not using StartTLS')
+        conf.title = '%s://%s' % (xmpp, conf.jid)
+    return conf

probes/xmpp/xmpp.py

@@ -0,0 +1,106 @@
+from typing import Optional
+
+import sentry_sdk
+import slixmpp
+from slixmpp import ClientXMPP
+from .config import Config
+
+
+class TracedClient(ClientXMPP):
+    def __init__(self, jid: str, password: str, starttls: bool, remote: Optional[str]):
+        super().__init__(jid, password)
+        self.starttls = starttls
+        self.remote = remote
+        if self.remote:
+            self.register_plugin('xep_0199')
+        self.spans = []
+        self.named_spans = {}
+        self.add_event_handler('session_start', self.started, disposable=True)
+        self.add_event_handler('message', self.message)
+        self.add_event_handler('presence_available', self.presence)
+        self.add_event_handler('connected', self.connected, disposable=True)
+        self.add_event_handler('connection_failed', self.failure, disposable=True)
+        self.add_event_handler('failed_auth', self.failure, disposable=True)
+        self.add_event_handler('tls_success', self.tls, disposable=True)
+        self.add_event_handler('auth_success', self.sasl, disposable=True)
+        self.add_event_handler('session_bind', self.bind, disposable=True)
+        self.spans.append(sentry_sdk.start_span(op='xmpp.session', description=jid))
+        self.spans.append(self.spans[-1].start_child(op='session_start', description=jid))
+        self.named_span('connect')
+
+    async def failure(self, event):
+        print("Failure", repr(event))
+        await self.disconnect()
+
+    def connected(self, event):
+        self.clear_span_maybe('connect')
+        if self.starttls:
+            self.named_span('tls')
+        else:
+            self.named_span('sasl')
+
+    def tls(self, event):
+        self.clear_span_maybe('tls')
+        self.named_span('sasl')
+
+    def sasl(self, event):
+        self.clear_span_maybe('sasl')
+        self.named_span('bind')
+
+    def bind(self, event):
+        self.clear_span_maybe('bind')
+
+    def named_span(self, name: str):
+        current = self.spans[-1]
+        self.named_spans[name] = current.start_child(op=name)
+
+    def clear_span_maybe(self, which: str):
+        span = self.named_spans.get(which)
+        if span:
+            span.finish()
+            del self.named_spans[which]
+
+    def clear_spans(self):
+        for span in self.named_spans.values():
+            span.finish()
+        self.named_spans = {}
+        while len(self.spans):
+            span = self.spans.pop()
+            span.finish()
+
+    async def started(self, event):
+        span = self.spans.pop()
+        span.finish()
+        self.named_span('presence')
+        self.send_presence()
+        with self.spans[-1].start_child(op='roster'):
+            await self.get_roster()
+        self.named_span('message')
+        self.send_message(self.boundjid, 'Test')
+
+    async def message(self, msg: slixmpp.Message):
+        if msg.get_from() != self.boundjid:
+            return
+        self.clear_span_maybe('message')
+        if self.remote:
+            with self.spans[-1].start_child(op='ping'):
+                await self.plugin['xep_0199'].ping(jid=self.remote)
+        with self.spans[-1].start_child(op='disconnect'):
+            await self.disconnect()
+
+    async def presence(self, presence: slixmpp.Presence):
+        if presence.get_from() != self.boundjid:
+            return
+        self.clear_span_maybe('presence')
+
+
+async def probe(config: Config):
+    print(config.title)
+    xmpp = TracedClient(config.jid, config.password, config.starttls, config.remote)
+    address = None
+    if config.address:
+        address = (config.address, config.port)
+    xmpp.connect(use_ssl=not config.starttls, force_starttls=config.starttls, address=address)
+    await xmpp.disconnected
+    await xmpp.disconnect(wait=0.1, ignore_send_queue=True)
+    xmpp.clear_spans()