Create indexes when creating database, also make it possible to search based on indexes.

3 files changed, 62 insertions, 19 deletions

diff --git a/src/db.py b/src/db.py
index 14a435d..4e7bcc5 100755
--- a/src/db.py
+++ b/src/db.py
@@ -9,6 +9,7 @@
 
 import time
 import couch
+import index
 
 
 class DictDB():
@@ -20,18 +21,11 @@ class DictDB():
             self.couchdb = self.server.database(database)
         except couch.exceptions.NotFound:
             print("Creating database and indexes.")
-            index = {
-                "index": {
-                    "fields": [
-                        "domain"
-                    ]
-                },
-                "name": "domain-json-index",
-                "type": "json"
-            }
-
             self.couchdb = self.server.create(database)
-            self.couchdb.index(index)
+            self.server.create('_users')
+
+            for i in index.indexes:
+                self.couchdb.index(i)
 
         self._ts = time.time()
 
@@ -67,10 +61,7 @@ class DictDB():
         selector = dict()
 
         if kwargs:
-            selector = {
-                "selector": {
-                }
-            }
+            selector = {"selector": {}}
 
             for key in kwargs:
                 if kwargs[key].isnumeric():
diff --git a/src/index.py b/src/index.py
new file mode 100644
index 0000000..7933824
--- /dev/null
+++ b/src/index.py
@@ -0,0 +1,47 @@
+indexes = [
+    {
+        "index": {
+            "fields": [
+                "domain"
+            ]
+        },
+        "name": "domain-json-index",
+        "type": "json"
+    },
+    {
+        "index": {
+            "fields": [
+                "ip"
+            ]
+        },
+        "name": "ip-json-index",
+        "type": "json"
+    },
+    {
+        "index": {
+            "fields": [
+                "port"
+            ]
+        },
+        "name": "port-json-index",
+        "type": "json"
+    },
+    {
+        "index": {
+            "fields": [
+                "asn"
+            ]
+        },
+        "name": "asn-json-index",
+        "type": "json"
+    },
+    {
+        "index": {
+            "fields": [
+                "asn_country_code"
+            ]
+        },
+        "name": "asn-country-code-json-index",
+        "type": "json"
+    }
+]
diff --git a/src/wsgi.py b/src/wsgi.py
index e1f645c..54d18d3 100755
--- a/src/wsgi.py
+++ b/src/wsgi.py
@@ -4,6 +4,7 @@ import os
 import sys
 import json
 import authn
+import index
 import falcon
 
 from db import DictDB
@@ -46,7 +47,9 @@ class CollectorResource():
 
 class EPGet(CollectorResource):
     def on_get(self, req, resp):
-        out = []
+        out = list()
+        selectors = dict()
+
         resp.status = falcon.HTTP_200
         resp.content_type = falcon.MEDIA_JSON
         orgs = self.user_auth(req.auth, self._users.read_perms)
@@ -57,9 +60,11 @@ class EPGet(CollectorResource):
                 {'status': 'error', 'message': 'Invalid username or password\n'})
             return
 
-        # We really shouldn't rely on req.params in its pure form since
-        # it might contain garbage.
-        selectors = req.params
+        for param in req.params:
+            for i in index.indexes:
+                for j in i['index']['fields']:
+                    if j == param:
+                        selectors[param] = req.params[param]
 
         for org in orgs:
             selectors['domain'] = org