1 files changed, 36 insertions, 12 deletions
diff --git a/src/main.py b/src/main.py
index 9de8eb8..a62d77c 100755
--- a/src/main.py
+++ b/src/main.py
@@ -11,14 +11,14 @@ from fastapi_jwt_auth import AuthJWT
 from fastapi_jwt_auth.exceptions import AuthJWTException
 from pydantic import BaseModel
 
-from db import DictDB
-from schema import get_index_keys, validate_collector_data
+from db.dictionary import DictDB
+from db.schema import get_index_keys
 
 app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["http://localhost:8001"],
+    allow_origins=["http://localhost:8000"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -37,9 +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)
@@ -90,25 +89,25 @@ class JWTConfig(BaseModel):
     authjwt_public_key: str = get_pubkey()
 
 
-@AuthJWT.load_config
+@ AuthJWT.load_config
 def jwt_config():
     return JWTConfig()
 
 
-@app.exception_handler(AuthJWTException)
+@ app.exception_handler(AuthJWTException)
 def authjwt_exception_handler(request: Request, exc: AuthJWTException):
     return JSONResponse(content={"status": "error", "message":
                                  exc.message}, status_code=400)
 
 
-@app.exception_handler(RuntimeError)
+@ app.exception_handler(RuntimeError)
 def app_exception_handler(request: Request, exc: RuntimeError):
     return JSONResponse(content={"status": "error", "message":
                                  str(exc.with_traceback(None))},
                         status_code=400)
 
 
-@app.get('/sc/v0/get')
+@ app.get('/sc/v0/get')
 async def get(key=None, limit=25, skip=0, ip=None, port=None,
               asn=None, Authorize: AuthJWT = Depends()):
 
@@ -134,7 +133,7 @@ async def get(key=None, limit=25, skip=0, ip=None, port=None,
     return JSONResponse(content={"status": "success", "docs": data})
 
 
-@app.get('/sc/v0/get/{key}')
+@ app.get('/sc/v0/get/{key}')
 async def get_key(key=None, Authorize: AuthJWT = Depends()):
 
     Authorize.jwt_required()
@@ -195,7 +194,7 @@ async def add(data: Request, Authorize: AuthJWT = Depends()):
     return JSONResponse(content={"status": "success", "docs": key})
 
 
-@app.delete('/sc/v0/delete/{key}')
+@ app.delete('/sc/v0/delete/{key}')
 async def delete(key, Authorize: AuthJWT = Depends()):
 
     Authorize.jwt_required()
@@ -232,6 +231,31 @@ async def delete(key, Authorize: AuthJWT = Depends()):
     return JSONResponse(content={"status": "success", "docs": data})
 
 
+@ app.get("/sc/v0/scanner/{name}")
+async def scanner_get(name, data: Request, Authorize: AuthJWT = Depends()):
+    Authorize.jwt_required()
+
+    scanners = Scanner.get(name)
+
+    return JSONResponse(content={"status": "success", "data": scanners})
+
+
+@ app.put("/sc/v0/scanner/{name}")
+async def scanner_put(name, data: Request, Authorize: AuthJWT = Depends()):
+    errors = None
+    Authorize.jwt_required()
+
+    json_data = await data.json()
+
+    if "active" in json_data and isinstance(json_data["active"], bool):
+        errors = Scanner.active(name, json_data["active"])
+
+    if errors:
+        return JSONResponse(content={"status": "error", "message": "\n".join(errors)}, status_code=400)
+
+    return JSONResponse(content={"status": "success", "data": Scanner.get(name)}, status_code=200)
+
+
 def main(standalone=False):
     if not standalone:
         return app