Added Token Validations

backend/api_handler.py

@@ -3,7 +3,6 @@
 import os
 from argparse import Namespace
 from time import perf_counter
-
 from ariadne import (
     gql,
     graphql_sync,

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

@@ -24,17 +24,19 @@ def generate_policy(verified_claims: dict, effect, incoming_resource_str: str):
         principal_id = verified_claims['sub']
 
         # Attach a claim called 'email'. This is needed by Api Handler
-        verified_claims['email'] = verified_claims[os.getenv('email')]
+        verified_claims['email'] = verified_claims[os.getenv('email', 'email')]
 
         for claim_name, claim_value in verified_claims.items():
             if isinstance(claim_value, list):
+                if claim_name == 'cognito:groups':
+                    verified_claims.update({claim_name: ','.join(claim_value)})
                 verified_claims.update({claim_name: json.dumps(claim_value)})
 
         context = {**verified_claims}
 
         context.update(
             {
-                'user_id': verified_claims[os.getenv('user_id')],
+                'user_id': verified_claims[os.getenv('user_id', 'email')],
                 'custom_authorizer': 'true',
             }
         )

deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py

@@ -28,6 +28,10 @@ def lambda_handler(incoming_event, context):
     if not verified_claims:
         raise Exception('Unauthorized. Token is not valid')
 
+    if os.getenv('provider') == 'Cognito':
+        access_token = incoming_event['headers']['AccessKeyId']
+        JWTServices.validate_access_token(access_token)
+
     effect = 'Allow'
     policy = AuthServices.generate_policy(verified_claims, effect, incoming_event['methodArn'])
     logger.debug('Generated policy is ', policy)
@@ -39,12 +43,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

@@ -99,3 +99,9 @@ def validate_jwt_token(jwt_token):
         except Exception as e:
             logger.error(f'Failed to validate token - {str(e)}')
             return None
+
+    @staticmethod
+    def validate_access_token(access_token):
+        user_info_url = os.getenv('user_info_url', '')
+        r = requests.get(user_info_url, headers={'Authorization': access_token})
+        r.raise_for_status()

deploy/stacks/backend_stack.py

@@ -35,6 +35,7 @@ def __init__(
         id,
         envname: str = 'dev',
         resource_prefix='dataall',
+        tooling_region=None,
         tooling_account_id=None,
         ecr_repository=None,
         image_tag=None,
@@ -195,6 +196,8 @@ 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,
+            user_pool_domain=cognito_stack.domain 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,