diff --git a/x-pack/packages/kbn-elastic-assistant-common/constants.ts b/x-pack/packages/kbn-elastic-assistant-common/constants.ts index d84d9d4cd6825..f7db791957435 100755 --- a/x-pack/packages/kbn-elastic-assistant-common/constants.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/constants.ts @@ -53,3 +53,7 @@ export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION = export const ELASTIC_AI_ASSISTANT_EVALUATE_URL = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/evaluate` as const; + +// 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/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 0000000000000..e070c3129e192 --- /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 0000000000000..5c27449c7d346 --- /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 0000000000000..fafaca8f48ead --- /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 0000000000000..2684bf53cf87b --- /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 0000000000000..0a2f3d618a869 --- /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 0000000000000..5d7e0b5358f81 --- /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 0000000000000..0518abdf6dcb7 --- /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 0000000000000..cc0ccfeea1980 --- /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 0000000000000..87c7cdbb81a8e --- /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 6304bfa4786cf..ff58ee4832225 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 ebef2dff8bdef..9b8007a9129bf 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 0000000000000..d25b8bb09b13d --- /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 9dc57bab25ef3..717747d1580d4 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, @@ -208,3 +214,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 d53ceaa586975..c4424e8410159 100644 --- a/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -27,6 +27,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'; export const createMockClients = () => { const core = coreMock.createRequestHandlerContext(); @@ -45,6 +46,7 @@ export const createMockClients = () => { getAIAssistantKnowledgeBaseDataClient: dataClientMock.create(), getAIAssistantPromptsDataClient: dataClientMock.create(), getAttackDiscoveryDataClient: attackDiscoveryDataClientMock.create(), + getDefendInsightsDataClient: dataClientMock.create(), getAIAssistantAnonymizationFieldsDataClient: dataClientMock.create(), getSpaceId: jest.fn(), getCurrentUser: jest.fn(), @@ -123,6 +125,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 0000000000000..5769ab4557102 --- /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 0000000000000..415487534a1b6 --- /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 0000000000000..4eeef2afd8738 --- /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 0000000000000..8e0793218154a --- /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 0000000000000..19a4cb66049db --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts @@ -0,0 +1,218 @@ +/* + * 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 { 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) { + 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 0000000000000..704ee9b962554 --- /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 0000000000000..b5cbbd6cd18a2 --- /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 0000000000000..f04c7ef505c2f --- /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 23a1a55564415..d187d8a926c3a 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 @@ -147,7 +147,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', @@ -155,6 +155,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); @@ -656,7 +657,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', @@ -666,6 +667,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); @@ -690,7 +692,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', @@ -699,6 +701,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); @@ -722,7 +725,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 () => { @@ -742,7 +745,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 15274f2323259..726eebc08ac33 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'; @@ -36,6 +37,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'; @@ -67,7 +69,8 @@ export type CreateDataStream = (params: { | 'conversations' | 'knowledgeBase' | 'prompts' - | 'attackDiscovery'; + | 'attackDiscovery' + | 'defendInsights'; fieldMap: FieldMap; kibanaVersion: string; spaceId?: string; @@ -82,6 +85,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; @@ -117,6 +121,11 @@ export class AIAssistantService { kibanaVersion: options.kibanaVersion, fieldMap: attackDiscoveryFieldMap, }); + this.defendInsightsDataStream = this.createDataStream({ + resource: 'defendInsights', + kibanaVersion: options.kibanaVersion, + fieldMap: defendInsightsFieldMap, + }); this.initPromise = this.initializeResources(); @@ -247,6 +256,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; @@ -265,6 +280,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'), @@ -272,6 +288,7 @@ export class AIAssistantService { prompts: getResourceName('prompts'), anonymizationFields: getResourceName('anonymization-fields'), attackDiscovery: getResourceName('attack-discovery'), + defendInsights: getResourceName('defend-insights'), }, indexPatterns: { conversations: getResourceName('conversations*'), @@ -279,6 +296,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'), @@ -286,6 +304,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'), @@ -428,6 +447,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 133419f45d175..9b2d444d643e4 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 5ff5ff894dffe..f7f751a826ed9 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 @@ -261,6 +261,100 @@ export const ATTACK_DISCOVERY_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, @@ -268,4 +362,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 0000000000000..6f29fb4913eb6 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts @@ -0,0 +1,139 @@ +/* + * 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 { 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); + }); + + 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 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 0000000000000..bfe4d7063d1a4 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.ts @@ -0,0 +1,87 @@ +/* + * 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 { 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 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 0000000000000..4b627d9d1671d --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts @@ -0,0 +1,139 @@ +/* + * 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 { 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); + }); + + 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 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 0000000000000..3826f71c00299 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.ts @@ -0,0 +1,87 @@ +/* + * 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, 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 { 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 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 0000000000000..cd47805fbe770 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts @@ -0,0 +1,251 @@ +/* + * 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 { 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' }]); + const result = getAssistantTool(getRegisteredTools, 'pluginName'); + expect(result).toEqual({ id: 'defend-insights' }); + }); + }); + + 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 0000000000000..e47b322c66308 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts @@ -0,0 +1,351 @@ +/* + * 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, + DefendInsightGenerationInterval, + DefendInsightsPostRequestBody, + DefendInsightsResponse, + ExecuteConnectorRequestBody, + 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 { + DefendInsightStatus, + DefendInsightType, + DefendInsightsGetRequestQuery, +} from '@kbn/elastic-assistant-common'; + +import type { GetRegisteredTools } from '../../services/app_context'; + +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'; + +const getDataFromJSON = (defendInsightStringified: string) => { + const { eventsContextCount, insights } = JSON.parse(defendInsightStringified); + return { eventsContextCount, insights }; +}; + +const 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 getAssistantTool(getRegisteredTools: GetRegisteredTools, pluginName: string) { + const assistantTools = getRegisteredTools(pluginName); + return assistantTools.find((tool) => tool.id === 'defend-insights'); +} + +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< + unknown, + unknown, + ExecuteConnectorRequestBody | DefendInsightsPostRequestBody + >; +}) { + 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 === 0 || insights === 0 + ? {} + : { + 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 const updateDefendInsightsLastViewedAt = async ({ + 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 0000000000000..a2835cb74c82d --- /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 0000000000000..57e76a2706dc9 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.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. + */ + +/* 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 } 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, + }); + + 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 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 0000000000000..1801d3dbfd190 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts @@ -0,0 +1,187 @@ +/* + * 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 { 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 { type IKibanaResponse, 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, +} 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 actions = (await context.elasticAssistant).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 { + 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 assistantTool = getAssistantTool( + (await context.elasticAssistant).getRegisteredTools, + pluginName + ); + + if (!assistantTool) { + return response.notFound(); + } + + 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/index.ts b/x-pack/plugins/elastic_assistant/server/routes/index.ts index a6d7a4298c2b7..a7b5e2b3ad975 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 { deleteKnowledgeBaseRoute } from './knowledge_base/delete_knowledge_base'; 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 7898629e15b5c..cdc04874103c6 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/register_routes.ts @@ -32,6 +32,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, @@ -87,4 +92,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 7d97029e7252a..be6aad2a44130 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 @@ -121,6 +121,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/types.ts b/x-pack/plugins/elastic_assistant/server/types.ts index e84b97ab43d7a..68570f1178f21 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; @@ -165,6 +168,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; indexTemplate: { conversations: string; @@ -172,6 +176,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; aliases: { conversations: string; @@ -179,6 +184,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; indexPatterns: { conversations: string; @@ -186,6 +192,7 @@ export interface AssistantResourceNames { prompts: string; anonymizationFields: string; attackDiscovery: string; + defendInsights: string; }; pipelines: { knowledgeBase: string; @@ -249,7 +256,7 @@ export interface AssistantToolParams { request: KibanaRequest< unknown, unknown, - ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody + ExecuteConnectorRequestBody | AttackDiscoveryPostRequestBody | DefendInsightsPostRequestBody >; size?: number; } diff --git a/x-pack/plugins/security_solution/common/endpoint/constants.ts b/x-pack/plugins/security_solution/common/endpoint/constants.ts index 19aa53eca6649..ba8a70d4c2c55 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 fileEventsIndexPattern = 'logs-endpoint.events.file-*'; export const alertsIndexPattern = 'logs-endpoint.alerts-*'; // metadata datastream 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 0000000000000..701808f1df6d7 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts @@ -0,0 +1,12 @@ +/* + * 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 class InvalidDefendInsightTypeError extends Error { + 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 0000000000000..a1f3674c4a74f --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts @@ -0,0 +1,47 @@ +/* + * 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 { fileEventsIndexPattern } from '../../../../../common/endpoint/constants'; + +const SIZE = 200; + +export const getFileEventsQuery = ({ endpointIds }: { endpointIds: string[] }) => ({ + allow_no_indices: true, + body: { + 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: [fileEventsIndexPattern], +}); 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 0000000000000..217233a2ef87e --- /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 0000000000000..793bc920edf97 --- /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[] }) { + 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 0000000000000..c34b0a37ae091 --- /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 { 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'); + 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 0000000000000..5ee4a55adf843 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/index.ts @@ -0,0 +1,113 @@ +/* + * 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 { 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 = { + id: 'defend-insights', + name: 'defendInsightsTool', + description: DEFEND_INSIGHTS_TOOL_DESCRIPTION, + sourceRegister: APP_UI_ID, + + isSupported: (params: AssistantToolParams) => { + const { llm, request } = params; + + return requestHasRequiredAnonymizationParams(request) && llm != null; + }, + + getTool(params: AssistantToolParams) { + 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'], + }); + }, +}; 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 0000000000000..db41494d927f7 --- /dev/null +++ b/x-pack/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts @@ -0,0 +1,27 @@ +/* + * 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 const getIncompatibleVirusOutputParser = () => + 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 0000000000000..1c3255de9abcb --- /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 const 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 0000000000000..aaaa5b1412b98 --- /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 const getIncompatibleAntivirusPrompt = ({ events }: { events: 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 0000000000000..1d92e0d530ecb --- /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 const getDefendInsightsPrompt = ({ + type, + events, +}: { + type: DefendInsightType; + events: 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 1b6e90eb7280f..da401c3fdb6ab 100644 --- a/x-pack/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/plugins/security_solution/server/assistant/tools/index.ts @@ -10,6 +10,7 @@ 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'; @@ -21,6 +22,7 @@ export const getAssistantTools = ({ }): AssistantTool[] => { const tools = [ ALERT_COUNTS_TOOL, + DEFEND_INSIGHTS_TOOL, NL_TO_ESQL_TOOL, KNOWLEDGE_BASE_RETRIEVAL_TOOL, KNOWLEDGE_BASE_WRITE_TOOL,