add pretoken generation lambda trigger

aws-samples · Nov 6, 2024 · 08a33ab · 08a33ab
1 parent bf85bd4
commit 08a33ab
Show file tree

Hide file tree

Showing 26 changed files with 527 additions and 640 deletions.
diff --git a/amplify/backend/api/team/schema.graphql b/amplify/backend/api/team/schema.graphql
@@ -265,9 +265,6 @@ type Query {
   getMgmtPermissions: MgmtPs
     @function(name: "teamgetMgmtAccountDetails-${env}")
     @auth(rules: [{ allow: private }])
-  getGroups: Groups
-    @function(name: "teamgetGroups-${env}")
-    @auth(rules: [{ allow: private }])
   getIdCGroups: [IdCGroups]
     @function(name: "teamgetIdCGroups-${env}")
     @auth(rules: [{ allow: private }])

diff --git a/amplify/backend/auth/team06dbb7fc/cli-inputs.json b/amplify/backend/auth/team06dbb7fc/cli-inputs.json
@@ -41,7 +41,11 @@
         "usernameAttributes": [
             "email"
         ],
-        "triggers": {},
+        "triggers": {
+        "PreTokenGeneration": [
+          "alter-claims"
+        ]
+      },
         "userPoolGroupList": [
             "Auditors",
             "Admin"
@@ -62,18 +66,29 @@
             ]
         },
         "breakCircularDependency": true,
-        "dependsOn": [],
+        "dependsOn": [
+        {
+          "category": "function",
+          "resourceName": "team06dbb7fcPreTokenGeneration",
+          "triggerProvider": "Cognito",
+          "attributes": [
+            "Arn",
+            "Name"
+          ]
+        }
+      ],
         "hostedUI": true,
         "parentStack": {
             "Ref": "AWS::StackId"
         },
         "permissions": [],
-        "hostedUIDomainName": "d11r99we6v7d2t",
+        "hostedUIDomainName": "d1s5ncogc730zk",
         "authProvidersUserPool": [],
         "hostedUIProviderMeta": "[]",
         "authProviders": [],
         "hostedUIProviderCreds": "[]",
         "adminQueryGroup": "Admin",
-        "oAuthMetadata": "{\"AllowedOAuthFlows\":[\"code\"],\"AllowedOAuthScopes\":[\"phone\",\"email\",\"openid\",\"profile\",\"aws.cognito.signin.user.admin\"],\"CallbackURLs\":[\"https://main.d11r99we6v7d2t.amplifyapp.com/\"],\"LogoutURLs\":[\"https://main.d11r99we6v7d2t.amplifyapp.com/\"]}"
+        "authTriggerConnections": "[\n  {\n    \"triggerType\": \"PreTokenGeneration\",\n    \"lambdaFunctionName\": \"team06dbb7fcPreTokenGeneration\"\n  }\n]",
+        "oAuthMetadata": "{\"AllowedOAuthFlows\":[\"code\"],\"AllowedOAuthScopes\":[\"phone\",\"email\",\"openid\",\"profile\",\"aws.cognito.signin.user.admin\"],\"CallbackURLs\":[\"https://main.d1s5ncogc730zk.amplifyapp.com/\"],\"LogoutURLs\":[\"https://main.d1s5ncogc730zk.amplifyapp.com/\"]}"
     }
 }
diff --git a/amplify/backend/backend-config.json b/amplify/backend/backend-config.json
@@ -36,7 +36,17 @@
   "auth": {
     "team06dbb7fc": {
       "customAuth": false,
-      "dependsOn": [],
+      "dependsOn": [
+        {
+          "attributes": [
+            "Arn",
+            "Name"
+          ],
+          "category": "function",
+          "resourceName": "team06dbb7fcPreTokenGeneration",
+          "triggerProvider": "Cognito"
+        }
+      ],
       "frontendAuthConfig": {
         "mfaConfiguration": "OFF",
         "mfaTypes": [
@@ -119,6 +129,21 @@
     }
   },
   "function": {
+    "team06dbb7fcPreTokenGeneration": {
+      "build": true,
+      "dependsOn": [
+        {
+          "attributes": [
+            "GraphQLAPIIdOutput",
+            "GraphQLAPIEndpointOutput"
+          ],
+          "category": "api",
+          "resourceName": "team"
+        }
+      ],
+      "providerPlugin": "awscloudformation",
+      "service": "Lambda"
+    },
     "teamGetPermissionSets": {
       "build": true,
       "dependsOn": [
@@ -264,35 +289,6 @@
       "providerPlugin": "awscloudformation",
       "service": "Lambda"
     },
-    "teamgetGroups": {
-      "build": true,
-      "dependsOn": [
-        {
-          "attributes": [
-            "GraphQLAPIIdOutput",
-            "GraphQLAPIEndpointOutput"
-          ],
-          "category": "api",
-          "resourceName": "team"
-        },
-        {
-          "attributes": [
-            "UserPoolId"
-          ],
-          "category": "auth",
-          "resourceName": "team06dbb7fc"
-        },
-        {
-          "attributes": [
-            "Arn"
-          ],
-          "category": "function",
-          "resourceName": "teamapplicationboto3layer"
-        }
-      ],
-      "providerPlugin": "awscloudformation",
-      "service": "Lambda"
-    },
     "teamgetIdCGroups": {
       "build": true,
       "providerPlugin": "awscloudformation",
@@ -410,6 +406,22 @@
     }
   },
   "parameters": {
+    "AMPLIFY_function_teamGetPermissionSets_deploymentBucketName": {
+      "usedBy": [
+        {
+          "category": "function",
+          "resourceName": "teamGetPermissionSets"
+        }
+      ]
+    },
+    "AMPLIFY_function_teamGetPermissionSets_s3Key": {
+      "usedBy": [
+        {
+          "category": "function",
+          "resourceName": "teamGetPermissionSets"
+        }
+      ]
+    },
     "AMPLIFY_function_teamListGroups_deploymentBucketName": {
       "usedBy": [
         {
@@ -442,6 +454,22 @@
         }
       ]
     },
+    "AMPLIFY_function_teamPublishOUs_deploymentBucketName": {
+      "usedBy": [
+        {
+          "category": "function",
+          "resourceName": "teamPublishOUs"
+        }
+      ]
+    },
+    "AMPLIFY_function_teamPublishOUs_s3Key": {
+      "usedBy": [
+        {
+          "category": "function",
+          "resourceName": "teamPublishOUs"
+        }
+      ]
+    },
     "AMPLIFY_function_teamRouter_deploymentBucketName": {
       "usedBy": [
         {
@@ -522,22 +550,6 @@
         }
       ]
     },
-    "AMPLIFY_function_teamgetGroups_deploymentBucketName": {
-      "usedBy": [
-        {
-          "category": "function",
-          "resourceName": "teamgetGroups"
-        }
-      ]
-    },
-    "AMPLIFY_function_teamgetGroups_s3Key": {
-      "usedBy": [
-        {
-          "category": "function",
-          "resourceName": "teamgetGroups"
-        }
-      ]
-    },
     "AMPLIFY_function_teamgetIdCGroups_deploymentBucketName": {
       "usedBy": [
         {

diff --git a/amplify/backend/custom/cloudtrailLake/parameters.json b/amplify/backend/custom/cloudtrailLake/parameters.json
@@ -1,3 +1,3 @@
 {
-    "CloudTrailAuditLogs": "arn:aws:cloudtrail:us-east-1:843551180572:eventdatastore/c2b1db81-d0ab-4857-a8d4-d85170209f9f"
+    "CloudTrailAuditLogs": "read_write"
 }
diff --git a/...fy/backend/function/teamgetGroups/Pipfile → ...on/team06dbb7fcPreTokenGeneration/Pipfile b/...fy/backend/function/teamgetGroups/Pipfile → ...on/team06dbb7fcPreTokenGeneration/Pipfile
@@ -9,4 +9,4 @@ verify_ssl = true
 src = {editable = true, path = "./src"}
 
 [requires]
-python_version = "3.9"
+python_version = "3.8"
diff --git a/...kend/function/teamgetGroups/amplify.state → ...m06dbb7fcPreTokenGeneration/amplify.state b/...kend/function/teamgetGroups/amplify.state → ...m06dbb7fcPreTokenGeneration/amplify.state
diff --git a/...nction/teamgetGroups/custom-policies.json → ...fcPreTokenGeneration/custom-policies.json b/...nction/teamgetGroups/custom-policies.json → ...fcPreTokenGeneration/custom-policies.json
diff --git a/amplify/backend/function/team06dbb7fcPreTokenGeneration/function-parameters.json b/amplify/backend/function/team06dbb7fcPreTokenGeneration/function-parameters.json
@@ -0,0 +1,10 @@
+{
+  "permissions": {
+    "api": {
+      "team": [
+        "Query"
+      ]
+    }
+  },
+  "lambdaLayers": []
+}
diff --git a/amplify/backend/function/team06dbb7fcPreTokenGeneration/parameters.json b/amplify/backend/function/team06dbb7fcPreTokenGeneration/parameters.json
@@ -0,0 +1,4 @@
+{
+    "teamAdminGroup": "TEAM-admins",
+    "teamAuditorGroup": "TEAM-auditors"
+}
diff --git a/...end/function/teamgetGroups/src/event.json → ...06dbb7fcPreTokenGeneration/src/event.json b/...end/function/teamgetGroups/src/event.json → ...06dbb7fcPreTokenGeneration/src/event.json
diff --git a/...ckend/function/teamgetGroups/src/index.py → ...am06dbb7fcPreTokenGeneration/src/index.py b/...ckend/function/teamgetGroups/src/index.py → ...am06dbb7fcPreTokenGeneration/src/index.py
@@ -1,12 +1,12 @@
-# © 2023 Amazon Web Services, Inc. or its affiliates. All Rights Reserved.
+# © 2024 Amazon Web Services, Inc. or its affiliates. All Rights Reserved.
 # This AWS Content is provided subject to the terms of the AWS Customer Agreement available at
 # http: // aws.amazon.com/agreement or other written agreement between Customer and either
 # Amazon Web Services, Inc. or Amazon Web Services EMEA SARL or both.
 import os
 from botocore.exceptions import ClientError
 import boto3
+import json
 
-user_pool_id = os.getenv("AUTH_AWSPIM06DBB7FC_USERPOOLID")
 team_admin_group = os.getenv("TEAM_ADMIN_GROUP")
 team_auditor_group = os.getenv("TEAM_AUDITOR_GROUP")
 settings_table_name = os.getenv("SETTINGS_TABLE_NAME")
@@ -31,31 +31,6 @@ def get_team_groups():
         print(f"Error retrieving TEAM settings from database: {e}")
     return team_admin_group, team_auditor_group
 
-def add_user_to_group(username, groupname):
-    client = boto3.client('cognito-idp')
-    try:
-        response = client.admin_add_user_to_group(
-            UserPoolId=user_pool_id,
-            Username=username,
-            GroupName=groupname
-        )
-        print(f"user {username} added to {groupname} group")
-    except ClientError as e:
-        print(e.response['Error']['Message'])
-
-
-def remove_user_from_group(username, groupname):
-    client = boto3.client('cognito-idp')
-    try:
-        response = client.admin_remove_user_from_group(
-            UserPoolId=user_pool_id,
-            Username=username,
-            GroupName=groupname
-        )
-        print(f"user {username} removed from {groupname} group")
-    except ClientError as e:
-        print(e.response['Error']['Message'])
-
 
 def get_identity_store_id():
     client = boto3.client('sso-admin')
@@ -123,29 +98,35 @@ def list_idc_group_membership(userId):
 def handler(event, context):
     team_admin_group, team_auditor_group = get_team_groups()
 
-    user = event["identity"]["username"]
-    # Strip idc prefix
-    username = user.removeprefix("idc_")
-    userId = get_user(username)
+    user = event["userName"].split("_", 1)[1]
+    userId = get_user(user)
     admin = get_group(team_admin_group)
     auditor = get_group(team_auditor_group)
     groups = []
-    groupIds = []
+    groupIds = str()
 
     groupData = list_idc_group_membership(userId)
-
+    
     for group in groupData:
-        groupIds.append(group["GroupId"])
-        if group['GroupId'] == admin:
-            add_user_to_group(user, "Admin")
+        groupIds += group["GroupId"] + ","
+        if group["GroupId"] == admin:
+            # add_user_to_group(user, "Admin")
             groups.append("Admin")
-        elif group['GroupId'] == auditor:
-            add_user_to_group(user, "Auditors")
+        elif group["GroupId"] == auditor:
+            # add_user_to_group(user, "Auditors")
             groups.append("Auditors")
 
-    if "Admin" not in groups:
-        remove_user_from_group(user, "Admin")
-    elif "Auditors" not in groups:
-        remove_user_from_group(user, "Auditors")
+    event["response"] = {
+        "claimsOverrideDetails": {
+            "claimsToAddOrOverride": {
+                "userId": userId,
+                "groupIds": groupIds,
+                "groups": ",".join(groups)
+            },
+            "groupOverrideDetails": {
+                "groupsToOverride": groups,
+            },
+        }
+    }
 
-    return {"groups": groups, "userId": userId, "groupIds": groupIds}
+    return event
diff --git a/...ckend/function/teamgetGroups/src/setup.py → ...am06dbb7fcPreTokenGeneration/src/setup.py b/...ckend/function/teamgetGroups/src/setup.py → ...am06dbb7fcPreTokenGeneration/src/setup.py