From 155db7af78232ab295830a1beb7f670340a3377d Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 11:07:10 -0500
Subject: [PATCH 01/32] Add checks to ensure access tokenn validity

---
 backend/api_handler.py                        | 18 +++++++++++++--
 .../dataall/base/utils/api_handler_utils.py   | 23 ++++++++++++++++++-
 frontend/src/authentication/hooks/useToken.js |  7 +++++-
 frontend/src/modules/Catalog/views/Catalog.js |  6 ++---
 .../contexts/RequestContext.js                |  2 +-
 frontend/src/services/hooks/useClient.js      |  9 ++++----
 6 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/backend/api_handler.py b/backend/api_handler.py
index 1e1645e9e..0799dd10b 100644
--- a/backend/api_handler.py
+++ b/backend/api_handler.py
@@ -3,7 +3,7 @@
 import os
 from argparse import Namespace
 from time import perf_counter
-
+import requests
 from ariadne import (
     gql,
     graphql_sync,
@@ -16,6 +16,7 @@
     check_reauth,
     validate_and_block_if_maintenance_window,
     redact_creds,
+    get_user_info,
 )
 from dataall.core.tasks.service_handlers import Worker
 from dataall.base.aws.sqs import SqsQueue
@@ -84,7 +85,20 @@ def handler(event, context):
 
         Return doc: https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html
     """
-
+    try:
+        r = get_user_info(event=event)
+    except Exception as e:
+        log.error(f'Error occured while getting user info due to - {e}')
+        return {
+            'statusCode': 400,
+            'headers': {
+                'content-type': 'application/json',
+                'Access-Control-Allow-Origin': ALLOWED_ORIGINS,
+                'Access-Control-Allow-Headers': '*',
+                'Access-Control-Allow-Methods': '*',
+            },
+            'body': str(e),
+        }
     event = redact_creds(event)
     log.info('Lambda Event %s', event)
     log.debug('Env name %s', ENVNAME)
diff --git a/backend/dataall/base/utils/api_handler_utils.py b/backend/dataall/base/utils/api_handler_utils.py
index df1a998ce..751a8f9bc 100644
--- a/backend/dataall/base/utils/api_handler_utils.py
+++ b/backend/dataall/base/utils/api_handler_utils.py
@@ -2,6 +2,7 @@
 import json
 import os
 import logging
+import requests
 
 from graphql import parse, utilities, OperationType, GraphQLSyntaxError
 from dataall.base.aws.parameter_store import ParameterStoreManager
@@ -24,6 +25,26 @@
 ]
 ENGINE = get_engine(envname=ENVNAME)
 ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS', '*')
+AWS_REGION = os.getenv('AWS_REGION')
+
+
+def get_user_info(event):
+    idp_domain = ParameterStoreManager.get_parameter_value(
+        region=AWS_REGION, parameter_path=f'/dataall/{ENVNAME}/cognito/domain'
+    )
+
+    accessToken = None
+    if event.get('headers', {}).get('accesskeyid'):
+        accessToken = event['headers']['accesskeyid']
+    else:
+        raise Exception('AccessKeyId not found in headers')
+
+    url = f'https://{idp_domain}.auth.{AWS_REGION}.amazoncognito.com/oauth2/userInfo'
+    r = requests.get(url, headers={'Authorization': accessToken})
+
+    r.raise_for_status()
+    # if r.status_code != 200:
+    #     raise Exception(r.json().get('error_description'))
 
 
 def redact_creds(event):
@@ -115,7 +136,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')
diff --git a/frontend/src/authentication/hooks/useToken.js b/frontend/src/authentication/hooks/useToken.js
index 08cfd5668..3ec75fd1f 100644
--- a/frontend/src/authentication/hooks/useToken.js
+++ b/frontend/src/authentication/hooks/useToken.js
@@ -7,12 +7,14 @@ export const useToken = () => {
   const dispatch = useDispatch();
   const auth = useAuth();
   const [token, setToken] = useState(null);
+  const [accessToken, setAccessToken] = useState(null);
   const fetchAuthToken = async () => {
     if (
       !process.env.REACT_APP_COGNITO_USER_POOL_ID &&
       process.env.REACT_APP_GRAPHQL_API.includes('localhost')
     ) {
       setToken('localToken');
+      setAccessToken('localAccessToken');
     } else {
       try {
         if (process.env.REACT_APP_CUSTOM_AUTH) {
@@ -22,6 +24,7 @@ export const useToken = () => {
             }
             const t = auth.user.id_token;
             setToken(t);
+            setAccessToken(t);
           } catch (error) {
             if (!auth) throw Error('User Token Not Found !');
           }
@@ -29,6 +32,8 @@ export const useToken = () => {
           const session = await Auth.currentSession();
           const t = await session.getIdToken().getJwtToken();
           setToken(t);
+          const at = await session.getAccessToken().getJwtToken();
+          setAccessToken(at);
         }
       } catch (error) {
         auth.dispatch({
@@ -45,5 +50,5 @@ export const useToken = () => {
       );
     }
   });
-  return token;
+  return { token, accessToken };
 };
diff --git a/frontend/src/modules/Catalog/views/Catalog.js b/frontend/src/modules/Catalog/views/Catalog.js
index c351c5da7..4a7bf733c 100644
--- a/frontend/src/modules/Catalog/views/Catalog.js
+++ b/frontend/src/modules/Catalog/views/Catalog.js
@@ -167,7 +167,7 @@ GlossaryFilter.propTypes = {
 };
 
 const Catalog = () => {
-  const token = useToken();
+  const { token, accessToken } = useToken();
   const { settings } = useSettings();
   const theme = useTheme();
   const classes = useStyles();
@@ -235,7 +235,7 @@ const Catalog = () => {
       credentials: { token },
       headers: {
         Authorization: token,
-        AccessKeyId: 'None',
+        AccessKeyId: accessToken ? `Bearer ${accessToken}` : '',
         SecretKey: 'None'
       }
     };
@@ -247,7 +247,7 @@ const Catalog = () => {
         : classes.darkListSearch
     );
   }, [settings.theme, classes]);
-  if (!token) {
+  if (!token || !accessToken) {
     return <CircularProgress />;
   }
 
diff --git a/frontend/src/reauthentication/contexts/RequestContext.js b/frontend/src/reauthentication/contexts/RequestContext.js
index 6986d5db3..ab6b55fb6 100644
--- a/frontend/src/reauthentication/contexts/RequestContext.js
+++ b/frontend/src/reauthentication/contexts/RequestContext.js
@@ -53,7 +53,7 @@ export const RequestContextProvider = (props) => {
   const [requestInfo, setRequestInfo] = useState(null);
   const navigate = useNavigate();
   const client = useClient();
-  const token = useToken();
+  const { token } = useToken();
   const auth = useAuth();
   const { enqueueSnackbar } = useSnackbar();
   const storeRequestInfo = (info) => {
diff --git a/frontend/src/services/hooks/useClient.js b/frontend/src/services/hooks/useClient.js
index bb2762ad4..7b41398ec 100644
--- a/frontend/src/services/hooks/useClient.js
+++ b/frontend/src/services/hooks/useClient.js
@@ -28,7 +28,7 @@ const defaultOptions = {
 export const useClient = () => {
   const dispatch = useDispatch();
   const [client, setClient] = useState(null);
-  const token = useToken();
+  const { token, accessToken } = useToken();
   const auth = useAuth();
 
   const setReAuth = useCallback(
@@ -47,6 +47,7 @@ export const useClient = () => {
   useEffect(() => {
     const initClient = async () => {
       const t = token;
+      const at = accessToken;
       const httpLink = new HttpLink({
         uri: process.env.REACT_APP_GRAPHQL_API
       });
@@ -55,7 +56,7 @@ export const useClient = () => {
         operation.setContext({
           headers: {
             Authorization: t ? `${t}` : '',
-            AccessKeyId: 'none',
+            AccessKeyId: at ? `Bearer ${at}` : '',
             SecretKey: 'none'
           }
         });
@@ -94,9 +95,9 @@ export const useClient = () => {
       });
       setClient(apolloClient);
     };
-    if (token) {
+    if (token && accessToken) {
       initClient().catch((e) => console.error(e));
     }
-  }, [token, dispatch]);
+  }, [token, accessToken, dispatch]);
   return client;
 };

From b70b55c67f6975dda456dbb7f249ef66995f0dda Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 12:03:52 -0500
Subject: [PATCH 02/32] store access token non cognito

---
 .../src/authentication/contexts/GenericAuthContext.js     | 8 ++++++--
 frontend/src/authentication/hooks/useToken.js             | 3 ++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/frontend/src/authentication/contexts/GenericAuthContext.js b/frontend/src/authentication/contexts/GenericAuthContext.js
index 07fbcc4df..aeb795c1b 100644
--- a/frontend/src/authentication/contexts/GenericAuthContext.js
+++ b/frontend/src/authentication/contexts/GenericAuthContext.js
@@ -84,7 +84,8 @@ export const GenericAuthProvider = (props) => {
               email: user.email,
               name: user.email,
               id_token: user.id_token,
-              short_id: user.short_id
+              short_id: user.short_id,
+              access_token: user.access_token
             }
           }
         });
@@ -129,7 +130,8 @@ export const GenericAuthProvider = (props) => {
               email: user.email,
               name: user.email,
               id_token: user.id_token,
-              short_id: user.short_id
+              short_id: user.short_id,
+              access_token: user.access_token
             }
           }
         });
@@ -178,6 +180,7 @@ export const GenericAuthProvider = (props) => {
             process.env.REACT_APP_CUSTOM_AUTH_EMAIL_CLAIM_MAPPING
           ],
         id_token: auth.user.id_token,
+        access_token: auth.user.access_token,
         short_id:
           auth.user.profile[
             process.env.REACT_APP_CUSTOM_AUTH_USERID_CLAIM_MAPPING
@@ -188,6 +191,7 @@ export const GenericAuthProvider = (props) => {
       return {
         email: user.attributes.email,
         id_token: user.signInUserSession.idToken.jwtToken,
+        access_token: user.signInUserSession.access_token.jwtToken,
         short_id: 'none'
       };
     }
diff --git a/frontend/src/authentication/hooks/useToken.js b/frontend/src/authentication/hooks/useToken.js
index 3ec75fd1f..267724060 100644
--- a/frontend/src/authentication/hooks/useToken.js
+++ b/frontend/src/authentication/hooks/useToken.js
@@ -23,8 +23,9 @@ export const useToken = () => {
               await auth.signinSilent();
             }
             const t = auth.user.id_token;
+            const at = auth.user.access_token;
             setToken(t);
-            setAccessToken(t);
+            setAccessToken(at);
           } catch (error) {
             if (!auth) throw Error('User Token Not Found !');
           }

From 88c03e51f9b82bffa90ad89eb06dcf5f00435c38 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 13:07:26 -0500
Subject: [PATCH 03/32] Enforce validation on Custom Auth Lambda

---
 backend/api_handler.py                        |  15 +-
 .../dataall/base/utils/api_handler_utils.py   |  29 +--
 .../custom_authorizer/auth_services.py        |   4 +-
 .../custom_authorizer_lambda.py               |   9 +-
 .../custom_authorizer/jwt_services.py         |  13 ++
 deploy/stacks/backend_stack.py                |   3 +
 deploy/stacks/lambda_api.py                   | 214 ++++++++++--------
 .../contexts/GenericAuthContext.js            |   2 +-
 8 files changed, 158 insertions(+), 131 deletions(-)

diff --git a/backend/api_handler.py b/backend/api_handler.py
index 0799dd10b..fd1ade4f4 100644
--- a/backend/api_handler.py
+++ b/backend/api_handler.py
@@ -85,20 +85,7 @@ def handler(event, context):
 
         Return doc: https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html
     """
-    try:
-        r = get_user_info(event=event)
-    except Exception as e:
-        log.error(f'Error occured while getting user info due to - {e}')
-        return {
-            'statusCode': 400,
-            'headers': {
-                'content-type': 'application/json',
-                'Access-Control-Allow-Origin': ALLOWED_ORIGINS,
-                'Access-Control-Allow-Headers': '*',
-                'Access-Control-Allow-Methods': '*',
-            },
-            'body': str(e),
-        }
+
     event = redact_creds(event)
     log.info('Lambda Event %s', event)
     log.debug('Env name %s', ENVNAME)
diff --git a/backend/dataall/base/utils/api_handler_utils.py b/backend/dataall/base/utils/api_handler_utils.py
index 751a8f9bc..4a16026ef 100644
--- a/backend/dataall/base/utils/api_handler_utils.py
+++ b/backend/dataall/base/utils/api_handler_utils.py
@@ -28,30 +28,19 @@
 AWS_REGION = os.getenv('AWS_REGION')
 
 
-def get_user_info(event):
-    idp_domain = ParameterStoreManager.get_parameter_value(
-        region=AWS_REGION, parameter_path=f'/dataall/{ENVNAME}/cognito/domain'
-    )
-
-    accessToken = None
-    if event.get('headers', {}).get('accesskeyid'):
-        accessToken = event['headers']['accesskeyid']
-    else:
-        raise Exception('AccessKeyId not found in headers')
+def redact_creds(event):
+    if event.get('headers', {}).get('Authorization'):
+        event['headers']['Authorization'] = 'XXXXXXXXXXXX'
 
-    url = f'https://{idp_domain}.auth.{AWS_REGION}.amazoncognito.com/oauth2/userInfo'
-    r = requests.get(url, headers={'Authorization': accessToken})
+    if event.get('multiValueHeaders', {}).get('Authorization'):
+        event['multiValueHeaders']['Authorization'] = 'XXXXXXXXXXXX'
 
-    r.raise_for_status()
-    # if r.status_code != 200:
-    #     raise Exception(r.json().get('error_description'))
+    if event.get('multiValueHeaders', {}).get('AccessKeyId'):
+        event['multiValueHeaders']['AccessKeyId'] = 'XXXXXXXXXXXX'
 
+    if event.get('headers', {}).get('AccessKeyId'):
+        event['headers']['AccessKeyId'] = 'XXXXXXXXXXXX'
 
-def redact_creds(event):
-    if 'headers' in event and 'Authorization' in event['headers']:
-        event['headers']['Authorization'] = 'XXXXXXXXXXXX'
-    if 'multiValueHeaders' in event and 'Authorization' in event['multiValueHeaders']:
-        event['multiValueHeaders']['Authorization'] = 'XXXXXXXXXXXX'
     return event
 
 
diff --git a/deploy/custom_resources/custom_authorizer/auth_services.py b/deploy/custom_resources/custom_authorizer/auth_services.py
index f5991a22e..a45e164ec 100644
--- a/deploy/custom_resources/custom_authorizer/auth_services.py
+++ b/deploy/custom_resources/custom_authorizer/auth_services.py
@@ -24,7 +24,7 @@ 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):
@@ -34,7 +34,7 @@ def generate_policy(verified_claims: dict, effect, incoming_resource_str: str):
 
         context.update(
             {
-                'user_id': verified_claims[os.getenv('user_id')],
+                'user_id': verified_claims[os.getenv('user_id', 'email')],
                 'custom_authorizer': 'true',
             }
         )
diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index ab710b1ae..053cc98d7 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/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)
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 812a03f01..d38ac3bdd 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -15,6 +15,8 @@
         'allowed_audiences': f'{os.environ.get("custom_auth_client")}',
     },
 }
+AWS_REGION = os.getenv('AWS_REGION')
+
 
 issuer_keys = {}
 
@@ -99,3 +101,14 @@ 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):
+        try:
+            idp_domain = os.getenv('auth_domain_name')
+            url = f'https://{idp_domain}.auth.{AWS_REGION}.amazoncognito.com/oauth2/userInfo'
+            r = requests.get(url, headers={'Authorization': access_token})
+            r.raise_for_status()
+        except Exception as e:
+            logger.error(f'Failed to validate access token - {str(e)}')
+            return None
diff --git a/deploy/stacks/backend_stack.py b/deploy/stacks/backend_stack.py
index 583069473..0a0273243 100644
--- a/deploy/stacks/backend_stack.py
+++ b/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,
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 1c5f788e5..fe168db01 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -30,6 +30,7 @@
 from .pyNestedStack import pyNestedClass
 from .solution_bundling import SolutionBundling
 from .waf_rules import get_waf_rules
+from .run_if import run_if
 
 
 class LambdaApiStack(pyNestedClass):
@@ -50,6 +51,8 @@ def __init__(
         apig_vpce=None,
         prod_sizing=False,
         user_pool=None,
+        user_pool_client=None,
+        user_pool_domain=None,
         pivot_role_name=None,
         reauth_ttl=5,
         email_notification_sender_email_id=None,
@@ -155,7 +158,9 @@ def __init__(
                 retention=getattr(logs.RetentionDays, self.log_retention_duration),
             ),
             description='dataall graphql function',
-            role=self.create_function_role(envname, resource_prefix, 'graphql', pivot_role_name, vpc),
+            role=self.create_function_role(
+                envname, resource_prefix, 'graphql', pivot_role_name, vpc, self._get_bedrock_policy_statement() or []
+            ),
             code=_lambda.DockerImageCode.from_ecr(
                 repository=ecr_repository, tag=image_tag, cmd=['api_handler.handler']
             ),
@@ -223,73 +228,79 @@ def __init__(
                 )
             )
 
-        if custom_auth is not None:
-            # Create the custom authorizer lambda
-            custom_authorizer_assets = os.path.realpath(
-                os.path.join(
-                    os.path.dirname(__file__),
-                    '..',
-                    'custom_resources',
-                    'custom_authorizer',
-                )
+        # if custom_auth is not None:
+        # Create the custom authorizer lambda
+        custom_authorizer_assets = os.path.realpath(
+            os.path.join(
+                os.path.dirname(__file__),
+                '..',
+                'custom_resources',
+                'custom_authorizer',
             )
+        )
+        ## GET CONGITO USERL POOL ID and APP CLIENT ID
 
-            if not os.path.isdir(custom_authorizer_assets):
-                raise Exception(f'Custom Authorizer Folder not found at {custom_authorizer_assets}')
+        if not os.path.isdir(custom_authorizer_assets):
+            raise Exception(f'Custom Authorizer Folder not found at {custom_authorizer_assets}')
 
-            custom_lambda_env = {
-                'envname': envname,
-                'LOG_LEVEL': 'DEBUG',
-                'custom_auth_provider': custom_auth.get('provider'),
-                'custom_auth_url': custom_auth.get('url'),
-                'custom_auth_client': custom_auth.get('client_id'),
-                'custom_auth_jwks_url': custom_auth.get('jwks_url'),
-            }
+        custom_lambda_env = {
+            'envname': envname,
+            'LOG_LEVEL': 'DEBUG',
+            'auth_domain_name': user_pool_domain.domain_name if user_pool_domain else '',
+            'custom_auth_provider': custom_auth.get('provider') or "Cognito",
+            'custom_auth_url': custom_auth.get('url') or f"https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}",
+            'custom_auth_client': custom_auth.get('client_id') or user_pool_client.user_pool_client_id,
+            'custom_auth_jwks_url': custom_auth.get('jwks_url') or f"https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}/.well-known/jwks.json",
+        }
 
+        if custom_auth.get('claims_mapping', {}):
             for claims_map in custom_auth.get('claims_mapping', {}):
                 custom_lambda_env[claims_map] = custom_auth.get('claims_mapping', '').get(claims_map, '')
+        else:
+            custom_lambda_env['email'] = 'email'
+            custom_lambda_env['user_id'] = 'email'
 
-            authorizer_fn_sg = self.create_lambda_sgs(envname, 'customauthorizer', resource_prefix, vpc)
-            self.authorizer_fn = _lambda.Function(
+        authorizer_fn_sg = self.create_lambda_sgs(envname, 'customauthorizer', resource_prefix, vpc)
+        self.authorizer_fn = _lambda.Function(
+            self,
+            f'CustomAuthorizerFunction-{envname}',
+            function_name=f'{resource_prefix}-{envname}-custom-authorizer',
+            log_group=logs.LogGroup(
                 self,
-                f'CustomAuthorizerFunction-{envname}',
-                function_name=f'{resource_prefix}-{envname}-custom-authorizer',
-                log_group=logs.LogGroup(
-                    self,
-                    'customauthorizerloggroup',
-                    log_group_name=f'/aws/lambda/{resource_prefix}-{envname}-custom-authorizer',
-                    retention=getattr(logs.RetentionDays, self.log_retention_duration),
-                ),
-                handler='custom_authorizer_lambda.lambda_handler',
-                code=_lambda.Code.from_asset(
-                    path=custom_authorizer_assets,
-                    bundling=BundlingOptions(
-                        image=_lambda.Runtime.PYTHON_3_9.bundling_image,
-                        local=SolutionBundling(source_path=custom_authorizer_assets),
-                    ),
+                'customauthorizerloggroup',
+                log_group_name=f'/aws/lambda/{resource_prefix}-{envname}-custom-authorizer',
+                retention=getattr(logs.RetentionDays, self.log_retention_duration),
+            ),
+            handler='custom_authorizer_lambda.lambda_handler',
+            code=_lambda.Code.from_asset(
+                path=custom_authorizer_assets,
+                bundling=BundlingOptions(
+                    image=_lambda.Runtime.PYTHON_3_9.bundling_image,
+                    local=SolutionBundling(source_path=custom_authorizer_assets),
                 ),
-                memory_size=512 if prod_sizing else 256,
-                description='dataall Custom authorizer replacing cognito authorizer',
-                timeout=Duration.seconds(20),
-                environment=custom_lambda_env,
-                environment_encryption=lambda_env_key,
-                vpc=vpc,
-                security_groups=[authorizer_fn_sg],
-                runtime=_lambda.Runtime.PYTHON_3_9,
-            )
+            ),
+            memory_size=512 if prod_sizing else 256,
+            description='dataall Custom authorizer replacing cognito authorizer',
+            timeout=Duration.seconds(20),
+            environment=custom_lambda_env,
+            environment_encryption=lambda_env_key,
+            vpc=vpc,
+            security_groups=[authorizer_fn_sg],
+            runtime=_lambda.Runtime.PYTHON_3_9,
+        )
 
-            # Add NAT Connectivity For Custom Authorizer Lambda
-            self.authorizer_fn.connections.allow_to(
-                ec2.Peer.any_ipv4(), ec2.Port.tcp(443), 'Allow NAT Internet Access SG Egress'
-            )
+        # Add NAT Connectivity For Custom Authorizer Lambda
+        self.authorizer_fn.connections.allow_to(
+            ec2.Peer.any_ipv4(), ec2.Port.tcp(443), 'Allow NAT Internet Access SG Egress'
+        )
 
-            # Store custom authorizer's ARN in ssm
-            ssm.StringParameter(
-                self,
-                f'{resource_prefix}-{envname}-custom-authorizer-arn',
-                parameter_name=f'/dataall/{envname}/customauth/customauthorizerarn',
-                string_value=self.authorizer_fn.function_arn,
-            )
+        # Store custom authorizer's ARN in ssm
+        ssm.StringParameter(
+            self,
+            f'{resource_prefix}-{envname}-custom-authorizer-arn',
+            parameter_name=f'/dataall/{envname}/customauth/customauthorizerarn',
+            string_value=self.authorizer_fn.function_arn,
+        )
 
         # Add VPC Endpoint Connectivity
         if vpce_connection:
@@ -346,7 +357,25 @@ def create_lambda_sgs(self, envname, name, resource_prefix, vpc):
         )
         return lambda_sg
 
-    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc):
+    @run_if(['modules.worksheets.features.nlq.active'])
+    def _get_bedrock_policy_statement(self):
+        return [
+            iam.PolicyStatement(
+                actions=[
+                    'bedrock:InvokeModel',
+                    'bedrock:GetPrompt',
+                    'bedrock:CreateFoundationModelAgreement',
+                    'bedrock:InvokeFlow',
+                ],
+                resources=[
+                    f'arn:aws:bedrock:{self.region}:{self.account}:flow/*',
+                    f'arn:aws:bedrock:{self.region}:{self.account}:prompt/*',
+                    f'arn:aws:bedrock:{self.region}::foundation-model/*',
+                ],
+            )
+        ]
+
+    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc, extra_statements=[]):
         role_name = f'{resource_prefix}-{envname}-{fn_name}-role'
 
         role_inline_policy = iam.Policy(
@@ -474,7 +503,8 @@ def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_nam
                     actions=['events:EnableRule', 'events:DisableRule'],
                     resources=[f'arn:aws:events:{self.region}:{self.account}:rule/dataall*'],
                 ),
-            ],
+            ]
+            + extra_statements,
         )
         role = iam.Role(
             self,
@@ -580,41 +610,41 @@ def set_up_graphql_api_gateway(
         user_pool,
         custom_auth,
     ):
-        if custom_auth is None:
-            cognito_authorizer = apigw.CognitoUserPoolsAuthorizer(
-                self,
-                'CognitoAuthorizer',
-                cognito_user_pools=[user_pool],
-                authorizer_name=f'{resource_prefix}-{envname}-cognito-authorizer',
-                identity_source='method.request.header.Authorization',
-                results_cache_ttl=Duration.minutes(60),
-            )
-        else:
-            # Create a custom Authorizer
-            custom_authorizer_role = iam.Role(
-                self,
-                f'{resource_prefix}-{envname}-custom-authorizer-role',
-                role_name=f'{resource_prefix}-{envname}-custom-authorizer-role',
-                assumed_by=iam.ServicePrincipal('apigateway.amazonaws.com'),
-                description='Allow Custom Authorizer to call custom auth lambda',
-            )
-            custom_authorizer_role.add_to_policy(
-                iam.PolicyStatement(
-                    effect=iam.Effect.ALLOW,
-                    actions=['lambda:InvokeFunction'],
-                    resources=[self.authorizer_fn.function_arn],
-                )
+        # if custom_auth is None:
+        #     cognito_authorizer = apigw.CognitoUserPoolsAuthorizer(
+        #         self,
+        #         'CognitoAuthorizer',
+        #         cognito_user_pools=[user_pool],
+        #         authorizer_name=f'{resource_prefix}-{envname}-cognito-authorizer',
+        #         identity_source='method.request.header.Authorization',
+        #         results_cache_ttl=Duration.minutes(60),
+        #     )
+        # else:
+        # Create a custom Authorizer
+        custom_authorizer_role = iam.Role(
+            self,
+            f'{resource_prefix}-{envname}-custom-authorizer-role',
+            role_name=f'{resource_prefix}-{envname}-custom-authorizer-role',
+            assumed_by=iam.ServicePrincipal('apigateway.amazonaws.com'),
+            description='Allow Custom Authorizer to call custom auth lambda',
+        )
+        custom_authorizer_role.add_to_policy(
+            iam.PolicyStatement(
+                effect=iam.Effect.ALLOW,
+                actions=['lambda:InvokeFunction'],
+                resources=[self.authorizer_fn.function_arn],
             )
+        )
 
-            custom_authorizer = apigw.RequestAuthorizer(
-                self,
-                'CustomAuthorizer',
-                handler=self.authorizer_fn,
-                identity_sources=[apigw.IdentitySource.header('Authorization')],
-                authorizer_name=f'{resource_prefix}-{envname}-custom-authorizer',
-                assume_role=custom_authorizer_role,
-                results_cache_ttl=Duration.minutes(60),
-            )
+        custom_authorizer = apigw.RequestAuthorizer(
+            self,
+            'CustomAuthorizer',
+            handler=self.authorizer_fn,
+            identity_sources=[apigw.IdentitySource.header('Authorization')],
+            authorizer_name=f'{resource_prefix}-{envname}-custom-authorizer',
+            assume_role=custom_authorizer_role,
+            results_cache_ttl=Duration.minutes(60),
+        )
         if not internet_facing:
             if apig_vpce:
                 api_vpc_endpoint = InterfaceVpcEndpoint.from_interface_vpc_endpoint_attributes(
diff --git a/frontend/src/authentication/contexts/GenericAuthContext.js b/frontend/src/authentication/contexts/GenericAuthContext.js
index aeb795c1b..2e2d095c5 100644
--- a/frontend/src/authentication/contexts/GenericAuthContext.js
+++ b/frontend/src/authentication/contexts/GenericAuthContext.js
@@ -191,7 +191,7 @@ export const GenericAuthProvider = (props) => {
       return {
         email: user.attributes.email,
         id_token: user.signInUserSession.idToken.jwtToken,
-        access_token: user.signInUserSession.access_token.jwtToken,
+        access_token: user.signInUserSession.accessToken.jwtToken,
         short_id: 'none'
       };
     }

From 5d1c2c9468d5de1e7138cab34eecfe1634d7be23 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 13:08:42 -0500
Subject: [PATCH 04/32] clean up

---
 deploy/stacks/backend_stack.py |  4 ++--
 deploy/stacks/lambda_api.py    | 12 +++++++-----
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/deploy/stacks/backend_stack.py b/deploy/stacks/backend_stack.py
index 0a0273243..c49e0017d 100644
--- a/deploy/stacks/backend_stack.py
+++ b/deploy/stacks/backend_stack.py
@@ -196,8 +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,
+            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,
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index fe168db01..0de45587f 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -247,10 +247,12 @@ def __init__(
             'envname': envname,
             'LOG_LEVEL': 'DEBUG',
             'auth_domain_name': user_pool_domain.domain_name if user_pool_domain else '',
-            'custom_auth_provider': custom_auth.get('provider') or "Cognito",
-            'custom_auth_url': custom_auth.get('url') or f"https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}",
+            'custom_auth_provider': custom_auth.get('provider') or 'Cognito',
+            'custom_auth_url': custom_auth.get('url')
+            or f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',
             'custom_auth_client': custom_auth.get('client_id') or user_pool_client.user_pool_client_id,
-            'custom_auth_jwks_url': custom_auth.get('jwks_url') or f"https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}/.well-known/jwks.json",
+            'custom_auth_jwks_url': custom_auth.get('jwks_url')
+            or f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}/.well-known/jwks.json',
         }
 
         if custom_auth.get('claims_mapping', {}):
@@ -751,7 +753,7 @@ def set_up_graphql_api_gateway(
         )
         graphql_proxy.add_method(
             'POST',
-            authorizer=cognito_authorizer if custom_auth is None else custom_authorizer,
+            authorizer=custom_authorizer,
             authorization_type=apigw.AuthorizationType.COGNITO
             if custom_auth is None
             else apigw.AuthorizationType.CUSTOM,
@@ -793,7 +795,7 @@ def set_up_graphql_api_gateway(
         )
         search_proxy.add_method(
             'POST',
-            authorizer=cognito_authorizer if custom_auth is None else custom_authorizer,
+            authorizer=custom_authorizer,
             authorization_type=apigw.AuthorizationType.COGNITO
             if custom_auth is None
             else apigw.AuthorizationType.CUSTOM,

From 02b5d36d96bf3d11280c05c8d2bda0b61d5059e5 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 14:29:35 -0500
Subject: [PATCH 05/32] env var user info endpoint

---
 deploy/custom_resources/custom_authorizer/jwt_services.py | 5 ++---
 deploy/stacks/lambda_api.py                               | 3 ++-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index d38ac3bdd..567aac7bc 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -105,9 +105,8 @@ def validate_jwt_token(jwt_token):
     @staticmethod
     def validate_access_token(access_token):
         try:
-            idp_domain = os.getenv('auth_domain_name')
-            url = f'https://{idp_domain}.auth.{AWS_REGION}.amazoncognito.com/oauth2/userInfo'
-            r = requests.get(url, headers={'Authorization': access_token})
+            user_info_url = os.getenv('user_info_url', '')
+            r = requests.get(user_info_url, headers={'Authorization': access_token})
             r.raise_for_status()
         except Exception as e:
             logger.error(f'Failed to validate access token - {str(e)}')
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 0de45587f..c4f14c3d9 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -246,7 +246,8 @@ def __init__(
         custom_lambda_env = {
             'envname': envname,
             'LOG_LEVEL': 'DEBUG',
-            'auth_domain_name': user_pool_domain.domain_name if user_pool_domain else '',
+            'user_info_url': custom_auth.get('user_info_url')
+            or f'https://{user_pool_domain.domain_name}.auth.{self.region}.amazoncognito.com/oauth2/userInfo',
             'custom_auth_provider': custom_auth.get('provider') or 'Cognito',
             'custom_auth_url': custom_auth.get('url')
             or f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',

From 7117d55fd81db2fbaf31a3c2fbbb7af42c193cef Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 14:43:56 -0500
Subject: [PATCH 06/32] fix lambda API CDK

---
 deploy/stacks/lambda_api.py | 55 ++++++++++++++++---------------------
 1 file changed, 23 insertions(+), 32 deletions(-)

diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index c4f14c3d9..97f9daf7d 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -228,7 +228,6 @@ def __init__(
                 )
             )
 
-        # if custom_auth is not None:
         # Create the custom authorizer lambda
         custom_authorizer_assets = os.path.realpath(
             os.path.join(
@@ -243,25 +242,31 @@ def __init__(
         if not os.path.isdir(custom_authorizer_assets):
             raise Exception(f'Custom Authorizer Folder not found at {custom_authorizer_assets}')
 
-        custom_lambda_env = {
-            'envname': envname,
-            'LOG_LEVEL': 'DEBUG',
-            'user_info_url': custom_auth.get('user_info_url')
-            or f'https://{user_pool_domain.domain_name}.auth.{self.region}.amazoncognito.com/oauth2/userInfo',
-            'custom_auth_provider': custom_auth.get('provider') or 'Cognito',
-            'custom_auth_url': custom_auth.get('url')
-            or f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',
-            'custom_auth_client': custom_auth.get('client_id') or user_pool_client.user_pool_client_id,
-            'custom_auth_jwks_url': custom_auth.get('jwks_url')
-            or f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}/.well-known/jwks.json',
-        }
+        if custom_auth:
+            custom_lambda_env = {
+                'envname': envname,
+                'LOG_LEVEL': 'DEBUG',
+                'user_info_url': custom_auth.get('user_info_url'),
+                'custom_auth_provider': custom_auth.get('provider'),
+                'custom_auth_url': custom_auth.get('url'),
+                'custom_auth_client': custom_auth.get('client_id'),
+                'custom_auth_jwks_url': custom_auth.get('jwks_url'),
+            }
 
-        if custom_auth.get('claims_mapping', {}):
             for claims_map in custom_auth.get('claims_mapping', {}):
                 custom_lambda_env[claims_map] = custom_auth.get('claims_mapping', '').get(claims_map, '')
         else:
-            custom_lambda_env['email'] = 'email'
-            custom_lambda_env['user_id'] = 'email'
+            custom_lambda_env = {
+                'envname': envname,
+                'LOG_LEVEL': 'DEBUG',
+                'user_info_url': f'https://{user_pool_domain.domain_name}.auth.{self.region}.amazoncognito.com/oauth2/userInfo',
+                'custom_auth_provider': 'Cognito',
+                'custom_auth_url': f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',
+                'custom_auth_client': user_pool_client.user_pool_client_id,
+                'custom_auth_jwks_url': f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}/.well-known/jwks.json',
+                'email': 'email',
+                'user_id': 'email',
+            }
 
         authorizer_fn_sg = self.create_lambda_sgs(envname, 'customauthorizer', resource_prefix, vpc)
         self.authorizer_fn = _lambda.Function(
@@ -613,16 +618,6 @@ def set_up_graphql_api_gateway(
         user_pool,
         custom_auth,
     ):
-        # if custom_auth is None:
-        #     cognito_authorizer = apigw.CognitoUserPoolsAuthorizer(
-        #         self,
-        #         'CognitoAuthorizer',
-        #         cognito_user_pools=[user_pool],
-        #         authorizer_name=f'{resource_prefix}-{envname}-cognito-authorizer',
-        #         identity_source='method.request.header.Authorization',
-        #         results_cache_ttl=Duration.minutes(60),
-        #     )
-        # else:
         # Create a custom Authorizer
         custom_authorizer_role = iam.Role(
             self,
@@ -755,9 +750,7 @@ def set_up_graphql_api_gateway(
         graphql_proxy.add_method(
             'POST',
             authorizer=custom_authorizer,
-            authorization_type=apigw.AuthorizationType.COGNITO
-            if custom_auth is None
-            else apigw.AuthorizationType.CUSTOM,
+            authorization_type=apigw.AuthorizationType.CUSTOM,
             request_validator=request_validator,
             request_models={'application/json': graphql_validation_model},
         )
@@ -797,9 +790,7 @@ def set_up_graphql_api_gateway(
         search_proxy.add_method(
             'POST',
             authorizer=custom_authorizer,
-            authorization_type=apigw.AuthorizationType.COGNITO
-            if custom_auth is None
-            else apigw.AuthorizationType.CUSTOM,
+            authorization_type=apigw.AuthorizationType.CUSTOM,
             request_validator=request_validator,
             request_models={'application/json': search_validation_model},
         )

From cba99e9da5052debfc7b4eabd6d3cbf91998c39b Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 17:38:25 -0500
Subject: [PATCH 07/32] remove unused import

---
 backend/api_handler.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/backend/api_handler.py b/backend/api_handler.py
index fd1ade4f4..9845aa245 100644
--- a/backend/api_handler.py
+++ b/backend/api_handler.py
@@ -16,7 +16,6 @@
     check_reauth,
     validate_and_block_if_maintenance_window,
     redact_creds,
-    get_user_info,
 )
 from dataall.core.tasks.service_handlers import Worker
 from dataall.base.aws.sqs import SqsQueue

From 09caacb1d63ce8d4ec24000417bacd3ed862c776 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 18:30:45 -0500
Subject: [PATCH 08/32] clean up and lint

---
 backend/api_handler.py                        |  1 -
 .../dataall/base/utils/api_handler_utils.py   |  9 ++++----
 .../custom_authorizer/jwt_services.py         |  2 --
 deploy/stacks/lambda_api.py                   | 23 ++-----------------
 .../contexts/GenericAuthContext.js            |  4 ++--
 5 files changed, 8 insertions(+), 31 deletions(-)

diff --git a/backend/api_handler.py b/backend/api_handler.py
index 9845aa245..3646c7e94 100644
--- a/backend/api_handler.py
+++ b/backend/api_handler.py
@@ -3,7 +3,6 @@
 import os
 from argparse import Namespace
 from time import perf_counter
-import requests
 from ariadne import (
     gql,
     graphql_sync,
diff --git a/backend/dataall/base/utils/api_handler_utils.py b/backend/dataall/base/utils/api_handler_utils.py
index 4a16026ef..1da2a6904 100644
--- a/backend/dataall/base/utils/api_handler_utils.py
+++ b/backend/dataall/base/utils/api_handler_utils.py
@@ -2,7 +2,6 @@
 import json
 import os
 import logging
-import requests
 
 from graphql import parse, utilities, OperationType, GraphQLSyntaxError
 from dataall.base.aws.parameter_store import ParameterStoreManager
@@ -35,11 +34,11 @@ def redact_creds(event):
     if event.get('multiValueHeaders', {}).get('Authorization'):
         event['multiValueHeaders']['Authorization'] = 'XXXXXXXXXXXX'
 
-    if event.get('multiValueHeaders', {}).get('AccessKeyId'):
-        event['multiValueHeaders']['AccessKeyId'] = 'XXXXXXXXXXXX'
+    if event.get('multiValueHeaders', {}).get('accesskeyid'):
+        event['multiValueHeaders']['accesskeyid'] = 'XXXXXXXXXXXX'
 
-    if event.get('headers', {}).get('AccessKeyId'):
-        event['headers']['AccessKeyId'] = 'XXXXXXXXXXXX'
+    if event.get('headers', {}).get('accesskeyid'):
+        event['headers']['accesskeyid'] = 'XXXXXXXXXXXX'
 
     return event
 
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 567aac7bc..40c9f839f 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -15,8 +15,6 @@
         'allowed_audiences': f'{os.environ.get("custom_auth_client")}',
     },
 }
-AWS_REGION = os.getenv('AWS_REGION')
-
 
 issuer_keys = {}
 
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 97f9daf7d..2576db5e3 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -365,25 +365,7 @@ def create_lambda_sgs(self, envname, name, resource_prefix, vpc):
         )
         return lambda_sg
 
-    @run_if(['modules.worksheets.features.nlq.active'])
-    def _get_bedrock_policy_statement(self):
-        return [
-            iam.PolicyStatement(
-                actions=[
-                    'bedrock:InvokeModel',
-                    'bedrock:GetPrompt',
-                    'bedrock:CreateFoundationModelAgreement',
-                    'bedrock:InvokeFlow',
-                ],
-                resources=[
-                    f'arn:aws:bedrock:{self.region}:{self.account}:flow/*',
-                    f'arn:aws:bedrock:{self.region}:{self.account}:prompt/*',
-                    f'arn:aws:bedrock:{self.region}::foundation-model/*',
-                ],
-            )
-        ]
-
-    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc, extra_statements=[]):
+    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc):
         role_name = f'{resource_prefix}-{envname}-{fn_name}-role'
 
         role_inline_policy = iam.Policy(
@@ -511,8 +493,7 @@ def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_nam
                     actions=['events:EnableRule', 'events:DisableRule'],
                     resources=[f'arn:aws:events:{self.region}:{self.account}:rule/dataall*'],
                 ),
-            ]
-            + extra_statements,
+            ],
         )
         role = iam.Role(
             self,
diff --git a/frontend/src/authentication/contexts/GenericAuthContext.js b/frontend/src/authentication/contexts/GenericAuthContext.js
index 2e2d095c5..d1fec575b 100644
--- a/frontend/src/authentication/contexts/GenericAuthContext.js
+++ b/frontend/src/authentication/contexts/GenericAuthContext.js
@@ -244,7 +244,7 @@ export const GenericAuthProvider = (props) => {
           }
         });
       } else {
-        await Auth.signOut();
+        await Auth.signOut({ global: true });
         dispatch({
           type: 'LOGOUT',
           payload: {
@@ -275,7 +275,7 @@ export const GenericAuthProvider = (props) => {
         console.error('Failed to ReAuth', error);
       }
     } else {
-      await Auth.signOut();
+      await Auth.signOut({ global: true });
       dispatch({
         type: 'REAUTH',
         payload: {

From 0e77776dc58f33ed0d708fb77ce74615aa8ac201 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 18:51:05 -0500
Subject: [PATCH 09/32] handle cognito groups header

---
 .../custom_authorizer/auth_services.py        |  2 ++
 deploy/stacks/lambda_api.py                   | 23 +++++++++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/auth_services.py b/deploy/custom_resources/custom_authorizer/auth_services.py
index a45e164ec..e4f1b771a 100644
--- a/deploy/custom_resources/custom_authorizer/auth_services.py
+++ b/deploy/custom_resources/custom_authorizer/auth_services.py
@@ -28,6 +28,8 @@ 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):
+                if claim_name == 'cognito:groups':
+                    verified_claims.update({'groups': ','.join(claim_value)})
                 verified_claims.update({claim_name: json.dumps(claim_value)})
 
         context = {**verified_claims}
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 2576db5e3..97f9daf7d 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -365,7 +365,25 @@ def create_lambda_sgs(self, envname, name, resource_prefix, vpc):
         )
         return lambda_sg
 
-    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc):
+    @run_if(['modules.worksheets.features.nlq.active'])
+    def _get_bedrock_policy_statement(self):
+        return [
+            iam.PolicyStatement(
+                actions=[
+                    'bedrock:InvokeModel',
+                    'bedrock:GetPrompt',
+                    'bedrock:CreateFoundationModelAgreement',
+                    'bedrock:InvokeFlow',
+                ],
+                resources=[
+                    f'arn:aws:bedrock:{self.region}:{self.account}:flow/*',
+                    f'arn:aws:bedrock:{self.region}:{self.account}:prompt/*',
+                    f'arn:aws:bedrock:{self.region}::foundation-model/*',
+                ],
+            )
+        ]
+
+    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc, extra_statements=[]):
         role_name = f'{resource_prefix}-{envname}-{fn_name}-role'
 
         role_inline_policy = iam.Policy(
@@ -493,7 +511,8 @@ def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_nam
                     actions=['events:EnableRule', 'events:DisableRule'],
                     resources=[f'arn:aws:events:{self.region}:{self.account}:rule/dataall*'],
                 ),
-            ],
+            ]
+            + extra_statements,
         )
         role = iam.Role(
             self,

From e6fcf8f4155a393471fc81b6ee1b34e163b3a594 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 19:08:52 -0500
Subject: [PATCH 10/32] update custom auth logic

---
 .../custom_resources/custom_authorizer/jwt_services.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 40c9f839f..2777c5cb7 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -102,10 +102,6 @@ def validate_jwt_token(jwt_token):
 
     @staticmethod
     def validate_access_token(access_token):
-        try:
-            user_info_url = os.getenv('user_info_url', '')
-            r = requests.get(user_info_url, headers={'Authorization': access_token})
-            r.raise_for_status()
-        except Exception as e:
-            logger.error(f'Failed to validate access token - {str(e)}')
-            return None
+        user_info_url = os.getenv('user_info_url', '')
+        r = requests.get(user_info_url, headers={'Authorization': access_token})
+        r.raise_for_status()

From 58a5440eeef3c490a5cde07f054b90fa38a8e572 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Mon, 4 Nov 2024 21:40:23 -0500
Subject: [PATCH 11/32] fix

---
 deploy/custom_resources/custom_authorizer/auth_services.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deploy/custom_resources/custom_authorizer/auth_services.py b/deploy/custom_resources/custom_authorizer/auth_services.py
index e4f1b771a..a2e4b1179 100644
--- a/deploy/custom_resources/custom_authorizer/auth_services.py
+++ b/deploy/custom_resources/custom_authorizer/auth_services.py
@@ -29,7 +29,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):
                 if claim_name == 'cognito:groups':
-                    verified_claims.update({'groups': ','.join(claim_value)})
+                    verified_claims.update({claim_name: ','.join(claim_value)})
                 verified_claims.update({claim_name: json.dumps(claim_value)})
 
         context = {**verified_claims}

From c04760aa51fe5fc9dcbac9fb28d15440f4920771 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 09:43:19 -0500
Subject: [PATCH 12/32] fix custom Auth logic

---
 .../custom_authorizer/custom_authorizer_lambda.py         | 8 ++++----
 deploy/custom_resources/custom_authorizer/jwt_services.py | 1 +
 deploy/stacks/lambda_api.py                               | 2 +-
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 053cc98d7..36e7590b1 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -19,6 +19,7 @@
 
 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')
@@ -28,9 +29,8 @@ 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)
+    access_token = incoming_event['headers']['accesskeyid']
+    JWTServices.validate_access_token(access_token)
 
     effect = 'Allow'
     policy = AuthServices.generate_policy(verified_claims, effect, incoming_event['methodArn'])
@@ -48,7 +48,7 @@ def lambda_handler(incoming_event, context):
     account_id = ''
     api_gw_id = ''
     event = {
-        'headers': {'Authorization': id_token, 'AccessKeyId': access_token},
+        'headers': {'Authorization': id_token, 'accesskeyid': access_token},
         'type': 'TOKEN',
         'methodArn': f'arn:aws:execute-api:us-east-1:{account_id}:{api_gw_id}/prod/POST/graphql/api',
     }
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 2777c5cb7..5bd43d262 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -104,4 +104,5 @@ def validate_jwt_token(jwt_token):
     def validate_access_token(access_token):
         user_info_url = os.getenv('user_info_url', '')
         r = requests.get(user_info_url, headers={'Authorization': access_token})
+        logger.debug(r.json())
         r.raise_for_status()
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 97f9daf7d..c210464ff 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -641,7 +641,7 @@ def set_up_graphql_api_gateway(
             identity_sources=[apigw.IdentitySource.header('Authorization')],
             authorizer_name=f'{resource_prefix}-{envname}-custom-authorizer',
             assume_role=custom_authorizer_role,
-            results_cache_ttl=Duration.minutes(60),
+            results_cache_ttl=Duration.minutes(1),
         )
         if not internet_facing:
             if apig_vpce:

From b5c0c082cb787beb555e51947a06d9a174de7417 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 09:55:12 -0500
Subject: [PATCH 13/32] fix cognito groups list handling

---
 deploy/custom_resources/custom_authorizer/auth_services.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/auth_services.py b/deploy/custom_resources/custom_authorizer/auth_services.py
index a2e4b1179..ae57ab455 100644
--- a/deploy/custom_resources/custom_authorizer/auth_services.py
+++ b/deploy/custom_resources/custom_authorizer/auth_services.py
@@ -28,9 +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):
-                if claim_name == 'cognito:groups':
-                    verified_claims.update({claim_name: ','.join(claim_value)})
-                verified_claims.update({claim_name: json.dumps(claim_value)})
+                verified_claims.update({claim_name: ','.join(claim_value)})
 
         context = {**verified_claims}
 

From 25334bc24aa8ebadcd5022ab3f919091fe04638d Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 10:03:34 -0500
Subject: [PATCH 14/32] lint

---
 backend/api_handler.py      |  1 +
 deploy/stacks/lambda_api.py | 27 +++------------------------
 2 files changed, 4 insertions(+), 24 deletions(-)

diff --git a/backend/api_handler.py b/backend/api_handler.py
index 3646c7e94..1e1645e9e 100644
--- a/backend/api_handler.py
+++ b/backend/api_handler.py
@@ -3,6 +3,7 @@
 import os
 from argparse import Namespace
 from time import perf_counter
+
 from ariadne import (
     gql,
     graphql_sync,
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index c210464ff..847c0d454 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -158,9 +158,7 @@ def __init__(
                 retention=getattr(logs.RetentionDays, self.log_retention_duration),
             ),
             description='dataall graphql function',
-            role=self.create_function_role(
-                envname, resource_prefix, 'graphql', pivot_role_name, vpc, self._get_bedrock_policy_statement() or []
-            ),
+            role=self.create_function_role(envname, resource_prefix, 'graphql', pivot_role_name, vpc),
             code=_lambda.DockerImageCode.from_ecr(
                 repository=ecr_repository, tag=image_tag, cmd=['api_handler.handler']
             ),
@@ -365,25 +363,7 @@ def create_lambda_sgs(self, envname, name, resource_prefix, vpc):
         )
         return lambda_sg
 
-    @run_if(['modules.worksheets.features.nlq.active'])
-    def _get_bedrock_policy_statement(self):
-        return [
-            iam.PolicyStatement(
-                actions=[
-                    'bedrock:InvokeModel',
-                    'bedrock:GetPrompt',
-                    'bedrock:CreateFoundationModelAgreement',
-                    'bedrock:InvokeFlow',
-                ],
-                resources=[
-                    f'arn:aws:bedrock:{self.region}:{self.account}:flow/*',
-                    f'arn:aws:bedrock:{self.region}:{self.account}:prompt/*',
-                    f'arn:aws:bedrock:{self.region}::foundation-model/*',
-                ],
-            )
-        ]
-
-    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc, extra_statements=[]):
+    def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_name, vpc):
         role_name = f'{resource_prefix}-{envname}-{fn_name}-role'
 
         role_inline_policy = iam.Policy(
@@ -511,8 +491,7 @@ def create_function_role(self, envname, resource_prefix, fn_name, pivot_role_nam
                     actions=['events:EnableRule', 'events:DisableRule'],
                     resources=[f'arn:aws:events:{self.region}:{self.account}:rule/dataall*'],
                 ),
-            ]
-            + extra_statements,
+            ],
         )
         role = iam.Role(
             self,

From fca14f2125aff411d69a10b083b2e66155ba1197 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 10:14:44 -0500
Subject: [PATCH 15/32] update checkov baseline

---
 .checkov.baseline           | 13 +++++++++++++
 deploy/stacks/lambda_api.py |  1 -
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/.checkov.baseline b/.checkov.baseline
index f80fb4c43..e0bbfce12 100644
--- a/.checkov.baseline
+++ b/.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": [
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 847c0d454..3708482bb 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -30,7 +30,6 @@
 from .pyNestedStack import pyNestedClass
 from .solution_bundling import SolutionBundling
 from .waf_rules import get_waf_rules
-from .run_if import run_if
 
 
 class LambdaApiStack(pyNestedClass):

From 64fa2f10d6367ecf5b8384baf1ba0e89e51dfe8a Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 12:26:57 -0500
Subject: [PATCH 16/32] use well-known/openid-config

---
 deploy/custom_resources/custom_authorizer/jwt_services.py | 8 +++++++-
 deploy/stacks/lambda_api.py                               | 2 --
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 5bd43d262..7ace6bf82 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -67,6 +67,12 @@ def fetch_public_keys():
 
 
 class JWTServices:
+    @staticmethod
+    def _fetch_openid_url(key):
+        response = requests.get(f'{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:
@@ -102,7 +108,7 @@ def validate_jwt_token(jwt_token):
 
     @staticmethod
     def validate_access_token(access_token):
-        user_info_url = os.getenv('user_info_url', '')
+        user_info_url = JWTServices._fetch_openid_url('userinfo_endpoint')
         r = requests.get(user_info_url, headers={'Authorization': access_token})
         logger.debug(r.json())
         r.raise_for_status()
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 3708482bb..29ecc7228 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -243,7 +243,6 @@ def __init__(
             custom_lambda_env = {
                 'envname': envname,
                 'LOG_LEVEL': 'DEBUG',
-                'user_info_url': custom_auth.get('user_info_url'),
                 'custom_auth_provider': custom_auth.get('provider'),
                 'custom_auth_url': custom_auth.get('url'),
                 'custom_auth_client': custom_auth.get('client_id'),
@@ -256,7 +255,6 @@ def __init__(
             custom_lambda_env = {
                 'envname': envname,
                 'LOG_LEVEL': 'DEBUG',
-                'user_info_url': f'https://{user_pool_domain.domain_name}.auth.{self.region}.amazoncognito.com/oauth2/userInfo',
                 'custom_auth_provider': 'Cognito',
                 'custom_auth_url': f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',
                 'custom_auth_client': user_pool_client.user_pool_client_id,

From d416a559df3a7f060fe087fb5451ed6f102f361c Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 17:43:15 -0500
Subject: [PATCH 17/32] Update deps custom authorizer

---
 .../custom_authorizer/auth_services.py        |  4 +-
 .../custom_authorizer_lambda.py               | 13 +--
 .../custom_authorizer/jwt_services.py         | 84 +++++--------------
 .../custom_authorizer/requirements.txt        |  5 +-
 4 files changed, 34 insertions(+), 72 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/auth_services.py b/deploy/custom_resources/custom_authorizer/auth_services.py
index ae57ab455..52efe14aa 100644
--- a/deploy/custom_resources/custom_authorizer/auth_services.py
+++ b/deploy/custom_resources/custom_authorizer/auth_services.py
@@ -24,7 +24,7 @@ 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', 'email')]
+        verified_claims['email'] = verified_claims[os.getenv('email')]
 
         for claim_name, claim_value in verified_claims.items():
             if isinstance(claim_value, list):
@@ -34,7 +34,7 @@ def generate_policy(verified_claims: dict, effect, incoming_resource_str: str):
 
         context.update(
             {
-                'user_id': verified_claims[os.getenv('user_id', 'email')],
+                'user_id': verified_claims[os.getenv('user_id')],
                 'custom_authorizer': 'true',
             }
         )
diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 36e7590b1..cf7498e4e 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -21,17 +21,20 @@ 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')
 
-    access_token = incoming_event['headers']['accesskeyid']
-    JWTServices.validate_access_token(access_token)
-
+    # Generate Allow Policy w/ Context
     effect = 'Allow'
     policy = AuthServices.generate_policy(verified_claims, effect, incoming_event['methodArn'])
     logger.debug('Generated policy is ', policy)
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 7ace6bf82..0432f9243 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/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,16 +20,8 @@ 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"
 }
 
 
@@ -76,30 +35,30 @@ def _fetch_openid_url(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:
@@ -108,6 +67,7 @@ def validate_jwt_token(jwt_token):
 
     @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})
         logger.debug(r.json())
diff --git a/deploy/custom_resources/custom_authorizer/requirements.txt b/deploy/custom_resources/custom_authorizer/requirements.txt
index 14e5c340e..31e7b3bbe 100644
--- a/deploy/custom_resources/custom_authorizer/requirements.txt
+++ b/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
\ No newline at end of file
+urllib3==1.26.19
+pyjwt[crypto]==2.9.0
\ No newline at end of file

From 913c1ced4d3161f3e1db1df3bc0a92febd39681e Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 17:52:48 -0500
Subject: [PATCH 18/32] clean up requirements

---
 .../custom_authorizer/requirements.txt                 | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/requirements.txt b/deploy/custom_resources/custom_authorizer/requirements.txt
index 31e7b3bbe..4ecfdcc46 100644
--- a/deploy/custom_resources/custom_authorizer/requirements.txt
+++ b/deploy/custom_resources/custom_authorizer/requirements.txt
@@ -1,9 +1,3 @@
-certifi==2024.7.4
-charset-normalizer==3.1.0
-idna==3.7
-pyasn1==0.5.0
+cryptography==43.0.3
+pyjwt[crypto]==2.9.0
 requests==2.32.2
-rsa==4.9
-six==1.16.0
-urllib3==1.26.19
-pyjwt[crypto]==2.9.0
\ No newline at end of file

From a8cfe08fb173ec0b426076add50243c79e25f74d Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Tue, 5 Nov 2024 20:00:37 -0500
Subject: [PATCH 19/32] update cryptography version

---
 deploy/custom_resources/custom_authorizer/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deploy/custom_resources/custom_authorizer/requirements.txt b/deploy/custom_resources/custom_authorizer/requirements.txt
index 4ecfdcc46..c85406261 100644
--- a/deploy/custom_resources/custom_authorizer/requirements.txt
+++ b/deploy/custom_resources/custom_authorizer/requirements.txt
@@ -1,3 +1,3 @@
-cryptography==43.0.3
+cryptography==41.0.7
 pyjwt[crypto]==2.9.0
 requests==2.32.2

From 647b6893e39318f88cf9603fb02e678079e68d40 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 00:06:39 -0500
Subject: [PATCH 20/32] Add cryptography lambda layer

---
 .../custom_authorizer/requirements.txt                 | 10 ++++++++--
 deploy/stacks/lambda_api.py                            |  7 +++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/requirements.txt b/deploy/custom_resources/custom_authorizer/requirements.txt
index c85406261..db3720bed 100644
--- a/deploy/custom_resources/custom_authorizer/requirements.txt
+++ b/deploy/custom_resources/custom_authorizer/requirements.txt
@@ -1,3 +1,9 @@
-cryptography==41.0.7
-pyjwt[crypto]==2.9.0
+certifi==2024.7.4
+charset-normalizer==3.1.0
+idna==3.7
+pyasn1==0.5.0
 requests==2.32.2
+rsa==4.9
+six==1.16.0
+urllib3==1.26.19
+pyjwt==2.9.0
\ No newline at end of file
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index 29ecc7228..fabbec0ce 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -263,6 +263,12 @@ def __init__(
                 'user_id': 'email',
             }
 
+        layer = _lambda.LayerVersion.from_layer_version_arn(
+            self,
+            'CryptographyLayerPy39',
+            f'arn:aws:lambda:{self.region}:770693421928:layer:Klayers-p39-cryptography:19',
+        )
+
         authorizer_fn_sg = self.create_lambda_sgs(envname, 'customauthorizer', resource_prefix, vpc)
         self.authorizer_fn = _lambda.Function(
             self,
@@ -290,6 +296,7 @@ def __init__(
             vpc=vpc,
             security_groups=[authorizer_fn_sg],
             runtime=_lambda.Runtime.PYTHON_3_9,
+            layers=[layer],
         )
 
         # Add NAT Connectivity For Custom Authorizer Lambda

From eddcb8688f0b75e5a9cc723f24e75d920c89f7c4 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 09:00:47 -0500
Subject: [PATCH 21/32] Reformat

---
 .../custom_authorizer/custom_authorizer_lambda.py            | 5 +++--
 deploy/custom_resources/custom_authorizer/jwt_services.py    | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index cf7498e4e..9676ef750 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/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
@@ -30,14 +31,14 @@ def lambda_handler(incoming_event, context):
 
     # 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
 
 
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 0432f9243..1068fa9cb 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -70,5 +70,5 @@ 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})
-        logger.debug(r.json())
         r.raise_for_status()
+        logger.debug(r.json())

From 3b4418d90f5bcfcf860c3556b09fb00802ec160a Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 10:48:51 -0500
Subject: [PATCH 22/32] Clean up and PR comments

---
 .../custom_authorizer/jwt_services.py         |  4 +-
 deploy/stacks/backend_stack.py                |  2 -
 deploy/stacks/lambda_api.py                   | 39 ++++++++++---------
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 1068fa9cb..bfdaac307 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -28,7 +28,9 @@
 class JWTServices:
     @staticmethod
     def _fetch_openid_url(key):
-        response = requests.get(f'{os.environ.get("custom_auth_url")}/.well-known/openid-configuration')
+        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)
 
diff --git a/deploy/stacks/backend_stack.py b/deploy/stacks/backend_stack.py
index c49e0017d..87183bb01 100644
--- a/deploy/stacks/backend_stack.py
+++ b/deploy/stacks/backend_stack.py
@@ -35,7 +35,6 @@ def __init__(
         id,
         envname: str = 'dev',
         resource_prefix='dataall',
-        tooling_region=None,
         tooling_account_id=None,
         ecr_repository=None,
         image_tag=None,
@@ -197,7 +196,6 @@ def __init__(
             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,
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index fabbec0ce..dc098e143 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -51,7 +51,6 @@ def __init__(
         prod_sizing=False,
         user_pool=None,
         user_pool_client=None,
-        user_pool_domain=None,
         pivot_role_name=None,
         reauth_ttl=5,
         email_notification_sender_email_id=None,
@@ -239,29 +238,31 @@ def __init__(
         if not os.path.isdir(custom_authorizer_assets):
             raise Exception(f'Custom Authorizer Folder not found at {custom_authorizer_assets}')
 
+        custom_lambda_env = {
+            'envname': envname,
+            'LOG_LEVEL': log_level,
+        }
         if custom_auth:
-            custom_lambda_env = {
-                'envname': envname,
-                'LOG_LEVEL': 'DEBUG',
-                'custom_auth_provider': custom_auth.get('provider'),
-                'custom_auth_url': custom_auth.get('url'),
-                'custom_auth_client': custom_auth.get('client_id'),
-                'custom_auth_jwks_url': custom_auth.get('jwks_url'),
-            }
+            custom_lambda_env.update(
+                {
+                    'custom_auth_provider': custom_auth.get('provider'),
+                    'custom_auth_url': custom_auth.get('url'),
+                    'custom_auth_client': custom_auth.get('client_id'),
+                }
+            )
 
             for claims_map in custom_auth.get('claims_mapping', {}):
                 custom_lambda_env[claims_map] = custom_auth.get('claims_mapping', '').get(claims_map, '')
         else:
-            custom_lambda_env = {
-                'envname': envname,
-                'LOG_LEVEL': 'DEBUG',
-                'custom_auth_provider': 'Cognito',
-                'custom_auth_url': f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',
-                'custom_auth_client': user_pool_client.user_pool_client_id,
-                'custom_auth_jwks_url': f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}/.well-known/jwks.json',
-                'email': 'email',
-                'user_id': 'email',
-            }
+            custom_lambda_env.update(
+                {
+                    'custom_auth_provider': 'Cognito',
+                    'custom_auth_url': f'https://cognito-idp.{self.region}.amazonaws.com/{user_pool.user_pool_id}',
+                    'custom_auth_client': user_pool_client.user_pool_client_id,
+                    'email': 'email',
+                    'user_id': 'email',
+                }
+            )
 
         layer = _lambda.LayerVersion.from_layer_version_arn(
             self,

From f30dd4f0a22949c1d4f61f9f1e405512f2518ff6 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 12:15:00 -0500
Subject: [PATCH 23/32] Fix Integration Tests to fetch ID and Access Tokens

---
 deploy/stacks/pipeline.py                    |  2 +
 tests_new/integration_tests/README.md        |  3 +
 tests_new/integration_tests/client.py        | 68 ++++++++++++++++----
 tests_new/integration_tests/requirements.txt |  3 +-
 4 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/deploy/stacks/pipeline.py b/deploy/stacks/pipeline.py
index 427a117c5..79ab7a5d4 100644
--- a/deploy/stacks/pipeline.py
+++ b/deploy/stacks/pipeline.py
@@ -715,6 +715,8 @@ def set_approval_tests_stage(
                                     'aws sts get-caller-identity --profile buildprofile',
                                     f'export COGNITO_CLIENT=$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/cognito/appclient --profile buildprofile --output text --query "Parameter.Value")',
                                     f'export API_ENDPOINT=$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/apiGateway/backendUrl --profile buildprofile --output text --query "Parameter.Value")',
+                                    f'export IDP_DOMAIN_URL=https://$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/cognito/domain --profile buildprofile --output text --query "Parameter.Value").auth.us-east-1.amazoncognito.com',
+                                    f'export DATAALL_DOMAIN_URL=https://$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/CloudfrontDistributionDomainName --profile buildprofile --output text --query "Parameter.Value")',
                                     f'export TESTDATA=$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/testdata --profile buildprofile --output text --query "Parameter.Value")',
                                     f'export ENVNAME={target_env["envname"]}',
                                     f'export AWS_REGION={target_env["region"]}',
diff --git a/tests_new/integration_tests/README.md b/tests_new/integration_tests/README.md
index 63493bc16..054b8fb31 100644
--- a/tests_new/integration_tests/README.md
+++ b/tests_new/integration_tests/README.md
@@ -173,6 +173,9 @@ You can also run the tests locally by...
   export AWS_REGION = "Introduce backend region"
   export COGNITO_CLIENT = "Introduce Cognito client id"
   export API_ENDPOINT = "Introduce API endpoint url"
+  export IDP_DOMAIN_URL = "Introduce your Identity Provider domain url"
+  export DATAALL_DOMAIN_URL = "Introduce data.all frontend domain url"
+
   echo "add your testdata here" > testdata.json 
   make integration-tests
   ```
diff --git a/tests_new/integration_tests/client.py b/tests_new/integration_tests/client.py
index 75fa0455e..2474d4161 100644
--- a/tests_new/integration_tests/client.py
+++ b/tests_new/integration_tests/client.py
@@ -1,9 +1,13 @@
 import requests
-import boto3
 import os
+import uuid
+from urllib.parse import parse_qs, urlparse
 from munch import DefaultMunch
 from retrying import retry
 from integration_tests.errors import GqlError
+import requests
+from oauthlib.oauth2 import WebApplicationClient
+from requests_oauthlib import OAuth2Session
 
 ENVNAME = os.getenv('ENVNAME', 'dev')
 
@@ -17,7 +21,7 @@ class Client:
     def __init__(self, username, password):
         self.username = username
         self.password = password
-        self.token = self._get_jwt_token()
+        self.id_token, self.access_token = self._get_jwt_tokens()
 
     @retry(
         retry_on_exception=_retry_if_connection_error,
@@ -27,7 +31,7 @@ def __init__(self, username, password):
     )
     def query(self, query: str):
         graphql_endpoint = os.path.join(os.environ['API_ENDPOINT'], 'graphql', 'api')
-        headers = {'AccessKeyId': 'none', 'SecretKey': 'none', 'authorization': self.token}
+        headers = {'accesskeyid': f'Bearer {self.access_token}', 'SecretKey': 'none', 'Authorization': self.id_token}
         r = requests.post(graphql_endpoint, json=query, headers=headers)
         if errors := r.json().get('errors'):
             raise GqlError(errors)
@@ -35,16 +39,52 @@ def query(self, query: str):
 
         return DefaultMunch.fromDict(r.json())
 
-    def _get_jwt_token(self):
-        cognito_client = boto3.client('cognito-idp', region_name=os.getenv('AWS_REGION', 'eu-west-1'))
-        kwargs = {
-            'ClientId': os.environ['COGNITO_CLIENT'],
-            'AuthFlow': 'USER_PASSWORD_AUTH',
-            'AuthParameters': {
-                'USERNAME': self.username,
-                'PASSWORD': self.password,
-            },
+    def _get_jwt_tokens(self):
+        token = uuid.uuid4()
+        scope = 'aws.cognito.signin.user.admin openid'
+
+        idp_domain_url = os.environ['IDP_DOMAIN_URL']
+
+        token_url = os.path.join(idp_domain_url, 'oauth2', 'token')
+        login_url = os.path.join(idp_domain_url, 'login')
+
+        client_id = os.environ['COGNITO_CLIENT']
+        redirect_uri = os.environ['DATAALL_DOMAIN_URL']
+
+        data = {
+            '_csrf': token,
+            'username': self.username,
+            'password': self.password,
+        }
+        params = {
+            'client_id': client_id,
+            'scope': scope,
+            'redirect_uri': redirect_uri,
+            'response_type': 'code',
         }
-        resp = cognito_client.initiate_auth(**kwargs)
 
-        return resp['AuthenticationResult']['IdToken']
+        headers = {'cookie': f'XSRF-TOKEN={token}; csrf-state=""; csrf-state-legacy=""'}
+        r = requests.post(
+            login_url,
+            params=params,
+            data=data,
+            headers=headers,
+            allow_redirects=False,
+        )
+        print(r)
+
+        r.raise_for_status()
+
+        code = parse_qs(urlparse(r.headers['location']).query)['code'][0]
+
+        client = WebApplicationClient(client_id=client_id)
+        oauth = OAuth2Session(client=client, redirect_uri=redirect_uri)
+        token = oauth.fetch_token(
+            token_url=token_url,
+            client_id=client_id,
+            code=code,
+            include_client_id=True,
+        )
+        print(token)
+
+        return token.get('id_token'), token.get('access_token')
diff --git a/tests_new/integration_tests/requirements.txt b/tests_new/integration_tests/requirements.txt
index 9c274748e..06deb022a 100644
--- a/tests_new/integration_tests/requirements.txt
+++ b/tests_new/integration_tests/requirements.txt
@@ -8,4 +8,5 @@ pytest-dependency==0.5.1
 requests==2.32.2
 dataclasses-json==0.6.6
 werkzeug==3.0.6
-retrying==1.3.4
\ No newline at end of file
+retrying==1.3.4
+requests-oauthlib==2.0.0
\ No newline at end of file

From 320f65216b2372ba2426c5d75dfc5c8821b23af0 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 12:28:14 -0500
Subject: [PATCH 24/32] use cdk klayers

---
 deploy/requirements.txt     |  3 ++-
 deploy/stacks/lambda_api.py | 16 +++++++++-------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/deploy/requirements.txt b/deploy/requirements.txt
index bf1d93556..157e91b1a 100644
--- a/deploy/requirements.txt
+++ b/deploy/requirements.txt
@@ -2,4 +2,5 @@ aws-cdk-lib==2.160.0
 boto3==1.35.26
 boto3-stubs==1.35.26
 cdk-nag==2.7.2
-typeguard==4.2.1
\ No newline at end of file
+typeguard==4.2.1
+cdk-klayers==0.3.0
\ No newline at end of file
diff --git a/deploy/stacks/lambda_api.py b/deploy/stacks/lambda_api.py
index dc098e143..64f6c6acd 100644
--- a/deploy/stacks/lambda_api.py
+++ b/deploy/stacks/lambda_api.py
@@ -20,6 +20,7 @@
     RemovalPolicy,
     BundlingOptions,
 )
+from cdk_klayers import Klayers
 from aws_cdk.aws_ec2 import (
     InterfaceVpcEndpoint,
     InterfaceVpcEndpointAwsService,
@@ -264,11 +265,12 @@ def __init__(
                 }
             )
 
-        layer = _lambda.LayerVersion.from_layer_version_arn(
-            self,
-            'CryptographyLayerPy39',
-            f'arn:aws:lambda:{self.region}:770693421928:layer:Klayers-p39-cryptography:19',
-        )
+        # Initialize Klayers
+        runtime = _lambda.Runtime.PYTHON_3_9
+        klayers = Klayers(self, python_version=runtime, region=self.region)
+
+        # get the latest layer version for the cryptography package
+        cryptography_layer = klayers.layer_version(self, 'cryptography')
 
         authorizer_fn_sg = self.create_lambda_sgs(envname, 'customauthorizer', resource_prefix, vpc)
         self.authorizer_fn = _lambda.Function(
@@ -296,8 +298,8 @@ def __init__(
             environment_encryption=lambda_env_key,
             vpc=vpc,
             security_groups=[authorizer_fn_sg],
-            runtime=_lambda.Runtime.PYTHON_3_9,
-            layers=[layer],
+            runtime=runtime,
+            layers=[cryptography_layer],
         )
 
         # Add NAT Connectivity For Custom Authorizer Lambda

From ff910c7a926653b0f6c3a2d27dd1706f3a6a7ebf Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 12:46:46 -0500
Subject: [PATCH 25/32] refactor jwt services

---
 .../custom_authorizer_lambda.py               |  7 ++--
 .../custom_authorizer/jwt_services.py         | 34 ++++++++-----------
 tests_new/integration_tests/client.py         |  1 -
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 9676ef750..7db7d7986 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -17,6 +17,9 @@
 Custom Lambda Authorizer is attached to the API Gateway. Check the deploy/stacks/lambda_api.py for more details on deployment
 """
 
+OPENID_CONFIG_PATH = os.path.join(os.environ.get('custom_auth_url', ''), '.well-known', 'openid-configuration')
+jwt_service = JWTServices(OPENID_CONFIG_PATH)
+
 
 def lambda_handler(incoming_event, context):
     # Get the Token which is sent in the Authorization Header
@@ -27,10 +30,10 @@ def lambda_handler(incoming_event, context):
         raise Exception('Unauthorized. Missing identity or access JWT')
 
     # Validate User is Active with Proper Access Token
-    JWTServices.validate_access_token(access_token)
+    jwt_service.validate_access_token(access_token)
 
     # Validate JWT
-    verified_claims = JWTServices.validate_jwt_token(auth_token)
+    verified_claims = jwt_service.validate_jwt_token(auth_token)
     if not verified_claims:
         raise Exception('Unauthorized. Token is not valid')
     logger.debug(verified_claims)
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index bfdaac307..131a0baf3 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -26,25 +26,22 @@
 
 
 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)
+    def __init__(self, openid_config_path):
+        # Get OpenID Config JSON
+        self.openid_config = self._fetch_openid_config(openid_config_path)
 
-    @staticmethod
-    def validate_jwt_token(jwt_token):
-        try:
-            # get JWK URI from OpenId Configuration
-            jwks_url = JWTServices._fetch_openid_url('jwks_uri')
+        # Init pyJWT.JWKClient with JWK URI
+        self.jwks_client = jwt.PyJWKClient(self.openid_config.get('jwks_uri'))
 
-            # Init pyJWT.JWKClient with JWK URI
-            jwks_client = jwt.PyJWKClient(jwks_url)
+    def _fetch_openid_config(self, openid_config_path):
+        response = requests.get(openid_config_path)
+        response.raise_for_status()
+        return response.json()
 
+    def validate_jwt_token(self, jwt_token):
+        try:
             # get signing_key from JWT
-            signing_key = jwks_client.get_signing_key_from_jwt(jwt_token)
+            signing_key = self.jwks_client.get_signing_key_from_jwt(jwt_token)
 
             # Decode and Verify JWT
             payload = jwt.decode(
@@ -67,10 +64,9 @@ def validate_jwt_token(jwt_token):
             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')
+    def validate_access_token(self, access_token):
+        # get UserInfo URI from OpenId Configuration
+        user_info_url = self.openid_config.get('userinfo_endpoint')
         r = requests.get(user_info_url, headers={'Authorization': access_token})
         r.raise_for_status()
         logger.debug(r.json())
diff --git a/tests_new/integration_tests/client.py b/tests_new/integration_tests/client.py
index 2474d4161..0992b4b8a 100644
--- a/tests_new/integration_tests/client.py
+++ b/tests_new/integration_tests/client.py
@@ -5,7 +5,6 @@
 from munch import DefaultMunch
 from retrying import retry
 from integration_tests.errors import GqlError
-import requests
 from oauthlib.oauth2 import WebApplicationClient
 from requests_oauthlib import OAuth2Session
 

From 9b9cbad6e87091c81861492ef4f5c3f8da42e5f2 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 13:57:58 -0500
Subject: [PATCH 26/32] fix string formatting

---
 .../custom_authorizer/custom_authorizer_lambda.py               | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 7db7d7986..08be6790a 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -41,7 +41,7 @@ def lambda_handler(incoming_event, context):
     # Generate Allow Policy w/ Context
     effect = 'Allow'
     policy = AuthServices.generate_policy(verified_claims, effect, incoming_event['methodArn'])
-    logger.debug('Generated policy is ', json.dumps(policy))
+    logger.debug(f'Generated policy is {json.dumps(policy)}')
     return policy
 
 

From 451b0d45ba948b391eb3852d591e7fcabf9fcc29 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 14:14:27 -0500
Subject: [PATCH 27/32] raise exceptions

---
 deploy/custom_resources/custom_authorizer/jwt_services.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 131a0baf3..f13167881 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -54,15 +54,15 @@ def validate_jwt_token(self, jwt_token):
                 options=jwt_options,
             )
             return payload
-        except jwt.exceptions.ExpiredSignatureError:
+        except jwt.exceptions.ExpiredSignatureError as e:
             logger.error('JWT token has expired.')
-            return None
+            raise e
         except jwt.exceptions.PyJWTError as e:
             logger.error(f'JWT token validation failed: {str(e)}')
-            return None
+            raise e
         except Exception as e:
             logger.error(f'Failed to validate token - {str(e)}')
-            return None
+            raise e
 
     def validate_access_token(self, access_token):
         # get UserInfo URI from OpenId Configuration

From 0e329a3943941ee4d9de31999578fb22c67d0bfa Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 15:50:42 -0500
Subject: [PATCH 28/32] only use access token, remove id token

---
 .../dataall/base/utils/api_handler_utils.py   |  7 -------
 .../custom_authorizer_lambda.py               | 15 ++++++++-------
 .../custom_authorizer/jwt_services.py         | 19 +++++++++++++------
 frontend/src/authentication/hooks/useToken.js | 12 +++---------
 frontend/src/modules/Catalog/views/Catalog.js |  8 ++++----
 .../contexts/RequestContext.js                |  2 +-
 frontend/src/services/hooks/useClient.js      | 11 +++++------
 tests_new/integration_tests/client.py         |  6 +++---
 8 files changed, 37 insertions(+), 43 deletions(-)

diff --git a/backend/dataall/base/utils/api_handler_utils.py b/backend/dataall/base/utils/api_handler_utils.py
index 1da2a6904..a43c100b8 100644
--- a/backend/dataall/base/utils/api_handler_utils.py
+++ b/backend/dataall/base/utils/api_handler_utils.py
@@ -33,13 +33,6 @@ def redact_creds(event):
 
     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
 
 
diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 08be6790a..57aba3def 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -25,23 +25,24 @@ 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']
-    access_token = incoming_event['headers']['accesskeyid']
-    if not auth_token or not access_token:
+    if not auth_token:
         raise Exception('Unauthorized. Missing identity or access JWT')
 
     # Validate User is Active with Proper Access Token
-    jwt_service.validate_access_token(access_token)
+    user_info = jwt_service.validate_access_token(auth_token)
 
     # Validate JWT
-    verified_claims = jwt_service.validate_jwt_token(auth_token)
+    verified_claims = jwt_service.validate_jwt_token(auth_token[7:])
     if not verified_claims:
         raise Exception('Unauthorized. Token is not valid')
     logger.debug(verified_claims)
 
     # Generate Allow Policy w/ Context
     effect = 'Allow'
+    verified_claims.update(user_info)
     policy = AuthServices.generate_policy(verified_claims, effect, incoming_event['methodArn'])
     logger.debug(f'Generated policy is {json.dumps(policy)}')
+    print(f'Generated policy is {json.dumps(policy)}')
     return policy
 
 
@@ -50,12 +51,12 @@ 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
-    id_token = ''
-    access_token = ''
+    #
+    access_token = 'Bearer eyJraWQiOiJtYTJ6SUxrbVMtQW1qZzZwVGtqZjhkN3JxY1FaNWE2eWtLS3dGQkFZckJBIiwidHlwIjoiYXBwbGljYXRpb25cL29rdGEtaW50ZXJuYWwtYXQrand0IiwiYWxnIjoiUlMyNTYifQ.eyJ2ZXIiOjEsImp0aSI6IkFULm9DSERGSHpVdGFUeTFDQXBENWF3amZRMEdyUzNPcEpyNE93czdnM3JKUXciLCJpc3MiOiJodHRwczovL2Rldi0zNzAxMTAxMC5va3RhLmNvbSIsImF1ZCI6Imh0dHBzOi8vZGV2LTM3MDExMDEwLm9rdGEuY29tIiwic3ViIjoibm9haHBhaWdAYW1hem9uLmNvbSIsImlhdCI6MTczMDkyNTQ1MiwiZXhwIjoxNzMwOTI5MDUyLCJjaWQiOiIwb2FkcndpcmVxcldoanFYaTVkNyIsInVpZCI6IjAwdWRydTNtNTZWS3hnWEtKNWQ3Iiwic2NwIjpbIm9wZW5pZCIsImVtYWlsIiwicHJvZmlsZSJdLCJhdXRoX3RpbWUiOjE3MzA5MjU0NTF9.uFZ123U7nbu6rN0L9WB2EZQTEZCnMcYOV_6uS4XRb8TAREcat-Kk88rLXONLwNWSaLaqGXOsr1tC1bd9FdTXyWG9WmVkihep8un_tmy1V410vEBtzXes6nqsr4-QZsx7csrWWtDetm4T7Smtl621z4isL8ePdYtkWe_2SELJjiOpr8qQ8pXMVEwMY8kiu-VuZHUXNnFGvrIRtNytsNzFVunbQxOX58uCq_J5eU7MRbj0tBAYqLXgXrj1iskb17uGHL4IqIWl1Te6qk05bLMZ9RrySEpyuCmYDPIgFpUZNiewLUNgPTNb4I8wrKycTpNfEEhTiLNxjo7QA5y2stTrFg'
     account_id = ''
     api_gw_id = ''
     event = {
-        'headers': {'Authorization': id_token, 'accesskeyid': access_token},
+        'headers': {'Authorization': access_token},
         'type': 'TOKEN',
         'methodArn': f'arn:aws:execute-api:us-east-1:{account_id}:{api_gw_id}/prod/POST/graphql/api',
     }
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index f13167881..10a57cdd4 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -10,18 +10,17 @@
 
 
 # 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
+# Only modification from default is to turn off verify_aud as Cognito Access Token does not provide this claim
 jwt_options = {
     'verify_signature': True,
-    'verify_aud': True,
+    'verify_aud': False,
     'verify_iat': True,
     'verify_exp': True,
     'verify_nbf': True,
     'verify_iss': True,
     'verify_sub': True,
     'verify_jti': True,
-    'verify_at_hash': True,
-    'require': ['aud', 'iat', 'exp', 'iss', 'sub', 'jti'],  # "nbf", "at_hash"
+    'require': ['iat', 'exp', 'iss', 'sub', 'jti'],
 }
 
 
@@ -38,7 +37,7 @@ def _fetch_openid_config(self, openid_config_path):
         response.raise_for_status()
         return response.json()
 
-    def validate_jwt_token(self, jwt_token):
+    def validate_jwt_token(self, jwt_token) -> dict:
         try:
             # get signing_key from JWT
             signing_key = self.jwks_client.get_signing_key_from_jwt(jwt_token)
@@ -53,6 +52,13 @@ def validate_jwt_token(self, jwt_token):
                 leeway=0,
                 options=jwt_options,
             )
+
+            # verify client_id if Cognito JWT
+            if os.environ['custom_auth_provider'] == 'Cognito' and payload['client_id'] != os.environ.get(
+                'custom_auth_client'
+            ):
+                raise Exception('Invalid Client ID in JWT Token')
+
             return payload
         except jwt.exceptions.ExpiredSignatureError as e:
             logger.error('JWT token has expired.')
@@ -64,9 +70,10 @@ def validate_jwt_token(self, jwt_token):
             logger.error(f'Failed to validate token - {str(e)}')
             raise e
 
-    def validate_access_token(self, access_token):
+    def validate_access_token(self, access_token) -> dict:
         # get UserInfo URI from OpenId Configuration
         user_info_url = self.openid_config.get('userinfo_endpoint')
         r = requests.get(user_info_url, headers={'Authorization': access_token})
         r.raise_for_status()
         logger.debug(r.json())
+        return r.json()
diff --git a/frontend/src/authentication/hooks/useToken.js b/frontend/src/authentication/hooks/useToken.js
index 267724060..8b17536db 100644
--- a/frontend/src/authentication/hooks/useToken.js
+++ b/frontend/src/authentication/hooks/useToken.js
@@ -7,14 +7,12 @@ export const useToken = () => {
   const dispatch = useDispatch();
   const auth = useAuth();
   const [token, setToken] = useState(null);
-  const [accessToken, setAccessToken] = useState(null);
   const fetchAuthToken = async () => {
     if (
       !process.env.REACT_APP_COGNITO_USER_POOL_ID &&
       process.env.REACT_APP_GRAPHQL_API.includes('localhost')
     ) {
       setToken('localToken');
-      setAccessToken('localAccessToken');
     } else {
       try {
         if (process.env.REACT_APP_CUSTOM_AUTH) {
@@ -22,19 +20,15 @@ export const useToken = () => {
             if (!auth.user) {
               await auth.signinSilent();
             }
-            const t = auth.user.id_token;
-            const at = auth.user.access_token;
+            const t = auth.user.access_token;
             setToken(t);
-            setAccessToken(at);
           } catch (error) {
             if (!auth) throw Error('User Token Not Found !');
           }
         } else {
           const session = await Auth.currentSession();
-          const t = await session.getIdToken().getJwtToken();
+          const t = await session.getAccessToken().getJwtToken();
           setToken(t);
-          const at = await session.getAccessToken().getJwtToken();
-          setAccessToken(at);
         }
       } catch (error) {
         auth.dispatch({
@@ -51,5 +45,5 @@ export const useToken = () => {
       );
     }
   });
-  return { token, accessToken };
+  return token;
 };
diff --git a/frontend/src/modules/Catalog/views/Catalog.js b/frontend/src/modules/Catalog/views/Catalog.js
index 4a7bf733c..d8d21b2a5 100644
--- a/frontend/src/modules/Catalog/views/Catalog.js
+++ b/frontend/src/modules/Catalog/views/Catalog.js
@@ -167,7 +167,7 @@ GlossaryFilter.propTypes = {
 };
 
 const Catalog = () => {
-  const { token, accessToken } = useToken();
+  const token = useToken();
   const { settings } = useSettings();
   const theme = useTheme();
   const classes = useStyles();
@@ -234,8 +234,8 @@ const Catalog = () => {
       url: transformedRequest.url,
       credentials: { token },
       headers: {
-        Authorization: token,
-        AccessKeyId: accessToken ? `Bearer ${accessToken}` : '',
+        Authorization: token ? `Bearer ${token}` : '',
+        AccessKeyId: 'None',
         SecretKey: 'None'
       }
     };
@@ -247,7 +247,7 @@ const Catalog = () => {
         : classes.darkListSearch
     );
   }, [settings.theme, classes]);
-  if (!token || !accessToken) {
+  if (!token) {
     return <CircularProgress />;
   }
 
diff --git a/frontend/src/reauthentication/contexts/RequestContext.js b/frontend/src/reauthentication/contexts/RequestContext.js
index ab6b55fb6..6986d5db3 100644
--- a/frontend/src/reauthentication/contexts/RequestContext.js
+++ b/frontend/src/reauthentication/contexts/RequestContext.js
@@ -53,7 +53,7 @@ export const RequestContextProvider = (props) => {
   const [requestInfo, setRequestInfo] = useState(null);
   const navigate = useNavigate();
   const client = useClient();
-  const { token } = useToken();
+  const token = useToken();
   const auth = useAuth();
   const { enqueueSnackbar } = useSnackbar();
   const storeRequestInfo = (info) => {
diff --git a/frontend/src/services/hooks/useClient.js b/frontend/src/services/hooks/useClient.js
index 7b41398ec..9e20d4619 100644
--- a/frontend/src/services/hooks/useClient.js
+++ b/frontend/src/services/hooks/useClient.js
@@ -28,7 +28,7 @@ const defaultOptions = {
 export const useClient = () => {
   const dispatch = useDispatch();
   const [client, setClient] = useState(null);
-  const { token, accessToken } = useToken();
+  const token = useToken();
   const auth = useAuth();
 
   const setReAuth = useCallback(
@@ -47,7 +47,6 @@ export const useClient = () => {
   useEffect(() => {
     const initClient = async () => {
       const t = token;
-      const at = accessToken;
       const httpLink = new HttpLink({
         uri: process.env.REACT_APP_GRAPHQL_API
       });
@@ -55,8 +54,8 @@ export const useClient = () => {
       const authLink = new ApolloLink((operation, forward) => {
         operation.setContext({
           headers: {
-            Authorization: t ? `${t}` : '',
-            AccessKeyId: at ? `Bearer ${at}` : '',
+            Authorization: t ? `Bearer ${t}` : '',
+            AccessKeyId: 'none',
             SecretKey: 'none'
           }
         });
@@ -95,9 +94,9 @@ export const useClient = () => {
       });
       setClient(apolloClient);
     };
-    if (token && accessToken) {
+    if (token) {
       initClient().catch((e) => console.error(e));
     }
-  }, [token, accessToken, dispatch]);
+  }, [token, dispatch]);
   return client;
 };
diff --git a/tests_new/integration_tests/client.py b/tests_new/integration_tests/client.py
index 0992b4b8a..5c8612519 100644
--- a/tests_new/integration_tests/client.py
+++ b/tests_new/integration_tests/client.py
@@ -20,7 +20,7 @@ class Client:
     def __init__(self, username, password):
         self.username = username
         self.password = password
-        self.id_token, self.access_token = self._get_jwt_tokens()
+        self.access_token = self._get_jwt_tokens()
 
     @retry(
         retry_on_exception=_retry_if_connection_error,
@@ -30,7 +30,7 @@ def __init__(self, username, password):
     )
     def query(self, query: str):
         graphql_endpoint = os.path.join(os.environ['API_ENDPOINT'], 'graphql', 'api')
-        headers = {'accesskeyid': f'Bearer {self.access_token}', 'SecretKey': 'none', 'Authorization': self.id_token}
+        headers = {'accesskeyid': 'none', 'SecretKey': 'none', 'Authorization': f'Bearer {self.access_token}'}
         r = requests.post(graphql_endpoint, json=query, headers=headers)
         if errors := r.json().get('errors'):
             raise GqlError(errors)
@@ -86,4 +86,4 @@ def _get_jwt_tokens(self):
         )
         print(token)
 
-        return token.get('id_token'), token.get('access_token')
+        return token.get('access_token')

From 5a0893be9e6505bad04addc3181268aad407691a Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 16:59:58 -0500
Subject: [PATCH 29/32] update exception str

---
 .../custom_authorizer/custom_authorizer_lambda.py             | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 57aba3def..b7c9e656e 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -26,7 +26,7 @@ def lambda_handler(incoming_event, context):
     logger.debug(incoming_event)
     auth_token = incoming_event['headers']['Authorization']
     if not auth_token:
-        raise Exception('Unauthorized. Missing identity or access JWT')
+        raise Exception('Unauthorized. Missing JWT')
 
     # Validate User is Active with Proper Access Token
     user_info = jwt_service.validate_access_token(auth_token)
@@ -52,7 +52,7 @@ def lambda_handler(incoming_event, context):
 if __name__ == '__main__':
     # for testing locally you can enter the JWT ID Token here
     #
-    access_token = 'Bearer eyJraWQiOiJtYTJ6SUxrbVMtQW1qZzZwVGtqZjhkN3JxY1FaNWE2eWtLS3dGQkFZckJBIiwidHlwIjoiYXBwbGljYXRpb25cL29rdGEtaW50ZXJuYWwtYXQrand0IiwiYWxnIjoiUlMyNTYifQ.eyJ2ZXIiOjEsImp0aSI6IkFULm9DSERGSHpVdGFUeTFDQXBENWF3amZRMEdyUzNPcEpyNE93czdnM3JKUXciLCJpc3MiOiJodHRwczovL2Rldi0zNzAxMTAxMC5va3RhLmNvbSIsImF1ZCI6Imh0dHBzOi8vZGV2LTM3MDExMDEwLm9rdGEuY29tIiwic3ViIjoibm9haHBhaWdAYW1hem9uLmNvbSIsImlhdCI6MTczMDkyNTQ1MiwiZXhwIjoxNzMwOTI5MDUyLCJjaWQiOiIwb2FkcndpcmVxcldoanFYaTVkNyIsInVpZCI6IjAwdWRydTNtNTZWS3hnWEtKNWQ3Iiwic2NwIjpbIm9wZW5pZCIsImVtYWlsIiwicHJvZmlsZSJdLCJhdXRoX3RpbWUiOjE3MzA5MjU0NTF9.uFZ123U7nbu6rN0L9WB2EZQTEZCnMcYOV_6uS4XRb8TAREcat-Kk88rLXONLwNWSaLaqGXOsr1tC1bd9FdTXyWG9WmVkihep8un_tmy1V410vEBtzXes6nqsr4-QZsx7csrWWtDetm4T7Smtl621z4isL8ePdYtkWe_2SELJjiOpr8qQ8pXMVEwMY8kiu-VuZHUXNnFGvrIRtNytsNzFVunbQxOX58uCq_J5eU7MRbj0tBAYqLXgXrj1iskb17uGHL4IqIWl1Te6qk05bLMZ9RrySEpyuCmYDPIgFpUZNiewLUNgPTNb4I8wrKycTpNfEEhTiLNxjo7QA5y2stTrFg'
+    access_token = ''
     account_id = ''
     api_gw_id = ''
     event = {

From 8b8eb1a74c0396775061c8ade092b3bcb36ed83c Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Wed, 6 Nov 2024 18:05:23 -0500
Subject: [PATCH 30/32] Clean Up

---
 .../custom_authorizer/custom_authorizer_lambda.py         | 1 +
 deploy/custom_resources/custom_authorizer/jwt_services.py | 8 +++++---
 tests_new/integration_tests/client.py                     | 2 --
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index b7c9e656e..210c11ddd 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -32,6 +32,7 @@ def lambda_handler(incoming_event, context):
     user_info = jwt_service.validate_access_token(auth_token)
 
     # Validate JWT
+    # Note: Removing the 7 Prefix Chars for 'Bearer ' from JWT
     verified_claims = jwt_service.validate_jwt_token(auth_token[7:])
     if not verified_claims:
         raise Exception('Unauthorized. Token is not valid')
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 10a57cdd4..080d3ee5b 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -54,9 +54,11 @@ def validate_jwt_token(self, jwt_token) -> dict:
             )
 
             # verify client_id if Cognito JWT
-            if os.environ['custom_auth_provider'] == 'Cognito' and payload['client_id'] != os.environ.get(
-                'custom_auth_client'
-            ):
+            if 'client_id' in payload and payload['client_id'] != os.environ.get('custom_auth_client'):
+                raise Exception('Invalid Client ID in JWT Token')
+
+            # verify cid for other IdPs
+            if 'cid' in payload and payload['cid'] != os.environ.get('custom_auth_client'):
                 raise Exception('Invalid Client ID in JWT Token')
 
             return payload
diff --git a/tests_new/integration_tests/client.py b/tests_new/integration_tests/client.py
index 5c8612519..bef02d39e 100644
--- a/tests_new/integration_tests/client.py
+++ b/tests_new/integration_tests/client.py
@@ -70,7 +70,6 @@ def _get_jwt_tokens(self):
             headers=headers,
             allow_redirects=False,
         )
-        print(r)
 
         r.raise_for_status()
 
@@ -84,6 +83,5 @@ def _get_jwt_tokens(self):
             code=code,
             include_client_id=True,
         )
-        print(token)
 
         return token.get('access_token')

From 86db3e8ddb6c48e787d647130168ce1be0a10504 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Thu, 7 Nov 2024 08:38:51 -0500
Subject: [PATCH 31/32] PR comments

---
 .../custom_authorizer/custom_authorizer_lambda.py         | 8 ++++----
 deploy/custom_resources/custom_authorizer/jwt_services.py | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
index 210c11ddd..47b9223e7 100644
--- a/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
+++ b/deploy/custom_resources/custom_authorizer/custom_authorizer_lambda.py
@@ -17,8 +17,8 @@
 Custom Lambda Authorizer is attached to the API Gateway. Check the deploy/stacks/lambda_api.py for more details on deployment
 """
 
-OPENID_CONFIG_PATH = os.path.join(os.environ.get('custom_auth_url', ''), '.well-known', 'openid-configuration')
-jwt_service = JWTServices(OPENID_CONFIG_PATH)
+OPENID_CONFIG_PATH = os.path.join(os.environ['custom_auth_url'], '.well-known', 'openid-configuration')
+JWT_SERVICE = JWTServices(OPENID_CONFIG_PATH)
 
 
 def lambda_handler(incoming_event, context):
@@ -29,11 +29,11 @@ def lambda_handler(incoming_event, context):
         raise Exception('Unauthorized. Missing JWT')
 
     # Validate User is Active with Proper Access Token
-    user_info = jwt_service.validate_access_token(auth_token)
+    user_info = JWT_SERVICE.validate_access_token(auth_token)
 
     # Validate JWT
     # Note: Removing the 7 Prefix Chars for 'Bearer ' from JWT
-    verified_claims = jwt_service.validate_jwt_token(auth_token[7:])
+    verified_claims = JWT_SERVICE.validate_jwt_token(auth_token[7:])
     if not verified_claims:
         raise Exception('Unauthorized. Token is not valid')
     logger.debug(verified_claims)
diff --git a/deploy/custom_resources/custom_authorizer/jwt_services.py b/deploy/custom_resources/custom_authorizer/jwt_services.py
index 080d3ee5b..c1f2f6a5c 100644
--- a/deploy/custom_resources/custom_authorizer/jwt_services.py
+++ b/deploy/custom_resources/custom_authorizer/jwt_services.py
@@ -47,7 +47,7 @@ def validate_jwt_token(self, jwt_token) -> dict:
                 jwt_token,
                 signing_key.key,
                 algorithms=['RS256', 'HS256'],
-                issuer=os.environ.get('custom_auth_url'),
+                issuer=os.environ['custom_auth_url'],
                 audience=os.environ.get('custom_auth_client'),
                 leeway=0,
                 options=jwt_options,

From 8db81a17944f2e2eb2a3dbefa2519e1475f6e448 Mon Sep 17 00:00:00 2001
From: Noah Paige <noahpaig@amazon.com>
Date: Thu, 7 Nov 2024 08:54:34 -0500
Subject: [PATCH 32/32] fix IDP_DOMAIN_URL setting

---
 deploy/stacks/pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deploy/stacks/pipeline.py b/deploy/stacks/pipeline.py
index a14ca98ba..a676f07d5 100644
--- a/deploy/stacks/pipeline.py
+++ b/deploy/stacks/pipeline.py
@@ -716,7 +716,7 @@ def set_approval_tests_stage(
                                     'aws sts get-caller-identity --profile buildprofile',
                                     f'export COGNITO_CLIENT=$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/cognito/appclient --profile buildprofile --output text --query "Parameter.Value")',
                                     f'export API_ENDPOINT=$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/apiGateway/backendUrl --profile buildprofile --output text --query "Parameter.Value")',
-                                    f'export IDP_DOMAIN_URL=https://$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/cognito/domain --profile buildprofile --output text --query "Parameter.Value").auth.us-east-1.amazoncognito.com',
+                                    f'export IDP_DOMAIN_URL=https://$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/cognito/domain --profile buildprofile --output text --query "Parameter.Value").auth.{target_env["region"]}.amazoncognito.com',
                                     f'export DATAALL_DOMAIN_URL=https://$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/CloudfrontDistributionDomainName --profile buildprofile --output text --query "Parameter.Value")',
                                     f'export TESTDATA=$(aws ssm get-parameter --name /dataall/{target_env["envname"]}/testdata --profile buildprofile --output text --query "Parameter.Value")',
                                     f'export ENVNAME={target_env["envname"]}',