auth4

Author: omertuc

src/app/endpoints/config.py

@@ -5,10 +5,9 @@
 
 from fastapi import APIRouter, Request
 
-from models.config import Configuration
-from configuration import configuration
 from authorization.middleware import authorize
-from models.config import Action
+from configuration import configuration
+from models.config import Action, Configuration
 from utils.endpoints import check_configuration_loaded
 
 logger = logging.getLogger(__name__)

src/app/endpoints/conversations.py

@@ -9,18 +9,18 @@
 
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
+from app.database import get_session
+from auth import get_auth_dependency
+from authorization.middleware import authorize
+from models.config import Action
+from models.database.conversations import UserConversation
 from models.responses import (
     ConversationResponse,
     ConversationDeleteResponse,
     ConversationsListResponse,
     ConversationDetails,
 )
-from models.database.conversations import UserConversation
-from auth import get_auth_dependency
-from app.database import get_session
 from utils.endpoints import check_configuration_loaded, validate_conversation_ownership
-from authorization.middleware import authorize
-from models.config import Action
 from utils.suid import check_suid
 
 logger = logging.getLogger("app.endpoints.handlers")

src/app/endpoints/feedback.py

@@ -10,16 +10,16 @@
 from auth import get_auth_dependency
 from auth.interface import AuthTuple
 from authorization.middleware import authorize
-from models.config import Action
 from configuration import configuration
+from models.config import Action
+from models.requests import FeedbackRequest
 from models.responses import (
     ErrorResponse,
     FeedbackResponse,
     StatusResponse,
     UnauthorizedResponse,
     ForbiddenResponse,
 )
-from models.requests import FeedbackRequest
 from utils.suid import get_suid
 
 logger = logging.getLogger(__name__)

src/app/endpoints/models.py

@@ -3,9 +3,9 @@
 import logging
 from typing import Any
 
+from fastapi import APIRouter, HTTPException, Request, status
 from fastapi.params import Depends
 from llama_stack_client import APIConnectionError
-from fastapi import APIRouter, HTTPException, Request, status
 
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
@@ -67,7 +67,6 @@ async def models_endpoint_handler(
     Returns:
         ModelsResponse: An object containing the list of available models.
     """
-
     # Used only by the middleware
     _ = auth
 

src/app/endpoints/query.py

@@ -24,18 +24,18 @@
 from configuration import configuration
 from app.database import get_session
 import metrics
+import constants
+from authorization.middleware import authorize
+from models.config import Action
 from models.database.conversations import UserConversation
-from models.responses import QueryResponse, UnauthorizedResponse, ForbiddenResponse
 from models.requests import QueryRequest, Attachment
-import constants
+from models.responses import QueryResponse, UnauthorizedResponse, ForbiddenResponse
 from utils.endpoints import (
     check_configuration_loaded,
     get_agent,
     get_system_prompt,
     validate_conversation_ownership,
 )
-from authorization.middleware import authorize
-from models.config import Action
 from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups
 from utils.suid import get_suid
 

src/app/endpoints/streaming_query.py

@@ -20,10 +20,10 @@
 from auth import get_auth_dependency
 from auth.interface import AuthTuple
 from authorization.middleware import authorize
-from models.config import Action
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
 import metrics
+from models.config import Action
 from models.requests import QueryRequest
 from models.database.conversations import UserConversation
 from utils.endpoints import check_configuration_loaded, get_agent, get_system_prompt

src/authorization/engine.py

@@ -28,7 +28,7 @@ class NoopAccessResolver(AccessResolver):  # pylint: disable=too-few-public-meth
     """No-op access resolver that does not perform any access checks."""
 
     async def check_access(self, action: Action, user_roles: UserRoles) -> bool:
-        """Always return True, indicating access is granted."""
+        """Return True always, indicating access is granted."""
         _ = action  # Unused
         _ = user_roles  # Unused
         return True
@@ -68,22 +68,19 @@ def __init__(self, role_rules: list[JwtRoleRule]):
 
     async def resolve_roles(self, auth: AuthTuple) -> UserRoles:
         """Extract roles from JWT claims using configured rules."""
-
         jwt_claims = self._get_claims(auth)
         return frozenset(
             role
             for rule in self.role_rules
-            for role in self._evaluate_role_rules(rule, jwt_claims)
+            for role in self.evaluate_role_rules(rule, jwt_claims)
         )
 
     @staticmethod
-    def _evaluate_role_rules(
-        rule: JwtRoleRule, jwt_claims: dict[str, Any]
-    ) -> UserRoles:
-        """Get roles from a JWT role rule if it matches the claims"""
+    def evaluate_role_rules(rule: JwtRoleRule, jwt_claims: dict[str, Any]) -> UserRoles:
+        """Get roles from a JWT role rule if it matches the claims."""
         return (
             frozenset(rule.roles)
-            if __class__._evaluate_operator(
+            if JwtRolesResolver._evaluate_operator(
                 rule.negate,
                 [match.value for match in parse(rule.jsonpath).find(jwt_claims)],
                 rule.operator,
@@ -147,6 +144,7 @@ def __init__(self, access_rules: list[AccessRule]):
             self._access_lookup[rule.role].update(rule.actions)
 
     async def check_access(self, action: Action, user_roles: UserRoles) -> bool:
+        """Check if the user has access to the specified action based on their roles."""
         if action != Action.ADMIN and self.check_access(action.ADMIN, user_roles):
             # Recurse to check if the roles allow the user to perform the admin action,
             # if they do, then we allow any action

src/authorization/middleware.py

@@ -83,11 +83,11 @@ async def _perform_authorization_check(action: Action, kwargs: dict[str, Any]) -
 
 
 def authorize(action: Action) -> Callable:
-    """Decorator to check authorization for an endpoint (async version)."""
+    """Check authorization for an endpoint (async version)."""
 
     def decorator(func: Callable) -> Callable:
         @wraps(func)
-        async def wrapper(*args, **kwargs):
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
             await _perform_authorization_check(action, kwargs)
             return await func(*args, **kwargs)
 

src/authorization/models.py

@@ -0,0 +1 @@
+"""Authorization models and data structures."""

src/models/config.py

@@ -215,9 +215,7 @@ class JwtRoleRule(BaseModel):
 
     @model_validator(mode="after")
     def check_jsonpath(self) -> Self:
-        """
-        Verify that the JSONPath expression is valid.
-        """
+        """Verify that the JSONPath expression is valid."""
         try:
             jsonpath_ng.parse(self.jsonpath)
             return self
@@ -228,9 +226,7 @@ def check_jsonpath(self) -> Self:
 
     @model_validator(mode="after")
     def check_roles(self) -> Self:
-        """
-        Ensure that at least one role is specified.
-        """
+        """Ensure that at least one role is specified."""
         if not self.roles:
             raise ValueError("At least one role must be specified in the rule")