add rest api support to ejabberdrpc.py

1 files changed, 28 insertions, 10 deletions

diff --git a/ejabberdrpc.py b/ejabberdrpc.py
index f1215e9..4cdaf88 100644
--- a/ejabberdrpc.py
+++ b/ejabberdrpc.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 from xmlrpc import client
+import requests
 import ipaddress
 
 # rfc6052: IPv6 Addressing of IPv4/IPv6 Translators
@@ -11,11 +12,28 @@ class EjabberdMetrics():
     """
     class to fetch metrics per xmlrpc
     """
-    def __init__(self, url, login = None):
-        self._server = client.ServerProxy(url)
+    def __init__(self, url, login = None, api = "rpc"):
         self._login = login
-
-    def _cmd(self, command, data):
+        if api == "rpc":
+            self._server = client.ServerProxy(url)
+            self._cmd = self._rpc
+        else:
+            self._url = url
+            self._cmd = self._rest
+    @property
+    def _auth(self):
+        if self._login is not None:
+            return ( "%s@%s" % (self._login['user'], self._login['server']),
+                self._login['password'])
+        return None
+
+    def _rest(self, command, data):
+        r = requests.post("%s/%s" % (self._url, command), auth=self._auth, json=data)
+        if r is not None:
+            return r.json()
+        return {}
+
+    def _rpc(self, command, data):
         fn = getattr(self._server, command)
         try:
             if self._login is not None:
@@ -60,7 +78,7 @@ class EjabberdMetrics():
     def fetch_onlineuser(self):
         tmp = self._cmd("connected_users_info", {})
         if "connected_users_info" not in tmp:
-            return None
+            return tmp
         data = []
         for c in tmp["connected_users_info"]:
             if "session" not in c:
@@ -75,7 +93,7 @@ class EjabberdMetrics():
     def fetch_nodes(self):
         result = self._cmd("list_cluster",{})
         if "nodes" not in result:
-            return None
+            return result
         data = []
         for node in result["nodes"]:
             data.append(node["node"])
@@ -84,7 +102,7 @@ class EjabberdMetrics():
     def fetch_vhosts(self):
         result = self._cmd("registered_vhosts",{})
         if "vhosts" not in result:
-            return None
+            return result
         data = []
         for vhost in result["vhosts"]:
             data.append(vhost["vhost"])
@@ -93,13 +111,13 @@ class EjabberdMetrics():
     def fetch_s2s_in(self):
         result = self._cmd("incoming_s2s_number",{})
         if "s2s_incoming" not in result:
-            return None
+            return result
         return result["s2s_incoming"]
 
     def fetch_s2s_out(self):
         result = self._cmd("outgoing_s2s_number",{})
         if "s2s_outgoing" not in result:
-            return None
+            return result
         return result["s2s_outgoing"]
 
     def fetch_registered(self, vhost=None):
@@ -302,7 +320,7 @@ class EjabberdMetrics():
 
 if __name__ == "__main__":
     from json import dumps
-    metric = EjabberdMetrics("http://[::1]:4560")
+    metric = EjabberdMetrics("http://[::1]:5280/api", api="rest")
    
     data = metric.get_all()
     print(dumps(data, indent=True))