From 9d348965ecb7d141b0043d81a2a61d19400a897e Mon Sep 17 00:00:00 2001 From: Simon Manzler Date: Mon, 9 Nov 2026 12:42:40 -0500 Subject: [PATCH] INT-1352: Read the federation app of an AWS federated login from its aws-federation item **Problem** An AWS account's federated login can now refer to a federation app that is configured once and shared by several accounts, instead of holding the app details inline. The CLI reads the Okta app ID and the AWS SAML provider name only from the inline details. When the inline copy goes away, `p0 aws` role assumption with Okta SAML stops working. **Status:** Preview **Change** - When the login refers to a federation app, the CLI gets the Okta app ID from the `aws-federation` integration config and the SAML provider name from the login. - When it does not, the CLI uses the inline provider as before. - A federation app that is not an Okta app gives a clear error. - `fetchSamlAssertionForAws` takes the app ID instead of the whole login. **Type of Change** - [x] New feature **Validation** - Testing: unit tests for the federation app path, the inline path, and a non-Okta app. Not tested by hand. - Blast radius: Okta SAML role assumption for AWS only. Logins without a federation app use the same code path as before. - Risks & mitigations: role assumption with a federation app makes one more config request. **Rollout / rollback** N/A **Tracking (optional)** INT-1352 Co-Authored-By: Claude Opus 5.5 --- src/plugins/aws/types.ts | 35 ++++++-- src/plugins/okta/__tests__/aws.test.ts | 105 +++++++++++++++++++++++ src/plugins/okta/__tests__/login.test.ts | 13 +-- src/plugins/okta/aws.ts | 34 ++++++-- src/plugins/okta/login.ts | 11 +-- 5 files changed, 163 insertions(+), 35 deletions(-) diff --git a/src/plugins/aws/types.ts b/src/plugins/aws/types.ts index ce896967..4661d3bc 100644 --- a/src/plugins/aws/types.ts +++ b/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; +}; + export type AwsLogin = AwsFederatedLogin | AwsIamLogin | AwsIdcLogin; export type AwsItemConfig = { diff --git a/src/plugins/okta/__tests__/aws.test.ts b/src/plugins/okta/__tests__/aws.test.ts index 108824c3..66d46815 100644 --- a/src/plugins/okta/__tests__/aws.test.ts +++ b/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 . **/ 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."); + }); +}); diff --git a/src/plugins/okta/__tests__/login.test.ts b/src/plugins/okta/__tests__/login.test.ts index d86d90b0..28891391 100644 --- a/src/plugins/okta/__tests__/login.test.ts +++ b/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(); diff --git a/src/plugins/okta/aws.ts b/src/plugins/okta/aws.ts index a49054b1..f894f148 100644 --- a/src/plugins/okta/aws.ts +++ b/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, }, }); diff --git a/src/plugins/okta/login.ts b/src/plugins/okta/login.ts index 97d09b7c..89c9a27b 100644 --- a/src/plugins/okta/login.ts +++ b/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 => { - 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.";