ejabberd-tools/ejabberdrpc.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import ipaddress
import re

from packaging import version

# rfc6052: IPv6 Addressing of IPv4/IPv6 Translators
nat64 = ipaddress.ip_network("64:ff9b::/96")


class EjabberdMetrics:
    """
    class to fetch metrics per xmlrpc
    """
    def __init__(self, url, login=None, api="rpc", muc_host: str = 'conference'):
        self._login = login
        self.muc_host = muc_host
        if api == "rpc":
            self.url = url
            self._cmd = self._rpc
        else:
            self._url = url
            self._cmd = self._rest

    @property
    def _auth(self):
        if self._login is not None:
            return f"{self._login['user']}@{self._login['server']}", self._login['password']
        return None

    @property
    def _verstring(self):
        if self._login is not None:
            ver_str = re.compile('([1-9][0-9.-]*)')
            status = self._cmd('status', {})

            # matches
            tmp = ver_str.findall(status)[0]

            # return parsed version string
            return version.parse(tmp)

        return None

    def _rest(self, command: str, data):
        import requests

        with requests.Session() as s:
            r = s.post(f'{self._url}/{command}', auth=self._auth, json=data)

            if r.status_code == 200:
                return r.json()
            return{}

    def _rpc(self, command: str, data):
        from xmlrpc import client

        with client.ServerProxy(self.url) as server:
            fn = getattr(server, command)
            try:
                if self._login is not None:
                    return fn(self._login, data)
                return fn(data)
            except:
                return {}

    def _client(self, resource):
        clientmap = {
            "Conv6ations for Sum7": ["Conversations with IPv6"],
            "Conversations": [],
            "Pix-Art Messenger": [],
            "Gajim": ["gajim"],
            "Psi+": [],
            "jitsi": [],
            "Dino": ["dino"],
            "poezio": [],
            "profanity": [],
            "Xabber": [],
            "ChatSecure": ["chatsecure"]
        }

        for client, names in clientmap.items():
            for c in names:
                if c in resource:
                    return client
            if client in resource:
                return client
        return "other"

    @staticmethod
    def _ipversion(ip):
        addr = ipaddress.ip_address(ip)
        if addr.version == 6:
            if addr.ipv4_mapped:
                return 4
            if addr in nat64:
                return 4
        return addr.version

    def fetch_onlineuser(self):
        tmp = self._cmd("connected_users_info", {})
        if "connected_users_info" not in tmp:
            return tmp
        data = []
        for c in tmp["connected_users_info"]:
            if "session" not in c:
                continue
            user = {}
            for attrs in c["session"]:
                for k, v in attrs.items():
                    user[k] = v
            data.append(user)
        return data

    def fetch_nodes(self):
        result = self._cmd("list_cluster", {})
        if "nodes" not in result:
            return result
        data = []
        for node in result["nodes"]:
            data.append(node["node"])
        return data

    def fetch_vhosts(self):
        result = self._cmd("registered_vhosts", {})
        if "vhosts" not in result:
            return result
        data = []
        for vhost in result["vhosts"]:
            data.append(vhost["vhost"])
        return data

    def fetch_s2s_in(self):
        result = self._cmd("incoming_s2s_number", {})
        if "s2s_incoming" not in result:
            return result
        return result["s2s_incoming"]

    def fetch_s2s_out(self):
        result = self._cmd("outgoing_s2s_number", {})
        if "s2s_outgoing" not in result:
            return result
        return result["s2s_outgoing"]

    def fetch_registered(self, vhost=None):
        if vhost is None:
            result = self._cmd("stats", {"name":"registeredusers"})
            if "stat" in result:
                return result["stat"]
        else:
            result = self._cmd("stats_host", {"name":"registeredusers", "host": vhost})
            if "stat" in result:
                return result["stat"]

    def fetch_muc(self, vhost=None):
        host = "global"
        if vhost is not None:
            if self._verstring.major >= 19:
                host = '.'.join([self.muc_host, vhost])
            else:
                host = vhost
        result = self._cmd("muc_online_rooms", {"host": host})
        if "rooms" in result:
            return len(result["rooms"])
        return len(result)

    def update(self):
        # nodes
        self._nodes = self.fetch_nodes()

        # vhosts
        self._vhosts = self.fetch_vhosts()

        # registered
        if not hasattr(self, "_registered"):
            self._registered = {}
        self._registered[None] = self.fetch_registered()

        # muc
        if not hasattr(self, "_muc"):
            self._muc = {}
        self._muc[None] = self.fetch_muc()

        # registered + muc
        for vhost in self._vhosts:
            self._registered[vhost] = self.fetch_registered(vhost)
            self._muc[vhost] = self.fetch_muc(vhost)

        # online user
        self._onlineuser = self.fetch_onlineuser()

        # s2s
        self._s2s_in = self.fetch_s2s_in()
        self._s2s_out = self.fetch_s2s_out()

    def get_online_by(self, by="node", parse=None, vhost=None, node=None):
        parser = parse or (lambda a: a)
        if not hasattr(self, "_onlineuser"):
            self._onlineuser = self.fetch_onlineuser()
        data = {}

        for conn in self._onlineuser:
            if vhost is not None and vhost not in conn["jid"]:
                continue
            if node is not None and node != conn["node"]:
                continue
            if by not in conn:
                continue
            value = parser(conn[by])
            if value not in data:
                data[value] = 1
            else:
                data[value] += 1
        return data

    def get_online_by_node(self, vhost=None):
        return self.get_online_by("node", vhost=vhost)

    def get_online_by_vhost(self, node=None):
        return self.get_online_by("jid", parse=lambda jid: jid[jid.find("@")+1:jid.find("/")], node=node)

    def get_online_by_status(self, vhost=None, node=None):
        return self.get_online_by("status", vhost=vhost, node=node)

    def get_online_by_connection(self, vhost=None, node=None):
        return self.get_online_by("connection", vhost=vhost, node=node)

    def get_online_by_client(self, vhost=None, node=None):
        return self.get_online_by("resource", parse=self._client, vhost=vhost, node=node)

    def get_online_by_ipversion(self, vhost=None, node=None):
        return self.get_online_by("ip", parse=self._ipversion, vhost=vhost, node=node)

    def get_online_client_by(self, by="ip", parse=None, vhost=None, node=None):
        parser = parse or self._ipversion
        if not hasattr(self, "_onlineuser"):
            self._onlineuser = self.fetch_onlineuser()
        data = {}

        for conn in self._onlineuser:
            client = "other"
            if "resource" in conn:
                client = self._client(conn["resource"])
            if client not in data:
                data[client] = {}
            if vhost is not None and vhost not in conn["jid"]:
                continue
            if node is not None and node != conn["node"]:
                continue
            if by not in conn:
                continue
            value = parser(conn[by])
            if value not in data[client]:
                data[client][value] = 1
            else:
                data[client][value] += 1
        return data

    def get_online_client_by_ipversion(self, vhost=None, node=None):
        return self.get_online_client_by("ip", parse=self._ipversion, vhost=vhost, node=node)

    def get_registered(self, vhost=None):
        if not hasattr(self, "_registered"):
            self._registered = {}
        if vhost not in self._registered:
            self._registered[vhost] = self.fetch_registered(vhost)
        return self._registered[vhost]

    def get_muc(self, vhost=None):
        if not hasattr(self, "_muc"):
            self._muc = {}
        if vhost not in self._muc:
            self._muc[vhost] = self.fetch_muc(vhost)
        return self._muc[vhost]

    def get_vhosts(self):
        if not hasattr(self, "_vhosts"):
            self._vhosts = self.fetch_vhosts()
        return self._vhosts

    def get_s2s_in(self):
        if not hasattr(self, "_s2s_in"):
            self._s2s_in = self.fetch_s2s_in()
        return self._s2s_in

    def get_s2s_out(self):
        if not hasattr(self, "_s2s_out"):
            self._s2s_out = self.fetch_s2s_out()
        return self._s2s_out

    def get_vhost_metrics(self, vhost): 
        data = {
            "registered": self.get_registered(vhost),
            "muc": self.get_muc(vhost),
            "online_by_status": self.get_online_by_status(vhost),
            "online_by_client": self.get_online_by_client(vhost),
            "online_by_ipversion": self.get_online_by_ipversion(vhost),
            "online_by_connection": self.get_online_by_connection(vhost),
            "online_by_node": self.get_online_by_node(vhost)
        }

        return data

    def get_nodes(self):
        if not hasattr(self, "_nodes"):
            self._nodes = self.fetch_nodes()
        return self._nodes

    def get_node_metrics(self, node): 
        data = {
            "online_by_status": self.get_online_by_status(node=node),
            "online_by_client": self.get_online_by_client(node=node),
            "online_by_ipversion": self.get_online_by_ipversion(node=node),
            "online_by_connection": self.get_online_by_connection(node=node),
            "online_by_vhost": self.get_online_by_vhost(node=node)
        }

        return data

    def get_all(self):
        data = {
            "registered": self.get_registered(),
            "muc": self.get_muc(),
            "online_by_status": self.get_online_by_status(),
            "online_by_client": self.get_online_by_client(),
            "online_by_ipversion": self.get_online_by_ipversion(),
            "online_by_connection": self.get_online_by_connection(),
            "online_by_node": self.get_online_by_node(),
            "online_by_vhost": self.get_online_by_vhost()
        }

        vhosts = {}
        for host in self.get_vhosts():
            vhosts[host] = self.get_vhost_metrics(host)
        data["vhosts"] = vhosts
        
        nodes = {}
        for node in self.get_nodes():
            nodes[node] = self.get_node_metrics(node)

        data["online_client_by_ipversion"] = self.get_online_client_by_ipversion()
        data["nodes"] = nodes
        
        data["s2s_in"] = self.get_s2s_in()
        data["s2s_out"] = self.get_s2s_out()
        return data


if __name__ == "__main__":
    import os
    import json

    # load config
    path = os.path.dirname(__file__)
    with open("/".join([path, "config.json"]), "r", encoding="utf-8") as f:
        config = json.load(f)

    url = config['url'] if "url" in config else "http://[::1]:5280/api"
    login = config['login'] if "login" in config else None
    api = config['api'] if "api" in config else "rest"

    metric = EjabberdMetrics(url, login=login, api=api)
   
    data = metric.get_all()
    print(json.dumps(data, indent=True))