Finish preliminary AWS IAM Auth method

2026-05-02 03:02:03 -04:00 · 2024-05-02 22:42:02 -07:00
parent 5dd144b97b
commit 9c002ad645
11 changed files with 575 additions and 69 deletions
--- a/frontend/src/hooks/api/identities/constants.tsx
+++ b/frontend/src/hooks/api/identities/constants.tsx
@@ -1,5 +1,6 @@
 import { IdentityAuthMethod } from "./enums";

 export const identityAuthToNameMap: { [I in IdentityAuthMethod]: string } = {
-  [IdentityAuthMethod.UNIVERSAL_AUTH]: "Universal Auth"
+  [IdentityAuthMethod.UNIVERSAL_AUTH]: "Universal Auth",
+  [IdentityAuthMethod.AWS_IAM_AUTH]: "AWS IAM Auth"
 };
--- a/frontend/src/hooks/api/identities/enums.tsx
+++ b/frontend/src/hooks/api/identities/enums.tsx
@@ -1,3 +1,4 @@
 export enum IdentityAuthMethod {
-  UNIVERSAL_AUTH = "universal-auth"
+  UNIVERSAL_AUTH = "universal-auth",
+  AWS_IAM_AUTH = "aws-iam-auth"
 }
--- a/frontend/src/hooks/api/identities/index.tsx
+++ b/frontend/src/hooks/api/identities/index.tsx
@@ -1,12 +1,16 @@
 export { identityAuthToNameMap } from "./constants";
 export { IdentityAuthMethod } from "./enums";
 export {
+  useAddIdentityAwsIamAuth,
  useAddIdentityUniversalAuth,
  useCreateIdentity,
  useCreateIdentityUniversalAuthClientSecret,
  useDeleteIdentity,
  useRevokeIdentityUniversalAuthClientSecret,
  useUpdateIdentity,
-  useUpdateIdentityUniversalAuth
-} from "./mutations";
-export { useGetIdentityUniversalAuth, useGetIdentityUniversalAuthClientSecrets } from "./queries";
+  useUpdateIdentityAwsIamAuth,
+  useUpdateIdentityUniversalAuth} from "./mutations";
+export {
+  useGetIdentityAwsIamAuth,
+  useGetIdentityUniversalAuth,
+  useGetIdentityUniversalAuthClientSecrets} from "./queries";
--- a/frontend/src/hooks/api/identities/mutations.tsx
+++ b/frontend/src/hooks/api/identities/mutations.tsx
@@ -5,6 +5,7 @@ import { apiRequest } from "@app/config/request";
 import { organizationKeys } from "../organization/queries";
 import { identitiesKeys } from "./queries";
 import {
+  AddIdentityAwsIamAuthDTO,
  AddIdentityUniversalAuthDTO,
  ClientSecretData,
  CreateIdentityDTO,
@@ -13,10 +14,11 @@ import {
  DeleteIdentityDTO,
  DeleteIdentityUniversalAuthClientSecretDTO,
  Identity,
+  IdentityAwsIamAuth,
  IdentityUniversalAuth,
+  UpdateIdentityAwsIamAuthDTO,
  UpdateIdentityDTO,
-  UpdateIdentityUniversalAuthDTO
-} from "./types";
+  UpdateIdentityUniversalAuthDTO} from "./types";

 export const useCreateIdentity = () => {
  const queryClient = useQueryClient();
@@ -169,3 +171,74 @@ export const useRevokeIdentityUniversalAuthClientSecret = () => {
    }
  });
 };
+
+export const useAddIdentityAwsIamAuth = () => {
+  const queryClient = useQueryClient();
+  return useMutation<IdentityAwsIamAuth, {}, AddIdentityAwsIamAuthDTO>({
+    mutationFn: async ({
+      identityId,
+      stsEndpoint,
+      allowedPrincipalArns,
+      allowedAccountIds,
+      accessTokenTTL,
+      accessTokenMaxTTL,
+      accessTokenNumUsesLimit,
+      accessTokenTrustedIps
+    }) => {
+      const {
+        data: { identityAwsIamAuth }
+      } = await apiRequest.post<{ identityAwsIamAuth: IdentityAwsIamAuth }>(
+        `/api/v1/auth/aws-iam-auth/identities/${identityId}`,
+        {
+          stsEndpoint,
+          allowedPrincipalArns,
+          allowedAccountIds,
+          accessTokenTTL,
+          accessTokenMaxTTL,
+          accessTokenNumUsesLimit,
+          accessTokenTrustedIps
+        }
+      );
+
+      return identityAwsIamAuth;
+    },
+    onSuccess: (_, { organizationId }) => {
+      queryClient.invalidateQueries(organizationKeys.getOrgIdentityMemberships(organizationId));
+    }
+  });
+};
+
+export const useUpdateIdentityAwsIamAuth = () => {
+  const queryClient = useQueryClient();
+  return useMutation<IdentityAwsIamAuth, {}, UpdateIdentityAwsIamAuthDTO>({
+    mutationFn: async ({
+      identityId,
+      stsEndpoint,
+      allowedPrincipalArns,
+      allowedAccountIds,
+      accessTokenTTL,
+      accessTokenMaxTTL,
+      accessTokenNumUsesLimit,
+      accessTokenTrustedIps
+    }) => {
+      const {
+        data: { identityAwsIamAuth }
+      } = await apiRequest.patch<{ identityAwsIamAuth: IdentityAwsIamAuth }>(
+        `/api/v1/auth/aws-iam-auth/identities/${identityId}`,
+        {
+          stsEndpoint,
+          allowedPrincipalArns,
+          allowedAccountIds,
+          accessTokenTTL,
+          accessTokenMaxTTL,
+          accessTokenNumUsesLimit,
+          accessTokenTrustedIps
+        }
+      );
+      return identityAwsIamAuth;
+    },
+    onSuccess: (_, { organizationId }) => {
+      queryClient.invalidateQueries(organizationKeys.getOrgIdentityMemberships(organizationId));
+    }
+  });
+};
--- a/frontend/src/hooks/api/identities/queries.tsx
+++ b/frontend/src/hooks/api/identities/queries.tsx
@@ -2,27 +2,26 @@ import { useQuery } from "@tanstack/react-query";

 import { apiRequest } from "@app/config/request";

-import { ClientSecretData, IdentityUniversalAuth } from "./types";
+import { ClientSecretData, IdentityAwsIamAuth,IdentityUniversalAuth } from "./types";

 export const identitiesKeys = {
  getIdentityUniversalAuth: (identityId: string) =>
    [{ identityId }, "identity-universal-auth"] as const,
  getIdentityUniversalAuthClientSecrets: (identityId: string) =>
-    [{ identityId }, "identity-universal-auth-client-secrets"] as const
+    [{ identityId }, "identity-universal-auth-client-secrets"] as const,
+  getIdentityAwsIamAuth: (identityId: string) => [{ identityId }, "identity-aws-iam-auth"] as const
 };

 export const useGetIdentityUniversalAuth = (identityId: string) => {
  return useQuery({
+    enabled: Boolean(identityId),
    queryKey: identitiesKeys.getIdentityUniversalAuth(identityId),
    queryFn: async () => {
-      if (identityId === "") throw new Error("Identity ID is required");
-
      const {
        data: { identityUniversalAuth }
      } = await apiRequest.get<{ identityUniversalAuth: IdentityUniversalAuth }>(
        `/api/v1/auth/universal-auth/identities/${identityId}`
      );
-
      return identityUniversalAuth;
    }
  });
@@ -30,17 +29,30 @@ export const useGetIdentityUniversalAuth = (identityId: string) => {

 export const useGetIdentityUniversalAuthClientSecrets = (identityId: string) => {
  return useQuery({
+    enabled: Boolean(identityId),
    queryKey: identitiesKeys.getIdentityUniversalAuthClientSecrets(identityId),
    queryFn: async () => {
-      if (identityId === "") return [];
-
      const {
        data: { clientSecretData }
      } = await apiRequest.get<{ clientSecretData: ClientSecretData[] }>(
        `/api/v1/auth/universal-auth/identities/${identityId}/client-secrets`
      );
-
      return clientSecretData;
    }
  });
 };
+
+export const useGetIdentityAwsIamAuth = (identityId: string) => {
+  return useQuery({
+    enabled: Boolean(identityId),
+    queryKey: identitiesKeys.getIdentityAwsIamAuth(identityId),
+    queryFn: async () => {
+      const {
+        data: { identityAwsIamAuth }
+      } = await apiRequest.get<{ identityAwsIamAuth: IdentityAwsIamAuth }>(
+        `/api/v1/auth/aws-iam-auth/identities/${identityId}`
+      );
+      return identityAwsIamAuth;
+    }
+  });
+};
--- a/frontend/src/hooks/api/identities/types.ts
+++ b/frontend/src/hooks/api/identities/types.ts
@@ -38,19 +38,19 @@ export type IdentityMembership = {
      customRoleSlug: string;
    } & (
      | {
-        isTemporary: false;
-        temporaryRange: null;
-        temporaryMode: null;
-        temporaryAccessEndTime: null;
-        temporaryAccessStartTime: null;
-      }
+          isTemporary: false;
+          temporaryRange: null;
+          temporaryMode: null;
+          temporaryAccessEndTime: null;
+          temporaryAccessStartTime: null;
+        }
      | {
-        isTemporary: true;
-        temporaryRange: string;
-        temporaryMode: string;
-        temporaryAccessEndTime: string;
-        temporaryAccessStartTime: string;
-      }
+          isTemporary: true;
+          temporaryRange: string;
+          temporaryMode: string;
+          temporaryAccessEndTime: string;
+          temporaryAccessStartTime: string;
+        }
    )
  >;
  createdAt: string;
@@ -113,6 +113,45 @@ export type UpdateIdentityUniversalAuthDTO = {
  }[];
 };

+export type IdentityAwsIamAuth = {
+  identityId: string;
+  stsEndpoint: string;
+  allowedPrincipalArns: string;
+  allowedAccountIds: string;
+  accessTokenTTL: number;
+  accessTokenMaxTTL: number;
+  accessTokenNumUsesLimit: number;
+  accessTokenTrustedIps: IdentityTrustedIp[];
+};
+
+export type AddIdentityAwsIamAuthDTO = {
+  organizationId: string;
+  identityId: string;
+  stsEndpoint: string;
+  allowedPrincipalArns: string;
+  allowedAccountIds: string;
+  accessTokenTTL: number;
+  accessTokenMaxTTL: number;
+  accessTokenNumUsesLimit: number;
+  accessTokenTrustedIps: {
+    ipAddress: string;
+  }[];
+};
+
+export type UpdateIdentityAwsIamAuthDTO = {
+  organizationId: string;
+  identityId: string;
+  stsEndpoint?: string;
+  allowedPrincipalArns?: string;
+  allowedAccountIds?: string;
+  accessTokenTTL?: number;
+  accessTokenMaxTTL?: number;
+  accessTokenNumUsesLimit?: number;
+  accessTokenTrustedIps?: {
+    ipAddress: string;
+  }[];
+};
+
 export type CreateIdentityUniversalAuthClientSecretDTO = {
  identityId: string;
  description?: string;
--- a/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityAuthMethodModal.tsx
+++ b/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityAuthMethodModal.tsx
@@ -1,3 +1,4 @@
+import { useEffect } from "react";
 import { Controller, useForm } from "react-hook-form";
 import { yupResolver } from "@hookform/resolvers/yup";
 import * as yup from "yup";
@@ -13,6 +14,7 @@ import {
 import { IdentityAuthMethod } from "@app/hooks/api/identities";
 import { UsePopUpState } from "@app/hooks/usePopUp";

+import { IdentityAwsIamAuthForm } from "./IdentityAwsIamAuthForm";
 import { IdentityUniversalAuthForm } from "./IdentityUniversalAuthForm";

 type Props = {
@@ -24,22 +26,25 @@ type Props = {
  ) => void;
 };

-const identityAuthMethods = [{ label: "Universal Auth", value: IdentityAuthMethod.UNIVERSAL_AUTH }];
+const identityAuthMethods = [
+  { label: "Universal Auth", value: IdentityAuthMethod.UNIVERSAL_AUTH },
+  { label: "AWS IAM Auth", value: IdentityAuthMethod.AWS_IAM_AUTH }
+];

 const schema = yup
  .object({
-    authMethod: yup.string().required("Auth method is required") // TODO: better enforcement here
+    authMethod: yup.string().required("Auth method is required")
  })
  .required();

 export type FormData = yup.InferType<typeof schema>;

 export const IdentityAuthMethodModal = ({ popUp, handlePopUpOpen, handlePopUpToggle }: Props) => {
-  const {
-    control
-    // watch,
-  } = useForm<FormData>({
-    resolver: yupResolver(schema)
+  const { control, watch, setValue } = useForm<FormData>({
+    resolver: yupResolver(schema),
+    defaultValues: {
+      authMethod: IdentityAuthMethod.UNIVERSAL_AUTH
+    }
  });

  const identityAuthMethodData = popUp?.identityAuthMethod?.data as {
@@ -48,16 +53,41 @@ export const IdentityAuthMethodModal = ({ popUp, handlePopUpOpen, handlePopUpTog
    authMethod?: IdentityAuthMethod;
  };

-  // const authMethod = watch("authMethod");
+  useEffect(() => {
+    if (identityAuthMethodData?.authMethod) {
+      setValue("authMethod", identityAuthMethodData.authMethod);
+      return;
+    }
+
+    setValue("authMethod", IdentityAuthMethod.UNIVERSAL_AUTH);
+  }, [identityAuthMethodData?.authMethod]);
+
+  const authMethod = watch("authMethod");

  const renderIdentityAuthForm = () => {
-    return (
-      <IdentityUniversalAuthForm
-        handlePopUpOpen={handlePopUpOpen}
-        handlePopUpToggle={handlePopUpToggle}
-        identityAuthMethodData={identityAuthMethodData}
-      />
-    );
+    switch (identityAuthMethodData?.authMethod ?? authMethod) {
+      case IdentityAuthMethod.AWS_IAM_AUTH: {
+        return (
+          <IdentityAwsIamAuthForm
+            handlePopUpOpen={handlePopUpOpen}
+            handlePopUpToggle={handlePopUpToggle}
+            identityAuthMethodData={identityAuthMethodData}
+          />
+        );
+      }
+      case IdentityAuthMethod.UNIVERSAL_AUTH: {
+        return (
+          <IdentityUniversalAuthForm
+            handlePopUpOpen={handlePopUpOpen}
+            handlePopUpToggle={handlePopUpToggle}
+            identityAuthMethodData={identityAuthMethodData}
+          />
+        );
+      }
+      default: {
+        return <div />;
+      }
+    }
  };

  return (
@@ -83,6 +113,7 @@ export const IdentityAuthMethodModal = ({ popUp, handlePopUpOpen, handlePopUpTog
                {...field}
                onValueChange={(e) => onChange(e)}
                className="w-full"
+                isDisabled={!!identityAuthMethodData?.authMethod}
              >
                {identityAuthMethods.map(({ label, value }) => (
                  <SelectItem value={String(value || "")} key={label}>
--- a/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityAwsIamAuthForm.tsx
+++ b/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityAwsIamAuthForm.tsx
@@ -0,0 +1,348 @@
+import { useEffect } from "react";
+import { Controller, useFieldArray, useForm } from "react-hook-form";
+import { faPlus, faXmark } from "@fortawesome/free-solid-svg-icons";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import { yupResolver } from "@hookform/resolvers/yup";
+import * as yup from "yup";
+
+import { createNotification } from "@app/components/notifications";
+import { Button, FormControl, IconButton, Input } from "@app/components/v2";
+import { useOrganization, useSubscription } from "@app/context";
+import {
+  useAddIdentityAwsIamAuth,
+  useGetIdentityAwsIamAuth,
+  useUpdateIdentityAwsIamAuth
+} from "@app/hooks/api";
+import { IdentityAuthMethod } from "@app/hooks/api/identities";
+import { IdentityTrustedIp } from "@app/hooks/api/identities/types";
+import { UsePopUpState } from "@app/hooks/usePopUp";
+
+const schema = yup
+  .object({
+    stsEndpoint: yup.string(),
+    allowedPrincipalArns: yup.string(),
+    allowedAccountIds: yup.string(),
+    accessTokenTTL: yup.string().required("Access Token TTL is required"),
+    accessTokenMaxTTL: yup.string().required("Access Max Token TTL is required"),
+    accessTokenNumUsesLimit: yup.string().required("Access Token Max Number of Uses is required"),
+    accessTokenTrustedIps: yup
+      .array(
+        yup.object({
+          ipAddress: yup.string().max(50).required().label("IP Address")
+        })
+      )
+      .min(1)
+      .required()
+      .label("Access Token Trusted IP")
+  })
+  .required();
+
+export type FormData = yup.InferType<typeof schema>;
+
+type Props = {
+  handlePopUpOpen: (popUpName: keyof UsePopUpState<["upgradePlan"]>) => void;
+  handlePopUpToggle: (
+    popUpName: keyof UsePopUpState<["identityAuthMethod"]>,
+    state?: boolean
+  ) => void;
+  identityAuthMethodData: {
+    identityId: string;
+    name: string;
+    authMethod?: IdentityAuthMethod;
+  };
+};
+
+export const IdentityAwsIamAuthForm = ({
+  handlePopUpOpen,
+  handlePopUpToggle,
+  identityAuthMethodData
+}: Props) => {
+  const { currentOrg } = useOrganization();
+  const orgId = currentOrg?.id || "";
+  const { subscription } = useSubscription();
+
+  const { mutateAsync: addMutateAsync } = useAddIdentityAwsIamAuth();
+  const { mutateAsync: updateMutateAsync } = useUpdateIdentityAwsIamAuth();
+
+  const { data } = useGetIdentityAwsIamAuth(identityAuthMethodData?.identityId ?? "");
+
+  const {
+    control,
+    handleSubmit,
+    reset,
+    formState: { isSubmitting }
+  } = useForm<FormData>({
+    resolver: yupResolver(schema),
+    defaultValues: {
+      stsEndpoint: "https://sts.amazonaws.com/",
+      allowedPrincipalArns: "",
+      allowedAccountIds: "",
+      accessTokenTTL: "2592000",
+      accessTokenMaxTTL: "2592000",
+      accessTokenNumUsesLimit: "0",
+      accessTokenTrustedIps: [{ ipAddress: "0.0.0.0/0" }, { ipAddress: "::/0" }]
+    }
+  });
+
+  const {
+    fields: accessTokenTrustedIpsFields,
+    append: appendAccessTokenTrustedIp,
+    remove: removeAccessTokenTrustedIp
+  } = useFieldArray({ control, name: "accessTokenTrustedIps" });
+
+  useEffect(() => {
+    if (data) {
+      reset({
+        stsEndpoint: data.stsEndpoint,
+        allowedPrincipalArns: data.allowedPrincipalArns,
+        allowedAccountIds: data.allowedAccountIds,
+        accessTokenTTL: String(data.accessTokenTTL),
+        accessTokenMaxTTL: String(data.accessTokenMaxTTL),
+        accessTokenNumUsesLimit: String(data.accessTokenNumUsesLimit),
+        accessTokenTrustedIps: data.accessTokenTrustedIps.map(
+          ({ ipAddress, prefix }: IdentityTrustedIp) => {
+            return {
+              ipAddress: `${ipAddress}${prefix !== undefined ? `/${prefix}` : ""}`
+            };
+          }
+        )
+      });
+    } else {
+      reset({
+        stsEndpoint: "https://sts.amazonaws.com/",
+        allowedPrincipalArns: "",
+        allowedAccountIds: "",
+        accessTokenTTL: "2592000",
+        accessTokenMaxTTL: "2592000",
+        accessTokenNumUsesLimit: "0",
+        accessTokenTrustedIps: [{ ipAddress: "0.0.0.0/0" }, { ipAddress: "::/0" }]
+      });
+    }
+  }, [data]);
+
+  const onFormSubmit = async ({
+    allowedPrincipalArns,
+    allowedAccountIds,
+    stsEndpoint,
+    accessTokenTTL,
+    accessTokenMaxTTL,
+    accessTokenNumUsesLimit,
+    accessTokenTrustedIps
+  }: FormData) => {
+    try {
+      if (!identityAuthMethodData) return;
+
+      if (data) {
+        await updateMutateAsync({
+          organizationId: orgId,
+          stsEndpoint,
+          allowedPrincipalArns,
+          allowedAccountIds,
+          identityId: identityAuthMethodData.identityId,
+          accessTokenTTL: Number(accessTokenTTL),
+          accessTokenMaxTTL: Number(accessTokenMaxTTL),
+          accessTokenNumUsesLimit: Number(accessTokenNumUsesLimit),
+          accessTokenTrustedIps
+        });
+      } else {
+        await addMutateAsync({
+          organizationId: orgId,
+          identityId: identityAuthMethodData.identityId,
+          stsEndpoint: stsEndpoint || "",
+          allowedPrincipalArns: allowedPrincipalArns || "",
+          allowedAccountIds: allowedAccountIds || "",
+          accessTokenTTL: Number(accessTokenTTL),
+          accessTokenMaxTTL: Number(accessTokenMaxTTL),
+          accessTokenNumUsesLimit: Number(accessTokenNumUsesLimit),
+          accessTokenTrustedIps
+        });
+      }
+
+      handlePopUpToggle("identityAuthMethod", false);
+
+      createNotification({
+        text: `Successfully ${
+          identityAuthMethodData?.authMethod ? "updated" : "configured"
+        } auth method`,
+        type: "success"
+      });
+
+      reset();
+    } catch (err) {
+      createNotification({
+        text: `Failed to ${identityAuthMethodData?.authMethod ? "update" : "configure"} identity`,
+        type: "error"
+      });
+    }
+  };
+
+  return (
+    <form onSubmit={handleSubmit(onFormSubmit)}>
+      <Controller
+        control={control}
+        defaultValue="2592000"
+        name="allowedPrincipalArns"
+        render={({ field, fieldState: { error } }) => (
+          <FormControl label="Allowed ARNs" isError={Boolean(error)} errorText={error?.message}>
+            <Input
+              {...field}
+              placeholder="arn:aws:iam::123456789012:role/MyRoleName, arn:aws:iam::123456789012:user/MyUserName..."
+              type="text"
+            />
+          </FormControl>
+        )}
+      />
+      <Controller
+        control={control}
+        name="allowedAccountIds"
+        render={({ field, fieldState: { error } }) => (
+          <FormControl
+            label="Allowed Account IDs"
+            isError={Boolean(error)}
+            errorText={error?.message}
+          >
+            <Input {...field} placeholder="123456789012, ..." />
+          </FormControl>
+        )}
+      />
+      <Controller
+        control={control}
+        defaultValue="https://sts.amazonaws.com/"
+        name="stsEndpoint"
+        render={({ field, fieldState: { error } }) => (
+          <FormControl label="STS Endpoint" isError={Boolean(error)} errorText={error?.message}>
+            <Input {...field} placeholder="https://sts.amazonaws.com/" type="text" />
+          </FormControl>
+        )}
+      />
+      <Controller
+        control={control}
+        defaultValue="2592000"
+        name="accessTokenTTL"
+        render={({ field, fieldState: { error } }) => (
+          <FormControl
+            label="Access Token TTL (seconds)"
+            isError={Boolean(error)}
+            errorText={error?.message}
+          >
+            <Input {...field} placeholder="2592000" type="number" min="1" step="1" />
+          </FormControl>
+        )}
+      />
+      <Controller
+        control={control}
+        defaultValue="2592000"
+        name="accessTokenMaxTTL"
+        render={({ field, fieldState: { error } }) => (
+          <FormControl
+            label="Access Token Max TTL (seconds)"
+            isError={Boolean(error)}
+            errorText={error?.message}
+          >
+            <Input {...field} placeholder="2592000" type="number" min="1" step="1" />
+          </FormControl>
+        )}
+      />
+      <Controller
+        control={control}
+        defaultValue="0"
+        name="accessTokenNumUsesLimit"
+        render={({ field, fieldState: { error } }) => (
+          <FormControl
+            label="Access Token Max Number of Uses"
+            isError={Boolean(error)}
+            errorText={error?.message}
+          >
+            <Input {...field} placeholder="0" type="number" min="0" step="1" />
+          </FormControl>
+        )}
+      />
+      {accessTokenTrustedIpsFields.map(({ id }, index) => (
+        <div className="mb-3 flex items-end space-x-2" key={id}>
+          <Controller
+            control={control}
+            name={`accessTokenTrustedIps.${index}.ipAddress`}
+            defaultValue="0.0.0.0/0"
+            render={({ field, fieldState: { error } }) => {
+              return (
+                <FormControl
+                  className="mb-0 flex-grow"
+                  label={index === 0 ? "Access Token Trusted IPs" : undefined}
+                  isError={Boolean(error)}
+                  errorText={error?.message}
+                >
+                  <Input
+                    value={field.value}
+                    onChange={(e) => {
+                      if (subscription?.ipAllowlisting) {
+                        field.onChange(e);
+                        return;
+                      }
+
+                      handlePopUpOpen("upgradePlan");
+                    }}
+                    placeholder="123.456.789.0"
+                  />
+                </FormControl>
+              );
+            }}
+          />
+          <IconButton
+            onClick={() => {
+              if (subscription?.ipAllowlisting) {
+                removeAccessTokenTrustedIp(index);
+                return;
+              }
+
+              handlePopUpOpen("upgradePlan");
+            }}
+            size="lg"
+            colorSchema="danger"
+            variant="plain"
+            ariaLabel="update"
+            className="p-3"
+          >
+            <FontAwesomeIcon icon={faXmark} />
+          </IconButton>
+        </div>
+      ))}
+      <div className="my-4 ml-1">
+        <Button
+          variant="outline_bg"
+          onClick={() => {
+            if (subscription?.ipAllowlisting) {
+              appendAccessTokenTrustedIp({
+                ipAddress: "0.0.0.0/0"
+              });
+              return;
+            }
+
+            handlePopUpOpen("upgradePlan");
+          }}
+          leftIcon={<FontAwesomeIcon icon={faPlus} />}
+          size="xs"
+        >
+          Add IP Address
+        </Button>
+      </div>
+      <div className="flex items-center">
+        <Button
+          className="mr-4"
+          size="sm"
+          type="submit"
+          isLoading={isSubmitting}
+          isDisabled={isSubmitting}
+        >
+          {identityAuthMethodData?.authMethod ? "Update" : "Configure"}
+        </Button>
+        <Button
+          colorSchema="secondary"
+          variant="plain"
+          onClick={() => handlePopUpToggle("identityAuthMethod", false)}
+        >
+          {identityAuthMethodData?.authMethod ? "Cancel" : "Skip"}
+        </Button>
+      </div>
+    </form>
+  );
+};
--- a/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityModal.tsx
+++ b/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityModal.tsx
@@ -15,7 +15,10 @@ import {
 } from "@app/components/v2";
 import { useOrganization } from "@app/context";
 import { useCreateIdentity, useGetOrgRoles, useUpdateIdentity } from "@app/hooks/api";
-import { IdentityAuthMethod, useAddIdentityUniversalAuth } from "@app/hooks/api/identities";
+import {
+  IdentityAuthMethod
+  // useAddIdentityUniversalAuth
+} from "@app/hooks/api/identities";
 import { UsePopUpState } from "@app/hooks/usePopUp";

 const schema = yup
@@ -40,9 +43,7 @@ type Props = {
  handlePopUpToggle: (popUpName: keyof UsePopUpState<["identity"]>, state?: boolean) => void;
 };

-export const IdentityModal = ({ popUp, /* handlePopUpOpen, */ handlePopUpToggle }: Props) => {
-  
-
+export const IdentityModal = ({ popUp, handlePopUpOpen, handlePopUpToggle }: Props) => {
  const { currentOrg } = useOrganization();
  const orgId = currentOrg?.id || "";

@@ -50,7 +51,7 @@ export const IdentityModal = ({ popUp, /* handlePopUpOpen, */ handlePopUpToggle

  const { mutateAsync: createMutateAsync } = useCreateIdentity();
  const { mutateAsync: updateMutateAsync } = useUpdateIdentity();
-  const { mutateAsync: addMutateAsync } = useAddIdentityUniversalAuth();
+  // const { mutateAsync: addMutateAsync } = useAddIdentityUniversalAuth();

  const {
    control,
@@ -113,31 +114,31 @@ export const IdentityModal = ({ popUp, /* handlePopUpOpen, */ handlePopUpToggle
        // create

        const {
-          id: createdId
-          // name: createdName,
-          // authMethod
+          id: createdId,
+          name: createdName,
+          authMethod
        } = await createMutateAsync({
          name,
          role: role || undefined,
          organizationId: orgId
        });

-        await addMutateAsync({
-          organizationId: orgId,
-          identityId: createdId,
-          clientSecretTrustedIps: [{ ipAddress: "0.0.0.0/0" }, { ipAddress: "::/0" }],
-          accessTokenTrustedIps: [{ ipAddress: "0.0.0.0/0" }, { ipAddress: "::/0" }],
-          accessTokenTTL: 2592000,
-          accessTokenMaxTTL: 2592000,
-          accessTokenNumUsesLimit: 0
-        });
+        // await addMutateAsync({
+        //   organizationId: orgId,
+        //   identityId: createdId,
+        //   clientSecretTrustedIps: [{ ipAddress: "0.0.0.0/0" }, { ipAddress: "::/0" }],
+        //   accessTokenTrustedIps: [{ ipAddress: "0.0.0.0/0" }, { ipAddress: "::/0" }],
+        //   accessTokenTTL: 2592000,
+        //   accessTokenMaxTTL: 2592000,
+        //   accessTokenNumUsesLimit: 0
+        // });

        handlePopUpToggle("identity", false);
-        // handlePopUpOpen("identityAuthMethod", {
-        //   identityId: createdId,
-        //   name: createdName,
-        //   authMethod
-        // });
+        handlePopUpOpen("identityAuthMethod", {
+          identityId: createdId,
+          name: createdName,
+          authMethod
+        });
      }

      createNotification({
--- a/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityTable.tsx
+++ b/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityTable.tsx
@@ -23,8 +23,6 @@ import { useGetIdentityMembershipOrgs, useGetOrgRoles, useUpdateIdentity } from
 import { IdentityAuthMethod, identityAuthToNameMap } from "@app/hooks/api/identities";
 import { UsePopUpState } from "@app/hooks/usePopUp";

-// TODO: some kind of map
-
 type Props = {
  handlePopUpOpen: (
    popUpName: keyof UsePopUpState<
@@ -44,7 +42,6 @@ type Props = {
 };

 export const IdentityTable = ({ handlePopUpOpen }: Props) => {
-  
  const { currentOrg } = useOrganization();
  const orgId = currentOrg?.id || "";

--- a/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityUniversalAuthForm.tsx
+++ b/frontend/src/views/Org/MembersPage/components/OrgIdentityTab/components/IdentitySection/IdentityUniversalAuthForm.tsx
@@ -63,7 +63,6 @@ export const IdentityUniversalAuthForm = ({
  handlePopUpToggle,
  identityAuthMethodData
 }: Props) => {
-  
  const { currentOrg } = useOrganization();
  const orgId = currentOrg?.id || "";
  const { subscription } = useSubscription();
@@ -384,7 +383,7 @@ export const IdentityUniversalAuthForm = ({
          variant="plain"
          onClick={() => handlePopUpToggle("identityAuthMethod", false)}
        >
-          Cancel
+          {identityAuthMethodData?.authMethod ? "Cancel" : "Skip"}
        </Button>
      </div>
    </form>