diff options
| -rwxr-xr-x | src/db.py | 10 | ||||
| -rw-r--r-- | src/index.py | 61 | ||||
| -rwxr-xr-x | src/main.py | 37 | ||||
| -rw-r--r-- | src/schema.py | 34 |
4 files changed, 64 insertions, 78 deletions
@@ -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)) |