From 8c9ec3ab9c61667ce949131ba2a5e08b61aaf342 Mon Sep 17 00:00:00 2001
From: Mark Veidemanis <m@zm.is>
Date: Thu, 21 Jul 2022 13:40:18 +0100
Subject: [PATCH] Implement getting number of channels and users

---
 api/views.py        | 38 ++++++++++++++++++++++++++++++++++++++
 modules/userinfo.py | 24 +++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/api/views.py b/api/views.py
index 98f49bc..84a5356 100644
--- a/api/views.py
+++ b/api/views.py
@@ -119,3 +119,41 @@ class API(object):
                 usermap[user] = False
 
         return dumps(usermap)
+
+    @app.route("/num_users/", methods=["POST"])
+    @login_required
+    def num_users(self, request):
+        try:
+            data = loads(request.content.read())
+        except JSONDecodeError:
+            return "Invalid JSON"
+        if "net" not in data:
+            return "No net provided"
+        if "query" not in data:
+            return "No users provided"
+        if not data["query"]:
+            warn(f"No chans provided: for online {data}")
+            return dumps({})
+        net = data["net"]
+        results = userinfo.getUserNum(net, data["query"])
+
+        return dumps(results)
+
+    @app.route("/num_chans/", methods=["POST"])
+    @login_required
+    def num_chans(self, request):
+        try:
+            data = loads(request.content.read())
+        except JSONDecodeError:
+            return "Invalid JSON"
+        if "net" not in data:
+            return "No net provided"
+        if "query" not in data:
+            return "No users provided"
+        if not data["query"]:
+            warn(f"No users provided: for online {data}")
+            return dumps({})
+        net = data["net"]
+        results = userinfo.getChanNum(net, data["query"])
+
+        return dumps(results)
diff --git a/modules/userinfo.py b/modules/userinfo.py
index d6b14e0..38246e3 100644
--- a/modules/userinfo.py
+++ b/modules/userinfo.py
@@ -23,7 +23,7 @@ def getWho(query):
 
 
 def getChansSingle(name, nick):
-    nick = ["live.chan." + name + "." + i for i in nick]
+    nick = ("live.chan." + name + "." + i for i in nick)
     result = main.r.sinter(*nick)
     if len(result) == 0:
         return None
@@ -38,6 +38,28 @@ def getChanList(name, nick):
     return (i.decode() for i in result)
 
 
+def getUserNum(name, channels):
+    """
+    Get the number of users on a list of channels.
+    """
+    chanspace = ("live.who." + name + "." + i for i in channels)
+    results = {}
+    for channel, space in zip(channels, chanspace):
+        results[channel] = main.r.scard(space)
+    return results
+
+
+def getChanNum(name, nicks):
+    """
+    Get the number of channels a list of users are on.
+    """
+    nickspace = ("live.chan." + name + "." + i for i in nicks)
+    results = {}
+    for nick, space in zip(nicks, nickspace):
+        results[nick] = main.r.scard(space)
+    return results
+
+
 def getChans(nick):
     result = {}
     for i in main.network.keys():