More database changes and endpoints.

1 files changed, 32 insertions, 13 deletions

diff --git a/src/main.py b/src/main.py
index c93029f..e6bb8e2 100755
--- a/src/main.py
+++ b/src/main.py
@@ -3,7 +3,10 @@ import sys
 import time
 
 import uvicorn
-from fastapi import Depends, FastAPI, Request
+
+from fastapi import Depends
+from fastapi import FastAPI
+from fastapi import Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi_jwt_auth import AuthJWT
@@ -12,13 +15,14 @@ from pydantic import BaseModel
 
 from db.dictionary import DictDB
 from db.index import CouchIindex
-from db.sql import Log, Scanner
+from db.sql import Log
+from db.sql import Scanner
 
 app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["http://localhost:8001"],
+    allow_origins=["http://localhost:8000"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -38,8 +42,7 @@ for i in range(10):
     try:
         db = DictDB()
     except Exception:
-        print(
-            f'Database not responding, will try again soon. Attempt {i + 1} of 10.')
+        print(f"Database not responding, will try again soon. Attempt {i + 1} of 10.")
     else:
         break
     time.sleep(1)
@@ -90,25 +93,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 +137,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()
@@ -166,7 +169,7 @@ 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
@@ -180,7 +183,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()
@@ -217,7 +220,7 @@ async def delete(key, Authorize: AuthJWT = Depends()):
     return JSONResponse(content={"status": "success", "docs": data})
 
 
-@app.get("/sc/v0/scanner/{name}")
+@ app.get("/sc/v0/scanner/{name}")
 async def scanner_get(name, data: Request, Authorize: AuthJWT = Depends()):
     Authorize.jwt_required()
 
@@ -226,6 +229,22 @@ async def scanner_get(name, data: Request, Authorize: AuthJWT = Depends()):
     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