[SecuritySolution] [UI] Re-score entity when asset criticality changes

x-pack/plugins/security_solution/common/api/entity_analytics/asset_criticality/common.gen.ts

@@ -40,6 +40,20 @@ export const CreateAssetCriticalityRecord = AssetCriticalityRecordIdParts.merge(
      * The criticality level of the asset.
      */
     criticality_level: z.enum(['low_impact', 'medium_impact', 'high_impact', 'extreme_impact']),
+    /**
+     * If 'wait_for' the request will wait for the index refresh.
+     */
+    refresh: z.literal('wait_for').optional(),
+  })
+);
+
+export type DeleteAssetCriticalityRecord = z.infer<typeof DeleteAssetCriticalityRecord>;
+export const DeleteAssetCriticalityRecord = AssetCriticalityRecordIdParts.merge(
+  z.object({
+    /**
+     * If 'wait_for' the request will wait for the index refresh.
+     */
+    refresh: z.literal('wait_for').optional(),
   })
 );
 

x-pack/plugins/security_solution/common/api/entity_analytics/asset_criticality/common.schema.yaml

@@ -50,8 +50,21 @@ components:
               type: string
               enum: [low_impact, medium_impact, high_impact, extreme_impact]
               description: The criticality level of the asset.
+            refresh:
+                type: string
+                enum: [wait_for]
+                description: If 'wait_for' the request will wait for the index refresh.
           required:
             - criticality_level
+    DeleteAssetCriticalityRecord:
+      allOf:
+        - $ref: '#/components/schemas/AssetCriticalityRecordIdParts'
+        - type: object
+          properties:
+            refresh:
+                type: string
+                enum: [wait_for]
+                description: If 'wait_for' the request will wait for the index refresh.
     AssetCriticalityRecord:
       allOf:
         - $ref: '#/components/schemas/CreateAssetCriticalityRecord'

x-pack/plugins/security_solution/common/api/entity_analytics/risk_engine/entity_calculation_route.gen.ts

@@ -28,6 +28,10 @@ export const RiskScoresEntityCalculationRequest = z.object({
    * Used to define the type of entity.
    */
   identifier_type: IdentifierType,
+  /**
+   * If 'wait_for' the request will wait for the index refresh.
+   */
+  refresh: z.literal('wait_for').optional(),
 });
 
 export type RiskScoresEntityCalculationResponse = z.infer<

x-pack/plugins/security_solution/common/api/entity_analytics/risk_engine/entity_calculation_route.schema.yaml

@@ -50,6 +50,10 @@ components:
         identifier_type:
           description: Used to define the type of entity.
           $ref: './common.schema.yaml#/components/schemas/IdentifierType'
+        refresh:
+            type: string
+            enum: [wait_for]
+            description: If 'wait_for' the request will wait for the index refresh.
 
     RiskScoresEntityCalculationResponse:
         type: object

x-pack/plugins/security_solution/public/entity_analytics/api/api.ts

@@ -6,6 +6,10 @@
  */
 
 import { useMemo } from 'react';
+import type {
+  RiskScoresEntityCalculationRequest,
+  RiskScoresEntityCalculationResponse,
+} from '../../../common/api/entity_analytics/risk_engine/entity_calculation_route.gen';
 import type { AssetCriticalityCsvUploadResponse } from '../../../common/entity_analytics/asset_criticality/types';
 import type { AssetCriticalityRecord } from '../../../common/api/entity_analytics/asset_criticality';
 import type { RiskScoreEntity } from '../../../common/search_strategy';
@@ -21,6 +25,7 @@ import {
   RISK_SCORE_INDEX_STATUS_API_URL,
   RISK_ENGINE_SETTINGS_URL,
   ASSET_CRITICALITY_CSV_UPLOAD_URL,
+  RISK_SCORE_ENTITY_CALCULATION_URL,
 } from '../../../common/constants';
 
 import type {
@@ -97,6 +102,17 @@ export const useEntityAnalyticsRoutes = () => {
         method: 'POST',
       });
 
+    /**
+     * Calculate and stores risk score for an entity
+     */
+    const calculateEntityRiskScore = (params: RiskScoresEntityCalculationRequest) => {
+      return http.fetch<RiskScoresEntityCalculationResponse>(RISK_SCORE_ENTITY_CALCULATION_URL, {
+        version: '1',
+        method: 'POST',
+        body: JSON.stringify(params),
+      });
+    };
+
     /**
      * Get risk engine privileges
      */
@@ -119,7 +135,9 @@ export const useEntityAnalyticsRoutes = () => {
      * Create asset criticality
      */
     const createAssetCriticality = async (
-      params: Pick<AssetCriticality, 'idField' | 'idValue' | 'criticalityLevel'>
+      params: Pick<AssetCriticality, 'idField' | 'idValue' | 'criticalityLevel'> & {
+        refresh?: 'wait_for';
+      }
     ): Promise<AssetCriticalityRecord> =>
       http.fetch<AssetCriticalityRecord>(ASSET_CRITICALITY_URL, {
         version: '1',
@@ -128,16 +146,17 @@ export const useEntityAnalyticsRoutes = () => {
           id_value: params.idValue,
           id_field: params.idField,
           criticality_level: params.criticalityLevel,
+          refresh: params.refresh,
         }),
       });
 
     const deleteAssetCriticality = async (
-      params: Pick<AssetCriticality, 'idField' | 'idValue'>
+      params: Pick<AssetCriticality, 'idField' | 'idValue'> & { refresh?: 'wait_for' }
     ): Promise<{ deleted: true }> => {
       await http.fetch(ASSET_CRITICALITY_URL, {
         version: '1',
         method: 'DELETE',
-        query: { id_value: params.idValue, id_field: params.idField },
+        query: { id_value: params.idValue, id_field: params.idField, refresh: params.refresh },
       });
 
       // spoof a response to allow us to better distnguish a delete from a create in use_asset_criticality.ts
@@ -219,6 +238,7 @@ export const useEntityAnalyticsRoutes = () => {
       uploadAssetCriticalityFile,
       getRiskScoreIndexStatus,
       fetchRiskEngineSettings,
+      calculateEntityRiskScore,
     };
   }, [http]);
 };

x-pack/plugins/security_solution/public/entity_analytics/api/hooks/use_calculate_entity_risk_score.test.ts

@@ -0,0 +1,103 @@
+/*
+ * 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 { act, renderHook } from '@testing-library/react-hooks';
+import { TestProviders } from '../../../common/mock';
+import { RiskScoreEntity } from '../../../../common/search_strategy';
+import { useCalculateEntityRiskScore } from './use_calculate_entity_risk_score';
+import { RiskEngineStatus } from '../../../../common/entity_analytics/risk_engine';
+import { waitFor } from '@testing-library/react';
+
+const enabledRiskEngineStatus = {
+  risk_engine_status: RiskEngineStatus.ENABLED,
+};
+const disabledRiskEngineStatus = {
+  risk_engine_status: RiskEngineStatus.DISABLED,
+};
+
+const mockUseRiskEngineStatus = jest.fn();
+jest.mock('./use_risk_engine_status', () => ({
+  useRiskEngineStatus: () => mockUseRiskEngineStatus(),
+}));
+
+const mockCalculateEntityRiskScore = jest.fn();
+jest.mock('../api', () => ({
+  useEntityAnalyticsRoutes: () => ({
+    calculateEntityRiskScore: mockCalculateEntityRiskScore,
+  }),
+}));
+
+const mockAddError = jest.fn();
+jest.mock('../../../common/hooks/use_app_toasts', () => ({
+  useAppToasts: jest.fn().mockReturnValue({
+    addError: () => mockAddError(),
+  }),
+}));
+
+const identifierType = RiskScoreEntity.user;
+const identifier = 'test-user';
+const options = {
+  onSuccess: jest.fn(),
+};
+
+describe('useRiskScoreData', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockUseRiskEngineStatus.mockReturnValue({ data: enabledRiskEngineStatus });
+    mockCalculateEntityRiskScore.mockResolvedValue({});
+  });
+
+  it('should call calculateEntityRiskScore API when the callback function is called', async () => {
+    const { result } = renderHook(
+      () => useCalculateEntityRiskScore(identifierType, identifier, options),
+      { wrapper: TestProviders }
+    );
+
+    await act(async () => {
+      result.current.calculateEntityRiskScore();
+
+      await waitFor(() =>
+        expect(mockCalculateEntityRiskScore).toHaveBeenCalledWith(
+          expect.objectContaining({
+            identifier_type: identifierType,
+            identifier,
+          })
+        )
+      );
+    });
+  });
+
+  it('should NOT call calculateEntityRiskScore API when risk engine is disabled', async () => {
+    mockUseRiskEngineStatus.mockReturnValue({
+      data: disabledRiskEngineStatus,
+    });
+    const { result } = renderHook(
+      () => useCalculateEntityRiskScore(identifierType, identifier, options),
+      { wrapper: TestProviders }
+    );
+
+    await act(async () => {
+      result.current.calculateEntityRiskScore();
+
+      await waitFor(() => expect(mockCalculateEntityRiskScore).not.toHaveBeenCalled());
+    });
+  });
+
+  it('should display a toast error when the API returns an error', async () => {
+    mockCalculateEntityRiskScore.mockRejectedValue({});
+    const { result } = renderHook(
+      () => useCalculateEntityRiskScore(identifierType, identifier, options),
+      { wrapper: TestProviders }
+    );
+
+    await act(async () => {
+      result.current.calculateEntityRiskScore();
+
+      await waitFor(() => expect(mockAddError).toHaveBeenCalled());
+    });
+  });
+});