Merged main and updated stuff.

1 files changed, 26 insertions, 15 deletions

diff --git a/src/main.py b/src/main.py
index e6bb8e2..a62d77c 100755
--- a/src/main.py
+++ b/src/main.py
@@ -1,12 +1,10 @@
+import json
 import os
 import sys
 import time
 
 import uvicorn
-
-from fastapi import Depends
-from fastapi import FastAPI
-from fastapi import Request
+from fastapi import Depends, FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi_jwt_auth import AuthJWT
@@ -14,9 +12,7 @@ from fastapi_jwt_auth.exceptions import AuthJWTException
 from pydantic import BaseModel
 
 from db.dictionary import DictDB
-from db.index import CouchIindex
-from db.sql import Log
-from db.sql import Scanner
+from db.schema import get_index_keys
 
 app = FastAPI()
 
@@ -41,8 +37,8 @@ async def mock_x_total_count_header(request: Request, call_next):
 for i in range(10):
     try:
         db = DictDB()
-    except Exception:
-        print(f"Database not responding, will try again soon. Attempt {i + 1} of 10.")
+    except Exception as e:
+        print(f"Database not responding, will try again soon: {e}")
     else:
         break
     time.sleep(1)
@@ -73,7 +69,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:
@@ -169,17 +165,32 @@ async def get_key(key=None, Authorize: AuthJWT = Depends()):
     return JSONResponse(content={"status": "success", "docs": data})
 
 
-@ app.post('/sc/v0/add')
+@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})