simplify to use a custom jwt data, also verify k8s object exists

Author: ikreymer

backend/btrixcloud/auth.py

@@ -4,7 +4,7 @@
 from uuid import UUID, uuid4
 import asyncio
 from datetime import timedelta
-from typing import Optional, Tuple, List, Callable
+from typing import Optional, Tuple, List
 from passlib import pwd
 from passlib.context import CryptContext
 
@@ -42,6 +42,7 @@
 PWD_CONTEXT = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 # Audiences
+CUSTOM_AUTH_AUD = "btrix:custom-auth"
 AUTH_AUD = "btrix:auth"
 RESET_AUD = "btrix:reset"
 VERIFY_AUD = "btrix:verify"
@@ -121,14 +122,22 @@ def create_access_token(user: User) -> str:
 
 
 # ============================================================================
-def create_internal_crawler_access_token(sub: str, role: str) -> str:
+def create_custom_jwt_token(sub: str, data: dict[str, str]) -> str:
     """create jwt token for internal crawler access"""
     return generate_jwt(
-        {"sub": sub, "internal_role": role, "aud": AUTH_AUD},
+        {**data, "sub": sub, "aud": CUSTOM_AUTH_AUD},
         INTERNAL_JWT_TOKEN_LIFETIME,
     )
 
 
+# ============================================================================
+def get_custom_jwt_token(request: Request) -> dict[str, str]:
+    """return data from custom jwt token"""
+    token = request.query_params.get("auth_bearer") or ""
+    payload = decode_jwt(token, [CUSTOM_AUTH_AUD])
+    return payload
+
+
 # ============================================================================
 def verify_password(plain_password: str, hashed_password: str) -> bool:
     """verify password by hash"""
@@ -157,7 +166,7 @@ def generate_password() -> str:
 
 # ============================================================================
 # pylint: disable=raise-missing-from
-def init_jwt_auth(user_manager) -> tuple[Callable, Callable, Callable, Callable]:
+def init_jwt_auth(user_manager):
     """init jwt auth router + current_active_user dependency"""
     oauth2_scheme = OA2BearerOrQuery(tokenUrl="/api/auth/jwt/login", auto_error=False)
 
@@ -167,8 +176,6 @@ async def get_current_user(
         try:
             payload = decode_jwt(token, AUTH_ALLOW_AUD)
             uid: Optional[str] = payload.get("sub") or payload.get("user_id")
-            # insure not an internal token
-            assert not payload.get("internal_role")
             user = await user_manager.get_by_id(UUID(uid))
             assert user
             return user
@@ -194,17 +201,6 @@ async def shared_secret_or_superuser(
 
         return user
 
-    def get_custom_access(role: str) -> Callable[[str], str]:
-        def get_access_dep(token: str = Depends(oauth2_scheme)) -> str:
-            payload = decode_jwt(token, AUTH_ALLOW_AUD)
-            sub = payload.get("sub")
-            if not sub or payload.get("internal_role") != role:
-                raise HTTPException(status_code=401, detail="invalid_credentials")
-
-            return sub
-
-        return get_access_dep
-
     current_active_user = get_current_user
 
     auth_jwt_router = APIRouter()
@@ -284,9 +280,4 @@ async def refresh_jwt(user=Depends(current_active_user)):
         user_info = await user_manager.get_user_info_with_orgs(user)
         return get_bearer_response(user, user_info)
 
-    return (
-        auth_jwt_router,
-        current_active_user,
-        shared_secret_or_superuser,
-        get_custom_access,
-    )
+    return auth_jwt_router, current_active_user, shared_secret_or_superuser

backend/btrixcloud/colls.py

@@ -63,6 +63,8 @@
     case_insensitive_collation,
 )
 
+from .auth import get_custom_jwt_token
+
 from .crawlmanager import CrawlManager
 
 if TYPE_CHECKING:
@@ -1274,7 +1276,6 @@ def init_collections_api(
     crawl_manager: CrawlManager,
     event_webhook_ops: EventWebhookOps,
     user_dep,
-    coll_dep,
 ) -> CollectionOps:
     """init collections api"""
     # pylint: disable=invalid-name, unused-argument, too-many-arguments
@@ -1288,10 +1289,20 @@ def init_collections_api(
     org_viewer_dep = orgs.org_viewer_dep
     org_public = orgs.org_public
 
-    def coll_access_dep(coll_id: UUID, coll_access_id=Depends(coll_dep)) -> UUID:
-        if coll_access_id == str(coll_id):
-            return coll_id
+    async def coll_access_dep(
+        coll_id: UUID, token_data: dict[str, str] = Depends(get_custom_jwt_token)
+    ) -> UUID:
+        # first, check subject match collection id and type is collection
+        if token_data.get("sub_type") == "coll" and token_data.get("sub") == str(
+            coll_id
+        ):
+            # second, check that the k8s object access is scoped to exists
+            if await crawl_manager.validate_k8s_obj_exists(
+                token_data.get("scope_type", ""), token_data.get("scope", "")
+            ):
+                return coll_id
 
+        # otherwise, deny access
         raise HTTPException(status_code=403, detail="access_denied")
 
     @app.post(

backend/btrixcloud/crawlmanager.py

@@ -10,7 +10,7 @@
 
 from .utils import dt_now, date_to_str, scale_from_browser_windows
 from .k8sapi import K8sAPI, ApiException
-from .auth import create_internal_crawler_access_token
+from .auth import create_custom_jwt_token
 
 from .models import (
     StorageRef,
@@ -245,7 +245,9 @@ async def run_index_import_job(
         )
 
         if job_type in ("purge", "import"):
-            auth_bearer = create_internal_crawler_access_token(coll_id, "coll")
+            auth_bearer = create_custom_jwt_token(
+                coll_id, {"sub_type": "coll", "scope_type": "job", "scope": name}
+            )
             import_source_url = (
                 f"{BACKEND_ORIGIN}/api/orgs/{oid}/collections/{coll_id}"
                 + f"/internal/replay.json?auth_bearer={auth_bearer}"
@@ -282,6 +284,13 @@ async def run_index_import_job(
 
         return name
 
+    async def validate_k8s_obj_exists(self, obj_type: str, name: str) -> bool:
+        """return true/false if specified k8s object exists"""
+        if obj_type == "job":
+            return await self.has_job(name)
+
+        return False
+
     async def delete_dedupe_index_resources(self, oid: str, coll_id: str) -> None:
         """Delete dedupe index-related jobs and index itself"""
         await self._delete_jobs(f"role=index-import-job,oid={oid},coll={coll_id}")

backend/btrixcloud/k8sapi.py

@@ -348,6 +348,17 @@ async def unsuspend_k8s_job(self, name) -> dict:
             traceback.print_exc()
             return {"error": str(exc)}
 
+    async def has_job(self, name) -> bool:
+        """return true/false if job exists"""
+        try:
+            await self.batch_api.read_namespaced_job(
+                name=name, namespace=self.namespace
+            )
+            return True
+        # pylint: disable=bare-except
+        except:
+            return False
+
     async def print_pod_logs(self, pod_names, lines=100):
         """print pod logs"""
         for pod in pod_names:

backend/btrixcloud/main.py

@@ -172,9 +172,7 @@ def main() -> None:
 
     user_manager = init_user_manager(mdb, email, invites)
 
-    current_active_user, shared_secret_or_superuser, custom_access = init_users_api(
-        app, user_manager
-    )
+    current_active_user, shared_secret_or_superuser = init_users_api(app, user_manager)
 
     org_ops = init_orgs_api(
         app,
@@ -247,7 +245,6 @@ def main() -> None:
         crawl_manager,
         event_webhook_ops,
         current_active_user,
-        custom_access("coll"),
     )
 
     base_crawl_init = (

backend/btrixcloud/users.py

@@ -593,13 +593,11 @@ def init_user_manager(mdb, emailsender, invites):
 
 # ============================================================================
 # pylint: disable=too-many-locals, raise-missing-from
-def init_users_api(
-    app, user_manager: UserManager
-) -> tuple[Callable, Callable, Callable]:
+def init_users_api(app, user_manager: UserManager):
     """init fastapi_users"""
 
-    auth_jwt_router, current_active_user, shared_secret_or_superuser, custom_access = (
-        init_jwt_auth(user_manager)
+    auth_jwt_router, current_active_user, shared_secret_or_superuser = init_jwt_auth(
+        user_manager
     )
 
     app.include_router(
@@ -620,7 +618,7 @@ def init_users_api(
         tags=["users"],
     )
 
-    return current_active_user, shared_secret_or_superuser, custom_access
+    return current_active_user, shared_secret_or_superuser
 
 
 # ============================================================================