Use the schema when creating indexes, also validate data before writing to CouchDB.

4 files changed, 64 insertions, 78 deletions

diff --git a/src/db.py b/src/db.py
index d186616..6f25ec3 100755
--- a/src/db.py
+++ b/src/db.py
@@ -12,7 +12,7 @@ import sys
 import time
 
 import couch
-import index
+from schema import as_index_list, validate_collector_data
 
 
 class DictDB():
@@ -47,7 +47,7 @@ class DictDB():
             print("Creating database and indexes.")
             self.couchdb = self.server.create(self.database)
 
-            for i in index.as_list():
+            for i in as_index_list():
                 self.couchdb.index(i)
 
         self._ts = time.time()
@@ -72,9 +72,15 @@ class DictDB():
 
         if type(data) is list:
             for item in data:
+                error = validate_collector_data(item)
+                if error != "":
+                    return error
                 item['_id'] = str(self.unique_key())
             ret = self.couchdb.save_bulk(data)
         else:
+            error = validate_collector_data(data)
+            if error != "":
+                return error
             data['_id'] = str(self.unique_key())
             ret = self.couchdb.save(data)
 
diff --git a/src/index.py b/src/index.py
deleted file mode 100644
index 688ceeb..0000000
--- a/src/index.py
+++ /dev/null
@@ -1,61 +0,0 @@
-from pydantic import BaseSettings
-
-
-class CouchIindex(BaseSettings):
-    domain: dict = {
-        "index": {
-            "fields": [
-                "domain",
-            ]
-        },
-        "name": "domain-json-index",
-        "type": "json"
-    }
-    ip: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "ip"
-            ]
-        },
-        "name": "ip-json-index",
-        "type": "json"
-    }
-    port: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "port"
-            ]
-        },
-        "name": "port-json-index",
-        "type": "json"
-    }
-    asn: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "asn"
-            ]
-        },
-        "name": "asn-json-index",
-        "type": "json"
-    }
-    asn_country_code: dict = {
-        "index": {
-            "fields": [
-                "domain",
-                "asn_country_code"
-            ]
-        },
-        "name": "asn-country-code-json-index",
-        "type": "json"
-    }
-
-
-def as_list():
-    index_list = list()
-    for item in CouchIindex().dict():
-        index_list.append(CouchIindex().dict()[item])
-
-    return index_list
diff --git a/src/main.py b/src/main.py
index fb359df..9de8eb8 100755
--- a/src/main.py
+++ b/src/main.py
@@ -1,16 +1,18 @@
+import json
 import os
 import sys
-import uvicorn
+import time
 
-from fastapi import FastAPI, Depends, Request
+import uvicorn
+from fastapi import Depends, FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi_jwt_auth import AuthJWT
 from fastapi_jwt_auth.exceptions import AuthJWTException
-from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from index import CouchIindex
-import time
+
 from db import DictDB
+from schema import get_index_keys, validate_collector_data
 
 app = FastAPI()
 
@@ -68,7 +70,7 @@ def get_data(key=None, limit=25, skip=0, ip=None,
         return db.get(key)
 
     selectors = dict()
-    indexes = CouchIindex().dict()
+    indexes = get_index_keys()
     selectors['domain'] = domain
 
     if ip and 'ip' in indexes:
@@ -166,15 +168,30 @@ async def get_key(key=None, Authorize: AuthJWT = Depends()):
 
 @app.post('/sc/v0/add')
 async def add(data: Request, Authorize: AuthJWT = Depends()):
-
-    # Maybe we should protect this enpoint too and let the scanner use
-    # a JWT token as well.
     # Authorize.jwt_required()
 
-    json_data = await data.json()
+    try:
+        json_data = await data.json()
+    except json.decoder.JSONDecodeError:
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": "Invalid JSON.",
+            },
+            status_code=400,
+        )
 
     key = db.add(json_data)
 
+    if isinstance(key, str):
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": key,
+            },
+            status_code=400,
+        )
+
     return JSONResponse(content={"status": "success", "docs": key})
 
 
diff --git a/src/schema.py b/src/schema.py
index f92a2ea..9bdf130 100644
--- a/src/schema.py
+++ b/src/schema.py
@@ -1,5 +1,6 @@
 import json
 import sys
+import traceback
 
 import jsonschema
 
@@ -94,18 +95,41 @@ schema = {
 # fmt:on
 
 
+def get_index_keys():
+    keys = list()
+    for key in schema["properties"]:
+        keys.append(key)
+    return keys
+
+
+def as_index_list():
+    index_list = list()
+    for key in schema["properties"]:
+        name = f"{key}-json-index"
+        index = {
+            "index": {
+                "fields": [
+                    key,
+                ]
+            },
+            "name": name,
+            "type": "json"
+        }
+        index_list.append(index)
+
+    return index_list
+
+
 def validate_collector_data(json_blob):
     try:
         jsonschema.validate(json_blob, schema)
     except jsonschema.exceptions.ValidationError as e:
-        print(f"Validation failed with error: {e}")
-        return False
-
-    return True
+        return f"Validation failed with error: {e.message}"
+    return ""
 
 
 if __name__ == "__main__":
     with open(sys.argv[1]) as fd:
         json_data = json.loads(fd.read())
 
-    validate_collector_data(json_data)
+    print(validate_collector_data(json_data))