#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from xmlrpc import client
import ipaddress

class EjabberdMetrics():
    """
    class to fetch metrics per xmlrpc
    """
    def __init__(self, url, login = None):
        self._server = client.ServerProxy(url)
        self._login = login

    def _cmd(self, command, data):
        fn = getattr(self._server, command)
        if self._login is not None:
            return fn(self._login, data)
        return fn(data)

    def fetch_onlineuser(self):
        data = None
        tmp = self._cmd("connected_users_info", {})
        if "connected_users_info" not in tmp:
            return None
        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 None
        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 None
        data = []
        for vhost in result["vhosts"]:
            data.append(vhost["vhost"])
        return data

    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 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()
        for vhost in self._vhosts:
            self._registered[vhost] = self.fetch_registered(vhost)
        # online user
        self._onlineuser = self.fetch_onlineuser()


    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):
        def client(r):
            clientmap = {
                "Conv6ations for Sum7": ["Conversations with IPv6"],
                "Conversations": [],
                "Pix-Art Messenger": [],
                "jitsi": [],
                "dino": [],
                "poezio": [],
            }
            for client,names in clientmap.items():
                for c in names:
                    if c in r:
                        return client
                if client in r:
                    return client
            return "other"
        return self.get_online_by("resource", parse=client, vhost=vhost, node=node)

    def get_online_by_ipversion(self, vhost=None, node=None):
        # rfc6052: IPv6 Addressing of IPv4/IPv6 Translators
        nat64 = ipaddress.ip_network("64:ff9b::/96")
        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
        return self.get_online_by("ip", parse=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_vhosts(self):
        if not hasattr(self, "_vhosts"):
            self._vhosts = self.fetch_vhosts()
        return self._vhosts

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

        data["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 = {}
        data["online_by_status"] = self.get_online_by_status(node=node)
        data["online_by_client"] = self.get_online_by_client(node=node)
        data["online_by_ipversion"] = self.get_online_by_ipversion(node=node)
        data["online_by_connection"] = self.get_online_by_connection(node=node)

        data["online_by_vhost"] = self.get_online_by_vhost(node=node)
        return data

    def get_all(self):
        data = {}
        data["registered"] = self.get_registered()
        data["online_by_status"] = self.get_online_by_status()
        data["online_by_client"] = self.get_online_by_client()
        data["online_by_ipversion"] = self.get_online_by_ipversion()
        data["online_by_connection"] = self.get_online_by_connection()

        data["online_by_node"] = self.get_online_by_node()
        data["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["nodes"] = nodes
        
        return data


if __name__ == "__main__":
    from json import dumps
    metric = EjabberdMetrics("http://[::1]:4560")
   
    data = metric.get_all()
    print(dumps(data, indent=True))