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

1 files changed, 27 insertions, 10 deletions

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})