aixtools 0.2.4__tar.gz → 0.2.6__tar.gz

{aixtools-0.2.4 → aixtools-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.4
+Version: 0.2.6
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.2.4 → aixtools-0.2.6}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.4'
-__version_tuple__ = version_tuple = (0, 2, 4)
+__version__ = version = '0.2.6'
+__version_tuple__ = version_tuple = (0, 2, 6)
 
-__commit_id__ = commit_id = 'g3a46ce295'
+__commit_id__ = commit_id = 'g9b369b8bb'

aixtools-0.2.6/aixtools/auth/auth.py

@@ -0,0 +1,149 @@
+"""
+Module that manages OAuth2 functions for authentication
+"""
+
+import enum
+import logging
+
+import jwt
+from fastapi import HTTPException
+from jwt import ExpiredSignatureError, InvalidAudienceError, InvalidIssuerError, InvalidSignatureError, PyJWKClient
+
+from aixtools.utils import config
+
+logger = logging.getLogger(__name__)
+
+
+class AuthTokenErrorCode(str, enum.Enum):
+    """Enum for error codes returned by the AuthTokenError exception."""
+
+    TOKEN_EXPIRED = "Token expired"
+    INVALID_AUDIENCE = "Token not for expected audience"
+    INVALID_ISSUER = "Token not for expected issuer"
+    INVALID_SIGNATURE = "Token signature error"
+    INVALID_TOKEN = "Invalid token"
+    JWT_ERROR = "Generic JWT error"
+    MISSING_GROUPS_ERROR = "Missing authorized groups"
+    INVALID_TOKEN_SCOPE = "Token scope does not match configured scope"
+
+
+class AuthTokenError(Exception):
+    """Exception raised for authentication token errors."""
+
+    def __init__(self, error_code: AuthTokenErrorCode, msg: str = None):
+        self.error_code = error_code
+        error_msg = error_code.value if msg is None else msg
+        super().__init__(error_msg)
+
+    def to_http_exception(self, required_scope: str = None, realm: str = "MCP") -> HTTPException:
+        """
+        Returns an HTTPException with 401 status for all AuthTokenErrorCode,
+        including MCP JSON body and WWW-Authenticate header.
+        """
+        status_code = 401
+        www_error = (
+            "insufficient_scope" if self.error_code == AuthTokenErrorCode.INVALID_TOKEN_SCOPE else "invalid_token"
+        )
+
+        header_value = f'Bearer realm="{realm}", error="{www_error}", error_description="{self.error_code.value}"'
+        if self.error_code == AuthTokenErrorCode.INVALID_TOKEN_SCOPE and required_scope:
+            header_value += f', scope="{required_scope}"'
+
+        detail = {"error": {"code": self.error_code.name, "message": self.error_code.value}}
+        if self.error_code == AuthTokenErrorCode.INVALID_TOKEN_SCOPE and required_scope:
+            detail["error"]["required_scope"] = required_scope
+
+        return HTTPException(status_code=status_code, detail=detail, headers={"WWW-Authenticate": header_value})
+
+
+class AccessTokenVerifier:
+    """
+    Verifies Microsoft SSO JWT token against the configured Tenant ID, Audience, API ID and Issuer URL.
+    """
+
+    def __init__(self):
+        tenant_id = config.APP_TENANT_ID
+        self.api_id = config.APP_API_ID
+        self.issuer_url = f"https://sts.windows.net/{tenant_id}/"
+
+        self.authorized_groups = set(config.APP_AUTHORIZED_GROUPS.split(",")) if config.APP_AUTHORIZED_GROUPS else set()
+        if not self.authorized_groups:
+            logger.warning("No authorized groups configured")
+
+        # Azure AD endpoints
+        jwks_url = f"https://login.microsoftonline.com/{tenant_id}/discovery/v2.0/keys"
+        self.jwks_client = PyJWKClient(
+            uri=jwks_url,
+            # cache keys url response to reduce SSO server network calls,
+            # as public keys are not expected to change frequently
+            cache_jwk_set=True,
+            # cache resolved public keys
+            cache_keys=True,
+            # cache url response for 10 hours
+            lifespan=36000,
+        )
+
+        logger.info("Using JWKS: %s", jwks_url)
+
+    def verify(self, token: str) -> dict:
+        """
+        Verifies The JWT access token and returns decoded claims as a dictionary if the token is
+        valid, otherwise raises an AuthTokenError
+        """
+        try:
+            signing_key = self.jwks_client.get_signing_key_from_jwt(token)
+            logger.info("Verifying JWT token")
+            claims = jwt.decode(
+                token,
+                signing_key.key,
+                algorithms=["RS256"],
+                audience=self.api_id,
+                issuer=self.issuer_url,
+                # ensure audience verification is carried out
+                options={"verify_aud": True},
+            )
+            logger.info("Verified JWT token")
+            return claims
+
+        except ExpiredSignatureError as e:
+            raise AuthTokenError(AuthTokenErrorCode.TOKEN_EXPIRED) from e
+        except InvalidAudienceError as e:
+            raise AuthTokenError(AuthTokenErrorCode.INVALID_AUDIENCE) from e
+        except InvalidIssuerError as e:
+            raise AuthTokenError(AuthTokenErrorCode.INVALID_ISSUER) from e
+        except InvalidSignatureError as e:
+            raise AuthTokenError(AuthTokenErrorCode.INVALID_SIGNATURE) from e
+        except jwt.exceptions.PyJWTError as e:
+            raise AuthTokenError(AuthTokenErrorCode.JWT_ERROR) from e
+
+    def authorize_claims(self, claims: dict, expected_scope: str):
+        """
+        Authorize claims based on token scope, expected scope and authorized groups
+        claims: decoded JWT claims
+        expected_scope: expected scope for the token
+        Raises AuthTokenError if authorization fails.
+        """
+        logger.info("Checking JWT token claims")
+        if expected_scope:
+            token_scopes = claims.get("scp", "").split()
+            if expected_scope not in token_scopes:
+                logger.error("Expected token scope: %s, got: %s", expected_scope, token_scopes)
+                raise AuthTokenError(
+                    AuthTokenErrorCode.INVALID_TOKEN_SCOPE,
+                    f"Expected token scope: {expected_scope}, got: {token_scopes}",
+                )
+
+        if not self.authorized_groups:
+            logger.info("Authorized JWT token, no authorized groups configured")
+            return
+
+        groups = claims.get("groups", [])
+        if self.authorized_groups & set(groups):
+            logger.info("Authorized JWT token, against %s", groups)
+            return
+
+        logger.error("Could not find any group in JWT token, matching: %s", self.authorized_groups)
+        raise AuthTokenError(
+            AuthTokenErrorCode.MISSING_GROUPS_ERROR,
+            f"Could not find any group in JWT token, matching: {self.authorized_groups}",
+        )

{aixtools-0.2.4 → aixtools-0.2.6}/aixtools/utils/config.py

@@ -135,5 +135,6 @@ APP_CLIENT_ID = get_variable_env("APP_CLIENT_ID")
 # used for token audience check
 APP_API_ID = get_variable_env("APP_API_ID")
 APP_TENANT_ID = get_variable_env("APP_TENANT_ID")
-# used for token issuer check
-APP_ISSUER_URL = get_variable_env("APP_ISSUER_URL")
+
+# used for token authorization check
+APP_AUTHORIZED_GROUPS = get_variable_env("APP_AUTHORIZED_GROUPS", allow_empty=True)

aixtools-0.2.4/aixtools/auth/auth.py

@@ -1,70 +0,0 @@
-"""
-Module that manages OAuth2 functions for authentication
-"""
-
-import logging
-
-import jwt
-from jwt import ExpiredSignatureError, InvalidAudienceError, InvalidIssuerError, PyJWKClient
-
-from aixtools.utils import config
-
-logger = logging.getLogger(__name__)
-
-
-class AuthTokenError(Exception):
-    """Exception raised for authentication token errors."""
-
-
-# pylint: disable=too-few-public-methods
-class AccessTokenVerifier:
-    """
-    Verifies Microsoft SSO JWT token against the configured Tenant ID, Audience, API ID and Issuer URL.
-    """
-
-    def __init__(self):
-        tenant_id = config.APP_TENANT_ID
-        self.api_id = config.APP_API_ID
-        self.issuer_url = config.APP_ISSUER_URL
-        # Azure AD endpoints
-        jwks_url = f"https://login.microsoftonline.com/{tenant_id}/discovery/v2.0/keys"
-        self.jwks_client = PyJWKClient(
-            uri=jwks_url,
-            # cache keys url response to reduce SSO server network calls,
-            # as public keys are not expected to change frequently
-            cache_jwk_set=True,
-            # cache resolved public keys
-            cache_keys=True,
-            # cache url response for 10 hours
-            lifespan=36000,
-        )
-
-        logger.info("Using JWKS: %s", jwks_url)
-
-    def verify(self, token: str) -> dict:
-        """
-        Verifies The JWT access token and returns decoded claims as a dictionary if the token is
-        valid, otherwise raises an AuthTokenError
-        """
-        try:
-            signing_key = self.jwks_client.get_signing_key_from_jwt(token)
-
-            claims = jwt.decode(
-                token,
-                signing_key.key,
-                algorithms=["RS256"],
-                audience=self.api_id,
-                issuer=self.issuer_url,
-                # ensure audience verification is carried out
-                options={"verify_aud": True},
-            )
-            return claims
-
-        except ExpiredSignatureError as e:
-            raise AuthTokenError("Token expired") from e
-        except InvalidAudienceError as e:
-            raise AuthTokenError(f"Token not for expected audience: {e}") from e
-        except InvalidIssuerError as e:
-            raise AuthTokenError(f"Token not for expected issuer: {e}") from e
-        except jwt.exceptions.PyJWTError as e:
-            raise AuthTokenError(f"Invalid token: {e}") from e