[Cloud Security] [Fleet] Get CloudFormation template accordingly with the current integration version to support auto-upgrade (#162206)

x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/components/post_install_cloud_formation_modal.tsx

@@ -22,7 +22,7 @@ import { useQuery } from '@tanstack/react-query';
 
 import type { AgentPolicy, PackagePolicy } from '../../../../../types';
 import { sendGetEnrollmentAPIKeys, useCreateCloudFormationUrl } from '../../../../../hooks';
-import { getCloudFormationTemplateUrlFromPackagePolicy } from '../../../../../services';
+import { getCloudFormationPropsFromPackagePolicy } from '../../../../../services';
 import { CloudFormationGuide } from '../../../../../components';
 
 export const PostInstallCloudFormationModal: React.FunctionComponent<{
@@ -39,13 +39,11 @@ export const PostInstallCloudFormationModal: React.FunctionComponent<{
     })
   );
 
-  const cloudFormationTemplateUrl =
-    getCloudFormationTemplateUrlFromPackagePolicy(packagePolicy) || '';
+  const cloudFormationProps = getCloudFormationPropsFromPackagePolicy(packagePolicy);
 
   const { cloudFormationUrl, error, isError, isLoading } = useCreateCloudFormationUrl({
-    cloudFormationTemplateUrl,
     enrollmentAPIKey: apyKeysData?.data?.items[0]?.api_key,
-    packagePolicy,
+    cloudFormationProps,
   });
 
   return (

x-pack/plugins/fleet/public/applications/fleet/sections/agent_policy/create_package_policy_page/single_page_layout/hooks/form.tsx

@@ -39,7 +39,7 @@ import type { PackagePolicyFormState } from '../../types';
 import { SelectedPolicyTab } from '../../components';
 import { useOnSaveNavigate } from '../../hooks';
 import { prepareInputPackagePolicyDataset } from '../../services/prepare_input_pkg_policy_dataset';
-import { getCloudFormationTemplateUrlFromPackagePolicy } from '../../../../../services';
+import { getCloudFormationPropsFromPackagePolicy } from '../../../../../services';
 
 async function createAgentPolicy({
   packagePolicy,
@@ -301,7 +301,7 @@ export function useOnSubmit({
       });
 
       const hasCloudFormation = data?.item
-        ? getCloudFormationTemplateUrlFromPackagePolicy(data.item)
+        ? getCloudFormationPropsFromPackagePolicy(data.item).templateUrl
         : false;
 
       if (hasCloudFormation) {

x-pack/plugins/fleet/public/components/agent_enrollment_flyout/cloud_formation_instructions.tsx

@@ -10,26 +10,23 @@ import { EuiButton, EuiSpacer, EuiCallOut, EuiSkeletonText } from '@elastic/eui'
 import { FormattedMessage } from '@kbn/i18n-react';
 import { i18n } from '@kbn/i18n';
 
-import type { PackagePolicy } from '../../../common';
-
 import { useCreateCloudFormationUrl } from '../../hooks';
 import { CloudFormationGuide } from '../cloud_formation_guide';
 
+import type { CloudSecurityIntegration } from './types';
+
 interface Props {
   enrollmentAPIKey?: string;
-  cloudFormationTemplateUrl: string;
-  packagePolicy?: PackagePolicy;
+  cloudSecurityIntegration: CloudSecurityIntegration;
 }
 
 export const CloudFormationInstructions: React.FunctionComponent<Props> = ({
   enrollmentAPIKey,
-  cloudFormationTemplateUrl,
-  packagePolicy,
+  cloudSecurityIntegration,
 }) => {
   const { isLoading, cloudFormationUrl, error, isError } = useCreateCloudFormationUrl({
     enrollmentAPIKey,
-    cloudFormationTemplateUrl,
-    packagePolicy,
+    cloudFormationProps: cloudSecurityIntegration?.cloudFormationProps,
   });
 
   if (error && isError) {
@@ -45,7 +42,7 @@ export const CloudFormationInstructions: React.FunctionComponent<Props> = ({
     <EuiSkeletonText
       lines={3}
       size="m"
-      isLoading={isLoading}
+      isLoading={isLoading || cloudSecurityIntegration?.isLoading}
       contentAriaLabel={i18n.translate(
         'xpack.fleet.agentEnrollment.cloudFormation.loadingAriaLabel',
         {

x-pack/plugins/fleet/public/components/agent_enrollment_flyout/hooks.tsx

@@ -8,15 +8,24 @@ import { useState, useEffect, useMemo } from 'react';
 import { i18n } from '@kbn/i18n';
 
 import type { PackagePolicy, AgentPolicy } from '../../types';
-import { sendGetOneAgentPolicy, useStartServices } from '../../hooks';
+import { sendGetOneAgentPolicy, useGetPackageInfoByKeyQuery, useStartServices } from '../../hooks';
 import {
   FLEET_KUBERNETES_PACKAGE,
   FLEET_CLOUD_SECURITY_POSTURE_PACKAGE,
   FLEET_CLOUD_DEFEND_PACKAGE,
 } from '../../../common';
-import { getCloudFormationTemplateUrlFromAgentPolicy } from '../../services';
 
-import type { K8sMode, CloudSecurityIntegrationType } from './types';
+import {
+  getCloudFormationTemplateUrlFromPackageInfo,
+  getCloudFormationTemplateUrlFromAgentPolicy,
+} from '../../services';
+
+import type {
+  K8sMode,
+  CloudSecurityIntegrationType,
+  CloudSecurityIntegrationAwsAccountType,
+  CloudSecurityIntegration,
+} from './types';
 
 // Packages that requires custom elastic-agent manifest
 const K8S_PACKAGES = new Set([FLEET_KUBERNETES_PACKAGE, FLEET_CLOUD_DEFEND_PACKAGE]);
@@ -74,19 +83,60 @@ export function useIsK8sPolicy(agentPolicy?: AgentPolicy) {
 }
 
 export function useCloudSecurityIntegration(agentPolicy?: AgentPolicy) {
-  const cloudSecurityIntegration = useMemo(() => {
-    if (!agentPolicy) {
+  const cloudSecurityPackagePolicy = useMemo(() => {
+    return getCloudSecurityPackagePolicyFromAgentPolicy(agentPolicy);
+  }, [agentPolicy]);
+
+  const integrationVersion = cloudSecurityPackagePolicy?.package?.version;
+
+  // Fetch the package info to get the CloudFormation template URL only
+  // if the package policy is a Cloud Security policy
+  const { data: packageInfoData, isLoading } = useGetPackageInfoByKeyQuery(
+    FLEET_CLOUD_SECURITY_POSTURE_PACKAGE,
+    integrationVersion,
+    { full: true },
+    { enabled: Boolean(cloudSecurityPackagePolicy) }
+  );
+
+  const cloudSecurityIntegration: CloudSecurityIntegration | undefined = useMemo(() => {
+    if (!agentPolicy || !cloudSecurityPackagePolicy) {
       return undefined;
     }
 
-    const integrationType = getCloudSecurityIntegrationTypeFromPackagePolicy(agentPolicy);
-    const cloudformationUrl = getCloudFormationTemplateUrlFromAgentPolicy(agentPolicy);
+    const integrationType = cloudSecurityPackagePolicy.inputs?.find((input) => input.enabled)
+      ?.policy_template as CloudSecurityIntegrationType;
+
+    if (!integrationType) return undefined;
+
+    const cloudFormationTemplateFromAgentPolicy =
+      getCloudFormationTemplateUrlFromAgentPolicy(agentPolicy);
+
+    // Use the latest CloudFormation template for the current version
+    // So it guarantee that the template version matches the integration version
+    // when the integration is upgraded.
+    // In case it can't find the template for the current version,
+    // it will fallback to the one from the agent policy.
+    const cloudFormationTemplateUrl = packageInfoData?.item
+      ? getCloudFormationTemplateUrlFromPackageInfo(packageInfoData.item, integrationType)
+      : cloudFormationTemplateFromAgentPolicy;
+
+    const AWS_ACCOUNT_TYPE = 'aws.account_type';
+
+    const cloudFormationAwsAccountType: CloudSecurityIntegrationAwsAccountType | undefined =
+      cloudSecurityPackagePolicy?.inputs?.find((input) => input.enabled)?.streams?.[0]?.vars?.[
+        AWS_ACCOUNT_TYPE
+      ]?.value;
 
     return {
+      isLoading,
       integrationType,
-      cloudformationUrl,
+      isCloudFormation: Boolean(cloudFormationTemplateFromAgentPolicy),
+      cloudFormationProps: {
+        awsAccountType: cloudFormationAwsAccountType,
+        templateUrl: cloudFormationTemplateUrl,
+      },
     };
-  }, [agentPolicy]);
+  }, [agentPolicy, packageInfoData?.item, isLoading, cloudSecurityPackagePolicy]);
 
   return { cloudSecurityIntegration };
 }
@@ -97,13 +147,10 @@ const isK8sPackage = (pkg: PackagePolicy) => {
   return K8S_PACKAGES.has(name);
 };
 
-const getCloudSecurityIntegrationTypeFromPackagePolicy = (
-  agentPolicy: AgentPolicy
-): CloudSecurityIntegrationType | undefined => {
-  const packagePolicy = agentPolicy?.package_policies?.find(
+const getCloudSecurityPackagePolicyFromAgentPolicy = (
+  agentPolicy?: AgentPolicy
+): PackagePolicy | undefined => {
+  return agentPolicy?.package_policies?.find(
     (input) => input.package?.name === FLEET_CLOUD_SECURITY_POSTURE_PACKAGE
   );
-  if (!packagePolicy) return undefined;
-  return packagePolicy?.inputs?.find((input) => input.enabled)
-    ?.policy_template as CloudSecurityIntegrationType;
 };

x-pack/plugins/fleet/public/components/agent_enrollment_flyout/instructions.tsx

@@ -80,8 +80,8 @@ export const Instructions = (props: InstructionProps) => {
       (fleetStatus.missingRequirements ?? []).some((r) => r === FLEET_SERVER_PACKAGE));
 
   useEffect(() => {
-    // If we have a cloudFormationTemplateUrl, we want to hide the selection type
-    if (props.cloudSecurityIntegration?.cloudformationUrl) {
+    // If we detect a CloudFormation integration, we want to hide the selection type
+    if (props.cloudSecurityIntegration?.isCloudFormation) {
       setSelectionType(undefined);
     } else if (!isIntegrationFlow && showAgentEnrollment) {
       setSelectionType('radio');
@@ -117,10 +117,7 @@ export const Instructions = (props: InstructionProps) => {
           {isFleetServerPolicySelected ? (
             <AdvancedTab selectedPolicyId={props.selectedPolicy?.id} onClose={() => undefined} />
           ) : (
-            <ManagedSteps
-              {...props}
-              cloudFormationTemplateUrl={props.cloudSecurityIntegration?.cloudformationUrl}
-            />
+            <ManagedSteps {...props} />
           )}
         </>
       );

x-pack/plugins/fleet/public/components/agent_enrollment_flyout/steps/compute_steps.tsx

@@ -200,7 +200,6 @@ export const ManagedSteps: React.FunctionComponent<InstructionProps> = ({
   isK8s,
   cloudSecurityIntegration,
   installedPackagePolicy,
-  cloudFormationTemplateUrl,
 }) => {
   const kibanaVersion = useKibanaVersion();
   const core = useStartServices();
@@ -247,14 +246,13 @@ export const ManagedSteps: React.FunctionComponent<InstructionProps> = ({
       );
     }
 
-    if (cloudFormationTemplateUrl) {
+    if (cloudSecurityIntegration?.isCloudFormation) {
       steps.push(
         InstallCloudFormationManagedAgentStep({
           apiKeyData,
           selectedApiKeyId,
           enrollToken,
-          cloudFormationTemplateUrl,
-          agentPolicy,
+          cloudSecurityIntegration,
         })
       );
     } else {
@@ -314,7 +312,6 @@ export const ManagedSteps: React.FunctionComponent<InstructionProps> = ({
     link,
     agentDataConfirmed,
     installedPackagePolicy,
-    cloudFormationTemplateUrl,
   ]);
 
   return <EuiSteps steps={instructionsSteps} />;

x-pack/plugins/fleet/public/components/agent_enrollment_flyout/steps/install_cloud_formation_managed_agent_step.tsx

@@ -11,46 +11,38 @@ import { i18n } from '@kbn/i18n';
 
 import type { EuiContainedStepProps } from '@elastic/eui/src/components/steps/steps';
 
-import type { AgentPolicy } from '../../../../common';
-
 import type { GetOneEnrollmentAPIKeyResponse } from '../../../../common/types/rest_spec/enrollment_api_key';
 
 import { CloudFormationInstructions } from '../cloud_formation_instructions';
-import { FLEET_CLOUD_SECURITY_POSTURE_PACKAGE } from '../../../../common';
+
+import type { CloudSecurityIntegration } from '../types';
 
 export const InstallCloudFormationManagedAgentStep = ({
   selectedApiKeyId,
   apiKeyData,
   enrollToken,
   isComplete,
-  cloudFormationTemplateUrl,
-  agentPolicy,
+  cloudSecurityIntegration,
 }: {
   selectedApiKeyId?: string;
   apiKeyData?: GetOneEnrollmentAPIKeyResponse | null;
   enrollToken?: string;
   isComplete?: boolean;
-  cloudFormationTemplateUrl: string;
-  agentPolicy?: AgentPolicy;
+  cloudSecurityIntegration?: CloudSecurityIntegration | undefined;
 }): EuiContainedStepProps => {
   const nonCompleteStatus = selectedApiKeyId ? undefined : 'disabled';
   const status = isComplete ? 'complete' : nonCompleteStatus;
 
-  const cloudSecurityPackagePolicy = agentPolicy?.package_policies?.find(
-    (p) => p.package?.name === FLEET_CLOUD_SECURITY_POSTURE_PACKAGE
-  );
-
   return {
     status,
     title: i18n.translate('xpack.fleet.agentEnrollment.cloudFormation.stepEnrollAndRunAgentTitle', {
       defaultMessage: 'Install Elastic Agent on your cloud',
     }),
     children:
-      selectedApiKeyId && apiKeyData ? (
+      selectedApiKeyId && apiKeyData && cloudSecurityIntegration ? (
         <CloudFormationInstructions
-          cloudFormationTemplateUrl={cloudFormationTemplateUrl}
+          cloudSecurityIntegration={cloudSecurityIntegration}
           enrollmentAPIKey={enrollToken}
-          packagePolicy={cloudSecurityPackagePolicy}
         />
       ) : (
         <React.Fragment />

x-pack/plugins/fleet/public/components/agent_enrollment_flyout/types.ts

@@ -16,13 +16,21 @@ export type K8sMode =
   | 'IS_KUBERNETES_MULTIPAGE';
 
 export type CloudSecurityIntegrationType = 'kspm' | 'vuln_mgmt' | 'cspm';
+export type CloudSecurityIntegrationAwsAccountType = 'single-account' | 'organization-account';
 
 export type FlyoutMode = 'managed' | 'standalone';
 export type SelectionType = 'tabs' | 'radio' | undefined;
 
+export interface CloudFormationProps {
+  templateUrl: string | undefined;
+  awsAccountType: CloudSecurityIntegrationAwsAccountType | undefined;
+}
+
 export interface CloudSecurityIntegration {
   integrationType: CloudSecurityIntegrationType | undefined;
-  cloudformationUrl: string | undefined;
+  isLoading: boolean;
+  isCloudFormation: boolean;
+  cloudFormationProps?: CloudFormationProps;
 }
 
 export interface BaseProps {
@@ -65,5 +73,4 @@ export interface InstructionProps extends BaseProps {
   setSelectedAPIKeyId: (key?: string) => void;
   fleetServerHosts: string[];
   fleetProxy?: FleetProxy;
-  cloudFormationTemplateUrl?: string;
 }

x-pack/plugins/fleet/public/hooks/use_create_cloud_formation_url.ts

@@ -7,34 +7,27 @@
 
 import { i18n } from '@kbn/i18n';
 
-import type { PackagePolicy, PackagePolicyInput } from '../../common';
+import type {
+  CloudFormationProps,
+  CloudSecurityIntegrationAwsAccountType,
+} from '../components/agent_enrollment_flyout/types';
 
 import { useKibanaVersion } from './use_kibana_version';
 import { useGetSettings } from './use_request';
 
-type AwsAccountType = 'single_account' | 'organization_account';
-
-const CLOUDBEAT_AWS = 'cloudbeat/cis_aws';
-
-const getAwsAccountType = (input?: PackagePolicyInput): AwsAccountType | undefined =>
-  input?.streams[0].vars?.['aws.account_type']?.value;
+const CLOUD_FORMATION_DEFAULT_ACCOUNT_TYPE = 'single-account';
 
 export const useCreateCloudFormationUrl = ({
   enrollmentAPIKey,
-  cloudFormationTemplateUrl,
-  packagePolicy,
+  cloudFormationProps,
 }: {
   enrollmentAPIKey: string | undefined;
-  cloudFormationTemplateUrl: string;
-  packagePolicy?: PackagePolicy;
+  cloudFormationProps: CloudFormationProps | undefined;
 }) => {
   const { data, isLoading } = useGetSettings();
 
   const kibanaVersion = useKibanaVersion();
 
-  const awsInput = packagePolicy?.inputs?.find((input) => input.type === CLOUDBEAT_AWS);
-  const awsAccountType = getAwsAccountType(awsInput) || '';
-
   let isError = false;
   let error: string | undefined;
 
@@ -56,13 +49,13 @@ export const useCreateCloudFormationUrl = ({
   }
 
   const cloudFormationUrl =
-    enrollmentAPIKey && fleetServerHost && cloudFormationTemplateUrl
+    enrollmentAPIKey && fleetServerHost && cloudFormationProps?.templateUrl
       ? createCloudFormationUrl(
-          cloudFormationTemplateUrl,
+          cloudFormationProps?.templateUrl,
           enrollmentAPIKey,
           fleetServerHost,
           kibanaVersion,
-          awsAccountType
+          cloudFormationProps?.awsAccountType
         )
       : undefined;
 
@@ -79,7 +72,7 @@ const createCloudFormationUrl = (
   enrollmentToken: string,
   fleetUrl: string,
   kibanaVersion: string,
-  awsAccountType: string
+  awsAccountType: CloudSecurityIntegrationAwsAccountType | undefined
 ) => {
   let cloudFormationUrl;
 
@@ -89,8 +82,15 @@ const createCloudFormationUrl = (
     .replace('KIBANA_VERSION', kibanaVersion);
 
   if (cloudFormationUrl.includes('ACCOUNT_TYPE')) {
-    cloudFormationUrl = cloudFormationUrl.replace('ACCOUNT_TYPE', awsAccountType);
+    cloudFormationUrl = cloudFormationUrl.replace(
+      'ACCOUNT_TYPE',
+      getAwsAccountType(awsAccountType)
+    );
   }
 
   return new URL(cloudFormationUrl).toString();
 };
+
+const getAwsAccountType = (awsAccountType: CloudSecurityIntegrationAwsAccountType | undefined) => {
+  return awsAccountType ? awsAccountType : CLOUD_FORMATION_DEFAULT_ACCOUNT_TYPE;
+};

x-pack/plugins/fleet/public/hooks/use_request/epm.ts

@@ -105,6 +105,13 @@ export const useGetPackageInfoByKeyQuery = (
     ignoreUnverified?: boolean;
     prerelease?: boolean;
     full?: boolean;
+  },
+  // Additional options for the useQuery hook
+  queryOptions: {
+    // If enabled is false, the query will not be fetched
+    enabled?: boolean;
+  } = {
+    enabled: true,
   }
 ) => {
   const confirmOpenUnverified = useConfirmOpenUnverified();
@@ -112,15 +119,18 @@ export const useGetPackageInfoByKeyQuery = (
     options?.ignoreUnverified
   );
 
-  const response = useQuery<GetInfoResponse, RequestError>([pkgName, pkgVersion, options], () =>
-    sendRequestForRq<GetInfoResponse>({
-      path: epmRouteService.getInfoPath(pkgName, pkgVersion),
-      method: 'get',
-      query: {
-        ...options,
-        ...(ignoreUnverifiedQueryParam && { ignoreUnverified: ignoreUnverifiedQueryParam }),
-      },
-    })
+  const response = useQuery<GetInfoResponse, RequestError>(
+    [pkgName, pkgVersion, options],
+    () =>
+      sendRequestForRq<GetInfoResponse>({
+        path: epmRouteService.getInfoPath(pkgName, pkgVersion),
+        method: 'get',
+        query: {
+          ...options,
+          ...(ignoreUnverifiedQueryParam && { ignoreUnverified: ignoreUnverifiedQueryParam }),
+        },
+      }),
+    { enabled: queryOptions.enabled }
   );
 
   const confirm = async () => {

x-pack/plugins/fleet/public/services/get_cloud_formation_props_from_package_policy.test.ts

@@ -0,0 +1,126 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { getCloudFormationPropsFromPackagePolicy } from './get_cloud_formation_props_from_package_policy';
+
+describe('getCloudFormationPropsFromPackagePolicy', () => {
+  test('returns empty CloudFormationProps when packagePolicy is undefined', () => {
+    const result = getCloudFormationPropsFromPackagePolicy(undefined);
+    expect(result).toEqual({
+      templateUrl: undefined,
+      awsAccountType: undefined,
+    });
+  });
+
+  test('returns empty CloudFormationProps when packagePolicy has no inputs', () => {
+    const packagePolicy = { otherProperty: 'value' };
+    // @ts-expect-error
+    const result = getCloudFormationPropsFromPackagePolicy(packagePolicy);
+    expect(result).toEqual({
+      templateUrl: undefined,
+      awsAccountType: undefined,
+    });
+  });
+
+  test('returns empty CloudFormationProps when no enabled input has a cloudFormationTemplateUrl', () => {
+    const packagePolicy = {
+      inputs: [
+        { enabled: false, config: { cloud_formation_template_url: { value: 'template1' } } },
+        { enabled: false, config: { cloud_formation_template_url: { value: 'template2' } } },
+      ],
+    };
+    // @ts-expect-error
+    const result = getCloudFormationPropsFromPackagePolicy(packagePolicy);
+    expect(result).toEqual({
+      templateUrl: undefined,
+      awsAccountType: undefined,
+    });
+  });
+
+  test('returns the cloudFormationTemplateUrl and awsAccountType when found in the enabled input', () => {
+    const packagePolicy = {
+      inputs: [
+        {
+          enabled: true,
+          config: { cloud_formation_template_url: { value: 'template1' } },
+          streams: [
+            {
+              vars: {
+                ['aws.account_type']: { value: 'aws_account_type_value' },
+              },
+            },
+          ],
+        },
+        {
+          enabled: false,
+          config: { cloud_formation_template_url: { value: 'template2' } },
+          streams: [
+            {
+              vars: {
+                ['aws.account_type']: { value: 'aws_account_type_value2' },
+              },
+            },
+          ],
+        },
+      ],
+    };
+    // @ts-expect-error
+    const result = getCloudFormationPropsFromPackagePolicy(packagePolicy);
+    expect(result).toEqual({
+      templateUrl: 'template1',
+      awsAccountType: 'aws_account_type_value',
+    });
+  });
+
+  test('returns the first cloudFormationTemplateUrl and awsAccountType when multiple enabled inputs have them', () => {
+    const packagePolicy = {
+      inputs: [
+        {
+          enabled: true,
+          config: {
+            cloud_formation_template_url: { value: 'template1' },
+          },
+          streams: [
+            {
+              vars: {
+                ['aws.account_type']: { value: 'aws_account_type_value1' },
+              },
+            },
+            {
+              vars: {
+                ['aws.account_type']: { value: 'aws_account_type_value2' },
+              },
+            },
+          ],
+        },
+        {
+          enabled: true,
+          config: {
+            cloud_formation_template_url: { value: 'template2' },
+          },
+          streams: [
+            {
+              vars: {
+                ['aws.account_type']: { value: 'aws_account_type_value1' },
+              },
+            },
+            {
+              vars: {
+                ['aws.account_type']: { value: 'aws_account_type_value2' },
+              },
+            },
+          ],
+        },
+      ],
+    };
+    // @ts-expect-error
+    const result = getCloudFormationPropsFromPackagePolicy(packagePolicy);
+    expect(result).toEqual({
+      templateUrl: 'template1',
+      awsAccountType: 'aws_account_type_value1',
+    });
+  });
+});

x-pack/plugins/fleet/public/services/get_cloud_formation_props_from_package_policy.ts

@@ -5,15 +5,23 @@
  * 2.0.
  */
 
+import type {
+  CloudFormationProps,
+  CloudSecurityIntegrationAwsAccountType,
+} from '../components/agent_enrollment_flyout/types';
 import type { PackagePolicy } from '../types';
 
+const AWS_ACCOUNT_TYPE = 'aws.account_type';
+
 /**
  * Get the cloud formation template url from a package policy
  * It looks for a config with a cloud_formation_template_url object present in
  * the enabled inputs of the package policy
  */
-export const getCloudFormationTemplateUrlFromPackagePolicy = (packagePolicy?: PackagePolicy) => {
-  const cloudFormationTemplateUrl = packagePolicy?.inputs?.reduce((accInput, input) => {
+export const getCloudFormationPropsFromPackagePolicy = (
+  packagePolicy?: PackagePolicy
+): CloudFormationProps => {
+  const templateUrl = packagePolicy?.inputs?.reduce((accInput, input) => {
     if (accInput !== '') {
       return accInput;
     }
@@ -23,5 +31,12 @@ export const getCloudFormationTemplateUrlFromPackagePolicy = (packagePolicy?: Pa
     return accInput;
   }, '');
 
-  return cloudFormationTemplateUrl !== '' ? cloudFormationTemplateUrl : undefined;
+  const awsAccountType: CloudSecurityIntegrationAwsAccountType | undefined =
+    packagePolicy?.inputs?.find((input) => input.enabled)?.streams?.[0]?.vars?.[AWS_ACCOUNT_TYPE]
+      ?.value;
+
+  return {
+    templateUrl: templateUrl !== '' ? templateUrl : undefined,
+    awsAccountType,
+  };
 };

x-pack/plugins/fleet/public/services/get_cloud_formation_template_url_from_package_info.test.ts

@@ -0,0 +1,66 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { getCloudFormationTemplateUrlFromPackageInfo } from './get_cloud_formation_template_url_from_package_info';
+
+describe('getCloudFormationTemplateUrlFromPackageInfo', () => {
+  test('returns undefined when packageInfo is undefined', () => {
+    const result = getCloudFormationTemplateUrlFromPackageInfo(undefined, 'test');
+    expect(result).toBeUndefined();
+  });
+
+  test('returns undefined when packageInfo has no policy_templates', () => {
+    const packageInfo = { inputs: [] };
+    // @ts-expect-error
+    const result = getCloudFormationTemplateUrlFromPackageInfo(packageInfo, 'test');
+    expect(result).toBeUndefined();
+  });
+
+  test('returns undefined when integrationType is not found in policy_templates', () => {
+    const packageInfo = { policy_templates: [{ name: 'template1' }, { name: 'template2' }] };
+    // @ts-expect-error
+    const result = getCloudFormationTemplateUrlFromPackageInfo(packageInfo, 'nonExistentTemplate');
+    expect(result).toBeUndefined();
+  });
+
+  test('returns undefined when no input in the policy template has a cloudFormationTemplate', () => {
+    const packageInfo = {
+      policy_templates: [
+        {
+          name: 'template1',
+          inputs: [
+            { name: 'input1', vars: [] },
+            { name: 'input2', vars: [{ name: 'var1', default: 'value1' }] },
+          ],
+        },
+      ],
+    };
+    // @ts-expect-error
+    const result = getCloudFormationTemplateUrlFromPackageInfo(packageInfo, 'template1');
+    expect(result).toBeUndefined();
+  });
+
+  test('returns the cloudFormationTemplate from the policy template', () => {
+    const packageInfo = {
+      policy_templates: [
+        {
+          name: 'template1',
+          inputs: [
+            { name: 'input1', vars: [] },
+            {
+              name: 'input2',
+              vars: [{ name: 'cloud_formation_template', default: 'cloud_formation_template_url' }],
+            },
+          ],
+        },
+      ],
+    };
+    // @ts-expect-error
+    const result = getCloudFormationTemplateUrlFromPackageInfo(packageInfo, 'template1');
+    expect(result).toBe('cloud_formation_template_url');
+  });
+});

x-pack/plugins/fleet/public/services/get_cloud_formation_template_url_from_package_info.ts

@@ -0,0 +1,32 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { PackageInfo } from '../types';
+
+/**
+ * Get the cloud formation template url from the PackageInfo
+ * It looks for a input var with a object containing cloud_formation_template_url present in
+ * the package_policies inputs of the given integration type
+ */
+export const getCloudFormationTemplateUrlFromPackageInfo = (
+  packageInfo: PackageInfo | undefined,
+  integrationType: string
+): string | undefined => {
+  if (!packageInfo?.policy_templates) return undefined;
+
+  const policyTemplate = packageInfo.policy_templates.find((p) => p.name === integrationType);
+  if (!policyTemplate) return undefined;
+
+  if ('inputs' in policyTemplate) {
+    const cloudFormationTemplate = policyTemplate.inputs?.reduce((acc, input): string => {
+      if (!input.vars) return acc;
+      const template = input.vars.find((v) => v.name === 'cloud_formation_template')?.default;
+      return template ? String(template) : acc;
+    }, '');
+    return cloudFormationTemplate !== '' ? cloudFormationTemplate : undefined;
+  }
+};

x-pack/plugins/fleet/public/services/get_cloud_formation_template_url_from_package_policy.test.ts

@@ -1,61 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-import { getCloudFormationTemplateUrlFromPackagePolicy } from './get_cloud_formation_template_url_from_package_policy';
-
-describe('getCloudFormationTemplateUrlFromPackagePolicy', () => {
-  test('returns undefined when packagePolicy is undefined', () => {
-    const result = getCloudFormationTemplateUrlFromPackagePolicy(undefined);
-    expect(result).toBeUndefined();
-  });
-
-  test('returns undefined when packagePolicy is defined but inputs are empty', () => {
-    const packagePolicy = { inputs: [] };
-    // @ts-expect-error
-    const result = getCloudFormationTemplateUrlFromPackagePolicy(packagePolicy);
-    expect(result).toBeUndefined();
-  });
-
-  test('returns undefined when no enabled input has a cloudFormationTemplateUrl', () => {
-    const packagePolicy = {
-      inputs: [
-        { enabled: false, config: { cloud_formation_template_url: { value: 'template1' } } },
-        { enabled: false, config: { cloud_formation_template_url: { value: 'template2' } } },
-      ],
-    };
-    // @ts-expect-error
-    const result = getCloudFormationTemplateUrlFromPackagePolicy(packagePolicy);
-    expect(result).toBeUndefined();
-  });
-
-  test('returns the cloudFormationTemplateUrl of the first enabled input', () => {
-    const packagePolicy = {
-      inputs: [
-        { enabled: false, config: { cloud_formation_template_url: { value: 'template1' } } },
-        { enabled: true, config: { cloud_formation_template_url: { value: 'template2' } } },
-        { enabled: true, config: { cloud_formation_template_url: { value: 'template3' } } },
-      ],
-    };
-    // @ts-expect-error
-    const result = getCloudFormationTemplateUrlFromPackagePolicy(packagePolicy);
-    expect(result).toBe('template2');
-  });
-
-  test('returns the cloudFormationTemplateUrl of the first enabled input and ignores subsequent inputs', () => {
-    const packagePolicy = {
-      inputs: [
-        { enabled: true, config: { cloud_formation_template_url: { value: 'template1' } } },
-        { enabled: true, config: { cloud_formation_template_url: { value: 'template2' } } },
-        { enabled: true, config: { cloud_formation_template_url: { value: 'template3' } } },
-      ],
-    };
-    // @ts-expect-error
-    const result = getCloudFormationTemplateUrlFromPackagePolicy(packagePolicy);
-    expect(result).toBe('template1');
-  });
-
-  // Add more test cases as needed
-});

x-pack/plugins/fleet/public/services/index.ts

@@ -48,5 +48,6 @@ export { isPackageUpdatable } from './is_package_updatable';
 export { pkgKeyFromPackageInfo } from './pkg_key_from_package_info';
 export { createExtensionRegistrationCallback } from './ui_extensions';
 export { incrementPolicyName } from './increment_policy_name';
-export { getCloudFormationTemplateUrlFromPackagePolicy } from './get_cloud_formation_template_url_from_package_policy';
+export { getCloudFormationPropsFromPackagePolicy } from './get_cloud_formation_props_from_package_policy';
 export { getCloudFormationTemplateUrlFromAgentPolicy } from './get_cloud_formation_template_url_from_agent_policy';
+export { getCloudFormationTemplateUrlFromPackageInfo } from './get_cloud_formation_template_url_from_package_info';