INT-1352: Read the federation app of an AWS federated login from its aws-federation item

src/plugins/aws/types.ts

@@ -41,18 +41,35 @@ export type AwsIdcLogin = {
   identityStoreId: string;
 };
 
-export type AwsFederatedLogin = {
-  type: "federated";
-  provider: {
-    type: "okta";
-    appId: string;
-    identityProvider: string;
-    method: {
-      type: "saml";
-    };
+type AwsFederatedOktaProvider = {
+  type: "okta";
+  appId: string;
+  identityProvider: string;
+  method: {
+    type: "saml";
   };
 };
 
+/** A login that holds its federation app inline */
+type AwsLegacyFederatedLogin = {
+  type: "federated";
+  federation?: undefined;
+  provider: AwsFederatedOktaProvider;
+};
+
+/** A login that refers to an `aws-federation` item */
+type AwsItemFederatedLogin = {
+  type: "federated";
+  federation: string;
+  samlProvider: string;
+};
+
+export type AwsFederatedLogin = AwsItemFederatedLogin | AwsLegacyFederatedLogin;
+
+export type AwsFederationConfig = {
+  federation: Record<string, { provider: { type: "okta"; appId: string } }>;
+};
+
 export type AwsLogin = AwsFederatedLogin | AwsIamLogin | AwsIdcLogin;
 
 export type AwsItemConfig = {

src/plugins/okta/__tests__/aws.test.ts

@@ -9,6 +9,7 @@ This file is part of @p0security/cli
 You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
 **/
 import { parseXml } from "../../../common/xml";
+import { fetchIntegrationConfig } from "../../../drivers/api";
 import { Authn } from "../../../types/identity";
 import { assumeRoleWithSaml } from "../../aws/assumeRole";
 import { getAwsConfig } from "../../aws/config";
@@ -18,6 +19,9 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 
 // Mock dependencies
 vi.mock("../../../drivers/auth");
+vi.mock("../../../drivers/api", () => ({
+  fetchIntegrationConfig: vi.fn(),
+}));
 vi.mock("../../aws/config", () => ({
   getAwsConfig: vi.fn(),
 }));
@@ -165,3 +169,104 @@ describe("assumeRoleWithOktaSaml retry logic", () => {
     );
   });
 });
+
+describe("assumeRoleWithOktaSaml federation app", () => {
+  const mockAuthn = {
+    identity: { org: { slug: "test-org" } },
+  } as unknown as Authn;
+
+  const useLogin = (login: object) =>
+    vi.mocked(getAwsConfig).mockResolvedValue({
+      identity: {} as any,
+      config: { id: "1234", login } as any,
+    });
+
+  const useFederation = (provider: object) =>
+    vi.mocked(fetchIntegrationConfig).mockResolvedValue({
+      config: { federation: { "okta-prod": { provider, state: "installed" } } },
+    });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.mocked(fetchSamlAssertionForAws).mockResolvedValue("base64-saml");
+    vi.mocked(parseXml).mockReturnValue({
+      "saml2p:Response": {
+        "saml2:Assertion": {
+          "saml2:AttributeStatement": {
+            "saml2:Attribute": [
+              {
+                _attributes: {
+                  Name: "https://aws.amazon.com/SAML/Attributes/Role",
+                },
+                "saml2:AttributeValue": [
+                  "arn:aws:iam::1234:saml-provider/okta-saml,arn:aws:iam::1234:role/Admin",
+                ],
+              },
+            ],
+          },
+        },
+      },
+    });
+  });
+
+  it("uses the app ID of the federation item and the SAML provider of the login", async () => {
+    useLogin({
+      type: "federated",
+      federation: "okta-prod",
+      samlProvider: "okta-saml",
+    });
+    useFederation({ type: "okta", appId: "0oa123" });
+
+    await assumeRoleWithOktaSaml(mockAuthn, { role: "Admin" }, false);
+
+    expect(fetchIntegrationConfig).toHaveBeenCalledWith(
+      mockAuthn,
+      "aws-federation",
+      false
+    );
+    expect(fetchSamlAssertionForAws).toHaveBeenCalledWith(
+      expect.anything(),
+      "0oa123",
+      false
+    );
+    expect(assumeRoleWithSaml).toHaveBeenCalledWith(
+      expect.objectContaining({
+        saml: { providerName: "okta-saml", response: "base64-saml" },
+      })
+    );
+  });
+
+  it("uses the inline provider of a login without a federation item", async () => {
+    useLogin({
+      type: "federated",
+      provider: { type: "okta", appId: "0oa456", identityProvider: "legacy" },
+    });
+
+    await assumeRoleWithOktaSaml(mockAuthn, { role: "Admin" }, false);
+
+    expect(fetchIntegrationConfig).not.toHaveBeenCalled();
+    expect(fetchSamlAssertionForAws).toHaveBeenCalledWith(
+      expect.anything(),
+      "0oa456",
+      false
+    );
+    expect(assumeRoleWithSaml).toHaveBeenCalledWith(
+      expect.objectContaining({
+        saml: { providerName: "legacy", response: "base64-saml" },
+      })
+    );
+  });
+
+  it("rejects a federation item that is not an Okta app", async () => {
+    useLogin({
+      type: "federated",
+      federation: "okta-prod",
+      samlProvider: "okta-saml",
+    });
+    useFederation({ type: "pingone" });
+
+    await expect(
+      assumeRoleWithOktaSaml(mockAuthn, { role: "Admin" }, false)
+    ).rejects.toBe("Federation app okta-prod is not an Okta app.");
+  });
+});

src/plugins/okta/__tests__/login.test.ts

@@ -14,7 +14,6 @@ import {
   getProviderType,
 } from "../../../types/authUtils";
 import { Identity } from "../../../types/identity";
-import { AwsFederatedLogin } from "../../aws/types";
 import { fetchSamlAssertionForAws } from "../login";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
@@ -32,16 +31,6 @@ const mockIdentity = {
   },
 } as unknown as Identity;
 
-const mockConfig: AwsFederatedLogin = {
-  type: "federated",
-  provider: {
-    type: "okta",
-    appId: "mock-app-id",
-    identityProvider: "mock-idp",
-    method: { type: "saml" },
-  },
-};
-
 describe("fetchSsoWebToken", () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -70,7 +59,7 @@ describe("fetchSsoWebToken", () => {
     );
 
     await expect(
-      fetchSamlAssertionForAws(mockIdentity, mockConfig)
+      fetchSamlAssertionForAws(mockIdentity, "mock-app-id")
     ).rejects.toThrow("400 Bad Request");
 
     expect(mockJson).toHaveBeenCalledOnce();

src/plugins/okta/aws.ts

@@ -10,11 +10,12 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { retryWithSleep } from "../../common/retry";
 import { parseXml } from "../../common/xml";
+import { fetchIntegrationConfig } from "../../drivers/api";
 import { cached } from "../../drivers/auth";
 import { Authn } from "../../types/identity";
 import { assumeRoleWithSaml } from "../aws/assumeRole";
 import { getAwsConfig } from "../aws/config";
-import { AwsFederatedLogin, AwsItem } from "../aws/types";
+import { AwsFederatedLogin, AwsFederationConfig, AwsItem } from "../aws/types";
 import { fetchSamlAssertionForAws } from "./login";
 import { flatten } from "lodash";
 
@@ -64,6 +65,26 @@ const isFederatedLogin = (
 ): config is AwsItem & { login: AwsFederatedLogin } =>
   config.login?.type === "federated";
 
+/** Gets the Okta app ID and the AWS SAML provider name of a federated login */
+const oktaFederationApp = async (
+  authn: Authn,
+  login: AwsFederatedLogin,
+  debug?: boolean
+) => {
+  if (login.federation === undefined)
+    return {
+      appId: login.provider.appId,
+      samlProvider: login.provider.identityProvider,
+    };
+  const { config } = await fetchIntegrationConfig<{
+    config: AwsFederationConfig;
+  }>(authn, "aws-federation", debug);
+  const provider = config?.federation?.[login.federation]?.provider;
+  if (provider?.type !== "okta")
+    throw `Federation app ${login.federation} is not an Okta app.`;
+  return { appId: provider.appId, samlProvider: login.samlProvider };
+};
+
 /** Retrieves the configured Okta SAML response for the specified account
  *
  * If no account is passed, and the organization only has one account configured,
@@ -77,14 +98,15 @@ const initOktaSaml = async (
   const { identity, config } = await getAwsConfig(authn, account, debug);
   if (!isFederatedLogin(config))
     throw `Account ${config.label ?? config.id} is not configured for Okta SAML login.`;
-  const samlResponse = await fetchSamlAssertionForAws(
-    identity,
+  const { appId, samlProvider } = await oktaFederationApp(
+    authn,
     config.login,
     debug
   );
+  const samlResponse = await fetchSamlAssertionForAws(identity, appId, debug);
   return {
     samlResponse,
-    config,
+    samlProvider,
     account: config.id,
   };
 };
@@ -102,7 +124,7 @@ export const assumeRoleWithOktaSaml = async (
       // Add retry logic to handle this race condition.
       return await retryWithSleep(
         async () => {
-          const { account, config, samlResponse } = await initOktaSaml(
+          const { account, samlProvider, samlResponse } = await initOktaSaml(
             authn,
             args.accountId,
             debug
@@ -116,7 +138,7 @@ export const assumeRoleWithOktaSaml = async (
             partition,
             role: args.role,
             saml: {
-              providerName: config.login.provider.identityProvider,
+              providerName: samlProvider,
               response: samlResponse,
             },
           });

src/plugins/okta/login.ts

@@ -24,7 +24,6 @@ import {
 import { Identity } from "../../types/identity";
 import { AuthorizeResponse, TokenResponse } from "../../types/oidc";
 import { OrgData } from "../../types/org";
-import { AwsFederatedLogin } from "../aws/types";
 import {
   oidcLogin,
   oidcLoginSteps,
@@ -184,21 +183,17 @@ export const oktaLogin = async (org: OrgData) =>
  *    SAML assertion from the resulting HTML response.
  *
  * @param identity - The user's P0 identity containing OIDC tokens from login
- * @param config - AWS federated login configuration with Okta app details
+ * @param appId - ID of the Okta AWS federation app
  * @returns Base64-encoded SAML assertion for AWS authentication
  * @throws Error if Okta session has expired or been terminated
  */
 // TODO: Inject Okta app
 export const fetchSamlAssertionForAws = async (
   identity: Identity,
-  config: AwsFederatedLogin,
+  appId: string,
   debug?: boolean
 ): Promise<string> => {
-  const webTokenResponse = await fetchSsoWebToken(
-    config.provider.appId,
-    identity,
-    debug
-  );
+  const webTokenResponse = await fetchSsoWebToken(appId, identity, debug);
   const samlResponse = await fetchSamlResponse(identity.org, webTokenResponse);
   if (!samlResponse) {
     throw "No SAML assertion obtained from Okta.";