[Cloud Posture] Create rules configuration when create new package

x-pack/plugins/cloud_security_posture/server/plugin.test.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { coreMock, httpServerMock, savedObjectsClientMock } from '@kbn/core/server/mocks';
+import { coreMock, httpServerMock } from '@kbn/core/server/mocks';
 import {
   createPackagePolicyServiceMock,
   createArtifactsClientMock,
@@ -19,7 +19,12 @@ import { createPackagePolicyMock } from '@kbn/fleet-plugin/common/mocks';
 import { dataPluginMock } from '@kbn/data-plugin/server/mocks';
 import { CspPlugin } from './plugin';
 import { CspServerPluginStartDeps } from './types';
-import { createFleetAuthzMock, Installation } from '@kbn/fleet-plugin/common';
+import {
+  createFleetAuthzMock,
+  Installation,
+  PackagePolicy,
+  UpdatePackagePolicy,
+} from '@kbn/fleet-plugin/common';
 import {
   ExternalCallback,
   FleetStartContract,
@@ -29,16 +34,16 @@ import { CLOUD_SECURITY_POSTURE_PACKAGE_NAME } from '../common/constants';
 import Chance from 'chance';
 import type { AwaitedProperties } from '@kbn/utility-types';
 import type { DeeplyMockedKeys } from '@kbn/utility-types/jest';
-import { RequestHandlerContext } from '@kbn/core/server';
+import {
+  ElasticsearchClient,
+  RequestHandlerContext,
+  SavedObjectsClientContract,
+} from '@kbn/core/server';
 
 const chance = new Chance();
 
 const mockRouteContext = {
-  core: {
-    savedObjects: {
-      client: savedObjectsClientMock.create(),
-    },
-  },
+  core: coreMock.createRequestHandlerContext(),
 } as unknown as AwaitedProperties<RequestHandlerContext>;
 
 const createMockFleetStartContract = (): DeeplyMockedKeys<FleetStartContract> => {
@@ -200,5 +205,64 @@ describe('Cloud Security Posture Plugin', () => {
 
       expect(spy).toHaveBeenCalledTimes(0);
     });
+
+    it('packagePolicyPostCreate should return the updated packagePolicy', async () => {
+      fleetMock.packageService.asInternalUser.getInstallation.mockImplementationOnce(
+        async (): Promise<Installation | undefined> => {
+          return;
+        }
+      );
+
+      fleetMock.packagePolicyService.update.mockImplementation(
+        (
+          soClient: SavedObjectsClientContract,
+          esClient: ElasticsearchClient,
+          id: string,
+          packagePolicyUpdate: UpdatePackagePolicy
+        ): Promise<PackagePolicy> => {
+          // @ts-expect-error 2322
+          return packagePolicyUpdate;
+        }
+      );
+
+      const packageMock = createPackagePolicyMock();
+      packageMock.package!.name = CLOUD_SECURITY_POSTURE_PACKAGE_NAME;
+      packageMock.vars = { dataYaml: { type: 'foo' } };
+
+      const packagePolicyPostCreateCallbacks: PostPackagePolicyPostCreateCallback[] = [];
+      fleetMock.registerExternalCallback.mockImplementation((...args) => {
+        if (args[0] === 'packagePolicyPostCreate') {
+          packagePolicyPostCreateCallbacks.push(args[1]);
+        }
+      });
+
+      const context = coreMock.createPluginInitializerContext<unknown>();
+      plugin = new CspPlugin(context);
+      const spy = jest.spyOn(plugin, 'initialize').mockImplementation();
+
+      // Act
+      await plugin.start(coreMock.createStart(), mockPlugins);
+      await mockPlugins.fleet.fleetSetupCompleted();
+
+      // Assert
+      expect(fleetMock.packageService.asInternalUser.getInstallation).toHaveBeenCalledTimes(1);
+      expect(spy).toHaveBeenCalledTimes(0);
+
+      expect(packagePolicyPostCreateCallbacks.length).toBeGreaterThan(0);
+
+      for (const cb of packagePolicyPostCreateCallbacks) {
+        const updatedPackagePolicy = await cb(
+          packageMock,
+          contextMock,
+          httpServerMock.createKibanaRequest()
+        );
+        if (fleetMock.packagePolicyService.update.mock.calls.length) {
+          expect(updatedPackagePolicy).toHaveProperty('vars');
+          expect(updatedPackagePolicy.vars).toHaveProperty('dataYaml');
+          expect(updatedPackagePolicy.vars!.dataYaml).toHaveProperty('value');
+        }
+      }
+      expect(fleetMock.packagePolicyService.update).toHaveBeenCalledTimes(1);
+    });
   });
 });

x-pack/plugins/cloud_security_posture/server/plugin.ts

@@ -40,6 +40,7 @@ import {
   removeCspRulesInstancesCallback,
 } from './fleet_integration/fleet_integration';
 import { CLOUD_SECURITY_POSTURE_PACKAGE_NAME } from '../common/constants';
+import { updateAgentConfiguration } from './routes/configuration/update_rules_configuration';
 
 import {
   removeFindingsStatsTask,
@@ -119,7 +120,16 @@ export class CspPlugin
             await this.initialize(core, plugins.taskManager);
 
             const soClient = (await context.core).savedObjects.client;
+            const esClient = (await context.core).elasticsearch.client.asCurrentUser;
             await onPackagePolicyPostCreateCallback(this.logger, packagePolicy, soClient);
+
+            const updatedPackagePolicy = await updateAgentConfiguration(
+              plugins.fleet.packagePolicyService,
+              packagePolicy,
+              esClient,
+              soClient
+            );
+            return updatedPackagePolicy;
           }
 
           return packagePolicy;

x-pack/plugins/cloud_security_posture/server/routes/configuration/update_rules_configuration.test.ts

@@ -17,7 +17,7 @@ import {
   defineUpdateRulesConfigRoute,
   getCspRules,
   setVarToPackagePolicy,
-  updatePackagePolicy,
+  updateAgentConfiguration,
 } from './update_rules_configuration';
 
 import { CspAppService } from '../../lib/csp_app_services';
@@ -35,6 +35,7 @@ import {
   SavedObjectsFindResponse,
 } from '@kbn/core/server';
 import { Chance } from 'chance';
+import { PackagePolicy, UpdatePackagePolicy } from '@kbn/fleet-plugin/common';
 
 describe('Update rules configuration API', () => {
   let logger: ReturnType<typeof loggingSystemMock.createLogger>;
@@ -173,39 +174,130 @@ describe('Update rules configuration API', () => {
     expect(cspConfig).toMatchObject({ data_yaml: { activated_rules: { cis_k8s: [] } } });
   });
 
-  it('validate adding new data.yaml to package policy instance', async () => {
+  it('validate adding new dataYaml to package policy instance', async () => {
+    const packagePolicy = createPackagePolicyMock();
+    packagePolicy.vars = { dataYaml: { type: 'yaml' } };
+
+    const dataYaml = 'data_yaml:\n  activated_rules:\n  cis_k8s:\n    - 1.1.1\n    - 1.1.2\n';
+    const updatedPackagePolicy = setVarToPackagePolicy(packagePolicy, dataYaml);
+    expect(updatedPackagePolicy.vars).toEqual({ dataYaml: { type: 'yaml', value: dataYaml } });
+  });
+
+  it('validate adding new datYaml to package policy instance when it not exists on source', async () => {
     const packagePolicy = createPackagePolicyMock();
 
     const dataYaml = 'data_yaml:\n  activated_rules:\n  cis_k8s:\n    - 1.1.1\n    - 1.1.2\n';
     const updatedPackagePolicy = setVarToPackagePolicy(packagePolicy, dataYaml);
-
-    expect(updatedPackagePolicy.vars).toEqual({ dataYaml: { type: 'config', value: dataYaml } });
+    expect(updatedPackagePolicy.vars).toEqual({ dataYaml: { type: 'yaml', value: dataYaml } });
   });
 
-  it('validate updatePackagePolicy is called with the right parameters', async () => {
+  it('verify that the API for updating package policy was invoked', async () => {
     mockEsClient = elasticsearchClientMock.createClusterClient().asScoped().asInternalUser;
     mockSoClient = savedObjectsClientMock.create();
     const mockPackagePolicyService = createPackagePolicyServiceMock();
 
-    const packagePolicyId1 = chance.guid();
-    const packagePolicyId2 = chance.guid();
-    const mockPackagePolicy1 = createPackagePolicyMock();
-    const mockPackagePolicy2 = createPackagePolicyMock();
-    mockPackagePolicy1.id = packagePolicyId1;
-    mockPackagePolicy2.id = packagePolicyId2;
-    const packagePolicies = mockPackagePolicy1;
-
-    const dataYaml = 'activated_rules:\n  cis_k8s:\n    - 1.1.1\n    - 1.1.2\n';
-
-    await updatePackagePolicy(
-      mockPackagePolicyService,
-      packagePolicies,
-      mockEsClient,
-      mockSoClient,
-      dataYaml
+    mockPackagePolicyService.update.mockImplementation(
+      (
+        soClient: SavedObjectsClientContract,
+        esClient: ElasticsearchClient,
+        id: string,
+        packagePolicyUpdate: UpdatePackagePolicy
+      ): Promise<PackagePolicy> => {
+        // @ts-expect-error 2322
+        return packagePolicyUpdate;
+      }
     );
 
+    mockSoClient.find.mockResolvedValueOnce({
+      page: 1,
+      per_page: 1000,
+      total: 2,
+      saved_objects: [
+        {
+          type: 'csp_rule',
+          rego_rule_id: '1.1.1',
+          attributes: { enabled: false, rego_rule_id: 'cis_1_1_1' },
+        },
+        {
+          type: 'csp_rule',
+          attributes: { enabled: false, rego_rule_id: 'cis_1_1_2' },
+        },
+        {
+          type: 'csp_rule',
+          attributes: { enabled: false, rego_rule_id: 'cis_1_1_3' },
+        },
+      ],
+    } as unknown as SavedObjectsFindResponse<CspRuleSchema>);
+
+    const mockPackagePolicy = createPackagePolicyMock();
+    mockPackagePolicy.vars = { dataYaml: { type: 'foo' } };
+    const packagePolicyId1 = chance.guid();
+    mockPackagePolicy.id = packagePolicyId1;
+
+    const updatePackagePolicy = await updateAgentConfiguration(
+      mockPackagePolicyService,
+      mockPackagePolicy,
+      mockEsClient,
+      mockSoClient
+    );
+
+    expect(updatePackagePolicy.vars!.dataYaml).toHaveProperty('value');
+    expect(updatePackagePolicy.vars!.dataYaml).toMatchObject({ type: 'yaml' });
     expect(mockPackagePolicyService.update).toBeCalledTimes(1);
     expect(mockPackagePolicyService.update.mock.calls[0][2]).toEqual(packagePolicyId1);
   });
+
+  it('validate updateAgentConfiguration not override vars', async () => {
+    mockEsClient = elasticsearchClientMock.createClusterClient().asScoped().asInternalUser;
+    mockSoClient = savedObjectsClientMock.create();
+    const mockPackagePolicyService = createPackagePolicyServiceMock();
+
+    mockSoClient.find.mockResolvedValueOnce({
+      page: 1,
+      per_page: 1000,
+      total: 2,
+      saved_objects: [
+        {
+          type: 'csp_rule',
+          rego_rule_id: '1.1.1',
+          attributes: { enabled: false, rego_rule_id: 'cis_1_1_1' },
+        },
+        {
+          type: 'csp_rule',
+          attributes: { enabled: false, rego_rule_id: 'cis_1_1_2' },
+        },
+        {
+          type: 'csp_rule',
+          attributes: { enabled: false, rego_rule_id: 'cis_1_1_3' },
+        },
+      ],
+    } as unknown as SavedObjectsFindResponse<CspRuleSchema>);
+
+    const mockPackagePolicy = createPackagePolicyMock();
+    const packagePolicyId1 = chance.guid();
+    mockPackagePolicy.id = packagePolicyId1;
+    mockPackagePolicy.vars = { foo: {}, dataYaml: { type: 'yaml' } };
+
+    mockPackagePolicyService.update.mockImplementation(
+      (
+        soClient: SavedObjectsClientContract,
+        esClient: ElasticsearchClient,
+        id: string,
+        packagePolicyUpdate: UpdatePackagePolicy
+      ): Promise<PackagePolicy> => {
+        // @ts-expect-error 2322
+        return packagePolicyUpdate;
+      }
+    );
+
+    const updatedPackagePolicy = await updateAgentConfiguration(
+      mockPackagePolicyService,
+      mockPackagePolicy,
+      mockEsClient,
+      mockSoClient
+    );
+
+    expect(mockPackagePolicyService.update).toBeCalledTimes(1);
+    expect(updatedPackagePolicy.vars).toHaveProperty('foo');
+  });
 });

x-pack/plugins/cloud_security_posture/server/routes/configuration/update_rules_configuration.ts

@@ -84,25 +84,30 @@ export const setVarToPackagePolicy = (
   dataYaml: string
 ): PackagePolicy => {
   const configFile: PackagePolicyConfigRecord = {
-    dataYaml: { type: 'config', value: dataYaml },
+    dataYaml: { type: 'yaml', value: dataYaml },
   };
   const updatedPackagePolicy = produce(packagePolicy, (draft) => {
     unset(draft, 'id');
-    draft.vars = configFile;
-    // TODO: disable comments after adding base config to integration
-    // draft.inputs[0].vars = configFile;
+    if (draft.vars) {
+      draft.vars.dataYaml = configFile.dataYaml;
+    } else {
+      draft.vars = configFile;
+    }
   });
   return updatedPackagePolicy;
 };
 
-export const updatePackagePolicy = (
+export const updateAgentConfiguration = async (
   packagePolicyService: PackagePolicyServiceInterface,
   packagePolicy: PackagePolicy,
   esClient: ElasticsearchClient,
-  soClient: SavedObjectsClientContract,
-  dataYaml: string
+  soClient: SavedObjectsClientContract
 ): Promise<PackagePolicy> => {
+  const cspRules = await getCspRules(soClient, packagePolicy);
+  const rulesConfig = createRulesConfig(cspRules);
+  const dataYaml = convertRulesConfigToYaml(rulesConfig);
   const updatedPackagePolicy = setVarToPackagePolicy(packagePolicy, dataYaml);
+
   return packagePolicyService.update(soClient, esClient, packagePolicy.id, updatedPackagePolicy);
 };
 
@@ -133,19 +138,14 @@ export const defineUpdateRulesConfigRoute = (router: CspRouter, cspContext: CspA
           packagePolicyId
         );
 
-        const cspRules = await getCspRules(soClient, packagePolicy);
-        const rulesConfig = createRulesConfig(cspRules);
-        const dataYaml = convertRulesConfigToYaml(rulesConfig);
-
-        const updatedPackagePolicies = await updatePackagePolicy(
-          packagePolicyService!,
+        const updatedPackagePolicy = await updateAgentConfiguration(
+          packagePolicyService,
           packagePolicy,
           esClient,
-          soClient,
-          dataYaml
+          soClient
         );
 
-        return response.ok({ body: updatedPackagePolicies });
+        return response.ok({ body: updatedPackagePolicy });
       } catch (err) {
         const error = transformError(err);
         cspContext.logger.error(