HandlerToken class created to manage token related operations.

Author: tmikula-dev

src/event_gate_lambda.py

@@ -15,23 +15,19 @@
 #
 
 """Event Gate Lambda function implementation."""
-import base64
 import json
 import logging
 import os
 import sys
-from typing import Any, Dict, cast
+from typing import Any, Dict
 
 import boto3
 import jwt
-import requests
 import urllib3
-from cryptography.exceptions import UnsupportedAlgorithm
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
 from jsonschema import validate
 from jsonschema.exceptions import ValidationError
 
+from src.handlers.handler_token import HandlerToken
 from src.writers import writer_eventbridge, writer_kafka, writer_postgres
 from src.utils.conf_path import CONF_DIR, INVALID_CONF_ENV
 
@@ -65,51 +61,28 @@
 logger.debug("Loaded TOPICS")
 
 with open(os.path.join(_CONF_DIR, "config.json"), "r", encoding="utf-8") as file:
-    CONFIG = json.load(file)
+    config = json.load(file)
 logger.debug("Loaded main CONFIG")
 
 aws_s3 = boto3.Session().resource("s3", verify=False)  # nosec Boto verify disabled intentionally
 logger.debug("Initialized AWS S3 Client")
 
-if CONFIG["access_config"].startswith("s3://"):
-    name_parts = CONFIG["access_config"].split("/")
+if config["access_config"].startswith("s3://"):
+    name_parts = config["access_config"].split("/")
     BUCKET_NAME = name_parts[2]
     BUCKET_OBJECT_KEY = "/".join(name_parts[3:])
     ACCESS = json.loads(aws_s3.Bucket(BUCKET_NAME).Object(BUCKET_OBJECT_KEY).get()["Body"].read().decode("utf-8"))
 else:
-    with open(CONFIG["access_config"], "r", encoding="utf-8") as file:
+    with open(config["access_config"], "r", encoding="utf-8") as file:
         ACCESS = json.load(file)
 logger.debug("Loaded ACCESS definitions")
 
-# Initialize token public keys
-TOKEN_PROVIDER_URL = CONFIG.get("token_provider_url")
-TOKEN_PUBLIC_KEYS_URL = CONFIG.get("token_public_keys_url") or CONFIG.get("token_public_key_url")
-
-try:
-    response_json = requests.get(TOKEN_PUBLIC_KEYS_URL, verify=False, timeout=5).json()
-    raw_keys: list[str] = []
-    if isinstance(response_json, dict):
-        if "keys" in response_json and isinstance(response_json["keys"], list):
-            for item in response_json["keys"]:
-                if "key" in item:
-                    raw_keys.append(item["key"].strip())
-        elif "key" in response_json:
-            raw_keys.append(response_json["key"].strip())
-
-    if not raw_keys:
-        raise KeyError(f"No public keys found in {TOKEN_PUBLIC_KEYS_URL} endpoint response")
-
-    TOKEN_PUBLIC_KEYS: list[RSAPublicKey] = [
-        cast(RSAPublicKey, serialization.load_der_public_key(base64.b64decode(raw_key))) for raw_key in raw_keys
-    ]
-    logger.debug("Loaded %d TOKEN_PUBLIC_KEYS", len(TOKEN_PUBLIC_KEYS))
-except (requests.RequestException, ValueError, KeyError, UnsupportedAlgorithm) as exc:
-    logger.exception("Failed to fetch or deserialize token public key from %s", TOKEN_PUBLIC_KEYS_URL)
-    raise RuntimeError("Token public key initialization failed") from exc
+# Initialize token handler and load token public keys
+handler_token = HandlerToken(config).load_public_keys()
 
 # Initialize EventGate writers
-writer_eventbridge.init(logger, CONFIG)
-writer_kafka.init(logger, CONFIG)
+writer_eventbridge.init(logger, config)
+writer_kafka.init(logger, config)
 writer_postgres.init(logger)
 
 
@@ -141,12 +114,6 @@ def get_api() -> Dict[str, Any]:
     return {"statusCode": 200, "body": API}
 
 
-def get_token() -> Dict[str, Any]:
-    """Return 303 redirect to token provider endpoint."""
-    logger.debug("Handling GET Token")
-    return {"statusCode": 303, "headers": {"Location": TOKEN_PROVIDER_URL}}
-
-
 def get_topics() -> Dict[str, Any]:
     """Return list of available topic names."""
     logger.debug("Handling GET Topics")
@@ -180,7 +147,7 @@ def post_topic_message(topic_name: str, topic_message: Dict[str, Any], token_enc
     """
     logger.debug("Handling POST %s", topic_name)
     try:
-        token = decode_jwt_all(token_encoded)
+        token: Dict[str, Any] = handler_token.decode_jwt(token_encoded)
     except jwt.PyJWTError:  # type: ignore[attr-defined]
         return _error_response(401, "auth", "Invalid or missing token")
 
@@ -222,55 +189,6 @@ def post_topic_message(topic_name: str, topic_message: Dict[str, Any], token_enc
     }
 
 
-def decode_jwt_all(token_encoded: str) -> Dict[str, Any]:
-    """Decode JWT using any of the loaded public keys.
-
-    Args:
-        token_encoded: Encoded bearer JWT token string.
-    """
-    for public_key in TOKEN_PUBLIC_KEYS:
-        try:
-            return jwt.decode(token_encoded, public_key, algorithms=["RS256"])
-        except jwt.PyJWTError:
-            continue
-    raise jwt.PyJWTError("Verification failed for all public keys")
-
-
-def extract_token(event_headers: Dict[str, str]) -> str:
-    """Extract bearer token from headers (case-insensitive).
-
-    Supports:
-      - Custom 'bearer' header (any casing) whose value is the raw token
-      - Standard 'Authorization: Bearer <token>' header (case-insensitive scheme & key)
-    Returns empty string if token not found or malformed.
-    """
-    if not event_headers:
-        return ""
-
-    # Normalize keys to lowercase for case-insensitive lookup
-    lowered = {str(k).lower(): v for k, v in event_headers.items()}
-
-    # Direct bearer header (raw token)
-    if "bearer" in lowered and isinstance(lowered["bearer"], str):
-        token_candidate = lowered["bearer"].strip()
-        if token_candidate:
-            return token_candidate
-
-    # Authorization header with Bearer scheme
-    auth_val = lowered.get("authorization", "")
-    if not isinstance(auth_val, str):  # defensive
-        return ""
-    auth_val = auth_val.strip()
-    if not auth_val:
-        return ""
-
-    # Case-insensitive match for 'Bearer ' prefix
-    if not auth_val.lower().startswith("bearer "):
-        return ""
-    token_part = auth_val[7:].strip()  # len('Bearer ')==7
-    return token_part
-
-
 def lambda_handler(event: Dict[str, Any], context: Any):  # pylint: disable=unused-argument,too-many-return-statements
     """AWS Lambda entry point.
 
@@ -281,7 +199,7 @@ def lambda_handler(event: Dict[str, Any], context: Any):  # pylint: disable=unus
         if resource == "/api":
             return get_api()
         if resource == "/token":
-            return get_token()
+            return handler_token.get_token()
         if resource == "/topics":
             return get_topics()
         if resource == "/topics/{topic_name}":
@@ -292,7 +210,7 @@ def lambda_handler(event: Dict[str, Any], context: Any):  # pylint: disable=unus
                 return post_topic_message(
                     event["pathParameters"]["topic_name"].lower(),
                     json.loads(event["body"]),
-                    extract_token(event.get("headers", {})),
+                    handler_token.extract_token(event.get("headers", {})),
                 )
         if resource == "/terminate":
             sys.exit("TERMINATING")  # pragma: no cover - deliberate termination path

src/handlers/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright 2025 ABSA Group Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#

src/handlers/handler_token.py

@@ -0,0 +1,142 @@
+#
+# Copyright 2025 ABSA Group Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+This module provides the HandlerToken class for managing the token related operations.
+"""
+
+import base64
+import logging
+import os
+from typing import Dict, Any, cast
+
+import jwt
+import requests
+from cryptography.exceptions import UnsupportedAlgorithm
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
+
+from src.utils.constants import TOKEN_PROVIDER_URL, TOKEN_PUBLIC_KEYS_URL, TOKEN_PUBLIC_KEY_URL
+
+logger = logging.getLogger(__name__)
+log_level = os.environ.get("LOG_LEVEL", "INFO")
+logger.setLevel(log_level)
+
+
+class HandlerToken:
+    """
+    HandlerToken manages token provider URL and public keys for JWT verification.
+    """
+
+    def __init__(self, config):
+        self.provider_url: str = config.get(TOKEN_PROVIDER_URL, "")
+        self.public_keys_url: str = config.get(TOKEN_PUBLIC_KEYS_URL) or config.get(TOKEN_PUBLIC_KEY_URL)
+        self.public_keys: list[RSAPublicKey] = []
+
+    def load_public_keys(self) -> "HandlerToken":
+        """
+        Load token public keys from the configured URL.
+        Returns:
+            HandlerToken: The current instance with loaded public keys.
+        Raises:
+            RuntimeError: If fetching or deserializing the public keys fails.
+        """
+        logger.debug("Loading token public keys from %s", self.public_keys_url)
+
+        try:
+            response_json = requests.get(self.public_keys_url, verify=False, timeout=5).json()
+            raw_keys: list[str] = []
+
+            if isinstance(response_json, dict):
+                if "keys" in response_json and isinstance(response_json["keys"], list):
+                    for item in response_json["keys"]:
+                        if "key" in item:
+                            raw_keys.append(item["key"].strip())
+                elif "key" in response_json:
+                    raw_keys.append(response_json["key"].strip())
+
+            if not raw_keys:
+                raise KeyError(f"No public keys found in {self.public_keys_url} endpoint response")
+
+            self.public_keys = [
+                cast(RSAPublicKey, serialization.load_der_public_key(base64.b64decode(raw_key))) for raw_key in raw_keys
+            ]
+            logger.debug("Loaded %d token public keys", len(self.public_keys))
+
+            return self
+        except (requests.RequestException, ValueError, KeyError, UnsupportedAlgorithm) as exc:
+            logger.exception("Failed to fetch or deserialize token public key from %s", self.public_keys_url)
+            raise RuntimeError("Token public key initialization failed") from exc
+
+    def decode_jwt(self, token_encoded: str) -> Dict[str, Any]:
+        """
+        Decode and verify a JWT using the loaded public keys.
+        Args:
+            token_encoded (str): The encoded JWT token.
+        Returns:
+            Dict[str, Any]: The decoded JWT payload.
+        Raises:
+            jwt.PyJWTError: If verification fails for all public keys.
+        """
+        logger.debug("Decoding JWT")
+        for public_key in self.public_keys:
+            try:
+                return jwt.decode(token_encoded, public_key, algorithms=["RS256"])
+            except jwt.PyJWTError:
+                continue
+        raise jwt.PyJWTError("Verification failed for all public keys")
+
+    def get_token(self) -> Dict[str, Any]:
+        """
+        Returns: A 303 redirect response to the token provider URL.
+        """
+        logger.debug("Handling GET Token")
+        return {"statusCode": 303, "headers": {"Location": self.provider_url}}
+
+    @staticmethod
+    def extract_token(event_headers: Dict[str, str]) -> str:
+        """
+        Extracts the bearer (custom/standard) token from event headers.
+        Args:
+            event_headers (Dict[str, str]): The event headers.
+        Returns:
+            str: The extracted bearer token, or an empty string if not found.
+        """
+        if not event_headers:
+            return ""
+
+        # Normalize keys to lowercase for case-insensitive lookup
+        lowered = {str(k).lower(): v for k, v in event_headers.items()}
+
+        # Direct bearer header (raw token)
+        if "bearer" in lowered and isinstance(lowered["bearer"], str):
+            token_candidate = lowered["bearer"].strip()
+            if token_candidate:
+                return token_candidate
+
+        # Authorization header with Bearer scheme
+        auth_val = lowered.get("authorization", "")
+        if not isinstance(auth_val, str):  # defensive
+            return ""
+        auth_val = auth_val.strip()
+        if not auth_val:
+            return ""
+
+        # Case-insensitive match for 'Bearer ' prefix
+        if not auth_val.lower().startswith("bearer "):
+            return ""
+        token_part = auth_val[7:].strip()  # len('Bearer ')==7
+        return token_part

src/utils/constants.py

@@ -0,0 +1,24 @@
+#
+# Copyright 2025 ABSA Group Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+This module contains all constants and enums used across the project.
+"""
+
+# Token related constants
+TOKEN_PROVIDER_URL = "token_provider_url"
+TOKEN_PUBLIC_KEY_URL = "token_public_key_url"
+TOKEN_PUBLIC_KEYS_URL = "token_public_keys_url"