[8.x] [Security Solution] [Attack discovery] Alerts filtering (#205070)…

… (#205137) # Backport This will backport the following commits from `main` to `8.x`: - [[Security Solution] [Attack discovery] Alerts filtering (#205070)](#205070)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  --------- Co-authored-by: Andrew Macri <andrew.macri@elastic.co>
elastic · Dec 24, 2024 · a8cf4eb · a8cf4eb
1 parent 8642be7
commit a8cf4eb
Show file tree

Hide file tree

Showing 61 changed files with 2,380 additions and 77 deletions.
diff --git a/...elastic-assistant-common/impl/alerts/get_open_and_acknowledged_alerts_query/index.test.ts b/...elastic-assistant-common/impl/alerts/get_open_and_acknowledged_alerts_query/index.test.ts
@@ -13,6 +13,7 @@ describe('getOpenAndAcknowledgedAlertsQuery', () => {
     const anonymizationFields = [
       { id: 'field1', field: 'field1', allowed: true, anonymized: false },
       { id: 'field2', field: 'field2', allowed: true, anonymized: false },
+      { id: 'field3', field: 'field3', allowed: false, anonymized: false },
     ];
     const size = 10;
 

diff --git a/.../kbn-elastic-assistant-common/impl/alerts/get_open_and_acknowledged_alerts_query/index.ts b/.../kbn-elastic-assistant-common/impl/alerts/get_open_and_acknowledged_alerts_query/index.ts
@@ -7,6 +7,34 @@
 
 import type { AnonymizationFieldResponse } from '../../schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen';
 
+export const DEFAULT_END = 'now';
+export const DEFAULT_START = 'now-24h';
+
+interface GetOpenAndAcknowledgedAlertsQuery {
+  allow_no_indices: boolean;
+  body: {
+    fields: Array<{
+      field: string;
+      include_unmapped: boolean;
+    }>;
+    query: {
+      bool: {
+        filter: Array<Record<string, unknown>>;
+      };
+    };
+    runtime_mappings: Record<string, unknown>;
+    size: number;
+    sort: Array<{
+      [key: string]: {
+        order: string;
+      };
+    }>;
+    _source: boolean;
+  };
+  ignore_unavailable: boolean;
+  index: string[];
+}
+
 /**
  * This query returns open and acknowledged (non-building block) alerts in the last 24 hours.
  *
@@ -15,12 +43,18 @@ import type { AnonymizationFieldResponse } from '../../schemas/anonymization_fie
 export const getOpenAndAcknowledgedAlertsQuery = ({
   alertsIndexPattern,
   anonymizationFields,
+  end,
+  filter,
   size,
+  start,
 }: {
   alertsIndexPattern: string;
   anonymizationFields: AnonymizationFieldResponse[];
+  end?: string | null;
+  filter?: Record<string, unknown> | null;
   size: number;
-}) => ({
+  start?: string | null;
+}): GetOpenAndAcknowledgedAlertsQuery => ({
   allow_no_indices: true,
   body: {
     fields: anonymizationFields
@@ -53,11 +87,12 @@ export const getOpenAndAcknowledgedAlertsQuery = ({
                     minimum_should_match: 1,
                   },
                 },
+                ...(filter != null ? [filter] : []),
                 {
                   range: {
                     '@timestamp': {
-                      gte: 'now-24h',
-                      lte: 'now',
+                      gte: start != null ? start : DEFAULT_START,
+                      lte: end != null ? end : DEFAULT_END,
                       format: 'strict_date_optional_time',
                     },
                   },

diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/capabilities/index.ts
@@ -20,5 +20,6 @@ export type AssistantFeatureKey = keyof AssistantFeatures;
  */
 export const defaultAssistantFeatures = Object.freeze({
   assistantModelEvaluation: false,
+  attackDiscoveryAlertFiltering: false,
   defendInsights: false,
 });
diff --git a/...elastic-assistant-common/impl/schemas/attack_discovery/post_attack_discovery_route.gen.ts b/...elastic-assistant-common/impl/schemas/attack_discovery/post_attack_discovery_route.gen.ts
@@ -28,11 +28,14 @@ export const AttackDiscoveryPostRequestBody = z.object({
    * LLM API configuration.
    */
   apiConfig: ApiConfig,
+  end: z.string().optional(),
+  filter: z.object({}).catchall(z.unknown()).optional(),
   langSmithProject: z.string().optional(),
   langSmithApiKey: z.string().optional(),
   model: z.string().optional(),
   replacements: Replacements.optional(),
   size: z.number(),
+  start: z.string().optional(),
   subAction: z.enum(['invokeAI', 'invokeStream']),
 });
 export type AttackDiscoveryPostRequestBodyInput = z.input<typeof AttackDiscoveryPostRequestBody>;

diff --git a/...ic-assistant-common/impl/schemas/attack_discovery/post_attack_discovery_route.schema.yaml b/...ic-assistant-common/impl/schemas/attack_discovery/post_attack_discovery_route.schema.yaml
@@ -38,6 +38,11 @@ paths:
                 apiConfig:
                   $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/ApiConfig'
                   description: LLM API configuration.
+                end:
+                  type: string
+                filter:
+                  type: object
+                  additionalProperties: true
                 langSmithProject:
                   type: string
                 langSmithApiKey:
@@ -48,6 +53,8 @@ paths:
                   $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements'
                 size:
                   type: number
+                start:
+                  type: string
                 subAction:
                   type: string
                   enum:

diff --git a/...ared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts b/...ared/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts
@@ -19,5 +19,6 @@ import { z } from '@kbn/zod';
 export type GetCapabilitiesResponse = z.infer<typeof GetCapabilitiesResponse>;
 export const GetCapabilitiesResponse = z.object({
   assistantModelEvaluation: z.boolean(),
+  attackDiscoveryAlertFiltering: z.boolean(),
   defendInsights: z.boolean(),
 });
diff --git a/...kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml b/...kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml
@@ -22,10 +22,13 @@ paths:
                 properties:
                   assistantModelEvaluation:
                     type: boolean
+                  attackDiscoveryAlertFiltering:
+                    type: boolean
                   defendInsights:
                     type: boolean
                 required:
                   - assistantModelEvaluation
+                  - attackDiscoveryAlertFiltering
                   - defendInsights
         '400':
           description: Generic Error

diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts
@@ -41,3 +41,10 @@ export { getRawDataOrDefault } from './impl/alerts/helpers/get_raw_data_or_defau
 
 /** Return true if the provided size is out of range */
 export { sizeIsOutOfRange } from './impl/alerts/helpers/size_is_out_of_range';
+
+export {
+  /** The default (relative) end of the date range (i.e. `now`) */
+  DEFAULT_END,
+  /** The default (relative) start of the date range (i.e. `now-24h`) */
+  DEFAULT_START,
+} from './impl/alerts/get_open_and_acknowledged_alerts_query';
diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/constants.tsx b/x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant_context/constants.tsx
@@ -10,10 +10,14 @@ import { KnowledgeBaseConfig } from '../assistant/types';
 export const ATTACK_DISCOVERY_STORAGE_KEY = 'attackDiscovery';
 export const DEFEND_INSIGHTS_STORAGE_KEY = 'defendInsights';
 export const DEFAULT_ASSISTANT_NAMESPACE = 'elasticAssistantDefault';
+export const END_LOCAL_STORAGE_KEY = 'end';
 export const LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY = 'lastConversationId';
+export const FILTERS_LOCAL_STORAGE_KEY = 'filters';
 export const MAX_ALERTS_LOCAL_STORAGE_KEY = 'maxAlerts';
 export const KNOWLEDGE_BASE_LOCAL_STORAGE_KEY = 'knowledgeBase';
+export const QUERY_LOCAL_STORAGE_KEY = 'query';
 export const SHOW_SETTINGS_TOUR_LOCAL_STORAGE_KEY = 'showSettingsTour';
+export const START_LOCAL_STORAGE_KEY = 'start';
 export const STREAMING_LOCAL_STORAGE_KEY = 'streaming';
 export const TRACE_OPTIONS_SESSION_STORAGE_KEY = 'traceOptions';
 export const CONVERSATION_TABLE_SESSION_STORAGE_KEY = 'conversationTable';

diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant/index.ts
@@ -84,11 +84,19 @@ export {
   DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS,
   DEFAULT_LATEST_ALERTS,
   DEFEND_INSIGHTS_STORAGE_KEY,
+  /** The end of the date range of alerts, sent as context to the LLM */
+  END_LOCAL_STORAGE_KEY,
+  /** Search bar filters that apply to the alerts sent as context to the LLM */
+  FILTERS_LOCAL_STORAGE_KEY,
   KNOWLEDGE_BASE_LOCAL_STORAGE_KEY,
   /** The local storage key that specifies the maximum number of alerts to send as context */
   MAX_ALERTS_LOCAL_STORAGE_KEY,
+  /** Search bar query that apply to the alerts sent as context to the LLM */
+  QUERY_LOCAL_STORAGE_KEY,
   /** The local storage key that specifies whether the settings tour should be shown */
   SHOW_SETTINGS_TOUR_LOCAL_STORAGE_KEY,
+  /** The start of the date range of alerts, sent as context to the LLM */
+  START_LOCAL_STORAGE_KEY,
 } from './impl/assistant_context/constants';
 
 export { useLoadConnectors } from './impl/connectorland/use_load_connectors';

diff --git a/...stic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/index.ts b/...stic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/index.ts
@@ -26,12 +26,15 @@ import type { GraphState } from './types';
 export interface GetDefaultAttackDiscoveryGraphParams {
   alertsIndexPattern?: string;
   anonymizationFields: AnonymizationFieldResponse[];
+  end?: string;
   esClient: ElasticsearchClient;
+  filter?: Record<string, unknown>;
   llm: ActionsClientLlm;
   logger?: Logger;
   onNewReplacements?: (replacements: Replacements) => void;
   replacements?: Replacements;
   size: number;
+  start?: string;
 }
 
 export type DefaultAttackDiscoveryGraph = ReturnType<typeof getDefaultAttackDiscoveryGraph>;
@@ -46,19 +49,22 @@ export type DefaultAttackDiscoveryGraph = ReturnType<typeof getDefaultAttackDisc
 export const getDefaultAttackDiscoveryGraph = ({
   alertsIndexPattern,
   anonymizationFields,
+  end,
   esClient,
+  filter,
   llm,
   logger,
   onNewReplacements,
   replacements,
   size,
+  start,
 }: GetDefaultAttackDiscoveryGraphParams): CompiledStateGraph<
   GraphState,
   Partial<GraphState>,
   'generate' | 'refine' | 'retrieve_anonymized_alerts' | '__start__'
 > => {
   try {
-    const graphState = getDefaultGraphState();
+    const graphState = getDefaultGraphState({ end, filter, start });
 
     // get nodes:
     const retrieveAnonymizedAlertsNode = getRetrieveAnonymizedAlertsNode({

diff --git a/...raphs/default_attack_discovery_graph/nodes/retriever/anonymized_alerts_retriever/index.ts b/...raphs/default_attack_discovery_graph/nodes/retriever/anonymized_alerts_retriever/index.ts
@@ -21,36 +21,48 @@ export class AnonymizedAlertsRetriever extends BaseRetriever {
 
   #alertsIndexPattern?: string;
   #anonymizationFields?: AnonymizationFieldResponse[];
+  #end?: string | null;
   #esClient: ElasticsearchClient;
+  #filter?: Record<string, unknown> | null;
   #onNewReplacements?: (newReplacements: Replacements) => void;
   #replacements?: Replacements;
   #size?: number;
+  #start?: string | null;
 
   constructor({
     alertsIndexPattern,
     anonymizationFields,
     fields,
+    end,
     esClient,
+    filter,
     onNewReplacements,
     replacements,
     size,
+    start,
   }: {
     alertsIndexPattern?: string;
     anonymizationFields?: AnonymizationFieldResponse[];
-    fields?: CustomRetrieverInput;
+    end?: string | null;
     esClient: ElasticsearchClient;
+    fields?: CustomRetrieverInput;
+    filter?: Record<string, unknown> | null;
     onNewReplacements?: (newReplacements: Replacements) => void;
     replacements?: Replacements;
     size?: number;
+    start?: string | null;
   }) {
     super(fields);
 
     this.#alertsIndexPattern = alertsIndexPattern;
     this.#anonymizationFields = anonymizationFields;
+    this.#end = end;
     this.#esClient = esClient;
+    this.#filter = filter;
     this.#onNewReplacements = onNewReplacements;
     this.#replacements = replacements;
     this.#size = size;
+    this.#start = start;
   }
 
   async _getRelevantDocuments(
@@ -60,10 +72,13 @@ export class AnonymizedAlertsRetriever extends BaseRetriever {
     const anonymizedAlerts = await getAnonymizedAlerts({
       alertsIndexPattern: this.#alertsIndexPattern,
       anonymizationFields: this.#anonymizationFields,
+      end: this.#end,
       esClient: this.#esClient,
+      filter: this.#filter,
       onNewReplacements: this.#onNewReplacements,
       replacements: this.#replacements,
       size: this.#size,
+      start: this.#start,
     });
 
     return anonymizedAlerts.map((alert) => ({

diff --git a/...phs/default_attack_discovery_graph/nodes/retriever/helpers/get_anonymized_alerts/index.ts b/...phs/default_attack_discovery_graph/nodes/retriever/helpers/get_anonymized_alerts/index.ts
@@ -21,17 +21,23 @@ import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/s
 export const getAnonymizedAlerts = async ({
   alertsIndexPattern,
   anonymizationFields,
+  end,
   esClient,
+  filter,
   onNewReplacements,
   replacements,
   size,
+  start,
 }: {
   alertsIndexPattern?: string;
   anonymizationFields?: AnonymizationFieldResponse[];
+  end?: string | null;
   esClient: ElasticsearchClient;
+  filter?: Record<string, unknown> | null;
   onNewReplacements?: (replacements: Replacements) => void;
   replacements?: Replacements;
   size?: number;
+  start?: string | null;
 }): Promise<string[]> => {
   if (alertsIndexPattern == null || size == null || sizeIsOutOfRange(size)) {
     return [];
@@ -40,7 +46,10 @@ export const getAnonymizedAlerts = async ({
   const query = getOpenAndAcknowledgedAlertsQuery({
     alertsIndexPattern,
     anonymizationFields: anonymizationFields ?? [],
+    end,
+    filter,
     size,
+    start,
   });
 
   const result = await esClient.search<SearchResponse>(query);

diff --git a/...erver/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.ts b/...erver/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.ts
@@ -36,17 +36,23 @@ export const getRetrieveAnonymizedAlertsNode = ({
     onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements
   };
 
-  const retriever = new AnonymizedAlertsRetriever({
-    alertsIndexPattern,
-    anonymizationFields,
-    esClient,
-    onNewReplacements: localOnNewReplacements,
-    replacements,
-    size,
-  });
-
   const retrieveAnonymizedAlerts = async (state: GraphState): Promise<GraphState> => {
     logger?.debug(() => '---RETRIEVE ANONYMIZED ALERTS---');
+
+    const { end, filter, start } = state;
+
+    const retriever = new AnonymizedAlertsRetriever({
+      alertsIndexPattern,
+      anonymizationFields,
+      end,
+      esClient,
+      filter,
+      onNewReplacements: localOnNewReplacements,
+      replacements,
+      size,
+      start,
+    });
+
     const documents = await retriever
       .withConfig({ runName: 'runAnonymizedAlertsRetriever' })
       .invoke('');