From 5530fd762c96c05a0a77cb92b93fae25a5dd5197 Mon Sep 17 00:00:00 2001
From: Mark Veidemanis <m@zm.is>
Date: Fri, 30 Sep 2022 07:22:22 +0100
Subject: [PATCH] Fix context menu

---
 core/db/druid.py           | 58 ++++++++++++++++++++++++++++++++++++++
 core/views/ui/drilldown.py | 12 +++++---
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/core/db/druid.py b/core/db/druid.py
index 0479407..38116a3 100644
--- a/core/db/druid.py
+++ b/core/db/druid.py
@@ -19,6 +19,64 @@ class DruidBackend(StorageBackend):
         # self.client = PyDruid("http://broker:8082", "druid/v2")
         pass  # we use requests
 
+    def construct_context_query(
+        self, index, net, channel, src, num, size, type=None, nicks=None
+    ):
+        search_query = self.construct_query(None, size, index, blank=True)
+        extra_must = []
+        extra_should = []
+        extra_should2 = []
+        if num:
+            extra_must.append({"num": num})
+        if net:
+            extra_must.append({"net": net})
+        if channel:
+            extra_must.append({"channel": channel})
+        if nicks:
+            for nick in nicks:
+                extra_should2.append({"nick": nick})
+        types = ["msg", "notice", "action", "kick", "topic", "mode"]
+
+        if index == "internal":
+            if channel == "*status" or type == "znc":
+                if {"channel": channel} in extra_must:
+                    extra_must.remove({"channel": channel})
+                extra_should2 = []
+                # Type is one of msg or notice
+                # extra_should.append({"match": {"mtype": "msg"}})
+                # extra_should.append({"match": {"mtype": "notice"}})
+                extra_should.append({"type": "znc"})
+                extra_should.append({"type": "self"})
+
+                extra_should2.append({"type": "znc"})
+                extra_should2.append({"nick": channel})
+            elif type == "auth":
+                if {"match": {"channel": channel}} in extra_must:
+                    extra_must.remove({"channel": channel})
+                extra_should2 = []
+                extra_should2.append({"nick": channel})
+                # extra_should2.append({"match": {"mtype": "msg"}})
+                # extra_should2.append({"match": {"mtype": "notice"}})
+
+                extra_should.append({"type": "query"})
+                extra_should2.append({"type": "self"})
+                extra_should.append({"nick": channel})
+            else:
+                for ctype in types:
+                    extra_should.append({"mtype": ctype})
+        else:
+            for ctype in types:
+                extra_should.append({"type": ctype})
+
+        if extra_must:
+            self.add_type("and", search_query, extra_must)
+
+        if extra_should:
+            self.add_type("or", search_query, extra_should)
+        if extra_should2:
+            self.add_type("or", search_query, extra_should2)
+        return search_query
+
     def construct_query(self, query, size, index, blank=False):
         search_query = {
             "limit": size,
diff --git a/core/views/ui/drilldown.py b/core/views/ui/drilldown.py
index 5465148..3e7d4f8 100644
--- a/core/views/ui/drilldown.py
+++ b/core/views/ui/drilldown.py
@@ -12,7 +12,6 @@ from rest_framework.parsers import FormParser
 from rest_framework.views import APIView
 
 from core.db.storage import db
-from core.lib.context import construct_query
 from core.lib.threshold import (
     annotate_num_chans,
     annotate_num_users,
@@ -361,9 +360,14 @@ class DrilldownContextModal(APIView):
             if query_params["type"] not in ["znc", "auth"]:
                 annotate = True
         # Create the query with the context helper
-        if query_params["num"].isdigit():
-            query_params["num"] = int(query_params["num"])
-        search_query = construct_query(
+        if "num" in query_params:
+            if query_params["num"]:
+                if query_params["num"].isdigit():
+                    query_params["num"] = int(query_params["num"])
+                else:
+                    return {"message": "Invalid num value", "class": "danger"}
+
+        search_query = db.construct_context_query(
             query_params["index"],
             query_params["net"],
             query_params["channel"],