feat(roles and states): Allow direct attachment of roles and states t…

…o federated cognito users (#479)

* Allow direct attachment of roles and states with no overwrite

* add ZZ state

src/libs/cognito-lib.ts

@@ -1,5 +1,13 @@
-import { CognitoIdentityProviderClient, AdminCreateUserCommand, AdminSetUserPasswordCommand, AdminUpdateUserAttributesCommand } from "@aws-sdk/client-cognito-identity-provider";
-const client = new CognitoIdentityProviderClient({ region: process.env.region });
+import {
+  CognitoIdentityProviderClient,
+  AdminCreateUserCommand,
+  AdminSetUserPasswordCommand,
+  AdminUpdateUserAttributesCommand,
+  AdminGetUserCommand,
+} from "@aws-sdk/client-cognito-identity-provider";
+const client = new CognitoIdentityProviderClient({
+  region: process.env.region,
+});
 
 export async function createUser(params: any): Promise<void> {
   try {
@@ -25,9 +33,91 @@ export async function setPassword(params: any): Promise<void> {
 
 export async function updateUserAttributes(params: any): Promise<void> {
   try {
+    // Fetch existing user attributes
+    const getUserCommand = new AdminGetUserCommand({
+      UserPoolId: params.UserPoolId,
+      Username: params.Username,
+    });
+    const user = await client.send(getUserCommand);
+
+    // Check for existing "custom:cms-roles"
+    const cmsRolesAttribute = user.UserAttributes?.find(
+      (attr) => attr.Name === "custom:cms-roles",
+    );
+    const existingRoles =
+      cmsRolesAttribute && cmsRolesAttribute.Value
+        ? cmsRolesAttribute.Value.split(",")
+        : [];
+
+    // Check for existing "custom:state"
+    const stateAttribute = user.UserAttributes?.find(
+      (attr) => attr.Name === "custom:state",
+    );
+    const existingStates =
+      stateAttribute && stateAttribute.Value
+        ? stateAttribute.Value.split(",")
+        : [];
+
+    // Prepare for updating user attributes
+    let attributeData: any = {
+      UserPoolId: params.UserPoolId,
+      Username: params.Username,
+      UserAttributes: params.UserAttributes,
+    };
+
+    // Ensure "onemac-micro-super" is preserved
+    if (existingRoles.includes("onemac-micro-super")) {
+      const rolesIndex = attributeData.UserAttributes.findIndex(
+        (attr: any) => attr.Name === "custom:cms-roles",
+      );
+      if (rolesIndex !== -1) {
+        // Only merge if new roles are not empty
+        let newRoles = attributeData.UserAttributes[rolesIndex].Value
+          ? new Set(
+              attributeData.UserAttributes[rolesIndex].Value.split(",").concat(
+                "onemac-micro-super",
+              ),
+            )
+          : new Set(["onemac-micro-super"]); // Ensure "onemac-micro-super" is always included
+        attributeData.UserAttributes[rolesIndex].Value =
+          Array.from(newRoles).join(",");
+      } else {
+        // Add "custom:cms-roles" with "onemac-micro-super"
+        attributeData.UserAttributes.push({
+          Name: "custom:cms-roles",
+          Value: "onemac-micro-super",
+        });
+      }
+    }
+
+    // Ensure "ZZ" state is preserved
+    if (existingStates.includes("ZZ")) {
+      const stateIndex = attributeData.UserAttributes.findIndex(
+        (attr: any) => attr.Name === "custom:state",
+      );
+      if (stateIndex !== -1) {
+        // Only merge if new states are not empty
+        let newStates = attributeData.UserAttributes[stateIndex].Value
+          ? new Set(
+              attributeData.UserAttributes[stateIndex].Value.split(",").concat(
+                "ZZ",
+              ),
+            )
+          : new Set(["ZZ"]); // Ensure "ZZ" is always included
+        attributeData.UserAttributes[stateIndex].Value =
+          Array.from(newStates).join(",");
+      } else {
+        // Add "custom:state" with "ZZ"
+        attributeData.UserAttributes.push({
+          Name: "custom:state",
+          Value: "ZZ",
+        });
+      }
+    }
+
     // Update the user's attributes
-    const command = new AdminUpdateUserAttributesCommand(params);
-    await client.send(command);
+    const updateCommand = new AdminUpdateUserAttributesCommand(attributeData);
+    await client.send(updateCommand);
     console.log(`Attributes for user ${params.Username} updated successfully.`);
   } catch (error) {
     console.error("Error updating user's attributes:", error);

src/packages/shared-types/states.ts

@@ -58,5 +58,6 @@ export const STATE_CODES = [
   "WV",
   "WI",
   "WY",
+  "ZZ",
 ] as const;
-export type StateCode = typeof STATE_CODES[number];
+export type StateCode = (typeof STATE_CODES)[number];

src/services/auth/handlers/postAuth.ts

@@ -1,6 +1,6 @@
 import { Handler } from "aws-lambda";
 import * as cognitolib from "../../../libs/cognito-lib";
-import { UserRoles, STATE_ROLES } from "shared-types"
+import { UserRoles, STATE_ROLES } from "shared-types";
 
 if (!process.env.apiKey) {
   throw "ERROR:  process.env.apiKey is required,";
@@ -13,68 +13,71 @@ if (!process.env.apiEndpoint) {
 const apiEndpoint: string = process.env.apiEndpoint;
 
 export const handler: Handler = async (event, context) => {
-  console.log(JSON.stringify(event,null,2));
+  console.log(JSON.stringify(event, null, 2));
   const { request, response } = event;
   const { userAttributes } = request;
 
-  if(!userAttributes.identities){
+  if (!userAttributes.identities) {
     console.log("User is not managed externally.  Nothing to do.");
   } else {
     console.log("Getting user attributes from external API");
-  
+
     try {
       const username = userAttributes["custom:username"]; // This is the four letter IDM username
-      const response = await fetch(`${apiEndpoint}/api/v1/authz/id/all?userId=${username}`, {
-        method: 'GET',
-        headers: {
-          'Content-Type': 'application/json',
-          'x-api-key': apiKey,
-        }
-      });
+      const response = await fetch(
+        `${apiEndpoint}/api/v1/authz/id/all?userId=${username}`,
+        {
+          method: "GET",
+          headers: {
+            "Content-Type": "application/json",
+            "x-api-key": apiKey,
+          },
+        },
+      );
       if (!response.ok) {
         console.log(response);
-        throw new Error(`Network response was not ok.  Response was ${response.status}: ${response.statusText}`);
+        throw new Error(
+          `Network response was not ok.  Response was ${response.status}: ${response.statusText}`,
+        );
       }
       let data = await response.json();
-      console.log(JSON.stringify(data, null, 2))
-      let roleArray :string[] = [];
-      let stateArray :string[] = [];
-      data.userProfileAppRoles.userRolesInfoList.forEach((element :any) => {
+      console.log(JSON.stringify(data, null, 2));
+      let roleArray: string[] = [];
+      let stateArray: string[] = [];
+      data.userProfileAppRoles.userRolesInfoList.forEach((element: any) => {
         let role = element.roleName;
-        if(Object.values(UserRoles).includes(role)){
+        if (Object.values(UserRoles).includes(role)) {
           roleArray.push(role);
-          if(STATE_ROLES.includes(role)){
-            element.roleAttributes.forEach((attr :any) => {
-              if(attr.name = "State/Territory"){
+          if (STATE_ROLES.includes(role)) {
+            element.roleAttributes.forEach((attr: any) => {
+              if ((attr.name = "State/Territory")) {
                 stateArray.push(attr.value);
               }
-            })
+            });
           }
         }
       });
 
-      let attributeData :any = {
+      let attributeData: any = {
         Username: event.userName,
         UserPoolId: event.userPoolId,
         UserAttributes: [
           {
-            "Name": "custom:cms-roles",
-            "Value": roleArray.join()
+            Name: "custom:cms-roles",
+            Value: roleArray.join(),
           },
           {
-            "Name": "custom:state",
-            "Value": stateArray.join()
-          }
+            Name: "custom:state",
+            Value: stateArray.join(),
+          },
         ],
       };
       console.log("Putting user attributes as: ");
-      console.log(JSON.stringify(attributeData,null,2));
+      console.log(JSON.stringify(attributeData, null, 2));
       await cognitolib.updateUserAttributes(attributeData);
-
     } catch (error) {
       console.error("Error performing post auth:", error);
     }
-
   }
   return event;
 };

src/services/auth/libs/users.json

@@ -144,6 +144,35 @@
       }
     ]
   },
+  {
+    "username": "[email protected]",
+    "attributes": [
+      {
+        "Name": "email",
+        "Value": "[email protected]"
+      },
+      {
+        "Name": "given_name",
+        "Value": "Superduper"
+      },
+      {
+        "Name": "family_name",
+        "Value": "Paratrooper"
+      },
+      {
+        "Name": "email_verified",
+        "Value": "true"
+      },
+      {
+        "Name": "custom:state",
+        "Value": ""
+      },
+      {
+        "Name": "custom:cms-roles",
+        "Value": "onemac-micro-reviewer,onemac-micro-super"
+      }
+    ]
+  },
   {
     "username": "[email protected]",
     "attributes": [
@@ -368,4 +397,4 @@
       }
     ]
   }
-]
+]

src/services/auth/serverless.yml

@@ -20,6 +20,7 @@ provider:
       statements:
         - Effect: "Allow"
           Action:
+            - cognito-idp:AdminGetUser
             - cognito-idp:AdminCreateUser
             - cognito-idp:AdminSetUserPassword
             - cognito-idp:AdminUpdateUserAttributes