From 0c69da8196c08a0469ed356568dbbd37fb70c8d6 Mon Sep 17 00:00:00 2001 From: "Joey F. Poon" Date: Tue, 19 Nov 2024 19:07:42 +0900 Subject: [PATCH] [Security Solution] add defend insights elastic assistant tool (#198676) ### Summary Adds the new Defend Insights Elastic Assistant tool. This assistant tool provides Elastic Defend configuration insights. For this initial PR, only incompatible antivirus detection is supported. Telemetry is collected for success and error events. For incompatible antivirus detection, Defend Insights will review the last 200 file events for the given endpoint and output suspected antiviruses. Improvements such as customizable event count and date range will come in the future. This PR does not include any UI, that will come in a separate PR. 3 internal APIs for interacting with Defend Insights are provided here: - `POST /defend_insights` for creating a new Defend Insight - `GET /defend_insights/{id}` for getting a Defend Insight - `GET /defend_insights` for getting multiple Defend Insights - available optional query params: - `size` - default 10 - `ids` - `connector_id` - `type` - `incompatible_antivirus` - `status` - `running`, `completed`, `failed`, `canceled` - `endpoint_ids` This initial implementation does not include the LangGraph/output chunking upgrades seen in Attack Discovery due to time constraints. We'll look to make this upgrade in a future PR. ### Checklist - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios ### For maintainers - [x] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#_add_your_labels) --- .github/CODEOWNERS | 5 + .../kbn-elastic-assistant-common/constants.ts | 5 + .../impl/capabilities/index.ts | 1 + .../get_capabilities_route.gen.ts | 1 + .../get_capabilities_route.schema.yaml | 3 + .../defend_insights/common_attributes.gen.ts | 217 +++++++++ .../common_attributes.schema.yaml | 224 ++++++++++ .../get_defend_insight_route.gen.ts | 34 ++ .../get_defend_insight_route.schema.yaml | 45 ++ .../get_defend_insights_route.gen.ts | 59 +++ .../get_defend_insights_route.schema.yaml | 82 ++++ .../impl/schemas/defend_insights/index.ts | 11 + .../post_defend_insights_route.gen.ts | 42 ++ .../post_defend_insights_route.schema.yaml | 77 ++++ .../impl/schemas/index.ts | 3 + .../common/anonymization/index.ts | 1 + .../__mocks__/defend_insights_schema.mock.ts | 109 +++++ .../server/__mocks__/request.ts | 27 ++ .../server/__mocks__/request_context.ts | 6 + .../field_maps_configuration.ts | 174 ++++++++ .../get_defend_insight.test.ts | 69 +++ .../defend_insights/get_defend_insight.ts | 78 ++++ .../defend_insights/helpers.test.ts | 64 +++ .../defend_insights/helpers.ts | 221 ++++++++++ .../defend_insights/index.test.ts | 410 ++++++++++++++++++ .../defend_insights/index.ts | 286 ++++++++++++ .../defend_insights/types.ts | 88 ++++ .../server/ai_assistant_service/index.test.ts | 13 +- .../server/ai_assistant_service/index.ts | 40 +- .../server/lib/langchain/helpers.ts | 4 +- .../lib/telemetry/event_based_telemetry.ts | 96 ++++ .../get_defend_insight.test.ts | 149 +++++++ .../defend_insights/get_defend_insight.ts | 96 ++++ .../get_defend_insights.test.ts | 149 +++++++ .../defend_insights/get_defend_insights.ts | 98 +++++ .../routes/defend_insights/helpers.test.ts | 255 +++++++++++ .../server/routes/defend_insights/helpers.ts | 387 +++++++++++++++++ .../server/routes/defend_insights/index.ts | 10 + .../post_defend_insights.test.ts | 184 ++++++++ .../defend_insights/post_defend_insights.ts | 196 +++++++++ .../server/routes/helpers.ts | 5 +- .../elastic_assistant/server/routes/index.ts | 5 + .../server/routes/register_routes.ts | 10 + .../server/routes/request_context_factory.ts | 10 + .../server/services/app_context.test.ts | 7 + .../plugins/elastic_assistant/server/types.ts | 9 +- .../common/endpoint/constants.ts | 1 + .../common/experimental_features.ts | 5 + .../assistant/tools/defend_insights/errors.ts | 14 + .../get_events/get_file_events_query.ts | 49 +++ .../defend_insights/get_events/index.test.ts | 84 ++++ .../tools/defend_insights/get_events/index.ts | 93 ++++ .../tools/defend_insights/index.test.ts | 69 +++ .../assistant/tools/defend_insights/index.ts | 114 +++++ .../output_parsers/incompatible_antivirus.ts | 28 ++ .../defend_insights/output_parsers/index.ts | 19 + .../prompts/incompatible_antivirus.ts | 16 + .../tools/defend_insights/prompts/index.ts | 25 ++ .../server/assistant/tools/index.ts | 2 + 59 files changed, 4574 insertions(+), 10 deletions(-) create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.schema.yaml create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.schema.yaml create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.schema.yaml create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/index.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.gen.ts create mode 100644 x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.schema.yaml create mode 100644 x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/field_maps_configuration.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts create mode 100644 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/index.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts create mode 100644 x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 030f888c3e9b..98f3aef6118c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2158,6 +2158,7 @@ x-pack/test/security_solution_cypress/cypress/tasks/expandable_flyout @elastic/ /x-pack/plugins/security_solution/public/flyout/document_details/isolate_host/ @elastic/security-defend-workflows /x-pack/plugins/security_solution/common/endpoint/ @elastic/security-defend-workflows /x-pack/plugins/security_solution/common/api/endpoint/ @elastic/security-defend-workflows +x-pack/plugins/security_solution/server/assistant/tools/defend_insights @elastic/security-defend-workflows /x-pack/plugins/security_solution/server/endpoint/ @elastic/security-defend-workflows /x-pack/plugins/security_solution/server/lists_integration/endpoint/ @elastic/security-defend-workflows /x-pack/plugins/security_solution/server/lib/license/ @elastic/security-defend-workflows @@ -2169,6 +2170,10 @@ x-pack/test/security_solution_cypress/cypress/tasks/expandable_flyout @elastic/ /x-pack/plugins/security_solution_serverless/public/upselling/sections/endpoint_management @elastic/security-defend-workflows /x-pack/plugins/security_solution_serverless/public/upselling/pages/endpoint_management @elastic/security-defend-workflows /x-pack/plugins/security_solution_serverless/server/endpoint @elastic/security-defend-workflows +x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights @elastic/security-defend-workflows +x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts @elastic/security-defend-workflows +x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights @elastic/security-defend-workflows +x-pack/plugins/elastic_assistant/server/routes/defend_insights @elastic/security-defend-workflows ## Security Solution sub teams - security-telemetry (Data Engineering) x-pack/plugins/security_solution/server/usage/ @elastic/security-data-analytics diff --git a/x-pack/packages/kbn-elastic-assistant-common/constants.ts b/x-pack/packages/kbn-elastic-assistant-common/constants.ts index 49db6c295a51..7a884936d04e 100755 --- a/x-pack/packages/kbn-elastic-assistant-common/constants.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/constants.ts @@ -55,3 +55,8 @@ export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_INDICES_URL = export const ELASTIC_AI_ASSISTANT_EVALUATE_URL = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/evaluate` as const; + +// Defend insights +export const DEFEND_INSIGHTS_TOOL_ID = 'defend-insights'; +export const DEFEND_INSIGHTS = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/defend_insights`; +export const DEFEND_INSIGHTS_BY_ID = `${DEFEND_INSIGHTS}/{id}`; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts index d883dfe98d56..0e204b4b949e 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/capabilities/index.ts @@ -20,4 +20,5 @@ export type AssistantFeatureKey = keyof AssistantFeatures; */ export const defaultAssistantFeatures = Object.freeze({ assistantModelEvaluation: false, + defendInsights: false, }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts index 0f8b6235d7dc..8777e8d72827 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.gen.ts @@ -19,4 +19,5 @@ import { z } from '@kbn/zod'; export type GetCapabilitiesResponse = z.infer; export const GetCapabilitiesResponse = z.object({ assistantModelEvaluation: z.boolean(), + defendInsights: z.boolean(), }); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml index a042abd39179..e9b6ca969725 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/capabilities/get_capabilities_route.schema.yaml @@ -22,8 +22,11 @@ paths: properties: assistantModelEvaluation: type: boolean + defendInsights: + type: boolean required: - assistantModelEvaluation + - defendInsights '400': description: Generic Error content: diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.gen.ts new file mode 100644 index 000000000000..e070c3129e19 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.gen.ts @@ -0,0 +1,217 @@ +/* + * 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. + */ + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Common Defend Insights Attributes + * version: not applicable + */ + +import { z } from '@kbn/zod'; + +import { NonEmptyString, User } from '../common_attributes.gen'; +import { Replacements, ApiConfig } from '../conversations/common_attributes.gen'; + +/** + * A Defend insight event + */ +export type DefendInsightEvent = z.infer; +export const DefendInsightEvent = z.object({ + /** + * The event's ID + */ + id: z.string(), + /** + * The endpoint's ID + */ + endpointId: z.string(), + /** + * The value of the event + */ + value: z.string(), +}); + +/** + * The insight type (ie. incompatible_antivirus) + */ +export type DefendInsightType = z.infer; +export const DefendInsightType = z.enum(['incompatible_antivirus', 'noisy_process_tree']); +export type DefendInsightTypeEnum = typeof DefendInsightType.enum; +export const DefendInsightTypeEnum = DefendInsightType.enum; + +/** + * A Defend insight generated from endpoint events + */ +export type DefendInsight = z.infer; +export const DefendInsight = z.object({ + /** + * The group category of the events (ie. Windows Defender) + */ + group: z.string(), + /** + * An array of event objects + */ + events: z.array(DefendInsightEvent).optional(), +}); + +/** + * Array of Defend insights + */ +export type DefendInsights = z.infer; +export const DefendInsights = z.array(DefendInsight); + +/** + * The status of the Defend insight. + */ +export type DefendInsightStatus = z.infer; +export const DefendInsightStatus = z.enum(['running', 'succeeded', 'failed', 'canceled']); +export type DefendInsightStatusEnum = typeof DefendInsightStatus.enum; +export const DefendInsightStatusEnum = DefendInsightStatus.enum; + +/** + * Run durations for the Defend insight + */ +export type DefendInsightGenerationInterval = z.infer; +export const DefendInsightGenerationInterval = z.object({ + /** + * The time the Defend insight was generated + */ + date: z.string(), + /** + * The duration of the Defend insight generation + */ + durationMs: z.number().int(), +}); + +export type DefendInsightsResponse = z.infer; +export const DefendInsightsResponse = z.object({ + id: NonEmptyString, + timestamp: NonEmptyString.optional(), + /** + * The last time the Defend insight was updated. + */ + updatedAt: z.string(), + /** + * The last time the Defend insight was viewed in the browser. + */ + lastViewedAt: z.string(), + /** + * The number of events in the context. + */ + eventsContextCount: z.number().int().optional(), + /** + * The time the Defend insight was created. + */ + createdAt: z.string(), + replacements: Replacements.optional(), + users: z.array(User), + /** + * The status of the Defend insight. + */ + status: DefendInsightStatus, + endpointIds: z.array(NonEmptyString), + insightType: DefendInsightType, + /** + * The Defend insights. + */ + insights: DefendInsights, + /** + * LLM API configuration. + */ + apiConfig: ApiConfig, + /** + * Kibana space + */ + namespace: z.string(), + /** + * The backing index required for update requests. + */ + backingIndex: z.string(), + /** + * The most 5 recent generation intervals + */ + generationIntervals: z.array(DefendInsightGenerationInterval), + /** + * The average generation interval in milliseconds + */ + averageIntervalMs: z.number().int(), + /** + * The reason for a status of failed. + */ + failureReason: z.string().optional(), +}); + +export type DefendInsightUpdateProps = z.infer; +export const DefendInsightUpdateProps = z.object({ + id: NonEmptyString, + /** + * LLM API configuration. + */ + apiConfig: ApiConfig.optional(), + /** + * The number of events in the context. + */ + eventsContextCount: z.number().int().optional(), + /** + * The Defend insights. + */ + insights: DefendInsights.optional(), + /** + * The status of the Defend insight. + */ + status: DefendInsightStatus.optional(), + replacements: Replacements.optional(), + /** + * The most 5 recent generation intervals + */ + generationIntervals: z.array(DefendInsightGenerationInterval).optional(), + /** + * The backing index required for update requests. + */ + backingIndex: z.string(), + /** + * The reason for a status of failed. + */ + failureReason: z.string().optional(), + /** + * The last time the Defend insight was viewed in the browser. + */ + lastViewedAt: z.string().optional(), +}); + +export type DefendInsightsUpdateProps = z.infer; +export const DefendInsightsUpdateProps = z.array(DefendInsightUpdateProps); + +export type DefendInsightCreateProps = z.infer; +export const DefendInsightCreateProps = z.object({ + /** + * The Defend insight id. + */ + id: z.string().optional(), + /** + * The status of the Defend insight. + */ + status: DefendInsightStatus, + /** + * The number of events in the context. + */ + eventsContextCount: z.number().int().optional(), + endpointIds: z.array(NonEmptyString), + insightType: DefendInsightType, + /** + * The Defend insights. + */ + insights: DefendInsights, + /** + * LLM API configuration. + */ + apiConfig: ApiConfig, + replacements: Replacements.optional(), +}); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.schema.yaml new file mode 100644 index 000000000000..5c27449c7d34 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/common_attributes.schema.yaml @@ -0,0 +1,224 @@ +openapi: 3.0.0 +info: + title: Common Defend Insights Attributes + version: 'not applicable' +paths: {} +components: + x-codegen-enabled: true + schemas: + DefendInsightEvent: + type: object + description: A Defend insight event + required: + - 'id' + - 'endpointId' + - 'value' + properties: + id: + description: The event's ID + type: string + endpointId: + description: The endpoint's ID + type: string + value: + description: The value of the event + type: string + + DefendInsightType: + description: The insight type (ie. incompatible_antivirus) + type: string + enum: + - incompatible_antivirus + - noisy_process_tree + + DefendInsight: + type: object + description: A Defend insight generated from endpoint events + required: + - 'group' + properties: + group: + description: The group category of the events (ie. Windows Defender) + type: string + events: + description: An array of event objects + type: array + items: + $ref: '#/components/schemas/DefendInsightEvent' + + DefendInsights: + type: array + description: Array of Defend insights + items: + $ref: '#/components/schemas/DefendInsight' + + DefendInsightsResponse: + type: object + required: + - apiConfig + - id + - createdAt + - updatedAt + - lastViewedAt + - users + - namespace + - endpointIds + - insightType + - insights + - status + - backingIndex + - generationIntervals + - averageIntervalMs + properties: + id: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + 'timestamp': + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + updatedAt: + description: The last time the Defend insight was updated. + type: string + lastViewedAt: + description: The last time the Defend insight was viewed in the browser. + type: string + eventsContextCount: + type: integer + description: The number of events in the context. + createdAt: + description: The time the Defend insight was created. + type: string + replacements: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements' + users: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/User' + status: + $ref: '#/components/schemas/DefendInsightStatus' + description: The status of the Defend insight. + endpointIds: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + insightType: + $ref: '#/components/schemas/DefendInsightType' + insights: + $ref: '#/components/schemas/DefendInsights' + description: The Defend insights. + apiConfig: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/ApiConfig' + description: LLM API configuration. + namespace: + type: string + description: Kibana space + backingIndex: + type: string + description: The backing index required for update requests. + generationIntervals: + type: array + description: The most 5 recent generation intervals + items: + $ref: '#/components/schemas/DefendInsightGenerationInterval' + averageIntervalMs: + type: integer + description: The average generation interval in milliseconds + failureReason: + type: string + description: The reason for a status of failed. + + DefendInsightGenerationInterval: + type: object + description: Run durations for the Defend insight + required: + - 'date' + - 'durationMs' + properties: + date: + description: The time the Defend insight was generated + type: string + durationMs: + description: The duration of the Defend insight generation + type: integer + + DefendInsightStatus: + type: string + description: The status of the Defend insight. + enum: + - running + - succeeded + - failed + - canceled + + DefendInsightUpdateProps: + type: object + required: + - id + - backingIndex + properties: + id: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + apiConfig: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/ApiConfig' + description: LLM API configuration. + eventsContextCount: + type: integer + description: The number of events in the context. + insights: + $ref: '#/components/schemas/DefendInsights' + description: The Defend insights. + status: + $ref: '#/components/schemas/DefendInsightStatus' + description: The status of the Defend insight. + replacements: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements' + generationIntervals: + type: array + description: The most 5 recent generation intervals + items: + $ref: '#/components/schemas/DefendInsightGenerationInterval' + backingIndex: + type: string + description: The backing index required for update requests. + failureReason: + type: string + description: The reason for a status of failed. + lastViewedAt: + description: The last time the Defend insight was viewed in the browser. + type: string + + DefendInsightsUpdateProps: + type: array + items: + $ref: '#/components/schemas/DefendInsightUpdateProps' + + DefendInsightCreateProps: + type: object + required: + - endpointIds + - insightType + - insights + - apiConfig + - status + properties: + id: + type: string + description: The Defend insight id. + status: + $ref: '#/components/schemas/DefendInsightStatus' + description: The status of the Defend insight. + eventsContextCount: + type: integer + description: The number of events in the context. + endpointIds: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + insightType: + $ref: '#/components/schemas/DefendInsightType' + insights: + $ref: '#/components/schemas/DefendInsights' + description: The Defend insights. + apiConfig: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/ApiConfig' + description: LLM API configuration. + replacements: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements' diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.gen.ts new file mode 100644 index 000000000000..fafaca8f48ea --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.gen.ts @@ -0,0 +1,34 @@ +/* + * 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. + */ + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Get Defend Insight API endpoint + * version: 1 + */ + +import { z } from '@kbn/zod'; + +import { NonEmptyString } from '../common_attributes.gen'; +import { DefendInsightsResponse } from './common_attributes.gen'; + +export type DefendInsightGetRequestParams = z.infer; +export const DefendInsightGetRequestParams = z.object({ + /** + * The Defend insight id + */ + id: NonEmptyString, +}); +export type DefendInsightGetRequestParamsInput = z.input; + +export type DefendInsightGetResponse = z.infer; +export const DefendInsightGetResponse = z.object({ + data: DefendInsightsResponse.nullable().optional(), +}); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.schema.yaml new file mode 100644 index 000000000000..2684bf53cf87 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insight_route.schema.yaml @@ -0,0 +1,45 @@ +openapi: 3.0.3 +info: + title: Get Defend Insight API endpoint + version: '1' +paths: + /internal/elastic_assistant/defend_insights/{id}: + get: + x-codegen-enabled: true + x-labels: [ess, serverless] + operationId: DefendInsightGet + description: Get Defend insight by id + summary: Get Defend insight data + tags: + - defend_insights + parameters: + - name: 'id' + in: path + required: true + description: The Defend insight id + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + responses: + '200': + description: Successful response + content: + application/json: + schema: + type: object + properties: + data: + $ref: './common_attributes.schema.yaml#/components/schemas/DefendInsightsResponse' + nullable: true + '400': + description: Generic Error + content: + application/json: + schema: + type: object + properties: + statusCode: + type: number + error: + type: string + message: + type: string diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.gen.ts new file mode 100644 index 000000000000..0a2f3d618a86 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.gen.ts @@ -0,0 +1,59 @@ +/* + * 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. + */ + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Get Defend Insights API endpoint + * version: 1 + */ + +import { z } from '@kbn/zod'; +import { ArrayFromString } from '@kbn/zod-helpers'; + +import { NonEmptyString } from '../common_attributes.gen'; +import { + DefendInsightType, + DefendInsightStatus, + DefendInsightsResponse, +} from './common_attributes.gen'; + +export type DefendInsightsGetRequestQuery = z.infer; +export const DefendInsightsGetRequestQuery = z.object({ + /** + * The insight ids for which to get Defend insights + */ + ids: ArrayFromString(NonEmptyString).optional(), + /** + * The connector id for which to get Defend insights + */ + connector_id: NonEmptyString.optional(), + /** + * The insight type for which to get Defend insights + */ + type: DefendInsightType.optional(), + /** + * The status for which to get Defend insights + */ + status: DefendInsightStatus.optional(), + /** + * The endpoint ids for which to get Defend insights + */ + endpoint_ids: ArrayFromString(NonEmptyString).optional(), + /** + * The number of Defend insights to return + */ + size: z.coerce.number().optional(), +}); +export type DefendInsightsGetRequestQueryInput = z.input; + +export type DefendInsightsGetResponse = z.infer; +export const DefendInsightsGetResponse = z.object({ + data: z.array(DefendInsightsResponse), +}); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.schema.yaml new file mode 100644 index 000000000000..5d7e0b5358f8 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/get_defend_insights_route.schema.yaml @@ -0,0 +1,82 @@ +openapi: 3.0.0 +info: + title: Get Defend Insights API endpoint + version: '1' +paths: + /internal/elastic_assistant/defend_insights: + get: + x-codegen-enabled: true + x-labels: [ess, serverless] + operationId: DefendInsightsGet + description: Get relevant data for Defend insights + summary: Get relevant data for Defend insights + tags: + - defend_insights + parameters: + - name: 'ids' + in: query + required: false + description: The insight ids for which to get Defend insights + schema: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + - name: 'connector_id' + in: query + required: false + description: The connector id for which to get Defend insights + schema: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + - name: 'type' + in: query + required: false + description: The insight type for which to get Defend insights + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/DefendInsightType' + - name: 'status' + in: query + required: false + description: The status for which to get Defend insights + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/DefendInsightStatus' + - name: 'endpoint_ids' + in: query + required: false + description: The endpoint ids for which to get Defend insights + schema: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + - name: 'size' + in: query + required: false + description: The number of Defend insights to return + schema: + type: number + responses: + '200': + description: Successful response + content: + application/json: + schema: + type: object + required: + - data + properties: + data: + type: array + items: + $ref: './common_attributes.schema.yaml#/components/schemas/DefendInsightsResponse' + '400': + description: Generic Error + content: + application/json: + schema: + type: object + properties: + statusCode: + type: number + error: + type: string + message: + type: string diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/index.ts new file mode 100644 index 000000000000..0518abdf6dcb --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/index.ts @@ -0,0 +1,11 @@ +/* + * 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. + */ + +export * from './common_attributes.gen'; +export * from './get_defend_insight_route.gen'; +export * from './get_defend_insights_route.gen'; +export * from './post_defend_insights_route.gen'; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.gen.ts new file mode 100644 index 000000000000..cc0ccfeea198 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.gen.ts @@ -0,0 +1,42 @@ +/* + * 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. + */ + +/* + * NOTICE: Do not edit this file manually. + * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. + * + * info: + * title: Post Defend Insights API endpoint + * version: 1 + */ + +import { z } from '@kbn/zod'; + +import { NonEmptyString } from '../common_attributes.gen'; +import { DefendInsightType, DefendInsightsResponse } from './common_attributes.gen'; +import { AnonymizationFieldResponse } from '../anonymization_fields/bulk_crud_anonymization_fields_route.gen'; +import { ApiConfig, Replacements } from '../conversations/common_attributes.gen'; + +export type DefendInsightsPostRequestBody = z.infer; +export const DefendInsightsPostRequestBody = z.object({ + endpointIds: z.array(NonEmptyString), + insightType: DefendInsightType, + anonymizationFields: z.array(AnonymizationFieldResponse), + /** + * LLM API configuration. + */ + apiConfig: ApiConfig, + langSmithProject: z.string().optional(), + langSmithApiKey: z.string().optional(), + model: z.string().optional(), + replacements: Replacements.optional(), + subAction: z.enum(['invokeAI', 'invokeStream']), +}); +export type DefendInsightsPostRequestBodyInput = z.input; + +export type DefendInsightsPostResponse = z.infer; +export const DefendInsightsPostResponse = DefendInsightsResponse; diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.schema.yaml new file mode 100644 index 000000000000..87c7cdbb81a8 --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/defend_insights/post_defend_insights_route.schema.yaml @@ -0,0 +1,77 @@ +openapi: 3.0.0 +info: + title: Post Defend Insights API endpoint + version: '1' +components: + x-codegen-enabled: true + +paths: + /internal/elastic_assistant/defend_insights: + post: + x-codegen-enabled: true + x-labels: [ess, serverless] + operationId: DefendInsightsPost + description: Generate Elastic Defend configuration insights + summary: Generate Elastic Defend configuration insights from endpoint events via the Elastic Assistant + tags: + - defend_insights + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - endpointIds + - insightType + - apiConfig + - anonymizationFields + - subAction + properties: + endpointIds: + type: array + items: + $ref: '../common_attributes.schema.yaml#/components/schemas/NonEmptyString' + insightType: + $ref: './common_attributes.schema.yaml#/components/schemas/DefendInsightType' + anonymizationFields: + type: array + items: + $ref: '../anonymization_fields/bulk_crud_anonymization_fields_route.schema.yaml#/components/schemas/AnonymizationFieldResponse' + apiConfig: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/ApiConfig' + description: LLM API configuration. + langSmithProject: + type: string + langSmithApiKey: + type: string + model: + type: string + replacements: + $ref: '../conversations/common_attributes.schema.yaml#/components/schemas/Replacements' + subAction: + type: string + enum: + - invokeAI + - invokeStream + responses: + '200': + description: Successful response + content: + application/json: + schema: + $ref: './common_attributes.schema.yaml#/components/schemas/DefendInsightsResponse' + '400': + description: Bad request + content: + application/json: + schema: + type: object + properties: + statusCode: + type: number + error: + type: string + message: + type: string + diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts index 9233791a870c..02ac9b7b1ba9 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/index.ts @@ -27,6 +27,9 @@ export * from './attack_discovery/get_attack_discovery_route.gen'; export * from './attack_discovery/post_attack_discovery_route.gen'; export * from './attack_discovery/cancel_attack_discovery_route.gen'; +// Defend insight Schemas +export * from './defend_insights'; + // Chat Schemas export * from './chat/post_chat_complete_route.gen'; diff --git a/x-pack/plugins/elastic_assistant/common/anonymization/index.ts b/x-pack/plugins/elastic_assistant/common/anonymization/index.ts index ebef2dff8bde..9b8007a9129b 100644 --- a/x-pack/plugins/elastic_assistant/common/anonymization/index.ts +++ b/x-pack/plugins/elastic_assistant/common/anonymization/index.ts @@ -9,6 +9,7 @@ export const DEFAULT_ALLOW = [ '_id', '@timestamp', + 'agent.id', 'cloud.availability_zone', 'cloud.provider', 'cloud.region', diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts new file mode 100644 index 000000000000..d25b8bb09b13 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts @@ -0,0 +1,109 @@ +/* + * 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 { estypes } from '@elastic/elasticsearch'; +import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; + +import type { EsDefendInsightSchema } from '../ai_assistant_data_clients/defend_insights/types'; + +export const getDefendInsightsSearchEsMock = () => { + const searchResponse: estypes.SearchResponse = { + took: 0, + timed_out: false, + _shards: { + total: 1, + successful: 1, + skipped: 0, + failed: 0, + }, + hits: { + total: { + value: 1, + relation: 'eq', + }, + max_score: 1, + hits: [ + { + _index: '.kibana-elastic-ai-assistant-defend-insights-default', + _id: '655c52ec-49ee-4d20-87e5-7edd6d8f84e8', + _score: 1, + _source: { + '@timestamp': '2024-09-24T10:48:46.847Z', + created_at: '2024-09-24T10:48:46.847Z', + users: [ + { + id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', + name: 'elastic', + }, + ], + status: DefendInsightStatus.Enum.succeeded, + api_config: { + action_type_id: '.bedrock', + connector_id: 'ac4e19d1-e2e2-49af-bf4b-59428473101c', + model: 'anthropic.claude-3-5-sonnet-20240620-v1:0', + }, + endpoint_ids: ['6e09ec1c-644c-4148-a02d-be451c35400d'], + insight_type: DefendInsightType.Enum.incompatible_antivirus, + insights: [ + { + group: 'windows_defenders', + events: [], + }, + ], + updated_at: '2024-09-24T10:48:59.952Z', + last_viewed_at: '2024-09-24T10:49:53.522Z', + namespace: 'default', + id: '655c52ec-49ee-4d20-87e5-7edd6d8f84e8', + generation_intervals: [ + { + date: '2024-09-24T10:48:59.952Z', + duration_ms: 13113, + }, + ], + average_interval_ms: 13113, + replacements: [ + { + uuid: '2009c67b-89b8-43d9-b502-2c32f71875a0', + value: 'root', + }, + { + uuid: '9f7f91b6-6853-48b7-bfb8-403f5efb2364', + value: 'joey-dev-default-3539', + }, + { + uuid: 'c08e4851-7234-408a-8083-7fd5740e4255', + value: 'syslog', + }, + { + uuid: '826c58bd-1466-42fd-af1f-9094c155811b', + value: 'messagebus', + }, + { + uuid: '1f8e3668-c7d7-4fdb-8195-3f337dfe10bf', + value: 'polkitd', + }, + { + uuid: 'e101d201-c675-47f3-b488-77bd0ce71920', + value: 'systemd-network', + }, + { + uuid: '0144102f-d69c-43a3-bf3b-04bde7d1b4e8', + value: 'systemd-resolve', + }, + { + uuid: '00c5a919-949e-4031-956e-3eeb071e9210', + value: 'systemd-timesync', + }, + ], + events_context_count: 100, + }, + }, + ], + }, + }; + return searchResponse; +}; diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts index b62cd24e938e..26db89124288 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request.ts @@ -11,10 +11,16 @@ import { ATTACK_DISCOVERY_CANCEL_BY_CONNECTOR_ID, CAPABILITIES, } from '../../common/constants'; +import type { + DefendInsightsGetRequestQuery, + DefendInsightsPostRequestBody, +} from '@kbn/elastic-assistant-common'; import { AttackDiscoveryPostRequestBody, ConversationCreateProps, ConversationUpdateProps, + DEFEND_INSIGHTS, + DEFEND_INSIGHTS_BY_ID, ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_BULK_ACTION, ELASTIC_AI_ASSISTANT_ANONYMIZATION_FIELDS_URL_FIND, ELASTIC_AI_ASSISTANT_CONVERSATIONS_URL, @@ -235,3 +241,24 @@ export const postAttackDiscoveryRequest = (body: AttackDiscoveryPostRequestBody) path: ATTACK_DISCOVERY, body, }); + +export const getDefendInsightRequest = (insightId: string) => + requestMock.create({ + method: 'get', + path: DEFEND_INSIGHTS_BY_ID, + params: { id: insightId }, + }); + +export const getDefendInsightsRequest = (queryParams: DefendInsightsGetRequestQuery) => + requestMock.create({ + method: 'get', + path: DEFEND_INSIGHTS, + query: queryParams, + }); + +export const postDefendInsightsRequest = (body: DefendInsightsPostRequestBody) => + requestMock.create({ + method: 'post', + path: DEFEND_INSIGHTS, + body, + }); diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts index 19d98633a83c..77bd6b00105b 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -28,6 +28,7 @@ import { } from '../ai_assistant_data_clients/knowledge_base'; import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; +import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights'; import { authenticatedUser } from './user'; export const createMockClients = () => { @@ -47,6 +48,7 @@ export const createMockClients = () => { getAIAssistantKnowledgeBaseDataClient: knowledgeBaseDataClientMock.create(), getAIAssistantPromptsDataClient: dataClientMock.create(), getAttackDiscoveryDataClient: attackDiscoveryDataClientMock.create(), + getDefendInsightsDataClient: dataClientMock.create(), getAIAssistantAnonymizationFieldsDataClient: dataClientMock.create(), getSpaceId: jest.fn(), getCurrentUser: jest.fn(), @@ -125,6 +127,10 @@ const createElasticAssistantRequestContextMock = ( () => clients.elasticAssistant.getAttackDiscoveryDataClient ) as unknown as jest.MockInstance, [], unknown> & (() => Promise), + getDefendInsightsDataClient: jest.fn( + () => clients.elasticAssistant.getDefendInsightsDataClient + ) as unknown as jest.MockInstance, [], unknown> & + (() => Promise), getAIAssistantKnowledgeBaseDataClient: jest.fn( () => clients.elasticAssistant.getAIAssistantKnowledgeBaseDataClient ) as unknown as jest.MockInstance< diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/field_maps_configuration.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/field_maps_configuration.ts new file mode 100644 index 000000000000..5769ab455710 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/field_maps_configuration.ts @@ -0,0 +1,174 @@ +/* + * 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 { FieldMap } from '@kbn/data-stream-adapter'; + +export const defendInsightsFieldMap: FieldMap = { + '@timestamp': { + type: 'date', + array: false, + required: false, + }, + users: { + type: 'nested', + array: true, + required: false, + }, + 'users.id': { + type: 'keyword', + array: false, + required: true, + }, + 'users.name': { + type: 'keyword', + array: false, + required: false, + }, + id: { + type: 'keyword', + array: false, + required: true, + }, + last_viewed_at: { + type: 'date', + array: false, + required: true, + }, + updated_at: { + type: 'date', + array: false, + required: true, + }, + created_at: { + type: 'date', + array: false, + required: true, + }, + endpoint_ids: { + type: 'keyword', + array: true, + required: false, + }, + insight_type: { + type: 'keyword', + required: true, + }, + insights: { + type: 'nested', + array: true, + required: false, + }, + 'insights.group': { + type: 'keyword', + array: true, + required: true, + }, + 'insights.events': { + type: 'nested', + array: true, + required: false, + }, + 'insights.events.endpoint_id': { + type: 'keyword', + array: false, + required: true, + }, + 'insights.events.id': { + type: 'keyword', + array: false, + required: true, + }, + 'insights.events.value': { + type: 'text', + array: false, + required: true, + }, + replacements: { + type: 'object', + array: false, + required: false, + }, + 'replacements.value': { + type: 'keyword', + array: false, + required: false, + }, + 'replacements.uuid': { + type: 'keyword', + array: false, + required: false, + }, + api_config: { + type: 'object', + array: false, + required: true, + }, + 'api_config.connector_id': { + type: 'keyword', + array: false, + required: true, + }, + 'api_config.action_type_id': { + type: 'keyword', + array: false, + required: false, + }, + 'api_config.default_system_prompt_id': { + type: 'keyword', + array: false, + required: false, + }, + 'api_config.provider': { + type: 'keyword', + array: false, + required: false, + }, + 'api_config.model': { + type: 'keyword', + array: false, + required: false, + }, + events_context_count: { + type: 'integer', + array: false, + required: false, + }, + status: { + type: 'keyword', + array: false, + required: true, + }, + namespace: { + type: 'keyword', + array: false, + required: true, + }, + average_interval_ms: { + type: 'integer', + array: false, + required: false, + }, + failure_reason: { + type: 'keyword', + array: false, + required: false, + }, + generation_intervals: { + type: 'nested', + array: true, + required: false, + }, + 'generation_intervals.date': { + type: 'date', + array: false, + required: true, + }, + 'generation_intervals.duration_ms': { + type: 'integer', + array: false, + required: true, + }, +} as const; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts new file mode 100644 index 000000000000..415487534a1b --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts @@ -0,0 +1,69 @@ +/* + * 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 { AuthenticatedUser } from '@kbn/core-security-common'; + +import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; +import { loggerMock } from '@kbn/logging-mocks'; + +import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import { getDefendInsight } from './get_defend_insight'; + +const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); +const mockLogger = loggerMock.create(); + +const mockResponse = getDefendInsightsSearchEsMock(); + +const user = { + username: 'test_user', + profile_uid: '1234', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, +} as AuthenticatedUser; +const mockRequest = { + esClient: mockEsClient, + index: 'defend-insights-index', + id: 'insight-id', + user, + logger: mockLogger, +}; +describe('getDefendInsight', () => { + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should get defend insight by id successfully', async () => { + mockEsClient.search.mockResolvedValueOnce(mockResponse); + + const response = await getDefendInsight(mockRequest); + + expect(response).not.toBeNull(); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).not.toHaveBeenCalled(); + }); + + it('should return null if no defend insights found', async () => { + mockEsClient.search.mockResolvedValueOnce({ ...mockResponse, hits: { hits: [] } }); + + const response = await getDefendInsight(mockRequest); + + expect(response).toBeNull(); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).not.toHaveBeenCalled(); + }); + + it('should throw error on elasticsearch search failure', async () => { + mockEsClient.search.mockRejectedValueOnce(new Error('Elasticsearch error')); + + await expect(getDefendInsight(mockRequest)).rejects.toThrowError('Elasticsearch error'); + + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).toHaveBeenCalledTimes(1); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.ts new file mode 100644 index 000000000000..4eeef2afd873 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.ts @@ -0,0 +1,78 @@ +/* + * 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 { AuthenticatedUser, ElasticsearchClient, Logger } from '@kbn/core/server'; +import { DefendInsightsResponse } from '@kbn/elastic-assistant-common'; + +import { EsDefendInsightSchema } from './types'; +import { transformESSearchToDefendInsights } from './helpers'; + +export interface GetDefendInsightParams { + esClient: ElasticsearchClient; + logger: Logger; + index: string; + id: string; + user: AuthenticatedUser; +} + +export const getDefendInsight = async ({ + esClient, + logger, + index, + id, + user, +}: GetDefendInsightParams): Promise => { + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + try { + const response = await esClient.search({ + query: { + bool: { + must: [ + { + bool: { + should: [ + { + term: { + _id: id, + }, + }, + ], + }, + }, + ...filterByUser, + ], + }, + }, + _source: true, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + }); + const insights = transformESSearchToDefendInsights(response); + return insights[0] ?? null; + } catch (err) { + logger.error(`Error fetching Defend insight: ${err} with id: ${id}`); + throw err; + } +}; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.test.ts new file mode 100644 index 000000000000..8e0793218154 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.test.ts @@ -0,0 +1,64 @@ +/* + * 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 { DefendInsightsGetRequestQuery } from '@kbn/elastic-assistant-common'; + +import { DefendInsightType, DefendInsightStatus } from '@kbn/elastic-assistant-common'; + +import { queryParamsToEsQuery } from './helpers'; + +describe('defend insights data client helpers', () => { + describe('queryParamsToEsQuery', () => { + let queryParams: DefendInsightsGetRequestQuery; + let expectedQuery: object[]; + + function getDefaultQueryParams(): DefendInsightsGetRequestQuery { + return { + ids: ['insight-id1', 'insight-id2'], + endpoint_ids: ['endpoint-id1', 'endpoint-id2'], + connector_id: 'connector-id1', + type: DefendInsightType.Enum.incompatible_antivirus, + status: DefendInsightStatus.Enum.succeeded, + }; + } + + function getDefaultExpectedQuery(): object[] { + return [ + { terms: { _id: queryParams.ids } }, + { terms: { endpoint_ids: queryParams.endpoint_ids } }, + { term: { 'api_config.connector_id': queryParams.connector_id } }, + { term: { insight_type: queryParams.type } }, + { term: { status: queryParams.status } }, + ]; + } + + beforeEach(() => { + queryParams = getDefaultQueryParams(); + expectedQuery = getDefaultExpectedQuery(); + }); + + it('should correctly convert valid query parameters to Elasticsearch query format', () => { + const result = queryParamsToEsQuery(queryParams); + expect(result).toEqual(expectedQuery); + }); + + it('should ignore invalid query parameters', () => { + const badParams = { + ...queryParams, + invalid_param: 'invalid value', + }; + + const result = queryParamsToEsQuery(badParams); + expect(result).toEqual(expectedQuery); + }); + + it('should handle empty query parameters', () => { + const result = queryParamsToEsQuery({}); + expect(result).toEqual([]); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts new file mode 100644 index 000000000000..b8164f53d981 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts @@ -0,0 +1,221 @@ +/* + * 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 { get as _get, isArray } from 'lodash'; + +import type { estypes } from '@elastic/elasticsearch'; +import type { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types'; +import type { AuthenticatedUser } from '@kbn/core/server'; +import type { + DefendInsightCreateProps, + DefendInsightUpdateProps, + DefendInsightsResponse, + DefendInsightsGetRequestQuery, +} from '@kbn/elastic-assistant-common'; + +import type { + CreateDefendInsightSchema, + EsDefendInsightSchema, + UpdateDefendInsightSchema, +} from './types'; + +export const transformESSearchToDefendInsights = ( + response: estypes.SearchResponse +): DefendInsightsResponse[] => { + return response.hits.hits + .filter((hit) => hit._source !== undefined) + .map((hit) => { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const insightSchema = hit._source!; + const defendInsight: DefendInsightsResponse = { + timestamp: insightSchema['@timestamp'], + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + id: hit._id!, + backingIndex: hit._index, + createdAt: insightSchema.created_at, + updatedAt: insightSchema.updated_at, + lastViewedAt: insightSchema.last_viewed_at, + users: + insightSchema.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + namespace: insightSchema.namespace, + status: insightSchema.status, + eventsContextCount: insightSchema.events_context_count, + apiConfig: { + connectorId: insightSchema.api_config.connector_id, + actionTypeId: insightSchema.api_config.action_type_id, + defaultSystemPromptId: insightSchema.api_config.default_system_prompt_id, + model: insightSchema.api_config.model, + provider: insightSchema.api_config.provider, + }, + endpointIds: insightSchema.endpoint_ids, + insightType: insightSchema.insight_type, + insights: insightSchema.insights.map((insight) => ({ + group: insight.group, + events: insight.events?.map((event) => ({ + id: event.id, + endpointId: event.endpoint_id, + value: event.value, + })), + })), + replacements: insightSchema.replacements?.reduce((acc: Record, r) => { + acc[r.uuid] = r.value; + return acc; + }, {}), + generationIntervals: + insightSchema.generation_intervals?.map((interval) => ({ + date: interval.date, + durationMs: interval.duration_ms, + })) ?? [], + averageIntervalMs: insightSchema.average_interval_ms ?? 0, + failureReason: insightSchema.failure_reason, + }; + + return defendInsight; + }); +}; + +export const transformToCreateScheme = ( + createdAt: string, + spaceId: string, + user: AuthenticatedUser, + { + endpointIds, + insightType, + insights, + apiConfig, + eventsContextCount, + replacements, + status, + }: DefendInsightCreateProps +): CreateDefendInsightSchema => { + return { + '@timestamp': createdAt, + created_at: createdAt, + users: [ + { + id: user.profile_uid, + name: user.username, + }, + ], + status, + api_config: { + action_type_id: apiConfig.actionTypeId, + connector_id: apiConfig.connectorId, + default_system_prompt_id: apiConfig.defaultSystemPromptId, + model: apiConfig.model, + provider: apiConfig.provider, + }, + events_context_count: eventsContextCount, + endpoint_ids: endpointIds, + insight_type: insightType, + insights: insights?.map((insight) => ({ + group: insight.group, + events: insight.events?.map((event) => ({ + id: event.id, + endpoint_id: event.endpointId, + value: event.value, + })), + })), + updated_at: createdAt, + last_viewed_at: createdAt, + replacements: replacements + ? Object.keys(replacements).map((key) => ({ + uuid: key, + value: replacements[key], + })) + : undefined, + namespace: spaceId, + }; +}; + +export const transformToUpdateScheme = ( + updatedAt: string, + { + eventsContextCount, + apiConfig, + insights, + failureReason, + generationIntervals, + id, + replacements, + lastViewedAt, + status, + }: DefendInsightUpdateProps +): UpdateDefendInsightSchema => { + const averageIntervalMsObj = + generationIntervals && generationIntervals.length > 0 + ? { + average_interval_ms: Math.trunc( + generationIntervals.reduce((acc, interval) => acc + interval.durationMs, 0) / + generationIntervals.length + ), + generation_intervals: generationIntervals.map((interval) => ({ + date: interval.date, + duration_ms: interval.durationMs, + })), + } + : {}; + return { + events_context_count: eventsContextCount, + ...(apiConfig + ? { + api_config: { + action_type_id: apiConfig.actionTypeId, + connector_id: apiConfig.connectorId, + default_system_prompt_id: apiConfig.defaultSystemPromptId, + model: apiConfig.model, + provider: apiConfig.provider, + }, + } + : {}), + ...(insights + ? { + insights: insights.map((insight) => ({ + group: insight.group, + events: insight.events?.map((event) => ({ + id: event.id, + endpoint_id: event.endpointId, + value: event.value, + })), + })), + } + : {}), + failure_reason: failureReason, + id, + replacements: replacements + ? Object.keys(replacements).map((key) => ({ + uuid: key, + value: replacements[key], + })) + : undefined, + ...(status ? { status } : {}), + // only update updated_at time if this is not an update to last_viewed_at + ...(lastViewedAt ? { last_viewed_at: lastViewedAt } : { updated_at: updatedAt }), + ...averageIntervalMsObj, + }; +}; + +const validParams = new Set(['ids', 'endpoint_ids', 'connector_id', 'type', 'status']); +const paramKeyMap = { ids: '_id', connector_id: 'api_config.connector_id', type: 'insight_type' }; +export function queryParamsToEsQuery( + queryParams: DefendInsightsGetRequestQuery +): QueryDslQueryContainer[] { + return Object.entries(queryParams).reduce((acc: object[], [k, v]) => { + if (!validParams.has(k)) { + return acc; + } + + const filterKey = isArray(v) ? 'terms' : 'term'; + const paramKey = _get(paramKeyMap, k, k); + const next = { [filterKey]: { [paramKey]: v } }; + + return [...acc, next]; + }, []); +} diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts new file mode 100644 index 000000000000..704ee9b96255 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts @@ -0,0 +1,410 @@ +/* + * 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. + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import type { AuthenticatedUser } from '@kbn/core-security-common'; +import type { + DefendInsightCreateProps, + DefendInsightsUpdateProps, + DefendInsightsGetRequestQuery, + DefendInsightsResponse, +} from '@kbn/elastic-assistant-common'; + +import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; +import { loggerMock } from '@kbn/logging-mocks'; +import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; + +import type { AIAssistantDataClientParams } from '..'; + +import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import { getDefendInsight } from './get_defend_insight'; +import { + queryParamsToEsQuery, + transformESSearchToDefendInsights, + transformToUpdateScheme, +} from './helpers'; +import { DefendInsightsDataClient } from '.'; + +jest.mock('./get_defend_insight'); +jest.mock('./helpers', () => { + const original = jest.requireActual('./helpers'); + return { + ...original, + queryParamsToEsQuery: jest.fn(), + }; +}); + +describe('DefendInsightsDataClient', () => { + const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); + const mockLogger = loggerMock.create(); + const mockGetDefendInsight = jest.mocked(getDefendInsight); + let user: AuthenticatedUser; + let dataClientParams: AIAssistantDataClientParams; + let dataClient: DefendInsightsDataClient; + + function getDefaultUser(): AuthenticatedUser { + return { + username: 'test_user', + profile_uid: '1234', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; + } + + function getDefaultDataClientParams(): AIAssistantDataClientParams { + return { + logger: mockLogger, + currentUser: user, + elasticsearchClientPromise: new Promise((resolve) => resolve(mockEsClient)), + indexPatternsResourceName: 'defend-insights-index', + kibanaVersion: '9.0.0', + spaceId: 'space-1', + } as AIAssistantDataClientParams; + } + + beforeEach(() => { + user = getDefaultUser(); + dataClientParams = getDefaultDataClientParams(); + dataClient = new DefendInsightsDataClient(dataClientParams); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + describe('getDefendInsight', () => { + it('should correctly get defend insight', async () => { + const id = 'some-id'; + mockGetDefendInsight.mockResolvedValueOnce({ id } as DefendInsightsResponse); + const response = await dataClient.getDefendInsight({ id, authenticatedUser: user }); + + expect(mockGetDefendInsight).toHaveBeenCalledTimes(1); + expect(response).not.toBeNull(); + expect(response!.id).toEqual(id); + }); + }); + + describe('createDefendInsight', () => { + const defendInsightCreate: DefendInsightCreateProps = { + endpointIds: [], + insightType: DefendInsightType.Enum.incompatible_antivirus, + insights: [], + apiConfig: { + actionTypeId: 'action-type-id', + connectorId: 'connector-id', + defaultSystemPromptId: 'default-prompt-id', + model: 'model-name', + provider: 'OpenAI', + }, + eventsContextCount: 10, + replacements: { key1: 'value1', key2: 'value2' }, + status: DefendInsightStatus.Enum.running, + }; + + it('should create defend insight successfully', async () => { + const id = 'created-id'; + // @ts-expect-error not full response interface + mockEsClient.create.mockResolvedValueOnce({ _id: id }); + mockGetDefendInsight.mockResolvedValueOnce({ id } as DefendInsightsResponse); + + const response = await dataClient.createDefendInsight({ + defendInsightCreate, + authenticatedUser: user, + }); + expect(mockEsClient.create).toHaveBeenCalledTimes(1); + expect(mockGetDefendInsight).toHaveBeenCalledTimes(1); + expect(response).not.toBeNull(); + expect(response!.id).toEqual(id); + }); + + it('should throw error on elasticsearch create failure', async () => { + mockEsClient.create.mockRejectedValueOnce(new Error('Elasticsearch error')); + const responsePromise = dataClient.createDefendInsight({ + defendInsightCreate, + authenticatedUser: user, + }); + await expect(responsePromise).rejects.toThrowError('Elasticsearch error'); + expect(mockEsClient.create).toHaveBeenCalledTimes(1); + expect(mockGetDefendInsight).not.toHaveBeenCalled(); + }); + }); + + describe('findDefendInsightsByParams', () => { + let mockQueryParamsToEsQuery: Function; + let queryParams: DefendInsightsGetRequestQuery; + let expectedTermFilters: object[]; + + function getDefaultQueryParams() { + return { + ids: ['insight-id1', 'insight-id2'], + endpoint_ids: ['endpoint-id1', 'endpoint-id2'], + connector_id: 'connector-id1', + type: DefendInsightType.Enum.incompatible_antivirus, + status: DefendInsightStatus.Enum.succeeded, + }; + } + + function getDefaultExpectedTermFilters() { + return [ + { terms: { _id: queryParams.ids } }, + { terms: { endpoint_ids: queryParams.endpoint_ids } }, + { term: { 'api_config.connector_id': queryParams.connector_id } }, + { term: { insight_type: queryParams.type } }, + { term: { status: queryParams.status } }, + ]; + } + + beforeEach(() => { + queryParams = getDefaultQueryParams(); + expectedTermFilters = getDefaultExpectedTermFilters(); + mockQueryParamsToEsQuery = jest + .mocked(queryParamsToEsQuery) + .mockReturnValueOnce(expectedTermFilters); + }); + + it('should return defend insights successfully', async () => { + const mockResponse = getDefendInsightsSearchEsMock(); + mockEsClient.search.mockResolvedValueOnce(mockResponse); + + const result = await dataClient.findDefendInsightsByParams({ + params: queryParams, + authenticatedUser: user, + }); + const expectedResult = transformESSearchToDefendInsights(mockResponse); + + expect(mockQueryParamsToEsQuery).toHaveBeenCalledTimes(1); + expect(mockQueryParamsToEsQuery).toHaveBeenCalledWith(queryParams); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockEsClient.search).toHaveBeenCalledWith( + expect.objectContaining({ + index: `${dataClientParams.indexPatternsResourceName}-${dataClientParams.spaceId}`, + size: 10, + sort: [ + { + '@timestamp': 'desc', + }, + ], + query: { + bool: { + must: [ + ...expectedTermFilters, + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: { 'users.id': user.profile_uid }, + }, + ], + }, + }, + }, + }, + ], + }, + }, + }) + ); + expect(result).toEqual(expectedResult); + }); + + it('should log and throw an error if search fails', async () => { + const mockError = new Error('Search failed'); + mockEsClient.search.mockRejectedValue(mockError); + + await expect( + dataClient.findDefendInsightsByParams({ + params: queryParams, + authenticatedUser: user, + }) + ).rejects.toThrow(mockError); + expect(mockLogger.error).toHaveBeenCalledWith( + `error fetching Defend insights: ${mockError} with params: ${JSON.stringify(queryParams)}` + ); + }); + }); + + describe('findAllDefendInsights', () => { + it('should correctly query ES', async () => { + const mockResponse = getDefendInsightsSearchEsMock(); + mockEsClient.search.mockResolvedValueOnce(mockResponse); + const searchParams = { + query: { + bool: { + must: [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: { 'users.id': user.profile_uid }, + }, + ], + }, + }, + }, + }, + ], + }, + }, + size: 10000, + _source: true, + ignore_unavailable: true, + index: `${dataClientParams.indexPatternsResourceName}-${dataClientParams.spaceId}`, + seq_no_primary_term: true, + }; + + const response = await dataClient.findAllDefendInsights({ + authenticatedUser: user, + }); + expect(response).not.toBeNull(); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockEsClient.search).toHaveBeenCalledWith(searchParams); + }); + + it('should throw error on elasticsearch search failure', async () => { + mockEsClient.search.mockRejectedValueOnce(new Error('Elasticsearch error')); + await expect( + dataClient.findAllDefendInsights({ + authenticatedUser: user, + }) + ).rejects.toThrowError('Elasticsearch error'); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).toHaveBeenCalledTimes(1); + }); + }); + + describe('updateDefendInsights', () => { + let defendInsightsUpdateProps: DefendInsightsUpdateProps; + + function getDefaultProps() { + return [ + { + id: 'insight-id1', + backingIndex: 'defend-insights-index', + status: DefendInsightStatus.Enum.succeeded, + insights: [ + { + group: 'windows_defender', + events: [ + { + id: 'event-id-1', + endpointId: 'endpoint-id-1', + value: '/windows/defender/scan.exe', + }, + ], + }, + ], + }, + ]; + } + + beforeEach(async () => { + defendInsightsUpdateProps = getDefaultProps(); + }); + + it('should update defend insights successfully', async () => { + // ensure startTime is before updatedAt timestamp + const startTime = new Date().getTime() - 1; + const mockResponse: DefendInsightsResponse[] = [ + { id: defendInsightsUpdateProps[0].id } as DefendInsightsResponse, + ]; + + const findDefendInsightsByParamsSpy = jest.spyOn(dataClient, 'findDefendInsightsByParams'); + findDefendInsightsByParamsSpy.mockResolvedValueOnce(mockResponse); + + const result = await dataClient.updateDefendInsights({ + defendInsightsUpdateProps, + authenticatedUser: user, + }); + const expectedDoc = transformToUpdateScheme('', defendInsightsUpdateProps[0]); + delete expectedDoc.updated_at; + + expect(mockEsClient.bulk).toHaveBeenCalledTimes(1); + expect(mockEsClient.bulk).toHaveBeenCalledWith({ + body: [ + { + update: { + _index: defendInsightsUpdateProps[0].backingIndex, + _id: defendInsightsUpdateProps[0].id, + }, + }, + { + doc: expect.objectContaining({ ...expectedDoc }), + }, + ], + refresh: 'wait_for', + }); + const updatedAt = (mockEsClient.bulk.mock.calls[0][0] as { body: any[] }).body[1].doc + .updated_at; + expect(new Date(updatedAt).getTime()).toBeGreaterThan(startTime); + expect(dataClient.findDefendInsightsByParams).toHaveBeenCalledTimes(1); + expect(dataClient.findDefendInsightsByParams).toHaveBeenCalledWith({ + params: { ids: [defendInsightsUpdateProps[0].id] }, + authenticatedUser: user, + }); + expect(result).toEqual(mockResponse); + }); + + it('should log a warning and throw an error if update fails', async () => { + const mockError = new Error('Update failed'); + mockEsClient.bulk.mockRejectedValue(mockError); + + await expect( + dataClient.updateDefendInsights({ + defendInsightsUpdateProps, + authenticatedUser: user, + }) + ).rejects.toThrow(mockError); + + expect(mockLogger.warn).toHaveBeenCalledWith( + `error updating Defend insights: ${mockError} for IDs: ${defendInsightsUpdateProps[0].id}` + ); + }); + }); + + describe('updateDefendInsight', () => { + it('correctly calls updateDefendInsights', async () => { + const defendInsightUpdateProps = { + id: 'insight-id1', + backingIndex: 'defend-insights-index', + status: DefendInsightStatus.Enum.succeeded, + insights: [ + { + group: 'windows_defender', + events: [ + { + id: 'event-id-1', + endpointId: 'endpoint-id-1', + value: '/windows/defender/scan.exe', + }, + ], + }, + ], + }; + const updateDefendInsightsSpy = jest.spyOn(dataClient, 'updateDefendInsights'); + updateDefendInsightsSpy.mockResolvedValueOnce([]); + await dataClient.updateDefendInsight({ + defendInsightUpdateProps, + authenticatedUser: user, + }); + + expect(updateDefendInsightsSpy).toHaveBeenCalledTimes(1); + expect(updateDefendInsightsSpy).toHaveBeenCalledWith({ + defendInsightsUpdateProps: [defendInsightUpdateProps], + authenticatedUser: user, + }); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts new file mode 100644 index 000000000000..b5cbbd6cd18a --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts @@ -0,0 +1,286 @@ +/* + * 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 { v4 as uuidv4 } from 'uuid'; + +import type { + DefendInsightCreateProps, + DefendInsightUpdateProps, + DefendInsightsUpdateProps, + DefendInsightsResponse, + DefendInsightsGetRequestQuery, +} from '@kbn/elastic-assistant-common'; +import type { AuthenticatedUser } from '@kbn/core-security-common'; + +import type { AIAssistantDataClientParams } from '..'; +import type { EsDefendInsightSchema } from './types'; + +import { AIAssistantDataClient } from '..'; +import { getDefendInsight } from './get_defend_insight'; +import { + queryParamsToEsQuery, + transformESSearchToDefendInsights, + transformToCreateScheme, + transformToUpdateScheme, +} from './helpers'; + +const DEFAULT_PAGE_SIZE = 10; + +export class DefendInsightsDataClient extends AIAssistantDataClient { + constructor(public readonly options: AIAssistantDataClientParams) { + super(options); + } + + /** + * Fetches a Defend insight + * @param options + * @param options.id The existing Defend insight id. + * @param options.authenticatedUser Current authenticated user. + * @returns The Defend insight response + */ + public getDefendInsight = async ({ + id, + authenticatedUser, + }: { + id: string; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + return getDefendInsight({ + esClient, + logger: this.options.logger, + index: this.indexTemplateAndPattern.alias, + id, + user: authenticatedUser, + }); + }; + + /** + * Creates a Defend insight, if given at least the "apiConfig" + * @param options + * @param options.defendInsightCreate + * @param options.authenticatedUser + * @returns The Defend insight created + */ + public createDefendInsight = async ({ + defendInsightCreate, + authenticatedUser, + }: { + defendInsightCreate: DefendInsightCreateProps; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const index = this.indexTemplateAndPattern.alias; + const user = authenticatedUser; + const id = defendInsightCreate?.id || uuidv4(); + const createdAt = new Date().toISOString(); + + const body = transformToCreateScheme(createdAt, this.spaceId, user, defendInsightCreate); + try { + const response = await esClient.create({ + body, + id, + index, + refresh: 'wait_for', + }); + + const createdDefendInsight = await getDefendInsight({ + esClient, + index, + id: response._id, + logger, + user, + }); + return createdDefendInsight; + } catch (err) { + logger.error(`error creating Defend insight: ${err} with id: ${id}`); + throw err; + } + }; + + /** + * Find Defend insights by params + * @param options + * @param options.params + * @param options.authenticatedUser + * @returns The Defend insights found + */ + public findDefendInsightsByParams = async ({ + params, + authenticatedUser, + }: { + params: DefendInsightsGetRequestQuery; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const index = this.indexTemplateAndPattern.alias; + const user = authenticatedUser; + const termFilters = queryParamsToEsQuery(params); + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + + try { + const query = { + bool: { + must: [...termFilters, ...filterByUser], + }, + }; + const response = await esClient.search({ + query, + _source: true, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + sort: [{ '@timestamp': 'desc' }], + size: params.size || DEFAULT_PAGE_SIZE, + }); + return transformESSearchToDefendInsights(response); + } catch (err) { + logger.error(`error fetching Defend insights: ${err} with params: ${JSON.stringify(params)}`); + throw err; + } + }; + + /** + * Finds all Defend insight for authenticated user + * @param options + * @param options.authenticatedUser + * @returns The Defend insight + */ + public findAllDefendInsights = async ({ + authenticatedUser, + }: { + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const index = this.indexTemplateAndPattern.alias; + const user = authenticatedUser; + const MAX_ITEMS = 10000; + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + + try { + const response = await esClient.search({ + query: { + bool: { + must: [...filterByUser], + }, + }, + size: MAX_ITEMS, + _source: true, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + }); + const insights = transformESSearchToDefendInsights(response); + return insights ?? []; + } catch (err) { + logger.error(`error fetching Defend insights: ${err}`); + throw err; + } + }; + + /** + * Updates Defend insights + * @param options + * @param options.defendInsightsUpdateProps + * @param options.authenticatedUser + */ + public updateDefendInsights = async ({ + defendInsightsUpdateProps, + authenticatedUser, + }: { + defendInsightsUpdateProps: DefendInsightsUpdateProps; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const updatedAt = new Date().toISOString(); + + let ids: string[] = []; + const bulkParams = defendInsightsUpdateProps.flatMap((updateProp) => { + const index = updateProp.backingIndex; + const params = transformToUpdateScheme(updatedAt, updateProp); + ids = [...ids, params.id]; + return [ + { + update: { + _index: index, + _id: params.id, + }, + }, + { + doc: params, + }, + ]; + }); + + try { + await esClient.bulk({ body: bulkParams, refresh: 'wait_for' }); + return this.findDefendInsightsByParams({ params: { ids }, authenticatedUser }); + } catch (err) { + logger.warn(`error updating Defend insights: ${err} for IDs: ${ids}`); + throw err; + } + }; + + /** + * Updates a Defend insight + * @param options + * @param options.defendInsightUpdateProps + * @param options.authenticatedUser + */ + public updateDefendInsight = async ({ + defendInsightUpdateProps, + authenticatedUser, + }: { + defendInsightUpdateProps: DefendInsightUpdateProps; + authenticatedUser: AuthenticatedUser; + }): Promise => { + return ( + await this.updateDefendInsights({ + defendInsightsUpdateProps: [defendInsightUpdateProps], + authenticatedUser, + }) + )[0]; + }; +} diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts new file mode 100644 index 000000000000..f04c7ef505c2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts @@ -0,0 +1,88 @@ +/* + * 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 { + DefendInsightStatus, + DefendInsightType, + Provider, + UUID, +} from '@kbn/elastic-assistant-common'; + +import type { EsReplacementSchema } from '../conversations/types'; + +interface DefendInsightInsightEventSchema { + id: string; + endpoint_id: string; + value: string; +} + +interface DefendInsightInsightSchema { + group: string; + events?: DefendInsightInsightEventSchema[]; +} + +interface BaseDefendInsightSchema { + '@timestamp': string; + created_at: string; + updated_at: string; + last_viewed_at: string; + status: DefendInsightStatus; + events_context_count?: number; + endpoint_ids: string[]; + insight_type: DefendInsightType; + insights: DefendInsightInsightSchema[]; + api_config: { + connector_id: string; + action_type_id: string; + default_system_prompt_id?: string; + provider?: Provider; + model?: string; + }; + replacements?: EsReplacementSchema[]; +} + +export interface EsDefendInsightSchema extends BaseDefendInsightSchema { + id: string; + namespace: string; + failure_reason?: string; + users?: Array<{ + id?: string; + name?: string; + }>; + average_interval_ms?: number; + generation_intervals?: Array<{ date: string; duration_ms: number }>; +} + +export interface CreateDefendInsightSchema extends BaseDefendInsightSchema { + id?: string | undefined; + users: Array<{ + id?: string; + name?: string; + }>; + namespace: string; +} + +export interface UpdateDefendInsightSchema { + id: UUID; + '@timestamp'?: string; + updated_at?: string; + last_viewed_at?: string; + status?: DefendInsightStatus; + events_context_count?: number; + insights?: DefendInsightInsightSchema[]; + api_config?: { + action_type_id?: string; + connector_id?: string; + default_system_prompt_id?: string; + provider?: Provider; + model?: string; + }; + replacements?: EsReplacementSchema[]; + average_interval_ms?: number; + generation_intervals?: Array<{ date: string; duration_ms: number }>; + failure_reason?: string; +} diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts index fb3ffe7442c1..4bfd4da6cfcb 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.test.ts @@ -141,7 +141,7 @@ describe('AI Assistant Service', () => { expect(assistantService.isInitialized()).toEqual(true); - expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(5); + expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(6); const expectedTemplates = [ '.kibana-elastic-ai-assistant-component-template-conversations', @@ -149,6 +149,7 @@ describe('AI Assistant Service', () => { '.kibana-elastic-ai-assistant-component-template-prompts', '.kibana-elastic-ai-assistant-component-template-anonymization-fields', '.kibana-elastic-ai-assistant-component-template-attack-discovery', + '.kibana-elastic-ai-assistant-component-template-defend-insights', ]; expectedTemplates.forEach((t, i) => { expect(clusterClient.cluster.putComponentTemplate.mock.calls[i][0].name).toEqual(t); @@ -650,7 +651,7 @@ describe('AI Assistant Service', () => { 'AI Assistant service initialized', async () => assistantService.isInitialized() === true ); - expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(7); + expect(clusterClient.cluster.putComponentTemplate).toHaveBeenCalledTimes(8); const expectedTemplates = [ '.kibana-elastic-ai-assistant-component-template-conversations', @@ -660,6 +661,7 @@ describe('AI Assistant Service', () => { '.kibana-elastic-ai-assistant-component-template-prompts', '.kibana-elastic-ai-assistant-component-template-anonymization-fields', '.kibana-elastic-ai-assistant-component-template-attack-discovery', + '.kibana-elastic-ai-assistant-component-template-defend-insights', ]; expectedTemplates.forEach((t, i) => { expect(clusterClient.cluster.putComponentTemplate.mock.calls[i][0].name).toEqual(t); @@ -684,7 +686,7 @@ describe('AI Assistant Service', () => { async () => (await getSpaceResourcesInitialized(assistantService)) === true ); - expect(clusterClient.indices.putIndexTemplate).toHaveBeenCalledTimes(7); + expect(clusterClient.indices.putIndexTemplate).toHaveBeenCalledTimes(8); const expectedTemplates = [ '.kibana-elastic-ai-assistant-index-template-conversations', '.kibana-elastic-ai-assistant-index-template-conversations', @@ -693,6 +695,7 @@ describe('AI Assistant Service', () => { '.kibana-elastic-ai-assistant-index-template-prompts', '.kibana-elastic-ai-assistant-index-template-anonymization-fields', '.kibana-elastic-ai-assistant-index-template-attack-discovery', + '.kibana-elastic-ai-assistant-index-template-defend-insights', ]; expectedTemplates.forEach((t, i) => { expect(clusterClient.indices.putIndexTemplate.mock.calls[i][0].name).toEqual(t); @@ -716,7 +719,7 @@ describe('AI Assistant Service', () => { async () => (await getSpaceResourcesInitialized(assistantService)) === true ); - expect(clusterClient.indices.putSettings).toHaveBeenCalledTimes(7); + expect(clusterClient.indices.putSettings).toHaveBeenCalledTimes(8); }); test('should retry updating index mappings for existing indices for transient ES errors', async () => { @@ -736,7 +739,7 @@ describe('AI Assistant Service', () => { async () => (await getSpaceResourcesInitialized(assistantService)) === true ); - expect(clusterClient.indices.putMapping).toHaveBeenCalledTimes(7); + expect(clusterClient.indices.putMapping).toHaveBeenCalledTimes(8); }); test('should retry creating concrete index for transient ES errors', async () => { diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index d7eff095b4be..81ddd69fb67d 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -13,6 +13,7 @@ import type { MlPluginSetup } from '@kbn/ml-plugin/server'; import { Subject } from 'rxjs'; import { LicensingApiRequestHandlerContext } from '@kbn/licensing-plugin/server'; import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration'; +import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; import { AssistantResourceNames, GetElser } from '../types'; import { AIAssistantConversationsDataClient } from '../ai_assistant_data_clients/conversations'; @@ -33,6 +34,7 @@ import { GetAIAssistantKnowledgeBaseDataClientParams, } from '../ai_assistant_data_clients/knowledge_base'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; +import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights'; import { createGetElserId, createPipeline, pipelineExists } from './helpers'; import { hasAIAssistantLicense } from '../routes/helpers'; @@ -64,7 +66,8 @@ export type CreateDataStream = (params: { | 'conversations' | 'knowledgeBase' | 'prompts' - | 'attackDiscovery'; + | 'attackDiscovery' + | 'defendInsights'; fieldMap: FieldMap; kibanaVersion: string; spaceId?: string; @@ -79,6 +82,7 @@ export class AIAssistantService { private promptsDataStream: DataStreamSpacesAdapter; private anonymizationFieldsDataStream: DataStreamSpacesAdapter; private attackDiscoveryDataStream: DataStreamSpacesAdapter; + private defendInsightsDataStream: DataStreamSpacesAdapter; private resourceInitializationHelper: ResourceInstallationHelper; private initPromise: Promise; private isKBSetupInProgress: boolean = false; @@ -112,6 +116,11 @@ export class AIAssistantService { kibanaVersion: options.kibanaVersion, fieldMap: attackDiscoveryFieldMap, }); + this.defendInsightsDataStream = this.createDataStream({ + resource: 'defendInsights', + kibanaVersion: options.kibanaVersion, + fieldMap: defendInsightsFieldMap, + }); this.initPromise = this.initializeResources(); @@ -222,6 +231,12 @@ export class AIAssistantService { logger: this.options.logger, pluginStop$: this.options.pluginStop$, }); + + await this.defendInsightsDataStream.install({ + esClient, + logger: this.options.logger, + pluginStop$: this.options.pluginStop$, + }); } catch (error) { this.options.logger.warn(`Error initializing AI assistant resources: ${error.message}`); this.initialized = false; @@ -240,6 +255,7 @@ export class AIAssistantService { prompts: getResourceName('component-template-prompts'), anonymizationFields: getResourceName('component-template-anonymization-fields'), attackDiscovery: getResourceName('component-template-attack-discovery'), + defendInsights: getResourceName('component-template-defend-insights'), }, aliases: { conversations: getResourceName('conversations'), @@ -247,6 +263,7 @@ export class AIAssistantService { prompts: getResourceName('prompts'), anonymizationFields: getResourceName('anonymization-fields'), attackDiscovery: getResourceName('attack-discovery'), + defendInsights: getResourceName('defend-insights'), }, indexPatterns: { conversations: getResourceName('conversations*'), @@ -254,6 +271,7 @@ export class AIAssistantService { prompts: getResourceName('prompts*'), anonymizationFields: getResourceName('anonymization-fields*'), attackDiscovery: getResourceName('attack-discovery*'), + defendInsights: getResourceName('defend-insights*'), }, indexTemplate: { conversations: getResourceName('index-template-conversations'), @@ -261,6 +279,7 @@ export class AIAssistantService { prompts: getResourceName('index-template-prompts'), anonymizationFields: getResourceName('index-template-anonymization-fields'), attackDiscovery: getResourceName('index-template-attack-discovery'), + defendInsights: getResourceName('index-template-defend-insights'), }, pipelines: { knowledgeBase: getResourceName('ingest-pipeline-knowledge-base'), @@ -393,6 +412,25 @@ export class AIAssistantService { }); } + public async createDefendInsightsDataClient( + opts: CreateAIAssistantClientParams + ): Promise { + const res = await this.checkResourcesInstallation(opts); + + if (res === null) { + return null; + } + + return new DefendInsightsDataClient({ + logger: this.options.logger.get('defendInsights'), + currentUser: opts.currentUser, + elasticsearchClientPromise: this.options.elasticsearchClientPromise, + indexPatternsResourceName: this.resourceNames.aliases.defendInsights, + kibanaVersion: this.options.kibanaVersion, + spaceId: opts.spaceId, + }); + } + public async createAIAssistantPromptsDataClient( opts: CreateAIAssistantClientParams ): Promise { diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts index 133419f45d17..9b2d444d643e 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/helpers.ts @@ -6,7 +6,7 @@ */ import { KibanaRequest } from '@kbn/core-http-server'; -import type { Message } from '@kbn/elastic-assistant-common'; +import type { DefendInsightsPostRequestBody, Message } from '@kbn/elastic-assistant-common'; import { AIMessage, BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages'; import { AttackDiscoveryPostRequestBody, @@ -36,7 +36,7 @@ export const requestHasRequiredAnonymizationParams = ( request: KibanaRequest< unknown, unknown, - ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody + ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody | DefendInsightsPostRequestBody > ): boolean => { const { replacements } = request?.body ?? {}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts b/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts index 1087703ba13a..92330b4960e7 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts @@ -411,6 +411,100 @@ export const CREATE_KNOWLEDGE_BASE_ENTRY_ERROR_EVENT: EventTypeOpts<{ }, }; +export const DEFEND_INSIGHT_SUCCESS_EVENT: EventTypeOpts<{ + actionTypeId: string; + eventsContextCount: number; + insightsGenerated: number; + durationMs: number; + model?: string; + provider?: string; +}> = { + eventType: 'defend_insight_success', + schema: { + actionTypeId: { + type: 'keyword', + _meta: { + description: 'Kibana connector type', + optional: false, + }, + }, + eventsContextCount: { + type: 'integer', + _meta: { + description: 'Number of events sent as context to the LLM', + optional: false, + }, + }, + insightsGenerated: { + type: 'integer', + _meta: { + description: 'Quantity of Defend insights generated', + optional: false, + }, + }, + durationMs: { + type: 'integer', + _meta: { + description: 'Duration of request in ms', + optional: false, + }, + }, + model: { + type: 'keyword', + _meta: { + description: 'LLM model', + optional: true, + }, + }, + provider: { + type: 'keyword', + _meta: { + description: 'OpenAI provider', + optional: true, + }, + }, + }, +}; + +export const DEFEND_INSIGHT_ERROR_EVENT: EventTypeOpts<{ + actionTypeId: string; + errorMessage: string; + model?: string; + provider?: string; +}> = { + eventType: 'defend_insight_error', + schema: { + actionTypeId: { + type: 'keyword', + _meta: { + description: 'Kibana connector type', + optional: false, + }, + }, + errorMessage: { + type: 'keyword', + _meta: { + description: 'Error message from Elasticsearch', + }, + }, + + model: { + type: 'keyword', + _meta: { + description: 'LLM model', + optional: true, + }, + }, + provider: { + type: 'keyword', + _meta: { + description: 'OpenAI provider', + optional: true, + }, + }, + }, +}; + export const events: Array> = [ KNOWLEDGE_BASE_EXECUTION_SUCCESS_EVENT, KNOWLEDGE_BASE_EXECUTION_ERROR_EVENT, @@ -420,4 +514,6 @@ export const events: Array> = [ INVOKE_ASSISTANT_ERROR_EVENT, ATTACK_DISCOVERY_SUCCESS_EVENT, ATTACK_DISCOVERY_ERROR_EVENT, + DEFEND_INSIGHT_SUCCESS_EVENT, + DEFEND_INSIGHT_ERROR_EVENT, ]; diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts new file mode 100644 index 000000000000..fa3ff15027e2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts @@ -0,0 +1,149 @@ +/* + * 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. + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import type { AuthenticatedUser } from '@kbn/core-security-common'; + +import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; + +import type { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; + +import { transformESSearchToDefendInsights } from '../../ai_assistant_data_clients/defend_insights/helpers'; +import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import { getDefendInsightRequest } from '../../__mocks__/request'; +import { + ElasticAssistantRequestHandlerContextMock, + requestContextMock, +} from '../../__mocks__/request_context'; +import { serverMock } from '../../__mocks__/server'; +import { isDefendInsightsEnabled, updateDefendInsightLastViewedAt } from './helpers'; +import { getDefendInsightRoute } from './get_defend_insight'; + +jest.mock('./helpers'); + +describe('getDefendInsightRoute', () => { + let server: ReturnType; + let context: ElasticAssistantRequestHandlerContextMock; + let mockUser: AuthenticatedUser; + let mockDataClient: DefendInsightsDataClient; + let mockCurrentInsight: any; + + function getDefaultUser(): AuthenticatedUser { + return { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; + } + + function getDefaultDataClient(): DefendInsightsDataClient { + return { + findDefendInsightByConnectorId: jest.fn(), + updateDefendInsight: jest.fn(), + createDefendInsight: jest.fn(), + getDefendInsight: jest.fn(), + } as unknown as DefendInsightsDataClient; + } + + beforeEach(() => { + const tools = requestContextMock.createTools(); + context = tools.context; + server = serverMock.create(); + tools.clients.core.elasticsearch.client = elasticsearchServiceMock.createScopedClusterClient(); + + mockUser = getDefaultUser(); + mockDataClient = getDefaultDataClient(); + mockCurrentInsight = transformESSearchToDefendInsights(getDefendInsightsSearchEsMock())[0]; + + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser); + context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValue(mockDataClient); + getDefendInsightRoute(server.router); + (updateDefendInsightLastViewedAt as jest.Mock).mockResolvedValue(mockCurrentInsight); + (isDefendInsightsEnabled as jest.Mock).mockReturnValue(true); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should handle successful request', async () => { + const response = await server.inject( + getDefendInsightRequest('insight-id1'), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(response.body).toEqual({ + data: mockCurrentInsight, + }); + }); + + it('should 404 if feature flag disabled', async () => { + (isDefendInsightsEnabled as jest.Mock).mockReturnValueOnce(false); + const response = await server.inject( + getDefendInsightRequest('insight-id1'), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(404); + }); + + it('should handle missing authenticated user', async () => { + context.elasticAssistant.getCurrentUser.mockReturnValueOnce(null); + const response = await server.inject( + getDefendInsightRequest('insight-id1'), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(401); + expect(response.body).toEqual({ + message: 'Authenticated user not found', + status_code: 401, + }); + }); + + it('should handle missing data client', async () => { + context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValueOnce(null); + const response = await server.inject( + getDefendInsightRequest('insight-id1'), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: 'Defend insights data client not initialized', + status_code: 500, + }); + }); + + it('should handle updateDefendInsightLastViewedAt empty array', async () => { + (updateDefendInsightLastViewedAt as jest.Mock).mockResolvedValueOnce([]); + const response = await server.inject( + getDefendInsightRequest('insight-id1'), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(response.body).toEqual({ data: [] }); + }); + + it('should handle updateDefendInsightLastViewedAt error', async () => { + (updateDefendInsightLastViewedAt as jest.Mock).mockRejectedValueOnce(new Error('Oh no!')); + const response = await server.inject( + getDefendInsightRequest('insight-id1'), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: { + error: 'Oh no!', + success: false, + }, + status_code: 500, + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.ts new file mode 100644 index 000000000000..5766b3d1b014 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.ts @@ -0,0 +1,96 @@ +/* + * 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 { IKibanaResponse } from '@kbn/core/server'; + +import { IRouter, Logger } from '@kbn/core/server'; +import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; +import { + DEFEND_INSIGHTS_BY_ID, + DefendInsightGetResponse, + DefendInsightGetRequestParams, + ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, +} from '@kbn/elastic-assistant-common'; +import { transformError } from '@kbn/securitysolution-es-utils'; + +import { buildResponse } from '../../lib/build_response'; +import { ElasticAssistantRequestHandlerContext } from '../../types'; +import { isDefendInsightsEnabled, updateDefendInsightLastViewedAt } from './helpers'; + +export const getDefendInsightRoute = (router: IRouter) => { + router.versioned + .get({ + access: 'internal', + path: DEFEND_INSIGHTS_BY_ID, + options: { + tags: ['access:elasticAssistant'], + }, + }) + .addVersion( + { + version: ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, + validate: { + request: { + params: buildRouteValidationWithZod(DefendInsightGetRequestParams), + }, + response: { + 200: { + body: { custom: buildRouteValidationWithZod(DefendInsightGetResponse) }, + }, + }, + }, + }, + async (context, request, response): Promise> => { + const resp = buildResponse(response); + const assistantContext = await context.elasticAssistant; + const logger: Logger = assistantContext.logger; + try { + const isEnabled = isDefendInsightsEnabled({ + request, + logger, + assistantContext, + }); + if (!isEnabled) { + return response.notFound(); + } + + const dataClient = await assistantContext.getDefendInsightsDataClient(); + const authenticatedUser = assistantContext.getCurrentUser(); + if (authenticatedUser == null) { + return resp.error({ + body: `Authenticated user not found`, + statusCode: 401, + }); + } + if (!dataClient) { + return resp.error({ + body: `Defend insights data client not initialized`, + statusCode: 500, + }); + } + + const defendInsight = await updateDefendInsightLastViewedAt({ + dataClient, + id: request.params.id, + authenticatedUser, + }); + + return response.ok({ + body: { data: defendInsight }, + }); + } catch (err) { + logger.error(err); + const error = transformError(err); + + return resp.error({ + body: { success: false, error: error.message }, + statusCode: error.statusCode, + }); + } + } + ); +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts new file mode 100644 index 000000000000..b27d71a690b6 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts @@ -0,0 +1,149 @@ +/* + * 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. + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import type { AuthenticatedUser } from '@kbn/core-security-common'; + +import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; + +import type { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; + +import { transformESSearchToDefendInsights } from '../../ai_assistant_data_clients/defend_insights/helpers'; +import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import { getDefendInsightsRequest } from '../../__mocks__/request'; +import { + ElasticAssistantRequestHandlerContextMock, + requestContextMock, +} from '../../__mocks__/request_context'; +import { serverMock } from '../../__mocks__/server'; +import { isDefendInsightsEnabled, updateDefendInsightsLastViewedAt } from './helpers'; +import { getDefendInsightsRoute } from './get_defend_insights'; + +jest.mock('./helpers'); + +describe('getDefendInsightsRoute', () => { + let server: ReturnType; + let context: ElasticAssistantRequestHandlerContextMock; + let mockUser: AuthenticatedUser; + let mockDataClient: DefendInsightsDataClient; + let mockCurrentInsights: any; + + function getDefaultUser(): AuthenticatedUser { + return { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; + } + + function getDefaultDataClient(): DefendInsightsDataClient { + return { + findDefendInsightByConnectorId: jest.fn(), + updateDefendInsight: jest.fn(), + createDefendInsight: jest.fn(), + getDefendInsight: jest.fn(), + } as unknown as DefendInsightsDataClient; + } + + beforeEach(() => { + const tools = requestContextMock.createTools(); + context = tools.context; + server = serverMock.create(); + tools.clients.core.elasticsearch.client = elasticsearchServiceMock.createScopedClusterClient(); + + mockUser = getDefaultUser(); + mockDataClient = getDefaultDataClient(); + mockCurrentInsights = transformESSearchToDefendInsights(getDefendInsightsSearchEsMock()); + + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser); + context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValue(mockDataClient); + getDefendInsightsRoute(server.router); + (updateDefendInsightsLastViewedAt as jest.Mock).mockResolvedValue(mockCurrentInsights); + (isDefendInsightsEnabled as jest.Mock).mockReturnValue(true); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should handle successful request', async () => { + const response = await server.inject( + getDefendInsightsRequest({ connector_id: 'connector-id1' }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(response.body).toEqual({ + data: mockCurrentInsights, + }); + }); + + it('should 404 if feature flag disabled', async () => { + (isDefendInsightsEnabled as jest.Mock).mockReturnValueOnce(false); + const response = await server.inject( + getDefendInsightsRequest({ connector_id: 'connector-id1' }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(404); + }); + + it('should handle missing authenticated user', async () => { + context.elasticAssistant.getCurrentUser.mockReturnValueOnce(null); + const response = await server.inject( + getDefendInsightsRequest({ connector_id: 'connector-id1' }), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(401); + expect(response.body).toEqual({ + message: 'Authenticated user not found', + status_code: 401, + }); + }); + + it('should handle missing data client', async () => { + context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValueOnce(null); + const response = await server.inject( + getDefendInsightsRequest({ connector_id: 'connector-id1' }), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: 'Defend insights data client not initialized', + status_code: 500, + }); + }); + + it('should handle updateDefendInsightsLastViewedAt empty array', async () => { + (updateDefendInsightsLastViewedAt as jest.Mock).mockResolvedValueOnce([]); + const response = await server.inject( + getDefendInsightsRequest({ connector_id: 'connector-id1' }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(response.body).toEqual({ data: [] }); + }); + + it('should handle updateDefendInsightsLastViewedAt error', async () => { + (updateDefendInsightsLastViewedAt as jest.Mock).mockRejectedValueOnce(new Error('Oh no!')); + const response = await server.inject( + getDefendInsightsRequest({ connector_id: 'connector-id1' }), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: { + error: 'Oh no!', + success: false, + }, + status_code: 500, + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.ts new file mode 100644 index 000000000000..e980c9be0915 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.ts @@ -0,0 +1,98 @@ +/* + * 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 { IKibanaResponse } from '@kbn/core/server'; + +import { IRouter, Logger } from '@kbn/core/server'; +import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; +import { + DEFEND_INSIGHTS, + DefendInsightsGetResponse, + DefendInsightsGetRequestQuery, + ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, +} from '@kbn/elastic-assistant-common'; +import { transformError } from '@kbn/securitysolution-es-utils'; + +import { buildResponse } from '../../lib/build_response'; +import { ElasticAssistantRequestHandlerContext } from '../../types'; +import { isDefendInsightsEnabled, updateDefendInsightsLastViewedAt } from './helpers'; + +export const getDefendInsightsRoute = (router: IRouter) => { + router.versioned + .get({ + access: 'internal', + path: DEFEND_INSIGHTS, + options: { + tags: ['access:elasticAssistant'], + }, + }) + .addVersion( + { + version: ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, + validate: { + request: { + query: buildRouteValidationWithZod(DefendInsightsGetRequestQuery), + }, + response: { + 200: { + body: { custom: buildRouteValidationWithZod(DefendInsightsGetResponse) }, + }, + }, + }, + }, + async (context, request, response): Promise> => { + const resp = buildResponse(response); + const assistantContext = await context.elasticAssistant; + const logger: Logger = assistantContext.logger; + try { + const isEnabled = isDefendInsightsEnabled({ + request, + logger, + assistantContext, + }); + if (!isEnabled) { + return response.notFound(); + } + + const dataClient = await assistantContext.getDefendInsightsDataClient(); + + const authenticatedUser = assistantContext.getCurrentUser(); + if (authenticatedUser == null) { + return resp.error({ + body: `Authenticated user not found`, + statusCode: 401, + }); + } + if (!dataClient) { + return resp.error({ + body: `Defend insights data client not initialized`, + statusCode: 500, + }); + } + + const defendInsights = await updateDefendInsightsLastViewedAt({ + dataClient, + params: request.query, + authenticatedUser, + }); + return response.ok({ + body: { + data: defendInsights, + }, + }); + } catch (err) { + logger.error(err); + const error = transformError(err); + + return resp.error({ + body: { success: false, error: error.message }, + statusCode: error.statusCode, + }); + } + } + ); +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts new file mode 100644 index 000000000000..22e89202e638 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts @@ -0,0 +1,255 @@ +/* + * 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. + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import moment from 'moment'; + +import { + DEFEND_INSIGHTS_TOOL_ID, + DefendInsightStatus, + DefendInsightType, +} from '@kbn/elastic-assistant-common'; +import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/openai/constants'; + +import { + DEFEND_INSIGHT_ERROR_EVENT, + DEFEND_INSIGHT_SUCCESS_EVENT, +} from '../../lib/telemetry/event_based_telemetry'; +import { + getAssistantTool, + getAssistantToolParams, + handleToolError, + updateDefendInsights, + updateDefendInsightLastViewedAt, +} from './helpers'; + +describe('defend insights route helpers', () => { + afterEach(() => { + jest.clearAllMocks(); + }); + + describe('getAssistantTool', () => { + it('should return the defend-insights tool', () => { + const getRegisteredTools = jest.fn().mockReturnValue([{ id: DEFEND_INSIGHTS_TOOL_ID }]); + const result = getAssistantTool(getRegisteredTools, 'pluginName'); + expect(result).toEqual({ id: DEFEND_INSIGHTS_TOOL_ID }); + }); + }); + + describe('getAssistantToolParams', () => { + it('should return the correct tool params', () => { + const params = { + endpointIds: ['endpoint-id1'], + insightType: DefendInsightType.Enum.incompatible_antivirus, + actionsClient: {} as any, + anonymizationFields: [], + apiConfig: { connectorId: 'connector-id1', actionTypeId: 'action-type-id1' }, + esClient: {} as any, + connectorTimeout: 1000, + langChainTimeout: 1000, + langSmithProject: 'project', + langSmithApiKey: 'apiKey', + logger: {} as any, + latestReplacements: {}, + onNewReplacements: jest.fn(), + request: {} as any, + }; + const result = getAssistantToolParams(params); + + expect(result).toHaveProperty('endpointIds', params.endpointIds); + expect(result).toHaveProperty('insightType', params.insightType); + expect(result).toHaveProperty('llm'); + }); + }); + + describe('handleToolError', () => { + it('should handle tool error and update defend insight', async () => { + const params = { + apiConfig: { + connectorId: 'connector-id1', + actionTypeId: 'action-type-id1', + model: 'model', + provider: OpenAiProviderType.OpenAi, + }, + defendInsightId: 'id', + authenticatedUser: {} as any, + dataClient: { + getDefendInsight: jest.fn().mockResolvedValueOnce({ + status: DefendInsightStatus.Enum.running, + backingIndex: 'index', + }), + updateDefendInsight: jest.fn(), + } as any, + err: new Error('error'), + latestReplacements: {}, + logger: { error: jest.fn() } as any, + telemetry: { reportEvent: jest.fn() } as any, + }; + await handleToolError(params); + + expect(params.dataClient.updateDefendInsight).toHaveBeenCalledTimes(1); + expect(params.telemetry.reportEvent).toHaveBeenCalledWith( + DEFEND_INSIGHT_ERROR_EVENT.eventType, + expect.any(Object) + ); + }); + }); + + describe('updateDefendInsights', () => { + it('should update defend insights', async () => { + const params = { + apiConfig: { + connectorId: 'connector-id1', + actionTypeId: 'action-type-id1', + model: 'model', + provider: OpenAiProviderType.OpenAi, + }, + defendInsightId: 'insight-id1', + authenticatedUser: {} as any, + dataClient: { + getDefendInsight: jest.fn().mockResolvedValueOnce({ + status: DefendInsightStatus.Enum.running, + backingIndex: 'backing-index-name', + generationIntervals: [], + }), + updateDefendInsight: jest.fn(), + } as any, + latestReplacements: {}, + logger: { error: jest.fn() } as any, + rawDefendInsights: '{"eventsContextCount": 5, "insights": ["insight1", "insight2"]}', + startTime: moment(), + telemetry: { reportEvent: jest.fn() } as any, + }; + await updateDefendInsights(params); + + expect(params.dataClient.getDefendInsight).toHaveBeenCalledTimes(1); + expect(params.dataClient.getDefendInsight).toHaveBeenCalledWith({ + id: params.defendInsightId, + authenticatedUser: params.authenticatedUser, + }); + expect(params.dataClient.updateDefendInsight).toHaveBeenCalledTimes(1); + expect(params.dataClient.updateDefendInsight).toHaveBeenCalledWith({ + defendInsightUpdateProps: { + eventsContextCount: 5, + insights: ['insight1', 'insight2'], + status: DefendInsightStatus.Enum.succeeded, + generationIntervals: expect.arrayContaining([ + expect.objectContaining({ + date: expect.any(String), + durationMs: expect.any(Number), + }), + ]), + id: params.defendInsightId, + replacements: params.latestReplacements, + backingIndex: 'backing-index-name', + }, + authenticatedUser: params.authenticatedUser, + }); + expect(params.telemetry.reportEvent).toHaveBeenCalledWith( + DEFEND_INSIGHT_SUCCESS_EVENT.eventType, + expect.any(Object) + ); + }); + + it('should handle error if rawDefendInsights is null', async () => { + const params = { + apiConfig: { + connectorId: 'connector-id1', + actionTypeId: 'action-type-id1', + model: 'model', + provider: OpenAiProviderType.OpenAi, + }, + defendInsightId: 'id', + authenticatedUser: {} as any, + dataClient: { + getDefendInsight: jest.fn().mockResolvedValueOnce({ + status: DefendInsightStatus.Enum.running, + backingIndex: 'index', + generationIntervals: [], + }), + updateDefendInsight: jest.fn(), + } as any, + latestReplacements: {}, + logger: { error: jest.fn() } as any, + rawDefendInsights: null, + startTime: moment(), + telemetry: { reportEvent: jest.fn() } as any, + }; + await updateDefendInsights(params); + + expect(params.logger.error).toHaveBeenCalledTimes(1); + expect(params.telemetry.reportEvent).toHaveBeenCalledTimes(1); + expect(params.telemetry.reportEvent).toHaveBeenCalledWith( + DEFEND_INSIGHT_ERROR_EVENT.eventType, + expect.any(Object) + ); + }); + }); + + describe('updateDefendInsightLastViewedAt', () => { + it('should update lastViewedAt time', async () => { + // ensure difference regardless of processing speed + const startTime = new Date().getTime() - 1; + const insightId = 'defend-insight-id1'; + const backingIndex = 'backing-index'; + const params = { + id: insightId, + authenticatedUser: {} as any, + dataClient: { + findDefendInsightsByParams: jest + .fn() + .mockResolvedValueOnce([{ id: insightId, backingIndex }]), + updateDefendInsights: jest.fn().mockResolvedValueOnce([{ id: insightId }]), + } as any, + }; + const result = await updateDefendInsightLastViewedAt(params); + + expect(params.dataClient.findDefendInsightsByParams).toHaveBeenCalledTimes(1); + expect(params.dataClient.findDefendInsightsByParams).toHaveBeenCalledWith({ + params: { ids: [insightId] }, + authenticatedUser: params.authenticatedUser, + }); + expect(params.dataClient.updateDefendInsights).toHaveBeenCalledTimes(1); + expect(params.dataClient.updateDefendInsights).toHaveBeenCalledWith({ + defendInsightsUpdateProps: [ + expect.objectContaining({ + id: insightId, + backingIndex, + }), + ], + authenticatedUser: params.authenticatedUser, + }); + expect( + new Date( + params.dataClient.updateDefendInsights.mock.calls[0][0].defendInsightsUpdateProps[0].lastViewedAt + ).getTime() + ).toBeGreaterThan(startTime); + expect(result).toEqual({ id: insightId }); + }); + + it('should return undefined if defend insight not found', async () => { + const insightId = 'defend-insight-id1'; + const params = { + id: insightId, + authenticatedUser: {} as any, + dataClient: { + findDefendInsightsByParams: jest.fn().mockResolvedValueOnce([]), + updateDefendInsight: jest.fn(), + } as any, + }; + const result = await updateDefendInsightLastViewedAt(params); + + expect(params.dataClient.findDefendInsightsByParams).toHaveBeenCalledTimes(1); + expect(params.dataClient.findDefendInsightsByParams).toHaveBeenCalledWith({ + params: { ids: [insightId] }, + authenticatedUser: params.authenticatedUser, + }); + expect(result).toBeUndefined(); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts new file mode 100644 index 000000000000..e67f00ef6514 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts @@ -0,0 +1,387 @@ +/* + * 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 moment, { Moment } from 'moment'; + +import type { + AnalyticsServiceSetup, + AuthenticatedUser, + KibanaRequest, + Logger, +} from '@kbn/core/server'; +import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; +import type { + ApiConfig, + DefendInsight, + DefendInsightGenerationInterval, + DefendInsightsPostRequestBody, + DefendInsightsResponse, + Replacements, +} from '@kbn/elastic-assistant-common'; +import type { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; +import type { ActionsClient } from '@kbn/actions-plugin/server'; + +import { ActionsClientLlm } from '@kbn/langchain/server'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { PublicMethodsOf } from '@kbn/utility-types'; +import { transformError } from '@kbn/securitysolution-es-utils'; +import { + DEFEND_INSIGHTS_TOOL_ID, + DefendInsightStatus, + DefendInsightType, + DefendInsightsGetRequestQuery, +} from '@kbn/elastic-assistant-common'; + +import type { GetRegisteredTools } from '../../services/app_context'; +import type { AssistantTool, ElasticAssistantApiRequestHandlerContext } from '../../types'; + +import { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; +import { + DEFEND_INSIGHT_ERROR_EVENT, + DEFEND_INSIGHT_SUCCESS_EVENT, +} from '../../lib/telemetry/event_based_telemetry'; +import { getLlmType } from '../utils'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; + +function getDataFromJSON(defendInsightStringified: string): { + eventsContextCount: number; + insights: DefendInsight[]; +} { + const { eventsContextCount, insights } = JSON.parse(defendInsightStringified); + return { eventsContextCount, insights }; +} + +function addGenerationInterval( + generationIntervals: DefendInsightGenerationInterval[], + generationInterval: DefendInsightGenerationInterval +): DefendInsightGenerationInterval[] { + const newGenerationIntervals = [generationInterval, ...generationIntervals]; + + const MAX_GENERATION_INTERVALS = 5; + if (newGenerationIntervals.length > MAX_GENERATION_INTERVALS) { + return newGenerationIntervals.slice(0, MAX_GENERATION_INTERVALS); // Return the first MAX_GENERATION_INTERVALS items + } + + return newGenerationIntervals; +} + +export function isDefendInsightsEnabled({ + request, + logger, + assistantContext, +}: { + request: KibanaRequest; + logger: Logger; + assistantContext: ElasticAssistantApiRequestHandlerContext; +}): boolean { + const pluginName = getPluginNameFromRequest({ + request, + logger, + defaultPluginName: DEFAULT_PLUGIN_NAME, + }); + + return assistantContext.getRegisteredFeatures(pluginName).defendInsights; +} + +export function getAssistantTool( + getRegisteredTools: GetRegisteredTools, + pluginName: string +): AssistantTool | undefined { + const assistantTools = getRegisteredTools(pluginName); + return assistantTools.find((tool) => tool.id === DEFEND_INSIGHTS_TOOL_ID); +} + +export function getAssistantToolParams({ + endpointIds, + insightType, + actionsClient, + anonymizationFields, + apiConfig, + esClient, + connectorTimeout, + langChainTimeout, + langSmithProject, + langSmithApiKey, + logger, + latestReplacements, + onNewReplacements, + request, +}: { + endpointIds: string[]; + insightType: DefendInsightType; + actionsClient: PublicMethodsOf; + anonymizationFields?: AnonymizationFieldResponse[]; + apiConfig: ApiConfig; + esClient: ElasticsearchClient; + connectorTimeout: number; + langChainTimeout: number; + langSmithProject?: string; + langSmithApiKey?: string; + logger: Logger; + latestReplacements: Replacements; + onNewReplacements: (newReplacements: Replacements) => void; + request: KibanaRequest; +}): { + endpointIds: string[]; + insightType: DefendInsightType; + anonymizationFields?: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + langChainTimeout: number; + llm: ActionsClientLlm; + logger: Logger; + replacements: Replacements; + onNewReplacements: (newReplacements: Replacements) => void; + request: KibanaRequest; + modelExists: boolean; + isEnabledKnowledgeBase: boolean; +} { + const traceOptions = { + projectName: langSmithProject, + tracers: [ + ...getLangSmithTracer({ + apiKey: langSmithApiKey, + projectName: langSmithProject, + logger, + }), + ], + }; + + const llm = new ActionsClientLlm({ + actionsClient, + connectorId: apiConfig.connectorId, + llmType: getLlmType(apiConfig.actionTypeId), + logger, + temperature: 0, // zero temperature because we want structured JSON output + timeout: connectorTimeout, + traceOptions, + }); + + return { + endpointIds, + insightType, + anonymizationFields, + esClient, + replacements: latestReplacements, + langChainTimeout, + llm, + logger, + onNewReplacements, + request, + modelExists: false, + isEnabledKnowledgeBase: false, + }; +} + +export async function handleToolError({ + apiConfig, + defendInsightId, + authenticatedUser, + dataClient, + err, + latestReplacements, + logger, + telemetry, +}: { + apiConfig: ApiConfig; + defendInsightId: string; + authenticatedUser: AuthenticatedUser; + dataClient: DefendInsightsDataClient; + err: Error; + latestReplacements: Replacements; + logger: Logger; + telemetry: AnalyticsServiceSetup; +}) { + try { + logger.error(err); + const error = transformError(err); + const currentInsight = await dataClient.getDefendInsight({ + id: defendInsightId, + authenticatedUser, + }); + + if (currentInsight === null || currentInsight?.status === DefendInsightStatus.Enum.canceled) { + return; + } + await dataClient.updateDefendInsight({ + defendInsightUpdateProps: { + insights: [], + status: DefendInsightStatus.Enum.failed, + id: defendInsightId, + replacements: latestReplacements, + backingIndex: currentInsight.backingIndex, + failureReason: error.message, + }, + authenticatedUser, + }); + telemetry.reportEvent(DEFEND_INSIGHT_ERROR_EVENT.eventType, { + actionTypeId: apiConfig.actionTypeId, + errorMessage: error.message, + model: apiConfig.model, + provider: apiConfig.provider, + }); + } catch (updateErr) { + const updateError = transformError(updateErr); + telemetry.reportEvent(DEFEND_INSIGHT_ERROR_EVENT.eventType, { + actionTypeId: apiConfig.actionTypeId, + errorMessage: updateError.message, + model: apiConfig.model, + provider: apiConfig.provider, + }); + } +} + +export async function createDefendInsight( + endpointIds: string[], + insightType: DefendInsightType, + dataClient: DefendInsightsDataClient, + authenticatedUser: AuthenticatedUser, + apiConfig: ApiConfig +): Promise<{ + currentInsight: DefendInsightsResponse; + defendInsightId: string; +}> { + const currentInsight = await dataClient?.createDefendInsight({ + defendInsightCreate: { + endpointIds, + insightType, + apiConfig, + insights: [], + status: DefendInsightStatus.Enum.running, + }, + authenticatedUser, + }); + + if (!currentInsight) { + throw new Error(`failed to create Defend insight for connectorId: ${apiConfig.connectorId}`); + } + + return { + defendInsightId: currentInsight.id, + currentInsight, + }; +} + +export async function updateDefendInsights({ + apiConfig, + defendInsightId, + authenticatedUser, + dataClient, + latestReplacements, + logger, + rawDefendInsights, + startTime, + telemetry, +}: { + apiConfig: ApiConfig; + defendInsightId: string; + authenticatedUser: AuthenticatedUser; + dataClient: DefendInsightsDataClient; + latestReplacements: Replacements; + logger: Logger; + rawDefendInsights: string | null; + startTime: Moment; + telemetry: AnalyticsServiceSetup; +}) { + try { + if (rawDefendInsights == null) { + throw new Error('tool returned no Defend insights'); + } + const currentInsight = await dataClient.getDefendInsight({ + id: defendInsightId, + authenticatedUser, + }); + if (currentInsight === null || currentInsight?.status === DefendInsightStatus.Enum.canceled) { + return; + } + const endTime = moment(); + const durationMs = endTime.diff(startTime); + const { eventsContextCount, insights } = getDataFromJSON(rawDefendInsights); + const updateProps = { + eventsContextCount, + insights, + status: DefendInsightStatus.Enum.succeeded, + ...(!eventsContextCount || !insights.length + ? {} + : { + generationIntervals: addGenerationInterval(currentInsight.generationIntervals, { + durationMs, + date: new Date().toISOString(), + }), + }), + id: defendInsightId, + replacements: latestReplacements, + backingIndex: currentInsight.backingIndex, + }; + + await dataClient.updateDefendInsight({ + defendInsightUpdateProps: updateProps, + authenticatedUser, + }); + telemetry.reportEvent(DEFEND_INSIGHT_SUCCESS_EVENT.eventType, { + actionTypeId: apiConfig.actionTypeId, + eventsContextCount: updateProps.eventsContextCount, + insightsGenerated: updateProps.insights.length, + durationMs, + model: apiConfig.model, + provider: apiConfig.provider, + }); + } catch (updateErr) { + logger.error(updateErr); + const updateError = transformError(updateErr); + telemetry.reportEvent(DEFEND_INSIGHT_ERROR_EVENT.eventType, { + actionTypeId: apiConfig.actionTypeId, + errorMessage: updateError.message, + model: apiConfig.model, + provider: apiConfig.provider, + }); + } +} + +export async function updateDefendInsightsLastViewedAt({ + params, + authenticatedUser, + dataClient, +}: { + params: DefendInsightsGetRequestQuery; + authenticatedUser: AuthenticatedUser; + dataClient: DefendInsightsDataClient; +}): Promise { + const defendInsights = await dataClient.findDefendInsightsByParams({ + params, + authenticatedUser, + }); + if (!defendInsights.length) { + return []; + } + + const defendInsightsUpdateProps = defendInsights.map((insight) => { + return { + id: insight.id, + lastViewedAt: new Date().toISOString(), + backingIndex: insight.backingIndex, + }; + }); + + return dataClient.updateDefendInsights({ + defendInsightsUpdateProps, + authenticatedUser, + }); +} + +export async function updateDefendInsightLastViewedAt({ + id, + authenticatedUser, + dataClient, +}: { + id: string; + authenticatedUser: AuthenticatedUser; + dataClient: DefendInsightsDataClient; +}): Promise { + return ( + await updateDefendInsightsLastViewedAt({ params: { ids: [id] }, authenticatedUser, dataClient }) + )[0]; +} diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/index.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/index.ts new file mode 100644 index 000000000000..a2835cb74c82 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/index.ts @@ -0,0 +1,10 @@ +/* + * 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. + */ + +export { getDefendInsightRoute } from './get_defend_insight'; +export { getDefendInsightsRoute } from './get_defend_insights'; +export { postDefendInsightsRoute } from './post_defend_insights'; diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts new file mode 100644 index 000000000000..95d6b521ed4b --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts @@ -0,0 +1,184 @@ +/* + * 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. + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import type { AuthenticatedUser } from '@kbn/core-security-common'; +import type { DefendInsightsPostRequestBody } from '@kbn/elastic-assistant-common'; + +import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; +import { actionsMock } from '@kbn/actions-plugin/server/mocks'; +import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/openai/constants'; +import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; + +import type { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; + +import { serverMock } from '../../__mocks__/server'; +import { + ElasticAssistantRequestHandlerContextMock, + requestContextMock, +} from '../../__mocks__/request_context'; +import { transformESSearchToDefendInsights } from '../../ai_assistant_data_clients/defend_insights/helpers'; +import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import { postDefendInsightsRequest } from '../../__mocks__/request'; +import { getAssistantTool, createDefendInsight, isDefendInsightsEnabled } from './helpers'; +import { postDefendInsightsRoute } from './post_defend_insights'; + +jest.mock('./helpers'); + +describe('postDefendInsightsRoute', () => { + let server: ReturnType; + let context: ElasticAssistantRequestHandlerContextMock; + let mockUser: AuthenticatedUser; + let mockDataClient: DefendInsightsDataClient; + let mockApiConfig: any; + let mockRequestBody: DefendInsightsPostRequestBody; + let mockCurrentInsight: any; + + function getDefaultUser(): AuthenticatedUser { + return { + username: 'my_username', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; + } + + function getDefaultDataClient(): DefendInsightsDataClient { + return { + findDefendInsightsByParams: jest.fn().mockResolvedValueOnce(mockCurrentInsight), + updateDefendInsight: jest.fn(), + createDefendInsight: jest.fn(), + } as unknown as DefendInsightsDataClient; + } + + function getDefaultApiConfig() { + return { + connectorId: 'connector-id', + actionTypeId: '.bedrock', + model: 'model', + provider: OpenAiProviderType.OpenAi, + }; + } + + function getDefaultRequestBody(): DefendInsightsPostRequestBody { + return { + endpointIds: [], + insightType: DefendInsightType.Enum.incompatible_antivirus, + subAction: 'invokeAI', + apiConfig: mockApiConfig, + anonymizationFields: [], + replacements: {}, + model: 'gpt-4', + langSmithProject: 'langSmithProject', + langSmithApiKey: 'langSmithApiKey', + }; + } + + beforeEach(() => { + const tools = requestContextMock.createTools(); + context = tools.context; + server = serverMock.create(); + tools.clients.core.elasticsearch.client = elasticsearchServiceMock.createScopedClusterClient(); + + mockCurrentInsight = transformESSearchToDefendInsights(getDefendInsightsSearchEsMock())[0]; + mockCurrentInsight.status = DefendInsightStatus.Enum.running; + + mockUser = getDefaultUser(); + mockDataClient = getDefaultDataClient(); + mockApiConfig = getDefaultApiConfig(); + mockRequestBody = getDefaultRequestBody(); + (getAssistantTool as jest.Mock).mockReturnValue({ getTool: jest.fn() }); + (createDefendInsight as jest.Mock).mockResolvedValue({ + currentInsight: mockCurrentInsight, + defendInsightId: mockCurrentInsight.id, + }); + (isDefendInsightsEnabled as jest.Mock).mockResolvedValue(true); + + context.elasticAssistant.getCurrentUser.mockReturnValue(mockUser); + context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValue(mockDataClient); + context.elasticAssistant.actions = actionsMock.createStart(); + + postDefendInsightsRoute(server.router); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should handle successful request', async () => { + const response = await server.inject( + postDefendInsightsRequest(mockRequestBody), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(200); + expect(response.body).toEqual(mockCurrentInsight); + }); + + it('should handle missing authenticated user', async () => { + context.elasticAssistant.getCurrentUser.mockReturnValueOnce(null); + const response = await server.inject( + postDefendInsightsRequest(mockRequestBody), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(401); + expect(response.body).toEqual({ + message: 'Authenticated user not found', + status_code: 401, + }); + }); + + it('should handle missing data client', async () => { + context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValueOnce(null); + const response = await server.inject( + postDefendInsightsRequest(mockRequestBody), + requestContextMock.convertContext(context) + ); + + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: 'Defend insights data client not initialized', + status_code: 500, + }); + }); + + it('should handle assistantTool null response', async () => { + (getAssistantTool as jest.Mock).mockReturnValueOnce(null); + const response = await server.inject( + postDefendInsightsRequest(mockRequestBody), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(404); + }); + + it('should 404 if feature flag disabled', async () => { + (isDefendInsightsEnabled as jest.Mock).mockReturnValueOnce(false); + const response = await server.inject( + postDefendInsightsRequest(mockRequestBody), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(404); + }); + + it('should handle createDefendInsight error', async () => { + (createDefendInsight as jest.Mock).mockRejectedValueOnce(new Error('Oh no!')); + const response = await server.inject( + postDefendInsightsRequest(mockRequestBody), + requestContextMock.convertContext(context) + ); + expect(response.status).toEqual(500); + expect(response.body).toEqual({ + message: { + error: 'Oh no!', + success: false, + }, + status_code: 500, + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts new file mode 100644 index 000000000000..d69b60a47880 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts @@ -0,0 +1,196 @@ +/* + * 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 moment from 'moment/moment'; + +import type { IKibanaResponse } from '@kbn/core/server'; + +import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; +import { + DEFEND_INSIGHTS, + DefendInsightsPostRequestBody, + DefendInsightsPostResponse, + ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, + Replacements, +} from '@kbn/elastic-assistant-common'; +import { transformError } from '@kbn/securitysolution-es-utils'; +import { IRouter, Logger } from '@kbn/core/server'; + +import { buildResponse } from '../../lib/build_response'; +import { ElasticAssistantRequestHandlerContext } from '../../types'; +import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; +import { + getAssistantTool, + getAssistantToolParams, + handleToolError, + createDefendInsight, + updateDefendInsights, + isDefendInsightsEnabled, +} from './helpers'; + +const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes +const LANG_CHAIN_TIMEOUT = ROUTE_HANDLER_TIMEOUT - 10_000; // 9 minutes 50 seconds +const CONNECTOR_TIMEOUT = LANG_CHAIN_TIMEOUT - 10_000; // 9 minutes 40 seconds + +export const postDefendInsightsRoute = (router: IRouter) => { + router.versioned + .post({ + access: 'internal', + path: DEFEND_INSIGHTS, + options: { + tags: ['access:elasticAssistant'], + timeout: { + idleSocket: ROUTE_HANDLER_TIMEOUT, + }, + }, + }) + .addVersion( + { + version: ELASTIC_AI_ASSISTANT_INTERNAL_API_VERSION, + validate: { + request: { + body: buildRouteValidationWithZod(DefendInsightsPostRequestBody), + }, + response: { + 200: { + body: { custom: buildRouteValidationWithZod(DefendInsightsPostResponse) }, + }, + }, + }, + }, + async (context, request, response): Promise> => { + const startTime = moment(); // start timing the generation + const resp = buildResponse(response); + const assistantContext = await context.elasticAssistant; + const logger: Logger = assistantContext.logger; + const telemetry = assistantContext.telemetry; + + try { + const isEnabled = isDefendInsightsEnabled({ + request, + logger, + assistantContext, + }); + if (!isEnabled) { + return response.notFound(); + } + + const actions = assistantContext.actions; + const actionsClient = await actions.getActionsClientWithRequest(request); + const dataClient = await assistantContext.getDefendInsightsDataClient(); + const authenticatedUser = assistantContext.getCurrentUser(); + if (authenticatedUser == null) { + return resp.error({ + body: `Authenticated user not found`, + statusCode: 401, + }); + } + if (!dataClient) { + return resp.error({ + body: `Defend insights data client not initialized`, + statusCode: 500, + }); + } + + const pluginName = getPluginNameFromRequest({ + request, + defaultPluginName: DEFAULT_PLUGIN_NAME, + logger, + }); + const assistantTool = getAssistantTool(assistantContext.getRegisteredTools, pluginName); + + if (!assistantTool) { + return response.notFound(); + } + + const { + endpointIds, + insightType, + apiConfig, + anonymizationFields, + langSmithApiKey, + langSmithProject, + replacements, + } = request.body; + + const esClient = (await context.core).elasticsearch.client.asCurrentUser; + + let latestReplacements: Replacements = { ...replacements }; + const onNewReplacements = (newReplacements: Replacements) => { + latestReplacements = { ...latestReplacements, ...newReplacements }; + }; + + const assistantToolParams = getAssistantToolParams({ + endpointIds, + insightType, + actionsClient, + anonymizationFields, + apiConfig, + esClient, + latestReplacements, + connectorTimeout: CONNECTOR_TIMEOUT, + langChainTimeout: LANG_CHAIN_TIMEOUT, + langSmithProject, + langSmithApiKey, + logger, + onNewReplacements, + request, + }); + + const toolInstance = assistantTool.getTool(assistantToolParams); + + const { currentInsight, defendInsightId } = await createDefendInsight( + endpointIds, + insightType, + dataClient, + authenticatedUser, + apiConfig + ); + + toolInstance + ?.invoke('') + .then((rawDefendInsights: string) => + updateDefendInsights({ + apiConfig, + defendInsightId, + authenticatedUser, + dataClient, + latestReplacements, + logger, + rawDefendInsights, + startTime, + telemetry, + }) + ) + .catch((err) => + handleToolError({ + apiConfig, + defendInsightId, + authenticatedUser, + dataClient, + err, + latestReplacements, + logger, + telemetry, + }) + ); + + return response.ok({ + body: currentInsight, + }); + } catch (err) { + logger.error(err); + const error = transformError(err); + + return resp.error({ + body: { success: false, error: error.message }, + statusCode: error.statusCode, + }); + } + } + ); +}; diff --git a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts index fcd051f1f215..23ec7011be5b 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/helpers.ts @@ -20,6 +20,7 @@ import { Message, Replacements, replaceAnonymizedValuesWithOriginalValues, + DEFEND_INSIGHTS_TOOL_ID, } from '@kbn/elastic-assistant-common'; import { ILicense } from '@kbn/licensing-plugin/server'; import { i18n } from '@kbn/i18n'; @@ -263,9 +264,11 @@ export const langChainExecute = async ({ logger, }); const assistantContext = context.elasticAssistant; + // We don't (yet) support invoking these tools interactively + const unsupportedTools = new Set(['attack-discovery', DEFEND_INSIGHTS_TOOL_ID]); const assistantTools = assistantContext .getRegisteredTools(pluginName) - .filter((x) => x.id !== 'attack-discovery'); // We don't (yet) support asking the assistant for NEW attack discoveries from a conversation + .filter((tool) => !unsupportedTools.has(tool.id)); // get a scoped esClient for assistant memory const esClient = context.core.elasticsearch.client.asCurrentUser; diff --git a/x-pack/plugins/elastic_assistant/server/routes/index.ts b/x-pack/plugins/elastic_assistant/server/routes/index.ts index c30a62872a82..ada5bf1c600d 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/index.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/index.ts @@ -12,6 +12,11 @@ export { postActionsConnectorExecuteRoute } from './post_actions_connector_execu export { postAttackDiscoveryRoute } from './attack_discovery/post/post_attack_discovery'; export { getAttackDiscoveryRoute } from './attack_discovery/get/get_attack_discovery'; +// Defend insights +export { postDefendInsightsRoute } from './defend_insights/post_defend_insights'; +export { getDefendInsightsRoute } from './defend_insights/get_defend_insights'; +export { getDefendInsightRoute } from './defend_insights/get_defend_insight'; + // Knowledge Base export { getKnowledgeBaseIndicesRoute } from './knowledge_base/get_knowledge_base_indices'; export { getKnowledgeBaseStatusRoute } from './knowledge_base/get_knowledge_base_status'; diff --git a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts index d722e31cb233..0124dfc7969c 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts @@ -33,6 +33,11 @@ import { postActionsConnectorExecuteRoute } from './post_actions_connector_execu import { bulkActionKnowledgeBaseEntriesRoute } from './knowledge_base/entries/bulk_actions_route'; import { createKnowledgeBaseEntryRoute } from './knowledge_base/entries/create_route'; import { findKnowledgeBaseEntriesRoute } from './knowledge_base/entries/find_route'; +import { + getDefendInsightRoute, + getDefendInsightsRoute, + postDefendInsightsRoute, +} from './defend_insights'; export const registerRoutes = ( router: ElasticAssistantPluginRouter, @@ -89,4 +94,9 @@ export const registerRoutes = ( getAttackDiscoveryRoute(router); postAttackDiscoveryRoute(router); cancelAttackDiscoveryRoute(router); + + // Defend insights + getDefendInsightRoute(router); + getDefendInsightsRoute(router); + postDefendInsightsRoute(router); }; diff --git a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts index 3f81763db49d..ef921d7c91a2 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/request_context_factory.ts @@ -116,6 +116,16 @@ export class RequestContextFactory implements IRequestContextFactory { }); }), + getDefendInsightsDataClient: memoize(() => { + const currentUser = getCurrentUser(); + return this.assistantService.createDefendInsightsDataClient({ + spaceId: getSpaceId(), + licensing: context.licensing, + logger: this.logger, + currentUser, + }); + }), + getAIAssistantPromptsDataClient: memoize(() => { const currentUser = getCurrentUser(); return this.assistantService.createAIAssistantPromptsDataClient({ diff --git a/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts index e91a0ec024c9..061e4e6f47af 100644 --- a/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts +++ b/x-pack/plugins/elastic_assistant/server/services/app_context.test.ts @@ -54,6 +54,7 @@ describe('AppContextService', () => { appContextService.start(mockAppContext); appContextService.registerFeatures('super', { assistantModelEvaluation: true, + defendInsights: true, }); appContextService.stop(); @@ -104,6 +105,7 @@ describe('AppContextService', () => { const features: AssistantFeatures = { ...defaultAssistantFeatures, assistantModelEvaluation: true, + defendInsights: true, }; appContextService.start(mockAppContext); @@ -119,11 +121,13 @@ describe('AppContextService', () => { const featuresOne: AssistantFeatures = { ...defaultAssistantFeatures, assistantModelEvaluation: true, + defendInsights: true, }; const pluginTwo = 'plugin2'; const featuresTwo: AssistantFeatures = { ...defaultAssistantFeatures, assistantModelEvaluation: false, + defendInsights: false, }; appContextService.start(mockAppContext); @@ -139,10 +143,12 @@ describe('AppContextService', () => { const featuresOne: AssistantFeatures = { ...defaultAssistantFeatures, assistantModelEvaluation: true, + defendInsights: true, }; const featuresTwo: AssistantFeatures = { ...defaultAssistantFeatures, assistantModelEvaluation: false, + defendInsights: false, }; appContextService.start(mockAppContext); @@ -164,6 +170,7 @@ describe('AppContextService', () => { const pluginName = 'pluginName'; const featuresSubset: Partial = { assistantModelEvaluation: true, + defendInsights: true, }; appContextService.start(mockAppContext); diff --git a/x-pack/plugins/elastic_assistant/server/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index d2dad4f9f998..d328001e86bb 100755 --- a/x-pack/plugins/elastic_assistant/server/types.ts +++ b/x-pack/plugins/elastic_assistant/server/types.ts @@ -27,6 +27,7 @@ import { TaskManagerSetupContract } from '@kbn/task-manager-plugin/server'; import { ElasticsearchClient } from '@kbn/core/server'; import { AttackDiscoveryPostRequestBody, + DefendInsightsPostRequestBody, AssistantFeatures, ExecuteConnectorRequestBody, Replacements, @@ -51,6 +52,7 @@ import { AIAssistantConversationsDataClient } from './ai_assistant_data_clients/ import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; import { AIAssistantDataClient } from './ai_assistant_data_clients'; import { AIAssistantKnowledgeBaseDataClient } from './ai_assistant_data_clients/knowledge_base'; +import type { DefendInsightsDataClient } from './ai_assistant_data_clients/defend_insights'; export const PLUGIN_ID = 'elasticAssistant' as const; @@ -129,6 +131,7 @@ export interface ElasticAssistantApiRequestHandlerContext { params?: GetAIAssistantKnowledgeBaseDataClientParams ) => Promise; getAttackDiscoveryDataClient: () => Promise; + getDefendInsightsDataClient: () => Promise; getAIAssistantPromptsDataClient: () => Promise; getAIAssistantAnonymizationFieldsDataClient: () => Promise; inference: InferenceServerStart; @@ -158,6 +161,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; indexTemplate: { conversations: string; @@ -165,6 +169,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; aliases: { conversations: string; @@ -172,6 +177,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; indexPatterns: { conversations: string; @@ -179,6 +185,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; pipelines: { knowledgeBase: string; @@ -230,7 +237,7 @@ export interface AssistantToolParams { request: KibanaRequest< unknown, unknown, - ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody + ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody | DefendInsightsPostRequestBody >; size?: number; telemetry?: AnalyticsServiceSetup; diff --git a/x-pack/plugins/security_solution/common/endpoint/constants.ts b/x-pack/plugins/security_solution/common/endpoint/constants.ts index 0ab749735d06..b53c7ae76154 100644 --- a/x-pack/plugins/security_solution/common/endpoint/constants.ts +++ b/x-pack/plugins/security_solution/common/endpoint/constants.ts @@ -18,6 +18,7 @@ export const ENDPOINT_ACTION_RESPONSES_INDEX = `${ENDPOINT_ACTION_RESPONSES_DS}- export const ENDPOINT_ACTION_RESPONSES_INDEX_PATTERN = `${ENDPOINT_ACTION_RESPONSES_DS}-*`; export const eventsIndexPattern = 'logs-endpoint.events.*'; +export const FILE_EVENTS_INDEX_PATTERN = 'logs-endpoint.events.file-*'; export const alertsIndexPattern = 'logs-endpoint.alerts-*'; // metadata datastream diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index dc6495e1d973..7fcdabad3b36 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -236,6 +236,11 @@ export const allowedExperimentalValues = Object.freeze({ * Enables the siem migrations feature */ siemMigrationsEnabled: false, + + /** + * Enables the Defend Insights feature + */ + defendInsights: false, }); type ExperimentalConfigKeys = Array; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts new file mode 100644 index 000000000000..03633d2ae1ee --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts @@ -0,0 +1,14 @@ +/* + * 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 { EndpointError } from '../../../../common/endpoint/errors'; + +export class InvalidDefendInsightTypeError extends EndpointError { + constructor() { + super('invalid defend insight type'); + } +} diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts new file mode 100644 index 000000000000..fa8f6fa1e33b --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts @@ -0,0 +1,49 @@ +/* + * 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 { SearchRequest } from '@elastic/elasticsearch/lib/api/types'; + +import { FILE_EVENTS_INDEX_PATTERN } from '../../../../../common/endpoint/constants'; + +const SIZE = 200; + +export function getFileEventsQuery({ endpointIds }: { endpointIds: string[] }): SearchRequest { + return { + allow_no_indices: true, + fields: ['_id', 'agent.id', 'process.executable'], + query: { + bool: { + must: [ + { + terms: { + 'agent.id': endpointIds, + }, + }, + { + range: { + '@timestamp': { + gte: 'now-24h', + lte: 'now', + }, + }, + }, + ], + }, + }, + size: SIZE, + sort: [ + { + '@timestamp': { + order: 'desc', + }, + }, + ], + _source: false, + ignore_unavailable: true, + index: [FILE_EVENTS_INDEX_PATTERN], + }; +} diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts new file mode 100644 index 000000000000..7c2fd9f61e25 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts @@ -0,0 +1,84 @@ +/* + * 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 { ElasticsearchClient } from '@kbn/core/server'; + +import { DefendInsightType, transformRawData } from '@kbn/elastic-assistant-common'; + +import { InvalidDefendInsightTypeError } from '../errors'; +import { getFileEventsQuery } from './get_file_events_query'; +import { getAnonymizedEvents } from '.'; + +jest.mock('@kbn/elastic-assistant-common', () => { + const originalModule = jest.requireActual('@kbn/elastic-assistant-common'); + return { + ...originalModule, + transformRawData: jest.fn(), + }; +}); + +jest.mock('./get_file_events_query', () => ({ + getFileEventsQuery: jest.fn(), +})); + +describe('getAnonymizedEvents', () => { + let mockEsClient: jest.Mocked; + + const mockHits = [ + { _index: 'test-index', fields: { field1: ['value1'] } }, + { _index: 'test-index', fields: { field2: ['value2'] } }, + ]; + + beforeEach(() => { + (getFileEventsQuery as jest.Mock).mockReturnValue({ index: 'test-index', body: {} }); + (transformRawData as jest.Mock).mockImplementation( + ({ rawData }) => `anonymized_${Object.values(rawData)[0]}` + ); + mockEsClient = { + search: jest.fn().mockResolvedValue({ + took: 1, + timed_out: false, + _shards: { + total: 1, + successful: 1, + skipped: 0, + failed: 0, + }, + hits: { + hits: mockHits, + }, + }), + } as unknown as jest.Mocked; + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should return anonymized events successfully', async () => { + const result = await getAnonymizedEvents({ + endpointIds: ['endpoint1'], + type: DefendInsightType.Enum.incompatible_antivirus, + esClient: mockEsClient, + }); + + expect(result).toEqual(['anonymized_value1', 'anonymized_value2']); + expect(getFileEventsQuery).toHaveBeenCalledWith({ endpointIds: ['endpoint1'] }); + expect(mockEsClient.search).toHaveBeenCalledWith({ index: 'test-index', body: {} }); + expect(transformRawData).toHaveBeenCalledTimes(2); + }); + + it('should throw InvalidDefendInsightTypeError for invalid type', async () => { + await expect( + getAnonymizedEvents({ + endpointIds: ['endpoint1'], + type: 'invalid_type' as DefendInsightType, + esClient: mockEsClient, + }) + ).rejects.toThrow(InvalidDefendInsightTypeError); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts new file mode 100644 index 000000000000..4d9fcaf89a34 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts @@ -0,0 +1,93 @@ +/* + * 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 { SearchRequest, SearchResponse } from '@elastic/elasticsearch/lib/api/types'; +import type { ElasticsearchClient } from '@kbn/core/server'; +import type { Replacements } from '@kbn/elastic-assistant-common'; +import type { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; + +import { + getAnonymizedValue, + transformRawData, + DefendInsightType, + getRawDataOrDefault, +} from '@kbn/elastic-assistant-common'; + +import { getFileEventsQuery } from './get_file_events_query'; +import { InvalidDefendInsightTypeError } from '../errors'; + +export async function getAnonymizedEvents({ + endpointIds, + type, + anonymizationFields, + esClient, + onNewReplacements, + replacements, +}: { + endpointIds: string[]; + type: DefendInsightType; + anonymizationFields?: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; +}): Promise { + const query = getQuery(type, { endpointIds }); + + return getAnonymized({ + query, + anonymizationFields, + esClient, + onNewReplacements, + replacements, + }); +} + +function getQuery(type: DefendInsightType, options: { endpointIds: string[] }): SearchRequest { + if (type === DefendInsightType.Enum.incompatible_antivirus) { + const { endpointIds } = options; + return getFileEventsQuery({ + endpointIds, + }); + } + + throw new InvalidDefendInsightTypeError(); +} + +const getAnonymized = async ({ + query, + anonymizationFields, + esClient, + onNewReplacements, + replacements, +}: { + query: SearchRequest; + anonymizationFields?: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; +}): Promise => { + const result = await esClient.search(query); + + // Accumulate replacements locally so we can, for example use the same + // replacement for a hostname when we see it in multiple alerts: + let localReplacements = { ...(replacements ?? {}) }; + const localOnNewReplacements = (newReplacements: Replacements) => { + localReplacements = { ...localReplacements, ...newReplacements }; + + onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements + }; + + return result.hits?.hits?.map((hit) => + transformRawData({ + anonymizationFields, + currentReplacements: localReplacements, // <-- the latest local replacements + getAnonymizedValue, + onNewReplacements: localOnNewReplacements, // <-- the local callback + rawData: getRawDataOrDefault(hit.fields), + }) + ); +}; diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts new file mode 100644 index 000000000000..5ef5aaeedf36 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts @@ -0,0 +1,69 @@ +/* + * 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 { DynamicTool } from '@langchain/core/tools'; + +import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; +import { DEFEND_INSIGHTS_TOOL_ID, DefendInsightType } from '@kbn/elastic-assistant-common'; +import { elasticsearchServiceMock } from '@kbn/core/server/mocks'; + +import type { DefendInsightsToolParams } from '.'; + +import { APP_UI_ID } from '../../../../common'; +import { DEFEND_INSIGHTS_TOOL, DEFEND_INSIGHTS_TOOL_DESCRIPTION } from '.'; + +jest.mock('@kbn/elastic-assistant-plugin/server/lib/langchain/helpers', () => ({ + requestHasRequiredAnonymizationParams: jest.fn(), +})); + +describe('DEFEND_INSIGHTS_TOOL', () => { + const mockLLM = {}; + const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); + const mockRequest = {}; + const mockParams: DefendInsightsToolParams = { + endpointIds: ['endpoint1'], + insightType: DefendInsightType.Enum.incompatible_antivirus, + anonymizationFields: [], + esClient: mockEsClient, + langChainTimeout: 1000, + llm: mockLLM, + onNewReplacements: jest.fn(), + replacements: {}, + request: mockRequest, + } as unknown as DefendInsightsToolParams; + + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should have correct properties', () => { + expect(DEFEND_INSIGHTS_TOOL.id).toBe(DEFEND_INSIGHTS_TOOL_ID); + expect(DEFEND_INSIGHTS_TOOL.name).toBe('defendInsightsTool'); + expect(DEFEND_INSIGHTS_TOOL.description).toBe(DEFEND_INSIGHTS_TOOL_DESCRIPTION); + expect(DEFEND_INSIGHTS_TOOL.sourceRegister).toBe(APP_UI_ID); + }); + + it('should return tool if supported', () => { + (requestHasRequiredAnonymizationParams as jest.Mock).mockReturnValue(true); + const tool = DEFEND_INSIGHTS_TOOL.getTool(mockParams); + expect(tool).toBeInstanceOf(DynamicTool); + }); + + it('should return null if not request missing anonymization params', () => { + (requestHasRequiredAnonymizationParams as jest.Mock).mockReturnValue(false); + const tool = DEFEND_INSIGHTS_TOOL.getTool(mockParams); + expect(tool).toBeNull(); + }); + + it('should return null if LLM is not provided', () => { + (requestHasRequiredAnonymizationParams as jest.Mock).mockReturnValue(true); + const paramsWithoutLLM = { ...mockParams, llm: undefined }; + const tool = DEFEND_INSIGHTS_TOOL.getTool(paramsWithoutLLM) as DynamicTool; + + expect(tool).toBeNull(); + }); +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.ts new file mode 100644 index 000000000000..1ea26b88a15c --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.ts @@ -0,0 +1,114 @@ +/* + * 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 { PromptTemplate } from '@langchain/core/prompts'; +import { DynamicTool } from '@langchain/core/tools'; +import { LLMChain } from 'langchain/chains'; +import { OutputFixingParser } from 'langchain/output_parsers'; + +import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; +import type { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; +import { DEFEND_INSIGHTS_TOOL_ID } from '@kbn/elastic-assistant-common'; + +import { APP_UI_ID } from '../../../../common'; +import { getAnonymizedEvents } from './get_events'; +import { getDefendInsightsOutputParser } from './output_parsers'; +import { getDefendInsightsPrompt } from './prompts'; + +export const DEFEND_INSIGHTS_TOOL_DESCRIPTION = 'Call this for Elastic Defend insights.'; + +export interface DefendInsightsToolParams extends AssistantToolParams { + endpointIds: string[]; + insightType: DefendInsightType; +} + +/** + * Returns a tool for generating Elastic Defend configuration insights + */ +export const DEFEND_INSIGHTS_TOOL: AssistantTool = Object.freeze({ + id: DEFEND_INSIGHTS_TOOL_ID, + name: 'defendInsightsTool', + description: DEFEND_INSIGHTS_TOOL_DESCRIPTION, + sourceRegister: APP_UI_ID, + + isSupported: (params: AssistantToolParams): boolean => { + const { llm, request } = params; + + return requestHasRequiredAnonymizationParams(request) && llm != null; + }, + + getTool(params: AssistantToolParams): DynamicTool | null { + if (!this.isSupported(params)) return null; + + const { + endpointIds, + insightType, + anonymizationFields, + esClient, + langChainTimeout, + llm, + onNewReplacements, + replacements, + } = params as DefendInsightsToolParams; + + return new DynamicTool({ + name: 'DefendInsightsTool', + description: DEFEND_INSIGHTS_TOOL_DESCRIPTION, + func: async () => { + if (llm == null) { + throw new Error('LLM is required for Defend Insights'); + } + + const anonymizedEvents = await getAnonymizedEvents({ + endpointIds, + type: insightType, + anonymizationFields, + esClient, + onNewReplacements, + replacements, + }); + + const eventsContextCount = anonymizedEvents.length; + if (eventsContextCount === 0) { + return JSON.stringify({ eventsContextCount, insights: [] }, null, 2); + } + + const outputParser = getDefendInsightsOutputParser({ type: insightType }); + const outputFixingParser = OutputFixingParser.fromLLM(llm, outputParser); + + const prompt = new PromptTemplate({ + template: `Answer the user's question as best you can:\n{format_instructions}\n{query}`, + inputVariables: ['query'], + partialVariables: { + format_instructions: outputFixingParser.getFormatInstructions(), + }, + }); + + const answerFormattingChain = new LLMChain({ + llm, + prompt, + outputKey: 'records', + outputParser: outputFixingParser, + }); + + const result = await answerFormattingChain.call({ + query: getDefendInsightsPrompt({ + type: insightType, + events: anonymizedEvents, + }), + timeout: langChainTimeout, + }); + const insights = result.records; + + return JSON.stringify({ eventsContextCount, insights }, null, 2); + }, + tags: [DEFEND_INSIGHTS_TOOL_ID], + }); + }, +}); diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts new file mode 100644 index 000000000000..b6430e440835 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts @@ -0,0 +1,28 @@ +/* + * 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 { StructuredOutputParser } from 'langchain/output_parsers'; + +import { z } from '@kbn/zod'; + +export function getIncompatibleVirusOutputParser() { + return StructuredOutputParser.fromZodSchema( + z.array( + z.object({ + group: z.string().describe('The program which is triggering the events'), + events: z + .object({ + id: z.string().describe('The event ID'), + endpointId: z.string().describe('The endpoint ID'), + value: z.string().describe('The process.executable value of the event'), + }) + .array() + .describe('The events that the insight is based on'), + }) + ) + ); +} diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts new file mode 100644 index 000000000000..78933b72702b --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts @@ -0,0 +1,19 @@ +/* + * 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 { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { InvalidDefendInsightTypeError } from '../errors'; +import { getIncompatibleVirusOutputParser } from './incompatible_antivirus'; + +export function getDefendInsightsOutputParser({ type }: { type: DefendInsightType }) { + if (type === DefendInsightType.Enum.incompatible_antivirus) { + return getIncompatibleVirusOutputParser(); + } + + throw new InvalidDefendInsightTypeError(); +} diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts new file mode 100644 index 000000000000..516de86a3097 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts @@ -0,0 +1,16 @@ +/* + * 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. + */ + +export function getIncompatibleAntivirusPrompt({ events }: { events: string[] }): string { + return `You are an Elastic Security user tasked with analyzing file events from Elastic Security to identify antivirus processes. Only focus on detecting antivirus processes. Ignore processes that belong to Elastic Agent or Elastic Defend, that are not antivirus processes, or are typical processes built into the operating system. Accuracy is of the utmost importance, try to minimize false positives. Group the processes by the antivirus program, keeping track of the agent.id and _id associated to each of the individual events as endpointId and eventId respectively. If there are no events, ignore the group field. Escape backslashes to respect JSON validation. New lines must always be escaped with double backslashes, i.e. \\\\n to ensure valid JSON. Only return JSON output, as described above. Do not add any additional text to describe your output. + + Use context from the following process events to provide insights: + """ + ${events.join('\n\n')} + """ + `; +} diff --git a/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts new file mode 100644 index 000000000000..d58778c3c544 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts @@ -0,0 +1,25 @@ +/* + * 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 { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { InvalidDefendInsightTypeError } from '../errors'; +import { getIncompatibleAntivirusPrompt } from './incompatible_antivirus'; + +export function getDefendInsightsPrompt({ + type, + events, +}: { + type: DefendInsightType; + events: string[]; +}): string { + if (type === DefendInsightType.Enum.incompatible_antivirus) { + return getIncompatibleAntivirusPrompt({ events }); + } + + throw new InvalidDefendInsightTypeError(); +} diff --git a/x-pack/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/plugins/security_solution/server/assistant/tools/index.ts index 9bb85f5beeda..f7824e688afe 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -10,12 +10,14 @@ import type { AssistantTool } from '@kbn/elastic-assistant-plugin/server'; import { NL_TO_ESQL_TOOL } from './esql/nl_to_esql_tool'; import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool'; +import { DEFEND_INSIGHTS_TOOL } from './defend_insights'; import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_retrieval_tool'; import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; import { SECURITY_LABS_KNOWLEDGE_BASE_TOOL } from './security_labs/security_labs_tool'; export const assistantTools: AssistantTool[] = [ ALERT_COUNTS_TOOL, + DEFEND_INSIGHTS_TOOL, NL_TO_ESQL_TOOL, KNOWLEDGE_BASE_RETRIEVAL_TOOL, KNOWLEDGE_BASE_WRITE_TOOL,