* Merge branch 'main' into feature.callhome

* New API endpoints
* Updated requirements

1 files changed, 59 insertions, 8 deletions

diff --git a/src/routers/collector.py b/src/routers/collector.py
index 3cda23a..7d91609 100644
--- a/src/routers/collector.py
+++ b/src/routers/collector.py
@@ -48,18 +48,22 @@ def get_data(key=None, limit=25, skip=0, ip=None,
 @router.get('/get')
 async def get(key=None, limit=25, skip=0, ip=None, port=None,
               asn=None, Authorize: AuthJWT = Depends()):
+
     Authorize.jwt_required()
 
     data = []
     raw_jwt = Authorize.get_raw_jwt()
 
-    if 'domains' not in raw_jwt:
-        return JSONResponse(content={"status": "error",
-                                     "message": "Could not find domains" +
-                                     "claim in JWT token"},
-                            status_code=400)
+    if "read" not in raw_jwt:
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": "Could not find read claim in JWT token",
+            },
+            status_code=400,
+        )
     else:
-        domains = raw_jwt['domains']
+        domains = raw_jwt["read"]
 
     for domain in domains:
         data.extend(get_data(key, limit, skip, ip, port, asn, domain))
@@ -69,17 +73,39 @@ async def get(key=None, limit=25, skip=0, ip=None, port=None,
 
 @router.get('/get/{key}')
 async def get_key(key=None, Authorize: AuthJWT = Depends()):
+
     Authorize.jwt_required()
 
-    # TODO: Use JWT authz and check e.g. domain here
+    raw_jwt = Authorize.get_raw_jwt()
+
+    if "read" not in raw_jwt:
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": "Could not find read claim in JWT token",
+            },
+            status_code=400,
+        )
+    else:
+        allowed_domains = raw_jwt["read"]
 
     data = get_data(key)
 
+    if data["domain"] not in allowed_domains:
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": "User not authorized to view this object",
+            },
+            status_code=400,
+        )
+
     return JSONResponse(content={"status": "success", "docs": data})
 
 
 @router.post('/add')
 async def add(data: Request, Authorize: AuthJWT = Depends()):
+
     Authorize.jwt_required()
 
     json_data = await data.json()
@@ -91,11 +117,36 @@ async def add(data: Request, Authorize: AuthJWT = Depends()):
 
 @router.delete('/delete/{key}')
 async def delete(key, Authorize: AuthJWT = Depends()):
+
     Authorize.jwt_required()
 
+    raw_jwt = Authorize.get_raw_jwt()
+
+    if "write" not in raw_jwt:
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": "Could not find write claim in JWT token",
+            },
+            status_code=400,
+        )
+    else:
+        allowed_domains = raw_jwt["write"]
+
+    data = get_data(key)
+
+    if data["domain"] not in allowed_domains:
+        return JSONResponse(
+            content={
+                "status": "error",
+                "message": "User not authorized to delete this object",
+            },
+            status_code=400,
+        )
+
     if db.delete(key) is None:
         return JSONResponse(content={"status": "error",
                                      "message": "Document not found"},
                             status_code=400)
 
-    return JSONResponse(content={"status": "success", "docs": {}})
+    return JSONResponse(content={"status": "success", "docs": data})