Add read/write permissions to JWTs based on YAML

- Uses Linus's YAML code, except with password stuff removed since
  auth-server-poc uses htpasswd.
- The collector checks JWT on API endpoints get, get/{key}, and
  delete/{key}, but not on add.

1 files changed, 55 insertions, 8 deletions

diff --git a/src/main.py b/src/main.py
index f95a09c..9beace0 100755
--- a/src/main.py
+++ b/src/main.py
@@ -116,13 +116,16 @@ async def get(key=None, limit=25, skip=0, ip=None, port=None,
     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))
@@ -135,10 +138,30 @@ 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})
 
 
@@ -161,12 +184,36 @@ 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})
 
 
 def main(standalone=False):