Added Token Validations

.checkov.baseline

@@ -192,6 +192,13 @@
                         "CKV_AWS_115"
                     ]
                 },
+                {
+                    "resource": "AWS::Lambda::Function.CustomAuthorizerFunctiondevB38B5CCB",
+                    "check_ids": [
+                        "CKV_AWS_115",
+                        "CKV_AWS_116"
+                    ]
+                },
                 {
                     "resource": "AWS::Lambda::Function.ElasticSearchProxyHandlerDBDE7574",
                     "check_ids": [
@@ -210,6 +217,12 @@
                         "CKV_AWS_158"
                     ]
                 },
+                {
+                    "resource": "AWS::Logs::LogGroup.customauthorizerloggroup8F3B5B9D",
+                    "check_ids": [
+                        "CKV_AWS_158"
+                    ]
+                },
                 {
                     "resource": "AWS::Logs::LogGroup.dataalldevapigateway2625FE76",
                     "check_ids": [

backend/dataall/base/utils/api_handler_utils.py

@@ -24,13 +24,22 @@
 ]
 ENGINE = get_engine(envname=ENVNAME)
 ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS', '*')
+AWS_REGION = os.getenv('AWS_REGION')
 
 
 def redact_creds(event):
-    if 'headers' in event and 'Authorization' in event['headers']:
+    if event.get('headers', {}).get('Authorization'):
         event['headers']['Authorization'] = 'XXXXXXXXXXXX'
-    if 'multiValueHeaders' in event and 'Authorization' in event['multiValueHeaders']:
+
+    if event.get('multiValueHeaders', {}).get('Authorization'):
         event['multiValueHeaders']['Authorization'] = 'XXXXXXXXXXXX'
+
+    if event.get('multiValueHeaders', {}).get('accesskeyid'):
+        event['multiValueHeaders']['accesskeyid'] = 'XXXXXXXXXXXX'
+
+    if event.get('headers', {}).get('accesskeyid'):
+        event['headers']['accesskeyid'] = 'XXXXXXXXXXXX'
+
     return event
 
 
@@ -115,7 +124,7 @@ def check_reauth(query, auth_time, username):
     # Determine if there are any Operations that Require ReAuth From SSM Parameter
     try:
         reauth_apis = ParameterStoreManager.get_parameter_value(
-            region=os.getenv('AWS_REGION', 'eu-west-1'), parameter_path=f'/dataall/{ENVNAME}/reauth/apis'
+            region=AWS_REGION, parameter_path=f'/dataall/{ENVNAME}/reauth/apis'
         ).split(',')
     except Exception:
         log.info('No ReAuth APIs Found in SSM')

deploy/custom_resources/custom_authorizer/auth_services.py

@@ -28,7 +28,7 @@ def generate_policy(verified_claims: dict, effect, incoming_resource_str: str):
 
         for claim_name, claim_value in verified_claims.items():
             if isinstance(claim_value, list):
-                verified_claims.update({claim_name: json.dumps(claim_value)})
+                verified_claims.update({claim_name: ','.join(claim_value)})
 
         context = {**verified_claims}
 

deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import json
 
 from auth_services import AuthServices
 from jwt_services import JWTServices
@@ -19,18 +20,25 @@
 
 def lambda_handler(incoming_event, context):
     # Get the Token which is sent in the Authorization Header
+    logger.debug(incoming_event)
     auth_token = incoming_event['headers']['Authorization']
-    if not auth_token:
-        raise Exception('Unauthorized . Token not found')
+    access_token = incoming_event['headers']['accesskeyid']
+    if not auth_token or not access_token:
+        raise Exception('Unauthorized. Missing identity or access JWT')
 
+    # Validate User is Active with Proper Access Token
+    JWTServices.validate_access_token(access_token)
+
+    # Validate JWT
     verified_claims = JWTServices.validate_jwt_token(auth_token)
-    logger.debug(verified_claims)
     if not verified_claims:
         raise Exception('Unauthorized. Token is not valid')
+    logger.debug(verified_claims)
 
+    # Generate Allow Policy w/ Context
     effect = 'Allow'
     policy = AuthServices.generate_policy(verified_claims, effect, incoming_event['methodArn'])
-    logger.debug('Generated policy is ', policy)
+    logger.debug('Generated policy is ', json.dumps(policy))
     return policy
 
 
@@ -39,12 +47,13 @@ def lambda_handler(incoming_event, context):
 # AWS Lambda and any other local environments
 if __name__ == '__main__':
     # for testing locally you can enter the JWT ID Token here
-    token = ''
+    id_token = ''
+    access_token = ''
     account_id = ''
     api_gw_id = ''
     event = {
+        'headers': {'Authorization': id_token, 'accesskeyid': access_token},
         'type': 'TOKEN',
-        'Authorization': token,
         'methodArn': f'arn:aws:execute-api:us-east-1:{account_id}:{api_gw_id}/prod/POST/graphql/api',
     }
     lambda_handler(event, None)

deploy/custom_resources/custom_authorizer/jwt_services.py

@@ -1,46 +1,13 @@
 import os
 
 import requests
-from jose import jwk
-from jose.jwt import get_unverified_header, decode, ExpiredSignatureError, JWTError
+import jwt
+
 import logging
 
 logger = logging.getLogger()
 logger.setLevel(os.environ.get('LOG_LEVEL', 'INFO'))
 
-# Configs required to fetch public keys from JWKS
-ISSUER_CONFIGS = {
-    f'{os.environ.get("custom_auth_url")}': {
-        'jwks_uri': f'{os.environ.get("custom_auth_jwks_url")}',
-        'allowed_audiences': f'{os.environ.get("custom_auth_client")}',
-    },
-}
-
-issuer_keys = {}
-
-
-# instead of re-downloading the public keys every time
-# we download them only on cold start
-# https://aws.amazon.com/blogs/compute/container-reuse-in-lambda/
-def fetch_public_keys():
-    try:
-        for issuer, issuer_config in ISSUER_CONFIGS.items():
-            jwks_response = requests.get(issuer_config['jwks_uri'])
-            jwks_response.raise_for_status()
-            jwks: dict = jwks_response.json()
-            for key in jwks['keys']:
-                value = {
-                    'issuer': issuer,
-                    'audience': issuer_config['allowed_audiences'],
-                    'jwk': jwk.construct(key),
-                    'public_key': jwk.construct(key).public_key(),
-                }
-                issuer_keys.update({key['kid']: value})
-    except Exception as e:
-        raise Exception(f'Unable to fetch public keys due to {str(e)}')
-
-
-fetch_public_keys()
 
 # Options to validate the JWT token
 # Only modification from default is to turn off verify_at_hash as we don't provide the access token for this validation
@@ -53,49 +20,57 @@ def fetch_public_keys():
     'verify_iss': True,
     'verify_sub': True,
     'verify_jti': True,
-    'verify_at_hash': False,
-    'require_aud': True,
-    'require_iat': True,
-    'require_exp': True,
-    'require_nbf': False,
-    'require_iss': True,
-    'require_sub': True,
-    'require_jti': True,
-    'require_at_hash': False,
-    'leeway': 0,
+    'verify_at_hash': True,
+    'require': ['aud', 'iat', 'exp', 'iss', 'sub', 'jti'],  # "nbf", "at_hash"
 }
 
 
 class JWTServices:
+    @staticmethod
+    def _fetch_openid_url(key):
+        response = requests.get(
+            os.path.join(os.environ.get('custom_auth_url', ''), '.well-known', 'openid-configuration')
+        )
+        response.raise_for_status()
+        return response.json().get(key)
+
     @staticmethod
     def validate_jwt_token(jwt_token):
         try:
-            # Decode and verify the JWT token
-            header = get_unverified_header(jwt_token)
-            kid = header['kid']
-            if kid not in issuer_keys:
-                logger.info('Public key not found in provided set of keys')
-                # Retry Fetching the public certificates again in case rotation occurs and lambda has cached the publicKeys
-                fetch_public_keys()
-                if kid not in issuer_keys:
-                    raise Exception('Unauthorized')
-            public_key = issuer_keys.get(kid)
-            payload = decode(
+            # get JWK URI from OpenId Configuration
+            jwks_url = JWTServices._fetch_openid_url('jwks_uri')
+
+            # Init pyJWT.JWKClient with JWK URI
+            jwks_client = jwt.PyJWKClient(jwks_url)
+
+            # get signing_key from JWT
+            signing_key = jwks_client.get_signing_key_from_jwt(jwt_token)
+
+            # Decode and Verify JWT
+            payload = jwt.decode(
                 jwt_token,
-                public_key.get('jwk'),
+                signing_key.key,
                 algorithms=['RS256', 'HS256'],
-                issuer=public_key.get('issuer'),
-                audience=public_key.get('audience'),
+                issuer=os.environ.get('custom_auth_url'),
+                audience=os.environ.get('custom_auth_client'),
+                leeway=0,
                 options=jwt_options,
             )
-
             return payload
-        except ExpiredSignatureError:
+        except jwt.exceptions.ExpiredSignatureError:
             logger.error('JWT token has expired.')
             return None
-        except JWTError as e:
+        except jwt.exceptions.PyJWTError as e:
             logger.error(f'JWT token validation failed: {str(e)}')
             return None
         except Exception as e:
             logger.error(f'Failed to validate token - {str(e)}')
             return None
+
+    @staticmethod
+    def validate_access_token(access_token):
+        # get JWK UserInfo URI from OpenId Configuration
+        user_info_url = JWTServices._fetch_openid_url('userinfo_endpoint')
+        r = requests.get(user_info_url, headers={'Authorization': access_token})
+        r.raise_for_status()
+        logger.debug(r.json())

deploy/custom_resources/custom_authorizer/requirements.txt

@@ -1,10 +1,9 @@
 certifi==2024.7.4
 charset-normalizer==3.1.0
-ecdsa==0.18.0
 idna==3.7
 pyasn1==0.5.0
-python-jose==3.3.0
 requests==2.32.2
 rsa==4.9
 six==1.16.0
-urllib3==1.26.19
+urllib3==1.26.19
+pyjwt==2.9.0

deploy/stacks/backend_stack.py

@@ -195,6 +195,7 @@ def __init__(
             apig_vpce=apig_vpce,
             prod_sizing=prod_sizing,
             user_pool=cognito_stack.user_pool if custom_auth is None else None,
+            user_pool_client=cognito_stack.client if custom_auth is None else None,
             pivot_role_name=self.pivot_role_name,
             reauth_ttl=reauth_config.get('ttl', 5) if reauth_config else 5,
             email_notification_sender_email_id=email_sender,