From 85e0a34eb3a1192ccc2bb118f47e5db17a8ffa07 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 24 Sep 2024 18:54:20 +0200 Subject: [PATCH 01/17] siem_migrations api scafolding --- .../common/api/siem_migrations/common.gen.ts | 38 +++++ .../api/siem_migrations/common.schema.yaml | 24 +++ .../common/api/siem_migrations/constants.ts | 8 + .../api/siem_migrations/splunk/constants.ts | 10 ++ .../siem_migrations/splunk/rules/constants.ts | 13 ++ .../splunk/rules/match_prebuilt_rule.gen.ts | 54 +++++++ .../rules/match_prebuilt_rule.schema.yaml | 65 +++++++++ .../splunk/rules/splunk_rule.gen.ts | 35 +++++ .../splunk/rules/splunk_rule.schema.yaml | 20 +++ .../server/lib/siem_migrations/routes.ts | 17 +++ .../lib/siem_migrations/splunk/rules/index.ts | 17 +++ .../rules/match_prebuilt_rule/agent/graph.ts | 82 +++++++++++ .../agent/match_prebuilt_rule.ts | 6 + .../api/match_prebuilt_rule.ts | 138 ++++++++++++++++++ .../server/lib/siem_migrations/util/llm.ts | 31 ++++ .../security_solution/server/routes/index.ts | 3 + 16 files changed, 561 insertions(+) create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/common.schema.yaml create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/constants.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/routes.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts new file mode 100644 index 0000000000000..cc1eb3e75c615 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts @@ -0,0 +1,38 @@ +/* + * 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 Migrations Attributes + * version: not applicable + */ + +import { z } from '@kbn/zod'; + +/** + * The GenAI connector id to use. + */ +export type ConnectorId = z.infer; +export const ConnectorId = z.string(); + +/** + * The LangSmith options object. + */ +export type LangSmithOptions = z.infer; +export const LangSmithOptions = z.object({ + /** + * The project name. + */ + projectName: z.string(), + /** + * The apiKey to use for tracing. + */ + apiKey: z.string(), +}); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/common.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/common.schema.yaml new file mode 100644 index 0000000000000..334a36eba91ba --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/common.schema.yaml @@ -0,0 +1,24 @@ +openapi: 3.0.3 +info: + title: Common SIEM Migrations Attributes + version: 'not applicable' +paths: {} +components: + x-codegen-enabled: true + schemas: + ConnectorId: + type: string + description: The GenAI connector id to use. + LangSmithOptions: + type: object + description: The LangSmith options object. + required: + - projectName + - apiKey + properties: + projectName: + type: string + description: The project name. + apiKey: + type: string + description: The apiKey to use for tracing. diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts new file mode 100644 index 0000000000000..2ffc1b8b42fb9 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts @@ -0,0 +1,8 @@ +/* + * 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 INTERNAL_SIEM_MIGRATIONS_PATH = '/internal/migrations' as const; diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/constants.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/constants.ts new file mode 100644 index 0000000000000..3cb4ff0083e4e --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/constants.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. + */ + +import { INTERNAL_SIEM_MIGRATIONS_PATH } from '../constants'; + +export const SPLUNK_MIGRATIONS_PATH = `${INTERNAL_SIEM_MIGRATIONS_PATH}/splunk` as const; diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts new file mode 100644 index 0000000000000..95d14208d673e --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts @@ -0,0 +1,13 @@ +/* + * 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 { SPLUNK_MIGRATIONS_PATH } from '../constants'; + +const SPLUNK_RULE_MIGRATIONS_PATH = `${SPLUNK_MIGRATIONS_PATH}/rules` as const; + +export const SPLUNK_MATCH_PREBUILT_RULE_PATH = + `${SPLUNK_RULE_MIGRATIONS_PATH}/match_prebuilt_rule` as const; diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts new file mode 100644 index 0000000000000..cd39d28603da0 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts @@ -0,0 +1,54 @@ +/* + * 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: Splunk Migration Match Prebuilt Rule API endpoint + * version: 1 + */ + +import { z } from '@kbn/zod'; + +import { Title, Description, Search } from './splunk_rule.gen'; +import { ConnectorId, LangSmithOptions } from '../../common.gen'; + +export type SplunkRuleMigrationMatchPrebuiltRuleRequestBody = z.infer< + typeof SplunkRuleMigrationMatchPrebuiltRuleRequestBody +>; +export const SplunkRuleMigrationMatchPrebuiltRuleRequestBody = z.object({ + rule: z.object({ + title: Title, + description: Description, + search: Search, + }), + connectorId: ConnectorId, + langSmithOptions: LangSmithOptions.optional(), +}); +export type SplunkRuleMigrationMatchPrebuiltRuleRequestBodyInput = z.input< + typeof SplunkRuleMigrationMatchPrebuiltRuleRequestBody +>; + +export type SplunkRuleMigrationMatchPrebuiltRuleResponse = z.infer< + typeof SplunkRuleMigrationMatchPrebuiltRuleResponse +>; +export const SplunkRuleMigrationMatchPrebuiltRuleResponse = z.object({ + /** + * The Elastic prebuilt rule id. + */ + id: z.string().min(1), + /** + * The Elastic prebuilt rule name. + */ + name: z.string().min(1), + /** + * Flag indicating if the rule is already installed. + */ + installed: z.boolean(), +}); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml new file mode 100644 index 0000000000000..f6ad8d7adb9b7 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml @@ -0,0 +1,65 @@ +openapi: 3.0.3 +info: + title: Splunk Migration Match Prebuilt Rule API endpoint + version: '1' +paths: + /internal/migrations/splunk/rules/match_prebuilt_rule: + post: + summary: Matches Splunk Security rules to Elastic prebuilt detection rules + operationId: SplunkRuleMigrationMatchPrebuiltRule + x-codegen-enabled: true + description: Perform Elastic prebuilt rule matching from Splunk Security rule + tags: + - Splunk Rule Migration API + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - rule + - connectorId + properties: + rule: + type: object + required: + - title + - description + - search + properties: + title: + $ref: './splunk_rule.schema.yaml#/components/schemas/Title' + description: + $ref: './splunk_rule.schema.yaml#/components/schemas/Description' + search: + $ref: './splunk_rule.schema.yaml#/components/schemas/Search' + connectorId: + $ref: '../../common.schema.yaml#/components/schemas/ConnectorId' + langSmithOptions: + $ref: '../../common.schema.yaml#/components/schemas/LangSmithOptions' + responses: + 200: + description: Indicates a successful match with a prebuilt rule. + content: + application/json: + schema: + type: object + required: + - id + - name + - installed + properties: + id: + type: string + minLength: 1 + description: The Elastic prebuilt rule id. + name: + type: string + minLength: 1 + description: The Elastic prebuilt rule name. + installed: + type: boolean + description: Flag indicating if the rule is already installed. + 204: + description: Indicates no prebuilt rule was matched. diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts new file mode 100644 index 0000000000000..d26ed7cd3a948 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts @@ -0,0 +1,35 @@ +/* + * 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 Splunk Rules Attributes + * version: not applicable + */ + +import { z } from '@kbn/zod'; + +/** + * The Splunk rule name. + */ +export type Title = z.infer; +export const Title = z.string().min(1); + +/** + * The Splunk rule search query. + */ +export type Search = z.infer; +export const Search = z.string().min(1); + +/** + * The Splunk rule description. + */ +export type Description = z.infer; +export const Description = z.string().min(1); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml new file mode 100644 index 0000000000000..54fd19ca6bd86 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml @@ -0,0 +1,20 @@ +openapi: 3.0.3 +info: + title: Common Splunk Rules Attributes + version: 'not applicable' +paths: {} +components: + x-codegen-enabled: true + schemas: + Title: + type: string + minLength: 1 + description: The Splunk rule name. + Search: + type: string + minLength: 1 + description: The Splunk rule search query. + Description: + type: string + minLength: 1 + description: The Splunk rule description. diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/routes.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/routes.ts new file mode 100644 index 0000000000000..f5e1be27e0e09 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/routes.ts @@ -0,0 +1,17 @@ +/* + * 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 { Logger } from '@kbn/core/server'; +import type { SecuritySolutionPluginRouter } from '../../types'; +import { registerSplunkMigrationRoutes } from './splunk/rules'; + +export const registerSiemMigrationsRoutes = ( + router: SecuritySolutionPluginRouter, + logger: Logger +) => { + registerSplunkMigrationRoutes(router, logger); +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts new file mode 100644 index 0000000000000..02a22222d443d --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts @@ -0,0 +1,17 @@ +/* + * 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 { Logger } from '@kbn/core/server'; +import type { SecuritySolutionPluginRouter } from '../../../../types'; +import { registerSplunkMatchPrebuiltRuleRoute } from './match_prebuilt_rule/api/match_prebuilt_rule'; + +export const registerSplunkMigrationRoutes = ( + router: SecuritySolutionPluginRouter, + logger: Logger +) => { + registerSplunkMatchPrebuiltRuleRoute(router, logger); +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts new file mode 100644 index 0000000000000..4b6fb08700bd5 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts @@ -0,0 +1,82 @@ +/* + * 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 { AIMessage } from '@langchain/core/messages'; +import type { RunnableConfig } from '@langchain/core/runnables'; +import { MessagesAnnotation, StateGraph } from '@langchain/langgraph'; +import { ToolNode } from '@langchain/langgraph/prebuilt'; + +import { ConfigurationSchema, ensureConfiguration } from './configuration.js'; +import { TOOLS } from './tools.js'; +import { loadChatModel } from './utils.js'; + +// Define the function that calls the model +async function callModel( + state: typeof MessagesAnnotation.State, + config: RunnableConfig +): Promise { + /** Call the LLM powering our agent. **/ + const configuration = ensureConfiguration(config); + + // Feel free to customize the prompt, model, and other logic! + const model = (await loadChatModel(configuration.model)).bindTools(TOOLS); + + const response = await model.invoke([ + { + role: 'system', + content: configuration.systemPromptTemplate.replace( + '{system_time}', + new Date().toISOString() + ), + }, + ...state.messages, + ]); + + // We return a list, because this will get added to the existing list + return { messages: [response] }; +} + +// Define the function that determines whether to continue or not +function routeModelOutput(state: typeof MessagesAnnotation.State): string { + const messages = state.messages; + const lastMessage = messages[messages.length - 1]; + // If the LLM is invoking tools, route there. + if ((lastMessage as AIMessage)?.tool_calls?.length || 0 > 0) { + return 'tools'; + } + // Otherwise end the graph. + else { + return '__end__'; + } +} + +// Define a new graph. We use the prebuilt MessagesAnnotation to define state: +// https://langchain-ai.github.io/langgraphjs/concepts/low_level/#messagesannotation +const workflow = new StateGraph(MessagesAnnotation, ConfigurationSchema) + // Define the two nodes we will cycle between + .addNode('callModel', callModel) + .addNode('tools', new ToolNode(TOOLS)) + // Set the entrypoint as `callModel` + // This means that this node is the first one called + .addEdge('__start__', 'callModel') + .addConditionalEdges( + // First, we define the edges' source node. We use `callModel`. + // This means these are the edges taken after the `callModel` node is called. + 'callModel', + // Next, we pass in the function that will determine the sink node(s), which + // will be called after the source node is called. + routeModelOutput + ) + // This means that after `tools` is called, `callModel` node is called next. + .addEdge('tools', 'callModel'); + +// Finally, we compile it! +// This compiles it into a graph you can invoke and deploy. +export const graph = workflow.compile({ + interruptBefore: [], // if you want to update the state before calling the tools + interruptAfter: [], +}); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts new file mode 100644 index 0000000000000..1fec1c76430eb --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts @@ -0,0 +1,6 @@ +/* + * 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. + */ diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts new file mode 100644 index 0000000000000..9ec79f443a8bd --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts @@ -0,0 +1,138 @@ +/* + * 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, Logger } from '@kbn/core/server'; +import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; +// import { getVersionBuckets } from '../../../../../detection_engine/prebuilt_rules/model/rule_versions/get_version_buckets'; +import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { APMTracer } from '@kbn/langchain/server/tracers/apm'; +import type { SplunkRuleMigrationMatchPrebuiltRuleResponse } from '../../../../../../../common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen'; +import { SplunkRuleMigrationMatchPrebuiltRuleRequestBody } from '../../../../../../../common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen'; +import { fetchRuleVersionsTriad } from '../../../../../detection_engine/prebuilt_rules/logic/rule_versions/fetch_rule_versions_triad'; +import { createPrebuiltRuleObjectsClient } from '../../../../../detection_engine/prebuilt_rules/logic/rule_objects/prebuilt_rule_objects_client'; +import { createPrebuiltRuleAssetsClient } from '../../../../../detection_engine/prebuilt_rules/logic/rule_assets/prebuilt_rule_assets_client'; +import type { SecuritySolutionPluginRouter } from '../../../../../../types'; +import { SPLUNK_MATCH_PREBUILT_RULE_PATH } from '../../../../../../../common/api/siem_migrations/splunk/rules/constants'; +import { getLLMType, getLLMClass } from '../../../../util/llm'; + +export const registerSplunkMatchPrebuiltRuleRoute = ( + router: SecuritySolutionPluginRouter, + logger: Logger +) => { + router.versioned + .post({ + path: SPLUNK_MATCH_PREBUILT_RULE_PATH, + access: 'internal', + options: { + tags: ['access:securitySolution'], + }, + }) + .addVersion( + { + version: '1', + validate: { + request: { + body: buildRouteValidationWithZod(SplunkRuleMigrationMatchPrebuiltRuleRequestBody), + }, + }, + }, + async ( + context, + request, + response + ): Promise> => { + const { langSmithOptions, rule: splunkRule, connectorId } = request.body; + try { + const ctx = await context.resolve(['core', 'actions', 'alerting']); + + const actionsClient = ctx.actions.getActionsClient(); + const connector = await actionsClient.get({ id: connectorId }); + + const abortSignal = getRequestAbortedSignal(request.events.aborted$); + + const soClient = ctx.core.savedObjects.client; + const rulesClient = ctx.alerting.getRulesClient(); + + const ruleAssetsClient = createPrebuiltRuleAssetsClient(soClient); + const ruleObjectsClient = createPrebuiltRuleObjectsClient(rulesClient); + + const prebuiltRulesMap = await fetchRuleVersionsTriad({ + ruleAssetsClient, + ruleObjectsClient, + }); + + // prebuiltRulesMap.forEach(({ current, target }) => { + // if (target != null) { + // // If this rule is available in the package + // totalAvailableRules.push(target); + // } + + // if (current != null) { + // // If this rule is installed + // currentRules.push(current); + // } + + // if (current == null && target != null) { + // // If this rule is not installed + // installableRules.push(target); + // } + + // if (current != null && target != null && current.version < target.version) { + // // If this rule is installed but outdated + // upgradeableRules.push({ + // current, + // target, + // }); + // } + // }); + + const actionTypeId = connector.actionTypeId; + const llmType = getLLMType(actionTypeId); + const llmClass = getLLMClass(llmType); + + const model = new llmClass({ + actionsClient, + connectorId: connector.id, + logger, + llmType, + model: connector.config?.defaultModel, + temperature: 0.05, + maxTokens: 16385, // 4096, + signal: abortSignal, + streaming: false, + }); + + const parameters = { + ruleTitle: splunkRule.title, + ruleDescription: splunkRule.description, + }; + + const options = { + callbacks: [ + new APMTracer({ projectName: langSmithOptions?.projectName ?? 'default' }, logger), + ...getLangSmithTracer({ ...langSmithOptions, logger }), + ], + }; + + const graph = await getMatchPrebuiltRuleGraph({ model }); + const results = await graph.invoke(parameters, options); + + const ruleMatched = false; + + if (!ruleMatched) { + return response.noContent(); + } + return response.ok({ body: { id: '1', name: 'dummy', installed: true } }); + } catch (err) { + return response.badRequest({ + body: err.message, + }); + } + } + ); +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts new file mode 100644 index 0000000000000..e380739c0b6cd --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts @@ -0,0 +1,31 @@ +/* + * 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 { + ActionsClientBedrockChatModel, + ActionsClientChatOpenAI, + ActionsClientGeminiChatModel, + ActionsClientSimpleChatModel, +} from '@kbn/langchain/server'; + +export const getLLMType = (actionTypeId: string): string | undefined => { + const llmTypeDictionary: Record = { + [`.gen-ai`]: `openai`, + [`.bedrock`]: `bedrock`, + [`.gemini`]: `gemini`, + }; + return llmTypeDictionary[actionTypeId]; +}; + +export const getLLMClass = (llmType?: string) => + llmType === 'openai' + ? ActionsClientChatOpenAI + : llmType === 'bedrock' + ? ActionsClientBedrockChatModel + : llmType === 'gemini' + ? ActionsClientGeminiChatModel + : ActionsClientSimpleChatModel; diff --git a/x-pack/plugins/security_solution/server/routes/index.ts b/x-pack/plugins/security_solution/server/routes/index.ts index 6f245bd04a02b..ab7c603aebc77 100644 --- a/x-pack/plugins/security_solution/server/routes/index.ts +++ b/x-pack/plugins/security_solution/server/routes/index.ts @@ -61,6 +61,7 @@ import { suggestUserProfilesRoute } from '../lib/detection_engine/routes/users/s import { registerTimelineRoutes } from '../lib/timeline/routes'; import { getFleetManagedIndexTemplatesRoute } from '../lib/security_integrations/cribl/routes'; import { registerEntityAnalyticsRoutes } from '../lib/entity_analytics/register_entity_analytics_routes'; +import { registerSiemMigrationsRoutes } from '../lib/siem_migrations/routes'; export const initRoutes = ( router: SecuritySolutionPluginRouter, @@ -145,6 +146,8 @@ export const initRoutes = ( } registerEntityAnalyticsRoutes({ router, config, getStartServices, logger }); + registerSiemMigrationsRoutes(router, logger); + // Security Integrations getFleetManagedIndexTemplatesRoute(router); }; From cbc2f67c608cc5463d3df8ed6704baf879b68907 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Fri, 27 Sep 2024 15:41:43 +0200 Subject: [PATCH 02/17] encapsulate chat model creation --- .../create_integration_assistant/state.ts | 6 +- .../common/api/siem_migrations/common.gen.ts | 2 +- .../common/api/siem_migrations/constants.ts | 2 +- .../splunk/rules/match_prebuilt_rule.gen.ts | 18 +-- .../rules/match_prebuilt_rule.schema.yaml | 35 ++---- .../splunk/rules/splunk_rule.gen.ts | 36 +++--- .../splunk/rules/splunk_rule.schema.yaml | 36 ++++-- .../siem_migrations/actions_client_chat.ts | 87 ++++++++++++++ .../rules/match_prebuilt_rule/agent/graph.ts | 112 +++++++----------- .../match_prebuilt_rule/agent/prompts.ts | 37 ++++++ .../rules/match_prebuilt_rule/agent/state.ts | 23 ++++ .../rules/match_prebuilt_rule/agent/types.ts | 20 ++++ .../api/match_prebuilt_rule.ts | 91 ++++---------- .../api/util/prebuilt_rules.ts | 68 +++++++++++ .../match_prebuilt_rule.ts => types.ts} | 9 ++ .../server/lib/siem_migrations/util/llm.ts | 31 ----- .../public/navigation/side_navigation.ts | 11 ++ 17 files changed, 381 insertions(+), 243 deletions(-) create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/prompts.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/types.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts rename x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/{agent/match_prebuilt_rule.ts => types.ts} (50%) delete mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts diff --git a/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts b/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts index bef5b35624df4..99f95a3eee039 100644 --- a/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts +++ b/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts @@ -51,7 +51,11 @@ export const reducer = (state: State, action: Action): State => { case 'SET_IS_GENERATING': return { ...state, isGenerating: action.payload }; case 'SET_GENERATED_RESULT': - return { ...state, result: action.payload }; + return { + ...state, + // keep original result as the samplesFormat is not always included in the payload + result: state.result ? { ...state.result, ...action.payload } : action.payload, + }; default: return state; } diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts index cc1eb3e75c615..db43041787722 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/common.gen.ts @@ -10,7 +10,7 @@ * This file is automatically generated by the OpenAPI Generator, @kbn/openapi-generator. * * info: - * title: Common Migrations Attributes + * title: Common SIEM Migrations Attributes * version: not applicable */ diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts index 2ffc1b8b42fb9..8dd2a22b859e5 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/constants.ts @@ -5,4 +5,4 @@ * 2.0. */ -export const INTERNAL_SIEM_MIGRATIONS_PATH = '/internal/migrations' as const; +export const INTERNAL_SIEM_MIGRATIONS_PATH = '/internal/siem_migrations' as const; diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts index cd39d28603da0..4e245673bd17d 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen.ts @@ -16,18 +16,14 @@ import { z } from '@kbn/zod'; -import { Title, Description, Search } from './splunk_rule.gen'; +import { SplunkRule } from './splunk_rule.gen'; import { ConnectorId, LangSmithOptions } from '../../common.gen'; export type SplunkRuleMigrationMatchPrebuiltRuleRequestBody = z.infer< typeof SplunkRuleMigrationMatchPrebuiltRuleRequestBody >; export const SplunkRuleMigrationMatchPrebuiltRuleRequestBody = z.object({ - rule: z.object({ - title: Title, - description: Description, - search: Search, - }), + splunkRule: SplunkRule, connectorId: ConnectorId, langSmithOptions: LangSmithOptions.optional(), }); @@ -40,15 +36,11 @@ export type SplunkRuleMigrationMatchPrebuiltRuleResponse = z.infer< >; export const SplunkRuleMigrationMatchPrebuiltRuleResponse = z.object({ /** - * The Elastic prebuilt rule id. + * The Elastic prebuilt rule information. */ - id: z.string().min(1), - /** - * The Elastic prebuilt rule name. - */ - name: z.string().min(1), + rule: z.object({}), /** * Flag indicating if the rule is already installed. */ - installed: z.boolean(), + isInstalled: z.boolean(), }); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml index f6ad8d7adb9b7..73325e8d13299 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.schema.yaml @@ -18,22 +18,11 @@ paths: schema: type: object required: - - rule + - splunkRule - connectorId properties: - rule: - type: object - required: - - title - - description - - search - properties: - title: - $ref: './splunk_rule.schema.yaml#/components/schemas/Title' - description: - $ref: './splunk_rule.schema.yaml#/components/schemas/Description' - search: - $ref: './splunk_rule.schema.yaml#/components/schemas/Search' + splunkRule: + $ref: './splunk_rule.schema.yaml#/components/schemas/SplunkRule' connectorId: $ref: '../../common.schema.yaml#/components/schemas/ConnectorId' langSmithOptions: @@ -46,19 +35,13 @@ paths: schema: type: object required: - - id - - name - - installed + - rule + - isInstalled properties: - id: - type: string - minLength: 1 - description: The Elastic prebuilt rule id. - name: - type: string - minLength: 1 - description: The Elastic prebuilt rule name. - installed: + rule: + type: object + description: The Elastic prebuilt rule information. + isInstalled: type: boolean description: Flag indicating if the rule is already installed. 204: diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts index d26ed7cd3a948..f0f2bc4ab520f 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts @@ -16,20 +16,22 @@ import { z } from '@kbn/zod'; -/** - * The Splunk rule name. - */ -export type Title = z.infer; -export const Title = z.string().min(1); - -/** - * The Splunk rule search query. - */ -export type Search = z.infer; -export const Search = z.string().min(1); - -/** - * The Splunk rule description. - */ -export type Description = z.infer; -export const Description = z.string().min(1); +export type SplunkRule = z.infer; +export const SplunkRule = z.object({ + /** + * The Splunk rule name. + */ + title: z.string().min(1), + /** + * The Splunk rule search query. + */ + search: z.string().min(1), + /** + * The Splunk rule description. + */ + description: z.string().min(1), + /** + * String array containing the rule Mitre Attack technique IDs. + */ + mitreAttackIds: z.array(z.string()).optional(), +}); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml index 54fd19ca6bd86..92d7b049f6179 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml @@ -6,15 +6,27 @@ paths: {} components: x-codegen-enabled: true schemas: - Title: - type: string - minLength: 1 - description: The Splunk rule name. - Search: - type: string - minLength: 1 - description: The Splunk rule search query. - Description: - type: string - minLength: 1 - description: The Splunk rule description. + SplunkRule: + type: object + required: + - title + - description + - search + properties: + title: + type: string + minLength: 1 + description: The Splunk rule name. + search: + type: string + minLength: 1 + description: The Splunk rule search query. + description: + type: string + minLength: 1 + description: The Splunk rule description. + mitreAttackIds: + type: array + items: + type: string + description: String array containing the rule Mitre Attack technique IDs. diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts new file mode 100644 index 0000000000000..94d98ee8078b1 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.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 { + ActionsClientBedrockChatModel, + ActionsClientChatOpenAI, + ActionsClientGeminiChatModel, + ActionsClientSimpleChatModel, +} from '@kbn/langchain/server'; +import type { Logger } from '@kbn/core/server'; +import type { ActionsClient } from '@kbn/actions-plugin/server'; +import type { ActionsClientChatOpenAIParams } from '@kbn/langchain/server/language_models/chat_openai'; +import type { CustomChatModelInput as ActionsClientBedrockChatModelParams } from '@kbn/langchain/server/language_models/bedrock_chat'; +import type { CustomChatModelInput as ActionsClientGeminiChatModelParams } from '@kbn/langchain/server/language_models/gemini_chat'; +import type { CustomChatModelInput as ActionsClientSimpleChatModelParams } from '@kbn/langchain/server/language_models/simple_chat_model'; + +export type ChatModel = + | ActionsClientSimpleChatModel + | ActionsClientChatOpenAI + | ActionsClientBedrockChatModel + | ActionsClientGeminiChatModel; + +export type ActionsClientChatModelClass = + | typeof ActionsClientSimpleChatModel + | typeof ActionsClientChatOpenAI + | typeof ActionsClientBedrockChatModel + | typeof ActionsClientGeminiChatModel; + +export type ChatModelParams = Partial & + Partial & + Partial & + Partial & { + /** Enables the streaming mode of the response, disabled by default */ + streaming?: boolean; + }; + +export class ActionsClientChat { + constructor( + private readonly connectorId: string, + private readonly actionsClient: ActionsClient, + private readonly logger: Logger + ) {} + + public async createModel(params?: ChatModelParams): Promise { + const connector = await this.actionsClient.get({ id: this.connectorId }); + if (!connector) { + throw new Error(`Connector not found: ${this.connectorId}`); + } + + const llmType = this.getLLMType(connector.actionTypeId); + const ChatModelClass = this.getLLMClass(llmType); + + const model = new ChatModelClass({ + actionsClient: this.actionsClient, + connectorId: this.connectorId, + logger: this.logger, + llmType, + model: connector.config?.defaultModel, + ...params, + streaming: params?.streaming ?? false, // disabling streaming by default, for some reason is enabled when omitted + }); + return model; + } + + private getLLMType(actionTypeId: string): string | undefined { + const llmTypeDictionary: Record = { + [`.gen-ai`]: `openai`, + [`.bedrock`]: `bedrock`, + [`.gemini`]: `gemini`, + }; + return llmTypeDictionary[actionTypeId]; + } + + private getLLMClass(llmType?: string): ActionsClientChatModelClass { + return llmType === 'openai' + ? ActionsClientChatOpenAI + : llmType === 'bedrock' + ? ActionsClientBedrockChatModel + : llmType === 'gemini' + ? ActionsClientGeminiChatModel + : ActionsClientSimpleChatModel; + } +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts index 4b6fb08700bd5..6e705e8c9ec28 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts @@ -5,78 +5,46 @@ * 2.0. */ -import type { AIMessage } from '@langchain/core/messages'; -import type { RunnableConfig } from '@langchain/core/runnables'; -import { MessagesAnnotation, StateGraph } from '@langchain/langgraph'; -import { ToolNode } from '@langchain/langgraph/prebuilt'; - -import { ConfigurationSchema, ensureConfiguration } from './configuration.js'; -import { TOOLS } from './tools.js'; -import { loadChatModel } from './utils.js'; - -// Define the function that calls the model -async function callModel( - state: typeof MessagesAnnotation.State, - config: RunnableConfig -): Promise { - /** Call the LLM powering our agent. **/ - const configuration = ensureConfiguration(config); - - // Feel free to customize the prompt, model, and other logic! - const model = (await loadChatModel(configuration.model)).bindTools(TOOLS); - - const response = await model.invoke([ - { - role: 'system', - content: configuration.systemPromptTemplate.replace( - '{system_time}', - new Date().toISOString() - ), - }, - ...state.messages, - ]); - - // We return a list, because this will get added to the existing list - return { messages: [response] }; -} - -// Define the function that determines whether to continue or not -function routeModelOutput(state: typeof MessagesAnnotation.State): string { - const messages = state.messages; - const lastMessage = messages[messages.length - 1]; - // If the LLM is invoking tools, route there. - if ((lastMessage as AIMessage)?.tool_calls?.length || 0 > 0) { - return 'tools'; +import { END, START, StateGraph } from '@langchain/langgraph'; +import { StringOutputParser } from '@langchain/core/output_parsers'; +import { matchPrebuiltRuleState } from './state'; +import type { + MatchPrebuiltRuleGraphParams, + MatchPrebuiltRuleNodeParams, + MatchPrebuiltRuleState, +} from './types'; +import { MATCH_PREBUILT_RULE_PROMPT } from './prompts'; + +const callModel = async ({ + state, + model, +}: MatchPrebuiltRuleNodeParams): Promise> => { + const outputParser = new StringOutputParser(); + const matchPrebuiltRule = MATCH_PREBUILT_RULE_PROMPT.pipe(model).pipe(outputParser); + + const elasticSecurityRules = Array(state.prebuiltRulesMap.keys()).join('\n'); + const result = await matchPrebuiltRule.invoke({ + elasticSecurityRules, + splunkRuleTitle: state.splunkRuleTitle, + splunkRuleDescription: state.splunkRuleDescription, + }); + const response = result.trim(); + if (response === 'no_match') { + return { matched: false }; } - // Otherwise end the graph. - else { - return '__end__'; + if (state.prebuiltRulesMap.get(response)) { + return { matched: true, rule: state.prebuiltRulesMap.get(response), response }; } + return { matched: false, response }; +}; + +export async function getMatchPrebuiltRuleGraph({ model }: MatchPrebuiltRuleGraphParams) { + const matchPrebuiltRuleGraph = new StateGraph(matchPrebuiltRuleState) + .addNode('callModel', (state: MatchPrebuiltRuleState) => callModel({ state, model })) + // .addNode('processModelResponse', processModelResponse) + .addEdge(START, 'callModel') + // .addEdge('callModel', 'processModelResponse') + .addEdge('callModel', END); + + return matchPrebuiltRuleGraph.compile(); } - -// Define a new graph. We use the prebuilt MessagesAnnotation to define state: -// https://langchain-ai.github.io/langgraphjs/concepts/low_level/#messagesannotation -const workflow = new StateGraph(MessagesAnnotation, ConfigurationSchema) - // Define the two nodes we will cycle between - .addNode('callModel', callModel) - .addNode('tools', new ToolNode(TOOLS)) - // Set the entrypoint as `callModel` - // This means that this node is the first one called - .addEdge('__start__', 'callModel') - .addConditionalEdges( - // First, we define the edges' source node. We use `callModel`. - // This means these are the edges taken after the `callModel` node is called. - 'callModel', - // Next, we pass in the function that will determine the sink node(s), which - // will be called after the source node is called. - routeModelOutput - ) - // This means that after `tools` is called, `callModel` node is called next. - .addEdge('tools', 'callModel'); - -// Finally, we compile it! -// This compiles it into a graph you can invoke and deploy. -export const graph = workflow.compile({ - interruptBefore: [], // if you want to update the state before calling the tools - interruptAfter: [], -}); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/prompts.ts new file mode 100644 index 0000000000000..4cae08c1da6ff --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/prompts.ts @@ -0,0 +1,37 @@ +/* + * 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 { ChatPromptTemplate } from '@langchain/core/prompts'; +export const MATCH_PREBUILT_RULE_PROMPT = ChatPromptTemplate.fromMessages([ + [ + 'system', + `You are an expert assistant in Cybersecurity, your task is to help migrating a SIEM detection rule, from Splunk Security to Elastic Security. +You will be provided with a Splunk Detection Rule name by the user, your goal is to try find an Elastic Detection Rule that covers the same threat, if any. +The list of Elastic Detection Rules suggested is provided in the context below. + + +If there is no Elastic rule in the list that covers the same threat, answer only with the string: no_match +If there is one Elastic rule in the list that covers the same threat, answer only with its name without any further explanation. +If there are multiple rules in the list that cover the same threat, answer with the most specific of them, for example: "Linux User Account Creation" is more specific than "User Account Creation". + + + + +{elasticSecurityRules} + +`, + ], + [ + 'human', + `The Splunk Detection Rule is: + +{splunkRuleTitle} + +`, + ], + ['ai', 'Please find the answer below:'], +]); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts new file mode 100644 index 0000000000000..dca4708a0ace4 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts @@ -0,0 +1,23 @@ +/* + * 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 { BaseMessage } from '@langchain/core/messages'; +import { Annotation, messagesStateReducer } from '@langchain/langgraph'; +import type { PrebuiltRuleMapped, PrebuiltRulesMapByName } from '../types'; + +export const matchPrebuiltRuleState = Annotation.Root({ + messages: Annotation({ + reducer: messagesStateReducer, + default: () => [], + }), + splunkRuleTitle: Annotation(), + splunkRuleDescription: Annotation, + matched: Annotation, + response: Annotation, + rule: Annotation, + prebuiltRulesMap: Annotation, +}); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/types.ts new file mode 100644 index 0000000000000..40e8ca6f9a0fd --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/types.ts @@ -0,0 +1,20 @@ +/* + * 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 { ChatModel } from '../../../../actions_client_chat'; +import type { matchPrebuiltRuleState } from './state'; + +export type MatchPrebuiltRuleState = typeof matchPrebuiltRuleState.State; + +export interface MatchPrebuiltRuleGraphParams { + model: ChatModel; +} + +export interface MatchPrebuiltRuleNodeParams { + model: ChatModel; + state: MatchPrebuiltRuleState; +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts index 9ec79f443a8bd..df07171ce123c 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts @@ -7,18 +7,17 @@ import type { IKibanaResponse, Logger } from '@kbn/core/server'; import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; -// import { getVersionBuckets } from '../../../../../detection_engine/prebuilt_rules/model/rule_versions/get_version_buckets'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import { APMTracer } from '@kbn/langchain/server/tracers/apm'; -import type { SplunkRuleMigrationMatchPrebuiltRuleResponse } from '../../../../../../../common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen'; import { SplunkRuleMigrationMatchPrebuiltRuleRequestBody } from '../../../../../../../common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen'; -import { fetchRuleVersionsTriad } from '../../../../../detection_engine/prebuilt_rules/logic/rule_versions/fetch_rule_versions_triad'; -import { createPrebuiltRuleObjectsClient } from '../../../../../detection_engine/prebuilt_rules/logic/rule_objects/prebuilt_rule_objects_client'; -import { createPrebuiltRuleAssetsClient } from '../../../../../detection_engine/prebuilt_rules/logic/rule_assets/prebuilt_rule_assets_client'; import type { SecuritySolutionPluginRouter } from '../../../../../../types'; import { SPLUNK_MATCH_PREBUILT_RULE_PATH } from '../../../../../../../common/api/siem_migrations/splunk/rules/constants'; -import { getLLMType, getLLMClass } from '../../../../util/llm'; +import { retrievePrebuiltRulesMap } from './util/prebuilt_rules'; +import { getMatchPrebuiltRuleGraph } from '../agent/graph'; +import type { MatchPrebuiltRuleState } from '../agent/types'; +import type { PrebuiltRuleMapped } from '../types'; +import { ActionsClientChat } from '../../../../actions_client_chat'; export const registerSplunkMatchPrebuiltRuleRoute = ( router: SecuritySolutionPluginRouter, @@ -41,75 +40,31 @@ export const registerSplunkMatchPrebuiltRuleRoute = ( }, }, }, - async ( - context, - request, - response - ): Promise> => { - const { langSmithOptions, rule: splunkRule, connectorId } = request.body; + async (context, request, response): Promise> => { + const { langSmithOptions, splunkRule, connectorId } = request.body; try { const ctx = await context.resolve(['core', 'actions', 'alerting']); const actionsClient = ctx.actions.getActionsClient(); - const connector = await actionsClient.get({ id: connectorId }); - - const abortSignal = getRequestAbortedSignal(request.events.aborted$); - const soClient = ctx.core.savedObjects.client; const rulesClient = ctx.alerting.getRulesClient(); - const ruleAssetsClient = createPrebuiltRuleAssetsClient(soClient); - const ruleObjectsClient = createPrebuiltRuleObjectsClient(rulesClient); - - const prebuiltRulesMap = await fetchRuleVersionsTriad({ - ruleAssetsClient, - ruleObjectsClient, + const prebuiltRulesMap = await retrievePrebuiltRulesMap({ + soClient, + rulesClient, + splunkRule, }); - // prebuiltRulesMap.forEach(({ current, target }) => { - // if (target != null) { - // // If this rule is available in the package - // totalAvailableRules.push(target); - // } - - // if (current != null) { - // // If this rule is installed - // currentRules.push(current); - // } - - // if (current == null && target != null) { - // // If this rule is not installed - // installableRules.push(target); - // } - - // if (current != null && target != null && current.version < target.version) { - // // If this rule is installed but outdated - // upgradeableRules.push({ - // current, - // target, - // }); - // } - // }); - - const actionTypeId = connector.actionTypeId; - const llmType = getLLMType(actionTypeId); - const llmClass = getLLMClass(llmType); - - const model = new llmClass({ - actionsClient, - connectorId: connector.id, - logger, - llmType, - model: connector.config?.defaultModel, + const actionsClientChat = new ActionsClientChat(connectorId, actionsClient, logger); + const model = await actionsClientChat.createModel({ + signal: getRequestAbortedSignal(request.events.aborted$), temperature: 0.05, - maxTokens: 16385, // 4096, - signal: abortSignal, - streaming: false, }); - const parameters = { - ruleTitle: splunkRule.title, - ruleDescription: splunkRule.description, + const parameters: Partial = { + splunkRuleTitle: splunkRule.title, + splunkRuleDescription: splunkRule.description, + prebuiltRulesMap, }; const options = { @@ -118,16 +73,14 @@ export const registerSplunkMatchPrebuiltRuleRoute = ( ...getLangSmithTracer({ ...langSmithOptions, logger }), ], }; - const graph = await getMatchPrebuiltRuleGraph({ model }); - const results = await graph.invoke(parameters, options); - - const ruleMatched = false; - if (!ruleMatched) { + const result = await graph.invoke(parameters, options); + const { matched, rule } = result as MatchPrebuiltRuleState; + if (!matched) { return response.noContent(); } - return response.ok({ body: { id: '1', name: 'dummy', installed: true } }); + return response.ok({ body: rule }); } catch (err) { return response.badRequest({ body: err.message, diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts new file mode 100644 index 0000000000000..b40efdf41259f --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts @@ -0,0 +1,68 @@ +/* + * 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 { RulesClient } from '@kbn/alerting-plugin/server'; +import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; +import type { SplunkRule } from '../../../../../../../../common/api/siem_migrations/splunk/rules/splunk_rule.gen'; +import { fetchRuleVersionsTriad } from '../../../../../../detection_engine/prebuilt_rules/logic/rule_versions/fetch_rule_versions_triad'; +import { createPrebuiltRuleObjectsClient } from '../../../../../../detection_engine/prebuilt_rules/logic/rule_objects/prebuilt_rule_objects_client'; +import { createPrebuiltRuleAssetsClient } from '../../../../../../detection_engine/prebuilt_rules/logic/rule_assets/prebuilt_rule_assets_client'; +import type { PrebuiltRulesMapByName } from '../../types'; +// import { getVersionBuckets } from '../../../../../detection_engine/prebuilt_rules/model/rule_versions/get_version_buckets'; + +interface RetrievePrebuiltRulesParams { + soClient: SavedObjectsClientContract; + rulesClient: RulesClient; + splunkRule: SplunkRule; +} + +export const retrievePrebuiltRulesMap = async ({ + soClient, + rulesClient, + splunkRule, +}: RetrievePrebuiltRulesParams): Promise => { + const ruleAssetsClient = createPrebuiltRuleAssetsClient(soClient); + const ruleObjectsClient = createPrebuiltRuleObjectsClient(rulesClient); + + const prebuiltRulesMap = await fetchRuleVersionsTriad({ + ruleAssetsClient, + ruleObjectsClient, + }); + const prebuiltRulesByName: PrebuiltRulesMapByName = new Map(); + prebuiltRulesMap.forEach((ruleVersions) => { + const rule = ruleVersions.target || ruleVersions.current; + if (rule) { + prebuiltRulesByName.set(rule.name, { + isInstalled: !!ruleVersions.current, + rule, + }); + } + }); + + if (splunkRule.mitreAttackIds?.length) { + // If this rule has MITRE ATT&CK IDs, remove unrelated prebuilt rules + prebuiltRulesByName.forEach(({ rule }, ruleName) => { + const mitreAttackThreat = rule.threat?.filter( + ({ framework }) => framework === 'MITRE ATT&CK' + ); + if (!mitreAttackThreat) { + // If this rule has no MITRE ATT&CK reference we can not ensure it is not related to the rule, keep it + return; + } + + const sameTechnique = mitreAttackThreat.find((threat) => + threat.technique?.some(({ id }) => splunkRule.mitreAttackIds?.includes(id)) + ); + + if (!sameTechnique) { + prebuiltRulesByName.delete(ruleName); + } + }); + } + + return prebuiltRulesByName; +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/types.ts similarity index 50% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/types.ts index 1fec1c76430eb..bffc7c2714f4f 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/match_prebuilt_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/types.ts @@ -4,3 +4,12 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ + +import type { PrebuiltRuleAsset } from '../../../../detection_engine/prebuilt_rules'; + +export interface PrebuiltRuleMapped { + isInstalled: boolean; + rule: PrebuiltRuleAsset; +} + +export type PrebuiltRulesMapByName = Map; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts deleted file mode 100644 index e380739c0b6cd..0000000000000 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/util/llm.ts +++ /dev/null @@ -1,31 +0,0 @@ -/* - * 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 { - ActionsClientBedrockChatModel, - ActionsClientChatOpenAI, - ActionsClientGeminiChatModel, - ActionsClientSimpleChatModel, -} from '@kbn/langchain/server'; - -export const getLLMType = (actionTypeId: string): string | undefined => { - const llmTypeDictionary: Record = { - [`.gen-ai`]: `openai`, - [`.bedrock`]: `bedrock`, - [`.gemini`]: `gemini`, - }; - return llmTypeDictionary[actionTypeId]; -}; - -export const getLLMClass = (llmType?: string) => - llmType === 'openai' - ? ActionsClientChatOpenAI - : llmType === 'bedrock' - ? ActionsClientBedrockChatModel - : llmType === 'gemini' - ? ActionsClientGeminiChatModel - : ActionsClientSimpleChatModel; diff --git a/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts b/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts index d72d22a7a79e0..2aeca2751f0a7 100644 --- a/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts +++ b/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts @@ -113,4 +113,15 @@ const stackManagementLinks: Array> title: 'Stack', children: [{ link: 'management:license_management' }, { link: 'management:upgrade_assistant' }], }, + { + title: 'Security Solution', + children: [ + { link: `${SECURITY_UI_APP_ID}:${SecurityPageName.entityAnalyticsManagement}` }, + { link: `${SECURITY_UI_APP_ID}:${SecurityPageName.entityAnalyticsAssetClassification}` }, + { link: 'maps' }, + { link: 'visualize' }, + { link: 'graph' }, + { link: 'canvas' }, + ], + }, ]; From d858f5f94c65f366fc6af6503ccfa762bb17203f Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Mon, 30 Sep 2024 10:35:04 +0200 Subject: [PATCH 03/17] cleaning --- .../create_integration_assistant/state.ts | 6 +----- .../match_prebuilt_rule/api/util/prebuilt_rules.ts | 1 - .../public/navigation/side_navigation.ts | 11 ----------- 3 files changed, 1 insertion(+), 17 deletions(-) diff --git a/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts b/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts index 99f95a3eee039..bef5b35624df4 100644 --- a/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts +++ b/x-pack/plugins/integration_assistant/public/components/create_integration/create_integration_assistant/state.ts @@ -51,11 +51,7 @@ export const reducer = (state: State, action: Action): State => { case 'SET_IS_GENERATING': return { ...state, isGenerating: action.payload }; case 'SET_GENERATED_RESULT': - return { - ...state, - // keep original result as the samplesFormat is not always included in the payload - result: state.result ? { ...state.result, ...action.payload } : action.payload, - }; + return { ...state, result: action.payload }; default: return state; } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts index b40efdf41259f..5d209928ec7fe 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.ts @@ -12,7 +12,6 @@ import { fetchRuleVersionsTriad } from '../../../../../../detection_engine/prebu import { createPrebuiltRuleObjectsClient } from '../../../../../../detection_engine/prebuilt_rules/logic/rule_objects/prebuilt_rule_objects_client'; import { createPrebuiltRuleAssetsClient } from '../../../../../../detection_engine/prebuilt_rules/logic/rule_assets/prebuilt_rule_assets_client'; import type { PrebuiltRulesMapByName } from '../../types'; -// import { getVersionBuckets } from '../../../../../detection_engine/prebuilt_rules/model/rule_versions/get_version_buckets'; interface RetrievePrebuiltRulesParams { soClient: SavedObjectsClientContract; diff --git a/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts b/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts index 2aeca2751f0a7..d72d22a7a79e0 100644 --- a/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts +++ b/x-pack/plugins/security_solution_ess/public/navigation/side_navigation.ts @@ -113,15 +113,4 @@ const stackManagementLinks: Array> title: 'Stack', children: [{ link: 'management:license_management' }, { link: 'management:upgrade_assistant' }], }, - { - title: 'Security Solution', - children: [ - { link: `${SECURITY_UI_APP_ID}:${SecurityPageName.entityAnalyticsManagement}` }, - { link: `${SECURITY_UI_APP_ID}:${SecurityPageName.entityAnalyticsAssetClassification}` }, - { link: 'maps' }, - { link: 'visualize' }, - { link: 'graph' }, - { link: 'canvas' }, - ], - }, ]; From 4084290eaf1ac1a3e3455e087a4b28ca4991fd8e Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Mon, 30 Sep 2024 18:28:26 +0200 Subject: [PATCH 04/17] add tests --- .../common/experimental_features.ts | 5 + .../api/util/prebuilt_rules.test.ts | 119 ++++++++++++++++++ .../security_solution/server/routes/index.ts | 9 +- 3 files changed, 130 insertions(+), 3 deletions(-) create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.test.ts diff --git a/x-pack/plugins/security_solution/common/experimental_features.ts b/x-pack/plugins/security_solution/common/experimental_features.ts index febfb6929333c..53af4309f4b3b 100644 --- a/x-pack/plugins/security_solution/common/experimental_features.ts +++ b/x-pack/plugins/security_solution/common/experimental_features.ts @@ -258,6 +258,11 @@ export const allowedExperimentalValues = Object.freeze({ * Enables the new Entity Store engine routes */ entityStoreEnabled: false, + + /** + * Enables the siem migrations feature + */ + siemMigrationsEnabled: false, }); type ExperimentalConfigKeys = Array; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.test.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.test.ts new file mode 100644 index 0000000000000..dd3801dde0e33 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/util/prebuilt_rules.test.ts @@ -0,0 +1,119 @@ +/* + * 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 { savedObjectsClientMock } from '@kbn/core-saved-objects-api-server-mocks'; +import { retrievePrebuiltRulesMap } from './prebuilt_rules'; +import { rulesClientMock } from '@kbn/alerting-plugin/server/mocks'; +import type { SplunkRule } from '../../../../../../../../common/api/siem_migrations/splunk/rules/splunk_rule.gen'; + +jest.mock( + '../../../../../../detection_engine/prebuilt_rules/logic/rule_objects/prebuilt_rule_objects_client', + () => ({ createPrebuiltRuleObjectsClient: jest.fn() }) +); +jest.mock( + '../../../../../../detection_engine/prebuilt_rules/logic/rule_assets/prebuilt_rule_assets_client', + () => ({ createPrebuiltRuleAssetsClient: jest.fn() }) +); + +const rule1 = { + name: 'rule one', + id: 'rule1', + threat: [ + { + framework: 'MITRE ATT&CK', + technique: [{ id: 'T1234', name: 'tactic one' }], + }, + ], +}; +const rule2 = { + name: 'rule two', + id: 'rule2', +}; + +const defaultRuleVersionsTriad = new Map([ + ['rule1', { target: rule1 }], + ['rule2', { target: rule2, current: rule2 }], +]); +const mockFetchRuleVersionsTriad = jest.fn().mockResolvedValue(defaultRuleVersionsTriad); +jest.mock( + '../../../../../../detection_engine/prebuilt_rules/logic/rule_versions/fetch_rule_versions_triad', + () => ({ + fetchRuleVersionsTriad: () => mockFetchRuleVersionsTriad(), + }) +); + +const splunkRule: SplunkRule = { + title: 'splunk rule', + description: 'splunk rule description', + search: 'index=*', + mitreAttackIds: ['T1234'], +}; + +const defaultParams = { + soClient: savedObjectsClientMock.create(), + rulesClient: rulesClientMock.create(), + splunkRule, +}; + +describe('retrievePrebuiltRulesMap', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('when prebuilt rule is installed', () => { + it('should return isInstalled flag', async () => { + const prebuiltRulesMap = await retrievePrebuiltRulesMap(defaultParams); + expect(prebuiltRulesMap.size).toBe(2); + expect(prebuiltRulesMap.get('rule one')).toEqual( + expect.objectContaining({ isInstalled: false }) + ); + expect(prebuiltRulesMap.get('rule two')).toEqual( + expect.objectContaining({ isInstalled: true }) + ); + }); + }); + + describe('when splunk rule does not contain mitreAttackIds', () => { + it('should return the full rules map', async () => { + const prebuiltRulesMap = await retrievePrebuiltRulesMap({ + ...defaultParams, + splunkRule: { ...splunkRule, mitreAttackIds: undefined }, + }); + expect(prebuiltRulesMap.size).toBe(2); + }); + }); + + describe('when splunk rule contains empty mitreAttackIds', () => { + it('should return the full rules map', async () => { + const prebuiltRulesMap = await retrievePrebuiltRulesMap({ + ...defaultParams, + splunkRule: { ...splunkRule, mitreAttackIds: [] }, + }); + expect(prebuiltRulesMap.size).toBe(2); + }); + }); + + describe('when splunk rule contains non matching mitreAttackIds', () => { + it('should return the full rules map', async () => { + const prebuiltRulesMap = await retrievePrebuiltRulesMap({ + ...defaultParams, + splunkRule: { ...splunkRule, mitreAttackIds: ['T2345'] }, + }); + expect(prebuiltRulesMap.size).toBe(1); + expect(prebuiltRulesMap.get('rule two')).toEqual(expect.objectContaining({ rule: rule2 })); + }); + }); + + describe('when splunk rule contains matching mitreAttackIds', () => { + it('should return the filtered rules map', async () => { + const prebuiltRulesMap = await retrievePrebuiltRulesMap(defaultParams); + expect(prebuiltRulesMap.size).toBe(2); + expect(prebuiltRulesMap.get('rule one')).toEqual(expect.objectContaining({ rule: rule1 })); + expect(prebuiltRulesMap.get('rule two')).toEqual(expect.objectContaining({ rule: rule2 })); + }); + }); +}); diff --git a/x-pack/plugins/security_solution/server/routes/index.ts b/x-pack/plugins/security_solution/server/routes/index.ts index ab7c603aebc77..85c4320a20325 100644 --- a/x-pack/plugins/security_solution/server/routes/index.ts +++ b/x-pack/plugins/security_solution/server/routes/index.ts @@ -12,7 +12,7 @@ import type { SecuritySolutionPluginRouter } from '../types'; import { registerFleetIntegrationsRoutes } from '../lib/detection_engine/fleet_integrations'; import { registerPrebuiltRulesRoutes } from '../lib/detection_engine/prebuilt_rules'; -// eslint-disable-next-line no-restricted-imports + import { registerLegacyRuleActionsRoutes } from '../lib/detection_engine/rule_actions_legacy'; import { registerRuleExceptionsRoutes } from '../lib/detection_engine/rule_exceptions'; import { registerRuleManagementRoutes } from '../lib/detection_engine/rule_management'; @@ -139,14 +139,17 @@ export const initRoutes = ( // Dashboards registerDashboardsRoutes(router, logger); registerTagsRoutes(router, logger); - const { previewTelemetryUrlEnabled } = config.experimentalFeatures; + const { previewTelemetryUrlEnabled, siemMigrationsEnabled } = config.experimentalFeatures; if (previewTelemetryUrlEnabled) { // telemetry preview endpoint for e2e integration tests only at the moment. telemetryDetectionRulesPreviewRoute(router, logger, previewTelemetryReceiver, telemetrySender); } registerEntityAnalyticsRoutes({ router, config, getStartServices, logger }); - registerSiemMigrationsRoutes(router, logger); + + if (siemMigrationsEnabled) { + registerSiemMigrationsRoutes(router, logger); + } // Security Integrations getFleetManagedIndexTemplatesRoute(router); From 454a2c41a3872340e04ac99cdea5b2f6d045b1bf Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 1 Oct 2024 11:56:03 +0200 Subject: [PATCH 05/17] clean state --- .../rules/match_prebuilt_rule/agent/graph.ts | 23 +++++++++++-------- .../rules/match_prebuilt_rule/agent/state.ts | 13 ++++------- .../api/match_prebuilt_rule.ts | 6 ++--- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts index 6e705e8c9ec28..cb17f46ea21b8 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/graph.ts @@ -23,28 +23,33 @@ const callModel = async ({ const matchPrebuiltRule = MATCH_PREBUILT_RULE_PROMPT.pipe(model).pipe(outputParser); const elasticSecurityRules = Array(state.prebuiltRulesMap.keys()).join('\n'); - const result = await matchPrebuiltRule.invoke({ + const response = await matchPrebuiltRule.invoke({ elasticSecurityRules, splunkRuleTitle: state.splunkRuleTitle, splunkRuleDescription: state.splunkRuleDescription, }); - const response = result.trim(); - if (response === 'no_match') { + return { response }; +}; + +const processResponse = (state: MatchPrebuiltRuleState): Partial => { + const cleanResponse = state.response.trim(); + if (cleanResponse === 'no_match') { return { matched: false }; } - if (state.prebuiltRulesMap.get(response)) { - return { matched: true, rule: state.prebuiltRulesMap.get(response), response }; + const result = state.prebuiltRulesMap.get(cleanResponse); + if (result != null) { + return { matched: true, result }; } - return { matched: false, response }; + return { matched: false }; }; export async function getMatchPrebuiltRuleGraph({ model }: MatchPrebuiltRuleGraphParams) { const matchPrebuiltRuleGraph = new StateGraph(matchPrebuiltRuleState) .addNode('callModel', (state: MatchPrebuiltRuleState) => callModel({ state, model })) - // .addNode('processModelResponse', processModelResponse) + .addNode('processResponse', processResponse) .addEdge(START, 'callModel') - // .addEdge('callModel', 'processModelResponse') - .addEdge('callModel', END); + .addEdge('callModel', 'processResponse') + .addEdge('processResponse', END); return matchPrebuiltRuleGraph.compile(); } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts index dca4708a0ace4..dc93281bac3c8 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/agent/state.ts @@ -5,19 +5,14 @@ * 2.0. */ -import type { BaseMessage } from '@langchain/core/messages'; -import { Annotation, messagesStateReducer } from '@langchain/langgraph'; +import { Annotation } from '@langchain/langgraph'; import type { PrebuiltRuleMapped, PrebuiltRulesMapByName } from '../types'; export const matchPrebuiltRuleState = Annotation.Root({ - messages: Annotation({ - reducer: messagesStateReducer, - default: () => [], - }), splunkRuleTitle: Annotation(), splunkRuleDescription: Annotation, - matched: Annotation, - response: Annotation, - rule: Annotation, prebuiltRulesMap: Annotation, + response: Annotation, + matched: Annotation, + result: Annotation, }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts index df07171ce123c..1c3f384eec30a 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/api/match_prebuilt_rule.ts @@ -74,13 +74,13 @@ export const registerSplunkMatchPrebuiltRuleRoute = ( ], }; const graph = await getMatchPrebuiltRuleGraph({ model }); + const matchPrebuiltRuleState = await graph.invoke(parameters, options); - const result = await graph.invoke(parameters, options); - const { matched, rule } = result as MatchPrebuiltRuleState; + const { matched, result } = matchPrebuiltRuleState as MatchPrebuiltRuleState; if (!matched) { return response.noContent(); } - return response.ok({ body: rule }); + return response.ok({ body: result }); } catch (err) { return response.badRequest({ body: err.message, From 7beb4ab5a892f252bd1d3e4cc58a02dd9a957ca4 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 1 Oct 2024 14:52:42 +0200 Subject: [PATCH 06/17] recover removed elint exception --- x-pack/plugins/security_solution/server/routes/index.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/security_solution/server/routes/index.ts b/x-pack/plugins/security_solution/server/routes/index.ts index 85c4320a20325..9a8e4bf157859 100644 --- a/x-pack/plugins/security_solution/server/routes/index.ts +++ b/x-pack/plugins/security_solution/server/routes/index.ts @@ -12,7 +12,7 @@ import type { SecuritySolutionPluginRouter } from '../types'; import { registerFleetIntegrationsRoutes } from '../lib/detection_engine/fleet_integrations'; import { registerPrebuiltRulesRoutes } from '../lib/detection_engine/prebuilt_rules'; - +// eslint-disable-next-line no-restricted-imports import { registerLegacyRuleActionsRoutes } from '../lib/detection_engine/rule_actions_legacy'; import { registerRuleExceptionsRoutes } from '../lib/detection_engine/rule_exceptions'; import { registerRuleManagementRoutes } from '../lib/detection_engine/rule_management'; From 42ad6e7275de144eb1891abc71c9d0f190dfa872 Mon Sep 17 00:00:00 2001 From: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Date: Tue, 1 Oct 2024 13:45:39 +0000 Subject: [PATCH 07/17] [CI] Auto-commit changed files from 'yarn openapi:generate' --- .../common/api/quickstart_client.gen.ts | 23 +++++++++++++++++++ .../services/security_solution_api.gen.ts | 20 ++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/x-pack/plugins/security_solution/common/api/quickstart_client.gen.ts b/x-pack/plugins/security_solution/common/api/quickstart_client.gen.ts index 255b56c485b9b..295ee46f7bbb8 100644 --- a/x-pack/plugins/security_solution/common/api/quickstart_client.gen.ts +++ b/x-pack/plugins/security_solution/common/api/quickstart_client.gen.ts @@ -292,6 +292,10 @@ import type { PreviewRiskScoreRequestBodyInput, PreviewRiskScoreResponse, } from './entity_analytics/risk_engine/preview_route.gen'; +import type { + SplunkRuleMigrationMatchPrebuiltRuleRequestBodyInput, + SplunkRuleMigrationMatchPrebuiltRuleResponse, +} from './siem_migrations/splunk/rules/match_prebuilt_rule.gen'; import type { CleanDraftTimelinesRequestBodyInput, CleanDraftTimelinesResponse, @@ -1899,6 +1903,22 @@ detection engine rules. }) .catch(catchAxiosErrorFormatAndThrow); } + /** + * Perform Elastic prebuilt rule matching from Splunk Security rule + */ + async splunkRuleMigrationMatchPrebuiltRule(props: SplunkRuleMigrationMatchPrebuiltRuleProps) { + this.log.info(`${new Date().toISOString()} Calling API SplunkRuleMigrationMatchPrebuiltRule`); + return this.kbnClient + .request({ + path: '/internal/migrations/splunk/rules/match_prebuilt_rule', + headers: { + [ELASTIC_HTTP_VERSION_HEADER]: '1', + }, + method: 'POST', + body: props.body, + }) + .catch(catchAxiosErrorFormatAndThrow); + } async startEntityEngine(props: StartEntityEngineProps) { this.log.info(`${new Date().toISOString()} Calling API StartEntityEngine`); return this.kbnClient @@ -2229,6 +2249,9 @@ export interface SetAlertsStatusProps { export interface SetAlertTagsProps { body: SetAlertTagsRequestBodyInput; } +export interface SplunkRuleMigrationMatchPrebuiltRuleProps { + body: SplunkRuleMigrationMatchPrebuiltRuleRequestBodyInput; +} export interface StartEntityEngineProps { params: StartEntityEngineRequestParamsInput; } diff --git a/x-pack/test/api_integration/services/security_solution_api.gen.ts b/x-pack/test/api_integration/services/security_solution_api.gen.ts index aba6550ea9c3b..cdb69327e1bf5 100644 --- a/x-pack/test/api_integration/services/security_solution_api.gen.ts +++ b/x-pack/test/api_integration/services/security_solution_api.gen.ts @@ -125,6 +125,7 @@ import { SearchAlertsRequestBodyInput } from '@kbn/security-solution-plugin/comm import { SetAlertAssigneesRequestBodyInput } from '@kbn/security-solution-plugin/common/api/detection_engine/alert_assignees/set_alert_assignees_route.gen'; import { SetAlertsStatusRequestBodyInput } from '@kbn/security-solution-plugin/common/api/detection_engine/signals/set_signal_status/set_signals_status_route.gen'; import { SetAlertTagsRequestBodyInput } from '@kbn/security-solution-plugin/common/api/detection_engine/alert_tags/set_alert_tags/set_alert_tags.gen'; +import { SplunkRuleMigrationMatchPrebuiltRuleRequestBodyInput } from '@kbn/security-solution-plugin/common/api/siem_migrations/splunk/rules/match_prebuilt_rule.gen'; import { StartEntityEngineRequestParamsInput } from '@kbn/security-solution-plugin/common/api/entity_analytics/entity_store/engine/start.gen'; import { StopEntityEngineRequestParamsInput } from '@kbn/security-solution-plugin/common/api/entity_analytics/entity_store/engine/stop.gen'; import { SuggestUserProfilesRequestQueryInput } from '@kbn/security-solution-plugin/common/api/detection_engine/users/suggest_user_profiles_route.gen'; @@ -1266,6 +1267,22 @@ detection engine rules. .set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana') .send(props.body as object); }, + /** + * Perform Elastic prebuilt rule matching from Splunk Security rule + */ + splunkRuleMigrationMatchPrebuiltRule( + props: SplunkRuleMigrationMatchPrebuiltRuleProps, + kibanaSpace: string = 'default' + ) { + return supertest + .post( + routeWithNamespace('/internal/migrations/splunk/rules/match_prebuilt_rule', kibanaSpace) + ) + .set('kbn-xsrf', 'true') + .set(ELASTIC_HTTP_VERSION_HEADER, '1') + .set(X_ELASTIC_INTERNAL_ORIGIN_REQUEST, 'kibana') + .send(props.body as object); + }, startEntityEngine(props: StartEntityEngineProps, kibanaSpace: string = 'default') { return supertest .post( @@ -1574,6 +1591,9 @@ export interface SetAlertsStatusProps { export interface SetAlertTagsProps { body: SetAlertTagsRequestBodyInput; } +export interface SplunkRuleMigrationMatchPrebuiltRuleProps { + body: SplunkRuleMigrationMatchPrebuiltRuleRequestBodyInput; +} export interface StartEntityEngineProps { params: StartEntityEngineRequestParamsInput; } From 68632f47fbbdc8ecbbae4e6be8ba571a42f44c7e Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Thu, 3 Oct 2024 18:40:53 +0200 Subject: [PATCH 08/17] route scafolding --- .../siem_migrations/splunk/rules/constants.ts | 2 + .../splunk/rules/splunk_rule.gen.ts | 82 ++++++++++++++++- .../splunk/rules/splunk_rule.schema.yaml | 92 +++++++++++++++++-- .../lib/siem_migrations/splunk/rules/index.ts | 4 +- 4 files changed, 169 insertions(+), 11 deletions(-) diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts index 95d14208d673e..67794c0eca73e 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/constants.ts @@ -11,3 +11,5 @@ const SPLUNK_RULE_MIGRATIONS_PATH = `${SPLUNK_MIGRATIONS_PATH}/rules` as const; export const SPLUNK_MATCH_PREBUILT_RULE_PATH = `${SPLUNK_RULE_MIGRATIONS_PATH}/match_prebuilt_rule` as const; + +export const SPLUNK_TRANSLATE_RULE_PATH = `${SPLUNK_RULE_MIGRATIONS_PATH}/translate` as const; diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts index f0f2bc4ab520f..219eb3601e138 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.gen.ts @@ -18,20 +18,94 @@ import { z } from '@kbn/zod'; export type SplunkRule = z.infer; export const SplunkRule = z.object({ + /** + * The Splunk rule id. + */ + id: z.string(), /** * The Splunk rule name. */ - title: z.string().min(1), + title: z.string(), /** - * The Splunk rule search query. + * The Splunk rule splSearch query. */ - search: z.string().min(1), + splSearch: z.string(), /** * The Splunk rule description. */ - description: z.string().min(1), + description: z.string(), /** * String array containing the rule Mitre Attack technique IDs. */ mitreAttackIds: z.array(z.string()).optional(), }); + +export type SplunkRuleMigration = z.infer; +export const SplunkRuleMigration = SplunkRule.merge( + z.object({ + /** + * The Splunk rule migration document uuid. + */ + uuid: z.string(), + /** + * The translated elastic query. + */ + elasticQuery: z.string(), + /** + * The translated elastic query language. + */ + elasticQueryLanguage: z.literal('esql').default('esql'), + /** + * The translation state. + */ + translationState: z.enum([ + 'untranslated', + 'matched', + 'translated:complete', + 'translated:partial', + 'translated:missing', + 'translated:failed', + ]), + /** + * The resources needed to complete the translation. + */ + resources: z + .array( + z.object({ + /** + * The resource type. + */ + type: z.enum(['macro', 'lookup']).optional(), + /** + * The resource name. + */ + name: z.string().optional(), + /** + * The resource value. + */ + value: z.string().optional(), + }) + ) + .optional(), + /** + * The missing resources needed to complete the translation. The format is `:`. E.g. `macro:my_macro`. + */ + missingResources: z.array(z.string()).optional(), + /** + * The status of the rule migration. + */ + status: z.enum(['pending', 'processing', 'finished', 'error']).default('pending'), + /** + * The migrated Elastic rule id. + */ + ruleId: z.string().optional(), + /** + * The summary of the migration in markdown. + */ + summary: z.string().optional(), + /** + * The agent messages list from the migration. + */ + messages: z.array(z.string()), + }) +); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml index 92d7b049f6179..e403c14bbe5ce 100644 --- a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/splunk_rule.schema.yaml @@ -9,24 +9,104 @@ components: SplunkRule: type: object required: + - id - title - description - - search + - splSearch properties: + id: + type: string + description: The Splunk rule id. title: type: string - minLength: 1 description: The Splunk rule name. - search: + splSearch: type: string - minLength: 1 - description: The Splunk rule search query. + description: The Splunk rule splSearch query. description: type: string - minLength: 1 description: The Splunk rule description. mitreAttackIds: type: array items: type: string description: String array containing the rule Mitre Attack technique IDs. + + SplunkRuleMigration: + allOf: # Combines the SplunkRule and the inline model + - $ref: '#/components/schemas/SplunkRule' + - type: object + description: The Splunk rule migration document object. + required: + - uuid + - elasticQuery + - elasticQueryLanguage + - translationState + - status + - messages + properties: + uuid: + type: string + description: The Splunk rule migration document uuid. + elasticQuery: + type: string + description: The translated elastic query. + elasticQueryLanguage: + type: string + description: The translated elastic query language. + enum: + - esql + default: esql + translationState: + type: string + description: The translation state. + enum: + - untranslated + - matched + - translated:complete + - translated:partial + - translated:missing + - translated:failed + resources: + type: array + description: The resources needed to complete the translation. + items: + type: object + properties: + type: + type: string + description: The resource type. + enum: + - macro + - lookup + name: + type: string + description: The resource name. + value: + type: string + description: The resource value. + missingResources: + type: array + description: The missing resources needed to complete the translation. The format is `:`. E.g. `macro:my_macro`. + items: + type: string + status: + type: string + description: The status of the rule migration. + enum: + - pending + - processing + - finished + - error + default: pending + ruleId: + type: string + description: The migrated Elastic rule id. + summary: + type: string + description: The summary of the migration in markdown. + messages: + type: array + description: The agent messages list from the migration. + items: + type: string diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts index 02a22222d443d..ed620623e9e5a 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/index.ts @@ -7,11 +7,13 @@ import type { Logger } from '@kbn/core/server'; import type { SecuritySolutionPluginRouter } from '../../../../types'; -import { registerSplunkMatchPrebuiltRuleRoute } from './match_prebuilt_rule/api/match_prebuilt_rule'; +import { registerSplunkMatchPrebuiltRuleRoute } from './match_prebuilt_rule'; +import { registerSplunkTranslateRuleRoute } from './translate_rule'; export const registerSplunkMigrationRoutes = ( router: SecuritySolutionPluginRouter, logger: Logger ) => { registerSplunkMatchPrebuiltRuleRoute(router, logger); + registerSplunkTranslateRuleRoute(router, logger); }; From 9a68f7be69572bdb855857b10853d81bca5450e0 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Thu, 3 Oct 2024 18:41:07 +0200 Subject: [PATCH 09/17] route directory scafolding --- .../splunk/rules/translate_rule.gen.ts | 39 ++++++++ .../splunk/rules/translate_rule.schema.yaml | 41 +++++++++ .../splunk/rules/match_prebuilt_rule/index.ts | 8 ++ .../rules/translate_rule/agent/graph.ts | 32 +++++++ .../rules/translate_rule/agent/prompts.ts | 37 ++++++++ .../rules/translate_rule/agent/state.ts | 14 +++ .../rules/translate_rule/agent/types.ts | 20 ++++ .../translate_rule/api/translate_rule.ts | 92 +++++++++++++++++++ .../api/util/find_missing_resources.ts | 26 ++++++ .../splunk/rules/translate_rule/index.ts | 8 ++ .../splunk/rules/translate_rule/types.ts | 6 ++ 11 files changed, 323 insertions(+) create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.gen.ts create mode 100644 x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.schema.yaml create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/index.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/util/find_missing_resources.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/index.ts create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/types.ts diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.gen.ts b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.gen.ts new file mode 100644 index 0000000000000..7b11a6fec34ab --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.gen.ts @@ -0,0 +1,39 @@ +/* + * 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: Splunk Migration Translation Rule API endpoint + * version: 1 + */ + +import { z } from '@kbn/zod'; + +import { SplunkRule, SplunkRuleMigration } from './splunk_rule.gen'; +import { ConnectorId, LangSmithOptions } from '../../common.gen'; + +export type SplunkRuleMigrationTranslateRuleRequestBody = z.infer< + typeof SplunkRuleMigrationTranslateRuleRequestBody +>; +export const SplunkRuleMigrationTranslateRuleRequestBody = z.object({ + splunkRule: SplunkRule, + connectorId: ConnectorId, + langSmithOptions: LangSmithOptions.optional(), +}); +export type SplunkRuleMigrationTranslateRuleRequestBodyInput = z.input< + typeof SplunkRuleMigrationTranslateRuleRequestBody +>; + +export type SplunkRuleMigrationTranslateRuleResponse = z.infer< + typeof SplunkRuleMigrationTranslateRuleResponse +>; +export const SplunkRuleMigrationTranslateRuleResponse = z.object({ + migration: SplunkRuleMigration, +}); diff --git a/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.schema.yaml b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.schema.yaml new file mode 100644 index 0000000000000..c0c736c3a50b2 --- /dev/null +++ b/x-pack/plugins/security_solution/common/api/siem_migrations/splunk/rules/translate_rule.schema.yaml @@ -0,0 +1,41 @@ +openapi: 3.0.3 +info: + title: Splunk Migration Translation Rule API endpoint + version: '1' +paths: + /internal/migrations/splunk/rules/translate: + post: + summary: Translates Splunk rules to Elastic custom rules + operationId: SplunkRuleMigrationTranslateRule + x-codegen-enabled: true + description: Perform Elastic custom rule rule translation from Splunk Security rule + tags: + - Splunk Rule Migration API + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - splunkRule + - connectorId + properties: + splunkRule: + $ref: './splunk_rule.schema.yaml#/components/schemas/SplunkRule' + connectorId: + $ref: '../../common.schema.yaml#/components/schemas/ConnectorId' + langSmithOptions: + $ref: '../../common.schema.yaml#/components/schemas/LangSmithOptions' + responses: + 200: + description: Indicates the translation of the Splunk rule to Elastic custom rule has been processed. + content: + application/json: + schema: + type: object + required: + - migration + properties: + migration: + $ref: './splunk_rule.schema.yaml#/components/schemas/SplunkRuleMigration' diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/index.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/index.ts new file mode 100644 index 0000000000000..beef0eef8c0fa --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/match_prebuilt_rule/index.ts @@ -0,0 +1,8 @@ +/* + * 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 { registerSplunkMatchPrebuiltRuleRoute } from './api/match_prebuilt_rule'; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts new file mode 100644 index 0000000000000..4b0ce26cc29e5 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts @@ -0,0 +1,32 @@ +/* + * 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 { END, START, StateGraph } from '@langchain/langgraph'; +// import { StringOutputParser } from '@langchain/core/output_parsers'; +import { translateRuleState } from './state'; +import type { + TranslateRuleGraphParams, + TranslateRuleNodeParams, + TranslateRuleState, +} from './types'; + +const callModel = async ({ + state, + model, +}: TranslateRuleNodeParams): Promise> => { + // TODO: + return state; +}; + +export async function getTranslateRuleGraph({ model }: TranslateRuleGraphParams) { + const matchPrebuiltRuleGraph = new StateGraph(translateRuleState) + .addNode('callModel', (state: TranslateRuleState) => callModel({ state, model })) + .addEdge(START, 'callModel') + .addEdge('callModel', END); + + return matchPrebuiltRuleGraph.compile(); +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts new file mode 100644 index 0000000000000..4cae08c1da6ff --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -0,0 +1,37 @@ +/* + * 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 { ChatPromptTemplate } from '@langchain/core/prompts'; +export const MATCH_PREBUILT_RULE_PROMPT = ChatPromptTemplate.fromMessages([ + [ + 'system', + `You are an expert assistant in Cybersecurity, your task is to help migrating a SIEM detection rule, from Splunk Security to Elastic Security. +You will be provided with a Splunk Detection Rule name by the user, your goal is to try find an Elastic Detection Rule that covers the same threat, if any. +The list of Elastic Detection Rules suggested is provided in the context below. + + +If there is no Elastic rule in the list that covers the same threat, answer only with the string: no_match +If there is one Elastic rule in the list that covers the same threat, answer only with its name without any further explanation. +If there are multiple rules in the list that cover the same threat, answer with the most specific of them, for example: "Linux User Account Creation" is more specific than "User Account Creation". + + + + +{elasticSecurityRules} + +`, + ], + [ + 'human', + `The Splunk Detection Rule is: + +{splunkRuleTitle} + +`, + ], + ['ai', 'Please find the answer below:'], +]); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts new file mode 100644 index 0000000000000..516094d07ae2b --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts @@ -0,0 +1,14 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { Annotation } from '@langchain/langgraph'; + +export const translateRuleState = Annotation.Root({ + splunkRuleTitle: Annotation(), + splunkRuleDescription: Annotation, + response: Annotation, +}); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts new file mode 100644 index 0000000000000..4e3dfac6b781d --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts @@ -0,0 +1,20 @@ +/* + * 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 { ChatModel } from '../../../../actions_client_chat'; +import type { translateRuleState } from './state'; + +export type TranslateRuleState = typeof translateRuleState.State; + +export interface TranslateRuleGraphParams { + model: ChatModel; +} + +export interface TranslateRuleNodeParams { + model: ChatModel; + state: TranslateRuleState; +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts new file mode 100644 index 0000000000000..dcefc33fe45fa --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts @@ -0,0 +1,92 @@ +/* + * 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, Logger } from '@kbn/core/server'; +import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; +import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { APMTracer } from '@kbn/langchain/server/tracers/apm'; +import type { SplunkRuleMigration } from '../../../../../../../common/api/siem_migrations/splunk/rules/splunk_rule.gen'; +import type { SplunkRuleMigrationTranslateRuleResponse } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; +import { SplunkRuleMigrationTranslateRuleRequestBody } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; +import type { SecuritySolutionPluginRouter } from '../../../../../../types'; +import { SPLUNK_TRANSLATE_RULE_PATH } from '../../../../../../../common/api/siem_migrations/splunk/rules/constants'; +import { getTranslateRuleGraph } from '../agent/graph'; +import type { TranslateRuleState } from '../agent/types'; +import { ActionsClientChat } from '../../../../actions_client_chat'; + +type SplunkTranslateRuleRouteResponse = IKibanaResponse; + +export const registerSplunkTranslateRuleRoute = ( + router: SecuritySolutionPluginRouter, + logger: Logger +) => { + router.versioned + .post({ + path: SPLUNK_TRANSLATE_RULE_PATH, + access: 'internal', + options: { + tags: ['access:securitySolution'], + }, + }) + .addVersion( + { + version: '1', + validate: { + request: { + body: buildRouteValidationWithZod(SplunkRuleMigrationTranslateRuleRequestBody), + }, + }, + }, + async (context, req, res): Promise => { + const { langSmithOptions, splunkRule, connectorId } = req.body; + try { + const ctx = await context.resolve(['core', 'actions']); + + const actionsClient = ctx.actions.getActionsClient(); + + const actionsClientChat = new ActionsClientChat(connectorId, actionsClient, logger); + const model = await actionsClientChat.createModel({ + signal: getRequestAbortedSignal(req.events.aborted$), + temperature: 0.05, + }); + + const parameters: Partial = { + splunkRuleTitle: splunkRule.title, + splunkRuleDescription: splunkRule.description, + }; + + const options = { + callbacks: [ + new APMTracer({ projectName: langSmithOptions?.projectName ?? 'default' }, logger), + ...getLangSmithTracer({ ...langSmithOptions, logger }), + ], + }; + const graph = await getTranslateRuleGraph({ model }); + const translateRuleState = await graph.invoke(parameters, options); + + const { response } = translateRuleState as TranslateRuleState; + const migration: SplunkRuleMigration = { + ...splunkRule, + uuid: '1234', + elasticQuery: response, + elasticQueryLanguage: 'esql', + status: 'finished', + translationState: 'translated:complete', + summary: 'This is a summary', + messages: [], + }; + + return res.ok({ body: { migration } }); + } catch (err) { + return res.badRequest({ + body: err.message, + }); + } + } + ); +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/util/find_missing_resources.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/util/find_missing_resources.ts new file mode 100644 index 0000000000000..7d918f0ac3929 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/util/find_missing_resources.ts @@ -0,0 +1,26 @@ +/* + * 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. + */ +const lookupRegex = /\b(?:lookup|inputlookup)\s+([^\s|]+)/g; +const macroRegex = /[\$`](\w+)[\$`]/g; + +export const findMissingResources = (splSearch: string): string[] => { + const missingResources: string[] = []; + + // lookups + let lookupMatch; + while ((lookupMatch = lookupRegex.exec(splSearch)) !== null) { + missingResources.push(`lookup:${lookupMatch[1]}`); + } + + // macros + let macroMatch; + while ((macroMatch = macroRegex.exec(splSearch)) !== null) { + missingResources.push(`macro:${macroMatch[0]}`); + } + + return missingResources; +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/index.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/index.ts new file mode 100644 index 0000000000000..ebf1ece565098 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/index.ts @@ -0,0 +1,8 @@ +/* + * 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 { registerSplunkTranslateRuleRoute } from './api/translate_rule'; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/types.ts new file mode 100644 index 0000000000000..1fec1c76430eb --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/types.ts @@ -0,0 +1,6 @@ +/* + * 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. + */ From 69df2f202c413523fdc28537b1174ffff7d2bd00 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 8 Oct 2024 20:36:26 +0200 Subject: [PATCH 10/17] initial version --- x-pack/plugins/security_solution/kibana.jsonc | 3 +- .../rules/translate_rule/agent/graph.ts | 52 ++++++++++++++++--- .../rules/translate_rule/agent/prompts.ts | 33 +++++------- .../rules/translate_rule/agent/state.ts | 10 +++- .../rules/translate_rule/agent/types.ts | 2 + .../translate_rule/api/translate_rule.ts | 41 +++++++++------ .../server/plugin_contract.ts | 2 + .../server/request_context_factory.ts | 1 + .../plugins/security_solution/server/types.ts | 2 + 9 files changed, 102 insertions(+), 44 deletions(-) diff --git a/x-pack/plugins/security_solution/kibana.jsonc b/x-pack/plugins/security_solution/kibana.jsonc index 075da90b44a0f..6fb931054e2d5 100644 --- a/x-pack/plugins/security_solution/kibana.jsonc +++ b/x-pack/plugins/security_solution/kibana.jsonc @@ -54,7 +54,8 @@ "savedSearch", "unifiedDocViewer", "charts", - "entityManager" + "entityManager", + "inference" ], "optionalPlugins": [ "encryptedSavedObjects", diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts index 4b0ce26cc29e5..f73166b4ca36e 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts @@ -6,27 +6,67 @@ */ import { END, START, StateGraph } from '@langchain/langgraph'; +import { ToolNode } from '@langchain/langgraph/prebuilt'; // import { StringOutputParser } from '@langchain/core/output_parsers'; +import type { StructuredTool } from '@langchain/core/tools'; +import type { AIMessage } from '@langchain/core/messages'; import { translateRuleState } from './state'; import type { TranslateRuleGraphParams, TranslateRuleNodeParams, TranslateRuleState, } from './types'; +import { TRANSLATE_RULE_MAIN_PROMPT } from './prompts'; const callModel = async ({ state, model, }: TranslateRuleNodeParams): Promise> => { - // TODO: - return state; + // const initialMessages = await TRANSLATE_RULE_MAIN_PROMPT.format({ + // splunkRuleTitle: state.splunkRuleTitle, + // splunkRuleDescription: state.splunkRuleDescription, + // splunkRuleQuery: state.splunkRuleQuery, + // }); + const translateCall = TRANSLATE_RULE_MAIN_PROMPT.pipe((initialMessages) => + model.invoke([...initialMessages.toChatMessages(), ...state.messages]) + ); + + const response = await translateCall.invoke({ + splunkRuleTitle: state.splunkRuleTitle, + splunkRuleDescription: state.splunkRuleDescription, + splunkRuleQuery: state.splunkRuleQuery, + }); + return { messages: [response] }; }; -export async function getTranslateRuleGraph({ model }: TranslateRuleGraphParams) { - const matchPrebuiltRuleGraph = new StateGraph(translateRuleState) +// Define the function that determines whether to continue or not +// We can extract the state typing via `StateAnnotation.State` +function toolConditionalEdge(state: TranslateRuleState) { + const messages = state.messages; + const lastMessage = messages[messages.length - 1] as AIMessage; + // If the LLM makes a tool call, then we route to the "tools" node + if (lastMessage.tool_calls?.length) { + return 'tools'; + } + return END; +} + +export async function getTranslateRuleGraph({ + model, + esqlKnowledgeBaseTool, +}: TranslateRuleGraphParams) { + if (model.bindTools === undefined) { + throw new Error(`The ${model.name} model does not support tools`); + } + const tools: StructuredTool[] = [esqlKnowledgeBaseTool]; + model.bindTools(tools); + + const translateRuleGraph = new StateGraph(translateRuleState) .addNode('callModel', (state: TranslateRuleState) => callModel({ state, model })) + .addNode('tools', new ToolNode(tools)) .addEdge(START, 'callModel') - .addEdge('callModel', END); + .addConditionalEdges('callModel', toolConditionalEdge) + .addEdge('tools', 'callModel'); - return matchPrebuiltRuleGraph.compile(); + return translateRuleGraph.compile(); } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts index 4cae08c1da6ff..d7ea53c866455 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -6,32 +6,27 @@ */ import { ChatPromptTemplate } from '@langchain/core/prompts'; -export const MATCH_PREBUILT_RULE_PROMPT = ChatPromptTemplate.fromMessages([ + +export const TRANSLATE_RULE_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ 'system', - `You are an expert assistant in Cybersecurity, your task is to help migrating a SIEM detection rule, from Splunk Security to Elastic Security. -You will be provided with a Splunk Detection Rule name by the user, your goal is to try find an Elastic Detection Rule that covers the same threat, if any. -The list of Elastic Detection Rules suggested is provided in the context below. + `You are a helpful assistant for translating queries from SPL (Splunk Search Processing Language) to Elastic ES|QL queries. +Your goal is to construct the equivalent ES|QL query given a SPL query. - -If there is no Elastic rule in the list that covers the same threat, answer only with the string: no_match -If there is one Elastic rule in the list that covers the same threat, answer only with its name without any further explanation. -If there are multiple rules in the list that cover the same threat, answer with the most specific of them, for example: "Linux User Account Creation" is more specific than "User Account Creation". - +VERY IMPORTANT: Use the provided tools to construct and validate the ES|QL query, do not make assumptions about the ES|QL queries. - - -{elasticSecurityRules} - -`, +The final response should be the ES|QL query inside a esql code block like: +\`\`\`esql + +\`\`\` +`, ], [ 'human', - `The Splunk Detection Rule is: - -{splunkRuleTitle} - + `The SPL query is: +\`\`\`spl +{splunkRuleQuery} +\`\`\` `, ], - ['ai', 'Please find the answer below:'], ]); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts index 516094d07ae2b..6f6023c9893cb 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts @@ -5,10 +5,16 @@ * 2.0. */ -import { Annotation } from '@langchain/langgraph'; +import type { BaseMessage } from '@langchain/core/messages'; +import { Annotation, messagesStateReducer } from '@langchain/langgraph'; export const translateRuleState = Annotation.Root({ - splunkRuleTitle: Annotation(), + messages: Annotation({ + reducer: messagesStateReducer, + default: () => [], + }), + splunkRuleTitle: Annotation, splunkRuleDescription: Annotation, + splunkRuleQuery: Annotation, response: Annotation, }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts index 4e3dfac6b781d..bf7a5eee44c6a 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts @@ -6,12 +6,14 @@ */ import type { ChatModel } from '../../../../actions_client_chat'; +import type { ESQLKnowledgeBaseTool } from '../../../../tools/esql_knowledge_base_tool'; import type { translateRuleState } from './state'; export type TranslateRuleState = typeof translateRuleState.State; export interface TranslateRuleGraphParams { model: ChatModel; + esqlKnowledgeBaseTool: ESQLKnowledgeBaseTool; } export interface TranslateRuleNodeParams { diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts index dcefc33fe45fa..2124b55bf1b15 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts @@ -10,6 +10,7 @@ import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import { APMTracer } from '@kbn/langchain/server/tracers/apm'; +import { HumanMessage, SystemMessage } from '@langchain/core/messages'; import type { SplunkRuleMigration } from '../../../../../../../common/api/siem_migrations/splunk/rules/splunk_rule.gen'; import type { SplunkRuleMigrationTranslateRuleResponse } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; import { SplunkRuleMigrationTranslateRuleRequestBody } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; @@ -18,6 +19,7 @@ import { SPLUNK_TRANSLATE_RULE_PATH } from '../../../../../../../common/api/siem import { getTranslateRuleGraph } from '../agent/graph'; import type { TranslateRuleState } from '../agent/types'; import { ActionsClientChat } from '../../../../actions_client_chat'; +import { getESQLKnowledgeBaseTool } from '../../../../tools/esql_knowledge_base_tool'; type SplunkTranslateRuleRouteResponse = IKibanaResponse; @@ -45,10 +47,16 @@ export const registerSplunkTranslateRuleRoute = ( async (context, req, res): Promise => { const { langSmithOptions, splunkRule, connectorId } = req.body; try { - const ctx = await context.resolve(['core', 'actions']); + const ctx = await context.resolve(['core', 'actions', 'securitySolution']); - const actionsClient = ctx.actions.getActionsClient(); + const inferenceClient = ctx.securitySolution.getInferenceClient({ request: req }); + const esqlKnowledgeBaseTool = getESQLKnowledgeBaseTool({ + inferenceClient, + connectorId, + logger, + }); + const actionsClient = ctx.actions.getActionsClient(); const actionsClientChat = new ActionsClientChat(connectorId, actionsClient, logger); const model = await actionsClientChat.createModel({ signal: getRequestAbortedSignal(req.events.aborted$), @@ -58,6 +66,7 @@ export const registerSplunkTranslateRuleRoute = ( const parameters: Partial = { splunkRuleTitle: splunkRule.title, splunkRuleDescription: splunkRule.description, + splunkRuleQuery: splunkRule.splSearch, }; const options = { @@ -66,22 +75,22 @@ export const registerSplunkTranslateRuleRoute = ( ...getLangSmithTracer({ ...langSmithOptions, logger }), ], }; - const graph = await getTranslateRuleGraph({ model }); + const graph = await getTranslateRuleGraph({ model, esqlKnowledgeBaseTool }); const translateRuleState = await graph.invoke(parameters, options); - const { response } = translateRuleState as TranslateRuleState; - const migration: SplunkRuleMigration = { - ...splunkRule, - uuid: '1234', - elasticQuery: response, - elasticQueryLanguage: 'esql', - status: 'finished', - translationState: 'translated:complete', - summary: 'This is a summary', - messages: [], - }; + // const { response, messages } = translateRuleState as TranslateRuleState; + // const migration: SplunkRuleMigration = { + // ...splunkRule, + // uuid: '1234', + // elasticQuery: response, + // elasticQueryLanguage: 'esql', + // status: 'finished', + // translationState: 'translated:complete', + // summary: 'This is a summary', + // messages, + // }; - return res.ok({ body: { migration } }); + return res.ok({ body: { messages: translateRuleState.messages } }); } catch (err) { return res.badRequest({ body: err.message, @@ -89,4 +98,4 @@ export const registerSplunkTranslateRuleRoute = ( } } ); -}; +}; \ No newline at end of file diff --git a/x-pack/plugins/security_solution/server/plugin_contract.ts b/x-pack/plugins/security_solution/server/plugin_contract.ts index c7ec67c1b07fc..c178f0654d9bd 100644 --- a/x-pack/plugins/security_solution/server/plugin_contract.ts +++ b/x-pack/plugins/security_solution/server/plugin_contract.ts @@ -45,6 +45,7 @@ import type { SharePluginStart } from '@kbn/share-plugin/server'; import type { GuidedOnboardingPluginSetup } from '@kbn/guided-onboarding-plugin/server'; import type { PluginSetup as UnifiedSearchServerPluginSetup } from '@kbn/unified-search-plugin/server'; import type { ElasticAssistantPluginStart } from '@kbn/elastic-assistant-plugin/server'; +import type { InferenceServerStart } from '@kbn/inference-plugin/server'; import type { ProductFeaturesService } from './lib/product_features_service/product_features_service'; import type { ExperimentalFeatures } from '../common'; @@ -88,6 +89,7 @@ export interface SecuritySolutionPluginStartDependencies { telemetry?: TelemetryPluginStart; share: SharePluginStart; actions: ActionsPluginStartContract; + inference: InferenceServerStart; } export interface SecuritySolutionPluginSetup { diff --git a/x-pack/plugins/security_solution/server/request_context_factory.ts b/x-pack/plugins/security_solution/server/request_context_factory.ts index 829509bb59817..597466a483613 100644 --- a/x-pack/plugins/security_solution/server/request_context_factory.ts +++ b/x-pack/plugins/security_solution/server/request_context_factory.ts @@ -214,6 +214,7 @@ export class RequestContextFactory implements IRequestContextFactory { }), }); }), + getInferenceClient: startPlugins.inference.getClient, }; } } diff --git a/x-pack/plugins/security_solution/server/types.ts b/x-pack/plugins/security_solution/server/types.ts index 31e10b70adbcf..4e4ecf4b446fa 100644 --- a/x-pack/plugins/security_solution/server/types.ts +++ b/x-pack/plugins/security_solution/server/types.ts @@ -20,6 +20,7 @@ import type { AlertsClient, IRuleDataService } from '@kbn/rule-registry-plugin/s import type { Readable } from 'stream'; import type { AuditLogger } from '@kbn/security-plugin-types-server'; +import type { InferenceServerStart } from '@kbn/inference-plugin/server'; import type { Immutable } from '../common/endpoint/types'; import { AppClient } from './client'; import type { ConfigType } from './config'; @@ -57,6 +58,7 @@ export interface SecuritySolutionApiRequestHandlerContext { getRiskScoreDataClient: () => RiskScoreDataClient; getAssetCriticalityDataClient: () => AssetCriticalityDataClient; getEntityStoreDataClient: () => EntityStoreDataClient; + getInferenceClient: InferenceServerStart['getClient']; } export type SecuritySolutionRequestHandlerContext = CustomRequestHandlerContext<{ From f9de8e45e7d2b6376240e9052d07cd205438408d Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 8 Oct 2024 20:40:09 +0200 Subject: [PATCH 11/17] esql KB tool wrapper --- .../translate_rule/agent/system_prompt.txt | 105 ++++++++++++++++++ .../tools/esql_knowledge_base_tool.ts | 73 ++++++++++++ 2 files changed, 178 insertions(+) create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/system_prompt.txt create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/system_prompt.txt b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/system_prompt.txt new file mode 100644 index 0000000000000..8691f4b732d3a --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/system_prompt.txt @@ -0,0 +1,105 @@ +You are a helpful assistant for translating queries from SPL (Splunk Search Processing Language) to Elastic ES|QL queries. +Your goal is to construct the equivalent ES|QL query given a SPL query. + +VERY IMPORTANT: Use the provided tools to construct and execute ES|QL query, do not make assumptions about the ES|QL queries. + +The final response should contain the ES|QL query inside a esql code block like: +```esql + +``` + +# Usage examples + +Here are some examples of ES|QL queries: + +```esql +FROM employees +| WHERE country == "NL" AND gender == "M" +| STATS COUNT(*) +``` + +```esql +FROM employees +| EVAL trunk_worked_seconds = avg_worked_seconds / 100000000 * 100000000 +| STATS c = count(languages.long) BY languages.long, trunk_worked_seconds +| SORT c desc, languages.long, trunk_worked_seconds +``` + +*Extracting structured data from logs using DISSECT* +```esql +ROW a = "2023-01-23T12:15:00.000Z - some text - 127.0.0.1" +| DISSECT a "%{date} - %{msg} - %{ip}" +| KEEP date, msg, ip +| EVAL date = TO_DATETIME(date) +``` + +```esql +FROM employees +| WHERE first_name LIKE "?b*" +| STATS doc_count = COUNT(*) by first_name, last_name +| SORT doc_count DESC +| KEEP first_name, last_name +``` + +**Returning average salary per hire date with 20 buckets** +```esql +FROM employees +| WHERE hire_date >= "1985-01-01T00:00:00Z" AND hire_date < "1986-01-01T00:00:00Z" +| STATS avg_salary = AVG(salary) BY date_bucket = BUCKET(hire_date, 20, "1985-01-01T00:00:00Z", "1986-01-01T00:00:00Z") +| SORT bucket +``` + +**Returning number of employees grouped by buckets of salary** +```esql +FROM employees +| WHERE hire_date >= "1985-01-01T00:00:00Z" AND hire_date < "1986-01-01T00:00:00Z" +| STATS c = COUNT(1) BY b = BUCKET(salary, 5000.) +| SORT b +``` + +```esql +FROM employees +| EVAL is_recent_hire = CASE(hire_date <= "2023-01-01T00:00:00Z", 1, 0) +| STATS total_recent_hires = SUM(is_recent_hire), total_hires = COUNT(*) BY country +| EVAL recent_hiring_rate = total_recent_hires / total_hires +``` + +```esql +FROM logs-* +| WHERE @timestamp <= NOW() - 24 hours +// convert a keyword field into a numeric field to aggregate over it +| EVAL is_5xx = CASE(http.response.status_code >= 500, 1, 0) +// count total events and failed events to calculate a rate +| STATS total_events = COUNT(*), total_failures = SUM(is_5xx) BY host.hostname, bucket = BUCKET(@timestamp, 1 hour) +| EVAL failure_rate_per_host = total_failures / total_events +| DROP total_events, total_failures +``` + +```esql +FROM logs-* +| WHERE @timestamp <= NOW() - 24 hours +| STATS count = COUNT(*) BY log.level +| SORT count DESC +``` + +**Returning all first names for each first letter** +```esql +FROM employees +| EVAL first_letter = SUBSTRING(first_name, 0, 1) +| STATS first_name = MV_SORT(VALUES(first_name)) BY first_letter +| SORT first_letter +``` + +```esql +FROM employees +| WHERE still_hired == true +| EVAL hired = DATE_FORMAT("YYYY", hire_date) +| STATS avg_salary = AVG(salary) BY languages +| EVAL avg_salary = ROUND(avg_salary) +| EVAL lang_code = TO_STRING(languages) +| ENRICH languages_policy ON lang_code WITH lang = language_name +| WHERE lang IS NOT NULL +| KEEP avg_salary, lang +| SORT avg_salary ASC +| LIMIT 3 +``` diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts new file mode 100644 index 0000000000000..638c070cd59ae --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts @@ -0,0 +1,73 @@ +/* + * 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 { Logger } from '@kbn/core/server'; +import { naturalLanguageToEsql, type InferenceClient } from '@kbn/inference-plugin/server'; +import type { StructuredTool } from '@langchain/core/tools'; +import { tool } from '@langchain/core/tools'; +import { lastValueFrom } from 'rxjs'; +import { z } from '@kbn/zod'; + +const TOOL_NAME = 'answer_ESQL_questions'; +const schema = z.object({ + question: z.string().describe(`The exact question about ES|QL`), +}); +type Schema = typeof schema; + +const toolParams = { + name: TOOL_NAME, + description: `You MUST use the "${TOOL_NAME}" function when the user wants to: +- run any arbitrary ES|QL query +- convert queries from another language to ES|QL +- asks general questions about ES|QL + +DO NOT UNDER ANY CIRCUMSTANCES generate ES|QL queries or explain anything about the ES|QL query language yourself. +DO NOT UNDER ANY CIRCUMSTANCES try to correct an ES|QL query yourself - always use the "${TOOL_NAME}" function for this. + +Even if the "${TOOL_NAME}" function was used before that, follow it up with the "${TOOL_NAME}" function. If a query fails, do not attempt to correct it yourself. Again you should call the "${TOOL_NAME}" function, +even if it has been called before.`, + schema, + tags: ['esql', 'query-generation', 'knowledge-base'], +}; + +export type ESQLKnowledgeBaseTool = StructuredTool; + +interface GetESQLKnowledgeBaseToolParams { + inferenceClient: InferenceClient; + connectorId: string; + logger: Logger; +} +export const getESQLKnowledgeBaseTool = ({ + inferenceClient: client, + connectorId, + logger, +}: GetESQLKnowledgeBaseToolParams): ESQLKnowledgeBaseTool => { + const callNaturalLanguageToEsql = async (question: string) => { + return lastValueFrom( + naturalLanguageToEsql({ + client, + connectorId, + input: question, + logger: { + debug: (source) => { + logger.debug(typeof source === 'function' ? source() : source); + }, + }, + }) + ); + }; + + const esqlKBTool = tool(async (input) => { + const generateEvent = await callNaturalLanguageToEsql(input.question); + const answer = generateEvent.content ?? 'An error occurred in the tool'; + + logger.debug(`Received response from NL to ESQL tool: ${answer}`); + return answer; + }, toolParams); + + return esqlKBTool; +}; From fee37a21cff6d5c4ea5a2878a44e88904efd76d5 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Fri, 11 Oct 2024 14:08:40 +0200 Subject: [PATCH 12/17] improve prompt --- .../rules/translate_rule/agent/graph.ts | 55 ++++++++++++++----- .../rules/translate_rule/agent/prompts.ts | 37 +++++++++---- .../rules/translate_rule/agent/state.ts | 5 +- .../translate_rule/api/translate_rule.ts | 4 +- 4 files changed, 72 insertions(+), 29 deletions(-) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts index f73166b4ca36e..e9c542a03077a 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts @@ -22,20 +22,16 @@ const callModel = async ({ state, model, }: TranslateRuleNodeParams): Promise> => { - // const initialMessages = await TRANSLATE_RULE_MAIN_PROMPT.format({ + const response = await model.invoke(state.messages); + // const translateCall = TRANSLATE_RULE_MAIN_PROMPT.pipe((initialMessages) => + // model.invoke([...initialMessages.toChatMessages(), ...state.messages]) + // ); + + // const response = await translateCall.invoke({ // splunkRuleTitle: state.splunkRuleTitle, // splunkRuleDescription: state.splunkRuleDescription, // splunkRuleQuery: state.splunkRuleQuery, // }); - const translateCall = TRANSLATE_RULE_MAIN_PROMPT.pipe((initialMessages) => - model.invoke([...initialMessages.toChatMessages(), ...state.messages]) - ); - - const response = await translateCall.invoke({ - splunkRuleTitle: state.splunkRuleTitle, - splunkRuleDescription: state.splunkRuleDescription, - splunkRuleQuery: state.splunkRuleQuery, - }); return { messages: [response] }; }; @@ -51,10 +47,27 @@ function toolConditionalEdge(state: TranslateRuleState) { return END; } -export async function getTranslateRuleGraph({ - model, - esqlKnowledgeBaseTool, -}: TranslateRuleGraphParams) { +// export async function getTranslateRuleGraph({ +// model, +// esqlKnowledgeBaseTool, +// }: TranslateRuleGraphParams) { +// if (model.bindTools === undefined) { +// throw new Error(`The ${model.name} model does not support tools`); +// } +// const tools: StructuredTool[] = [esqlKnowledgeBaseTool]; +// model.bindTools(tools); + +// const translateRuleGraph = new StateGraph(translateRuleState) +// .addNode('callModel', (state: TranslateRuleState) => callModel({ state, model })) +// .addNode('tools', new ToolNode(tools)) +// .addEdge(START, 'callModel') +// .addConditionalEdges('callModel', toolConditionalEdge) +// .addEdge('tools', 'callModel'); + +// return translateRuleGraph.compile(); +// } + +export function getTranslateRuleGraph({ model, esqlKnowledgeBaseTool }: TranslateRuleGraphParams) { if (model.bindTools === undefined) { throw new Error(`The ${model.name} model does not support tools`); } @@ -68,5 +81,17 @@ export async function getTranslateRuleGraph({ .addConditionalEdges('callModel', toolConditionalEdge) .addEdge('tools', 'callModel'); - return translateRuleGraph.compile(); + const graph = translateRuleGraph.compile(); + + const invoke: typeof graph.invoke = async (state, options) => { + const mainPrompts = await TRANSLATE_RULE_MAIN_PROMPT.formatMessages({ + splunkRuleTitle: state.splunkRuleTitle, + splunkRuleDescription: state.splunkRuleDescription, + splunkRuleQuery: state.splunkRuleQuery, + }); + + return graph.invoke({ ...state, messages: mainPrompts }, options); + }; + + return { invoke }; } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts index d7ea53c866455..3913330c03204 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -10,23 +10,40 @@ import { ChatPromptTemplate } from '@langchain/core/prompts'; export const TRANSLATE_RULE_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ 'system', - `You are a helpful assistant for translating queries from SPL (Splunk Search Processing Language) to Elastic ES|QL queries. -Your goal is to construct the equivalent ES|QL query given a SPL query. + `You are a helpful cybersecurity (SIEM) expert. Your task is to translate "detection rules" from Splunk to Elastic Security. +You will be given a Splunk rule information, such as the title, description and the SPL (Search Processing Language) query, and you will need to construct the equivalent ES|QL (Elasticsearch Query Language) query. -VERY IMPORTANT: Use the provided tools to construct and validate the ES|QL query, do not make assumptions about the ES|QL queries. +VERY IMPORTANT: Use the provided tools to convert and validate the ES|QL query, do not make assumptions about the ES|QL queries. -The final response should be the ES|QL query inside a esql code block like: -\`\`\`esql - -\`\`\` +The final response should contain two parts: + +1- A summary of the translation process in a human-readable format, including any description and the goal of the original rule. And why it was translated in a certain way. +The summary should be inside a block like: +<> +[the summary goes here] +<> + +2- The translated ES|QL query inside a esql code block like: +<> +[the translated query goes here] +<> `, ], [ 'human', - `The SPL query is: -\`\`\`spl + `Translate this Splunk rule into an ES|QL query: +<> +{splunkRuleTitle} +<> + +<> +{splunkRuleDescription} +<> + +<> {splunkRuleQuery} -\`\`\` +<> + `, ], ]); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts index 6f6023c9893cb..ac143d80ac5a5 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts @@ -13,7 +13,10 @@ export const translateRuleState = Annotation.Root({ reducer: messagesStateReducer, default: () => [], }), - splunkRuleTitle: Annotation, + splunkRuleTitle: Annotation({ + reducer: (a, b) => b ?? a, + default: () => '', + }), splunkRuleDescription: Annotation, splunkRuleQuery: Annotation, response: Annotation, diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts index 2124b55bf1b15..27c9df341d312 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts @@ -10,8 +10,6 @@ import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import { APMTracer } from '@kbn/langchain/server/tracers/apm'; -import { HumanMessage, SystemMessage } from '@langchain/core/messages'; -import type { SplunkRuleMigration } from '../../../../../../../common/api/siem_migrations/splunk/rules/splunk_rule.gen'; import type { SplunkRuleMigrationTranslateRuleResponse } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; import { SplunkRuleMigrationTranslateRuleRequestBody } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; import type { SecuritySolutionPluginRouter } from '../../../../../../types'; @@ -75,7 +73,7 @@ export const registerSplunkTranslateRuleRoute = ( ...getLangSmithTracer({ ...langSmithOptions, logger }), ], }; - const graph = await getTranslateRuleGraph({ model, esqlKnowledgeBaseTool }); + const graph = getTranslateRuleGraph({ model, esqlKnowledgeBaseTool }); const translateRuleState = await graph.invoke(parameters, options); // const { response, messages } = translateRuleState as TranslateRuleState; From 7000aec3165462e228c46a314f8afe842086e211 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 15 Oct 2024 11:06:50 +0200 Subject: [PATCH 13/17] improve tool --- .../rules/translate_rule/agent/prompts.ts | 18 +-- .../agent/tools/esql_knowledge_base_tool.ts | 107 ++++++++++++++++++ .../rules/translate_rule/agent/types.ts | 2 +- .../translate_rule/api/translate_rule.ts | 2 +- .../tools/esql_knowledge_base_tool.ts | 73 ------------ 5 files changed, 112 insertions(+), 90 deletions(-) create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts delete mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts index 3913330c03204..3f30275562677 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -11,22 +11,10 @@ export const TRANSLATE_RULE_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ 'system', `You are a helpful cybersecurity (SIEM) expert. Your task is to translate "detection rules" from Splunk to Elastic Security. -You will be given a Splunk rule information, such as the title, description and the SPL (Search Processing Language) query, and you will need to construct the equivalent ES|QL (Elasticsearch Query Language) query. +You will be provided with a Splunk rule information: the title, description and the SPL (Search Processing Language) query. +You will also be provided with the tools necessary to translate the ES|QL query and summarize the translation, do not make assumptions about the ES|QL queries. -VERY IMPORTANT: Use the provided tools to convert and validate the ES|QL query, do not make assumptions about the ES|QL queries. - -The final response should contain two parts: - -1- A summary of the translation process in a human-readable format, including any description and the goal of the original rule. And why it was translated in a certain way. -The summary should be inside a block like: -<> -[the summary goes here] -<> - -2- The translated ES|QL query inside a esql code block like: -<> -[the translated query goes here] -<> +The output should be an ES|QL query that is equivalent to the provided Splunk SPL query and a markdown summary of the translation process followed. `, ], [ diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts new file mode 100644 index 0000000000000..e442f3d8d48e5 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts @@ -0,0 +1,107 @@ +/* + * 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 { Logger } from '@kbn/core/server'; +import { naturalLanguageToEsql, type InferenceClient } from '@kbn/inference-plugin/server'; +import type { StructuredTool } from '@langchain/core/tools'; +import { tool } from '@langchain/core/tools'; +import { lastValueFrom } from 'rxjs'; +import { z } from '@kbn/zod'; + +const TOOL_NAME = 'esql_translator'; +const schema = z.object({ + splQuery: z.string().describe(`The exact SPL query to translate to ES|QL`), + title: z.string().describe(`The title of the splunk rule`), + description: z.string().describe(`The description of the splunk rule`), +}); +type Schema = typeof schema; +type SchemaInput = z.output; + +const toolParams = { + name: TOOL_NAME, + description: `You MUST use the "${TOOL_NAME}" function to convert queries from SPL language to ES|QL. +The output will contain the ES|QL equivalent inside a \`\`\`esql code block and a markdown summary of the translation process inside a \`\`\`markdown code block. + +IMPORTANT: The SPL query must be passed directly without any wrapping or modification. +`, + schema, + tags: ['esql', 'query-translation', 'knowledge-base'], +}; + +const createInputPrompt = ({ + splQuery, + title, + description, +}: SchemaInput) => `Translate the following SPL (Search Processing Language) query to ES|QL in order to be used in a Security detection rule: + +\`\`\`spl +${splQuery} +\`\`\` + +This SPL query is part of a Splunk rule with the following title and description: + + +${title} + + + +${description} + + +Along with the translated ES|QL query, you should also provide a summary of the translation process you followed, in markdown format. +The output will be parsed using a regular expression, please format the output using the following guidelines: + +The summary of the translation process must be placed in a "markdown" code block, example: + +\`\`\`markdown +[the summary goes here] +\`\`\` + +The ES|QL translated query must be placed in a "esql" code block, example: + +\`\`\`esql +[the translated query goes here] +\`\`\` +`; + +export type ESQLKnowledgeBaseTool = StructuredTool; + +interface GetESQLKnowledgeBaseToolParams { + inferenceClient: InferenceClient; + connectorId: string; + logger: Logger; +} +export const getESQLKnowledgeBaseTool = ({ + inferenceClient: client, + connectorId, + logger, +}: GetESQLKnowledgeBaseToolParams): ESQLKnowledgeBaseTool => { + const callNaturalLanguageToEsql = async (input: SchemaInput) => { + return lastValueFrom( + naturalLanguageToEsql({ + client, + connectorId, + input: createInputPrompt(input), + logger: { + debug: (source) => { + logger.debug(typeof source === 'function' ? source() : source); + }, + }, + }) + ); + }; + + const esqlKBTool = tool(async (input) => { + const generateEvent = await callNaturalLanguageToEsql(input); + const answer = generateEvent.content ?? 'An error occurred in the tool'; + + logger.debug(`Received response from NL to ESQL tool: ${answer}`); + return answer; + }, toolParams); + + return esqlKBTool; +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts index bf7a5eee44c6a..4c83598b95b0f 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts @@ -6,7 +6,7 @@ */ import type { ChatModel } from '../../../../actions_client_chat'; -import type { ESQLKnowledgeBaseTool } from '../../../../tools/esql_knowledge_base_tool'; +import type { ESQLKnowledgeBaseTool } from './tools/esql_knowledge_base_tool'; import type { translateRuleState } from './state'; export type TranslateRuleState = typeof translateRuleState.State; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts index 27c9df341d312..3346164851a8e 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts @@ -17,7 +17,7 @@ import { SPLUNK_TRANSLATE_RULE_PATH } from '../../../../../../../common/api/siem import { getTranslateRuleGraph } from '../agent/graph'; import type { TranslateRuleState } from '../agent/types'; import { ActionsClientChat } from '../../../../actions_client_chat'; -import { getESQLKnowledgeBaseTool } from '../../../../tools/esql_knowledge_base_tool'; +import { getESQLKnowledgeBaseTool } from '../agent/tools/esql_knowledge_base_tool'; type SplunkTranslateRuleRouteResponse = IKibanaResponse; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts deleted file mode 100644 index 638c070cd59ae..0000000000000 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/tools/esql_knowledge_base_tool.ts +++ /dev/null @@ -1,73 +0,0 @@ -/* - * 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 { Logger } from '@kbn/core/server'; -import { naturalLanguageToEsql, type InferenceClient } from '@kbn/inference-plugin/server'; -import type { StructuredTool } from '@langchain/core/tools'; -import { tool } from '@langchain/core/tools'; -import { lastValueFrom } from 'rxjs'; -import { z } from '@kbn/zod'; - -const TOOL_NAME = 'answer_ESQL_questions'; -const schema = z.object({ - question: z.string().describe(`The exact question about ES|QL`), -}); -type Schema = typeof schema; - -const toolParams = { - name: TOOL_NAME, - description: `You MUST use the "${TOOL_NAME}" function when the user wants to: -- run any arbitrary ES|QL query -- convert queries from another language to ES|QL -- asks general questions about ES|QL - -DO NOT UNDER ANY CIRCUMSTANCES generate ES|QL queries or explain anything about the ES|QL query language yourself. -DO NOT UNDER ANY CIRCUMSTANCES try to correct an ES|QL query yourself - always use the "${TOOL_NAME}" function for this. - -Even if the "${TOOL_NAME}" function was used before that, follow it up with the "${TOOL_NAME}" function. If a query fails, do not attempt to correct it yourself. Again you should call the "${TOOL_NAME}" function, -even if it has been called before.`, - schema, - tags: ['esql', 'query-generation', 'knowledge-base'], -}; - -export type ESQLKnowledgeBaseTool = StructuredTool; - -interface GetESQLKnowledgeBaseToolParams { - inferenceClient: InferenceClient; - connectorId: string; - logger: Logger; -} -export const getESQLKnowledgeBaseTool = ({ - inferenceClient: client, - connectorId, - logger, -}: GetESQLKnowledgeBaseToolParams): ESQLKnowledgeBaseTool => { - const callNaturalLanguageToEsql = async (question: string) => { - return lastValueFrom( - naturalLanguageToEsql({ - client, - connectorId, - input: question, - logger: { - debug: (source) => { - logger.debug(typeof source === 'function' ? source() : source); - }, - }, - }) - ); - }; - - const esqlKBTool = tool(async (input) => { - const generateEvent = await callNaturalLanguageToEsql(input.question); - const answer = generateEvent.content ?? 'An error occurred in the tool'; - - logger.debug(`Received response from NL to ESQL tool: ${answer}`); - return answer; - }, toolParams); - - return esqlKBTool; -}; From dde06d0dcb2b0fede69353181fe116d6fec4e632 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Tue, 15 Oct 2024 18:20:44 +0200 Subject: [PATCH 14/17] code cleaning --- .../siem_migrations/actions_client_chat.ts | 44 +++++++++------- .../rules/translate_rule/agent/graph.ts | 33 +----------- .../rules/translate_rule/agent/prompts.ts | 12 +++-- .../rules/translate_rule/agent/state.ts | 5 +- ...e_base_tool.ts => esql_translator_tool.ts} | 51 ++++++------------- .../rules/translate_rule/agent/types.ts | 4 +- .../translate_rule/api/translate_rule.ts | 8 +-- 7 files changed, 57 insertions(+), 100 deletions(-) rename x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/{esql_knowledge_base_tool.ts => esql_translator_tool.ts} (59%) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts index 94d98ee8078b1..204978c901df6 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/actions_client_chat.ts @@ -5,39 +5,45 @@ * 2.0. */ +import type { ActionsClientSimpleChatModel } from '@kbn/langchain/server'; import { ActionsClientBedrockChatModel, ActionsClientChatOpenAI, - ActionsClientGeminiChatModel, - ActionsClientSimpleChatModel, + ActionsClientChatVertexAI, } from '@kbn/langchain/server'; import type { Logger } from '@kbn/core/server'; import type { ActionsClient } from '@kbn/actions-plugin/server'; import type { ActionsClientChatOpenAIParams } from '@kbn/langchain/server/language_models/chat_openai'; import type { CustomChatModelInput as ActionsClientBedrockChatModelParams } from '@kbn/langchain/server/language_models/bedrock_chat'; -import type { CustomChatModelInput as ActionsClientGeminiChatModelParams } from '@kbn/langchain/server/language_models/gemini_chat'; +import type { CustomChatModelInput as ActionsClientChatVertexAIParams } from '@kbn/langchain/server/language_models/gemini_chat'; import type { CustomChatModelInput as ActionsClientSimpleChatModelParams } from '@kbn/langchain/server/language_models/simple_chat_model'; export type ChatModel = | ActionsClientSimpleChatModel | ActionsClientChatOpenAI | ActionsClientBedrockChatModel - | ActionsClientGeminiChatModel; + | ActionsClientChatVertexAI; export type ActionsClientChatModelClass = | typeof ActionsClientSimpleChatModel | typeof ActionsClientChatOpenAI | typeof ActionsClientBedrockChatModel - | typeof ActionsClientGeminiChatModel; + | typeof ActionsClientChatVertexAI; export type ChatModelParams = Partial & Partial & Partial & - Partial & { + Partial & { /** Enables the streaming mode of the response, disabled by default */ streaming?: boolean; }; +const llmTypeDictionary: Record = { + [`.gen-ai`]: `openai`, + [`.bedrock`]: `bedrock`, + [`.gemini`]: `gemini`, +}; + export class ActionsClientChat { constructor( private readonly connectorId: string, @@ -67,21 +73,21 @@ export class ActionsClientChat { } private getLLMType(actionTypeId: string): string | undefined { - const llmTypeDictionary: Record = { - [`.gen-ai`]: `openai`, - [`.bedrock`]: `bedrock`, - [`.gemini`]: `gemini`, - }; - return llmTypeDictionary[actionTypeId]; + if (llmTypeDictionary[actionTypeId]) { + return llmTypeDictionary[actionTypeId]; + } + throw new Error(`Unknown LLM type for action type ID: ${actionTypeId}`); } private getLLMClass(llmType?: string): ActionsClientChatModelClass { - return llmType === 'openai' - ? ActionsClientChatOpenAI - : llmType === 'bedrock' - ? ActionsClientBedrockChatModel - : llmType === 'gemini' - ? ActionsClientGeminiChatModel - : ActionsClientSimpleChatModel; + switch (llmType) { + case 'bedrock': + return ActionsClientBedrockChatModel; + case 'gemini': + return ActionsClientChatVertexAI; + case 'openai': + default: + return ActionsClientChatOpenAI; + } } } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts index e9c542a03077a..868dabbbf6b7e 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts @@ -8,7 +8,6 @@ import { END, START, StateGraph } from '@langchain/langgraph'; import { ToolNode } from '@langchain/langgraph/prebuilt'; // import { StringOutputParser } from '@langchain/core/output_parsers'; -import type { StructuredTool } from '@langchain/core/tools'; import type { AIMessage } from '@langchain/core/messages'; import { translateRuleState } from './state'; import type { @@ -23,15 +22,6 @@ const callModel = async ({ model, }: TranslateRuleNodeParams): Promise> => { const response = await model.invoke(state.messages); - // const translateCall = TRANSLATE_RULE_MAIN_PROMPT.pipe((initialMessages) => - // model.invoke([...initialMessages.toChatMessages(), ...state.messages]) - // ); - - // const response = await translateCall.invoke({ - // splunkRuleTitle: state.splunkRuleTitle, - // splunkRuleDescription: state.splunkRuleDescription, - // splunkRuleQuery: state.splunkRuleQuery, - // }); return { messages: [response] }; }; @@ -47,31 +37,10 @@ function toolConditionalEdge(state: TranslateRuleState) { return END; } -// export async function getTranslateRuleGraph({ -// model, -// esqlKnowledgeBaseTool, -// }: TranslateRuleGraphParams) { -// if (model.bindTools === undefined) { -// throw new Error(`The ${model.name} model does not support tools`); -// } -// const tools: StructuredTool[] = [esqlKnowledgeBaseTool]; -// model.bindTools(tools); - -// const translateRuleGraph = new StateGraph(translateRuleState) -// .addNode('callModel', (state: TranslateRuleState) => callModel({ state, model })) -// .addNode('tools', new ToolNode(tools)) -// .addEdge(START, 'callModel') -// .addConditionalEdges('callModel', toolConditionalEdge) -// .addEdge('tools', 'callModel'); - -// return translateRuleGraph.compile(); -// } - -export function getTranslateRuleGraph({ model, esqlKnowledgeBaseTool }: TranslateRuleGraphParams) { +export function getTranslateRuleGraph({ model, tools }: TranslateRuleGraphParams) { if (model.bindTools === undefined) { throw new Error(`The ${model.name} model does not support tools`); } - const tools: StructuredTool[] = [esqlKnowledgeBaseTool]; model.bindTools(tools); const translateRuleGraph = new StateGraph(translateRuleState) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts index 3f30275562677..b8bd1d8fc43e4 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -10,16 +10,20 @@ import { ChatPromptTemplate } from '@langchain/core/prompts'; export const TRANSLATE_RULE_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ 'system', - `You are a helpful cybersecurity (SIEM) expert. Your task is to translate "detection rules" from Splunk to Elastic Security. + `You are a helpful cybersecurity (SIEM) expert agent. Your task is to translate "detection rules" from Splunk to Elastic Security. You will be provided with a Splunk rule information: the title, description and the SPL (Search Processing Language) query. -You will also be provided with the tools necessary to translate the ES|QL query and summarize the translation, do not make assumptions about the ES|QL queries. +Your goal is to translate the SPL query into an equivalent Elastic Security Query Language (ES|QL) query. -The output should be an ES|QL query that is equivalent to the provided Splunk SPL query and a markdown summary of the translation process followed. +IMPORTANT: Use the tools provided to translate the ES|QL query and summarize the translation, rather than making assumptions about how the ES|QL language works. + +The output should contain: +- First, the ES|QL query inside an \`\`\`esql code block. +- At the end, the summary of the translation process followed in markdown, starting with "## Translation Summary". `, ], [ 'human', - `Translate this Splunk rule into an ES|QL query: + `Translate this Splunk rule into an Elastic ES|QL query rule: <> {splunkRuleTitle} <> diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts index ac143d80ac5a5..6f6023c9893cb 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts @@ -13,10 +13,7 @@ export const translateRuleState = Annotation.Root({ reducer: messagesStateReducer, default: () => [], }), - splunkRuleTitle: Annotation({ - reducer: (a, b) => b ?? a, - default: () => '', - }), + splunkRuleTitle: Annotation, splunkRuleDescription: Annotation, splunkRuleQuery: Annotation, response: Annotation, diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts similarity index 59% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts index e442f3d8d48e5..9b022f93d280d 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_knowledge_base_tool.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts @@ -23,11 +23,7 @@ type SchemaInput = z.output; const toolParams = { name: TOOL_NAME, - description: `You MUST use the "${TOOL_NAME}" function to convert queries from SPL language to ES|QL. -The output will contain the ES|QL equivalent inside a \`\`\`esql code block and a markdown summary of the translation process inside a \`\`\`markdown code block. - -IMPORTANT: The SPL query must be passed directly without any wrapping or modification. -`, + description: `ALWAYS use the "${TOOL_NAME}" tool to convert the Splunk detection rule (SPL) to an Elastic ES|QL query.`, schema, tags: ['esql', 'query-translation', 'knowledge-base'], }; @@ -36,50 +32,35 @@ const createInputPrompt = ({ splQuery, title, description, -}: SchemaInput) => `Translate the following SPL (Search Processing Language) query to ES|QL in order to be used in a Security detection rule: +}: SchemaInput) => `Translate the following Splunk SPL (Search Processing Language) query rule to an ES|QL query, in order to be used as an Elastic Security detection rule: + +Splunk rule title: ${title} + +Splunk rule description: ${description} +Splunk rule SPL query: \`\`\`spl ${splQuery} \`\`\` -This SPL query is part of a Splunk rule with the following title and description: - - -${title} - - - -${description} - - Along with the translated ES|QL query, you should also provide a summary of the translation process you followed, in markdown format. -The output will be parsed using a regular expression, please format the output using the following guidelines: - -The summary of the translation process must be placed in a "markdown" code block, example: - -\`\`\`markdown -[the summary goes here] -\`\`\` - -The ES|QL translated query must be placed in a "esql" code block, example: - -\`\`\`esql -[the translated query goes here] -\`\`\` +The output should contain: +- First, the ES|QL query inside an \`\`\`esql code block. +- At the end, the summary of the translation process followed in markdown, starting with "## Translation Summary". `; -export type ESQLKnowledgeBaseTool = StructuredTool; +export type EsqlTranslatorTool = StructuredTool; -interface GetESQLKnowledgeBaseToolParams { +interface GetEsqlTranslatorToolParams { inferenceClient: InferenceClient; connectorId: string; logger: Logger; } -export const getESQLKnowledgeBaseTool = ({ +export const getEsqlTranslatorTool = ({ inferenceClient: client, connectorId, logger, -}: GetESQLKnowledgeBaseToolParams): ESQLKnowledgeBaseTool => { +}: GetEsqlTranslatorToolParams): EsqlTranslatorTool => { const callNaturalLanguageToEsql = async (input: SchemaInput) => { return lastValueFrom( naturalLanguageToEsql({ @@ -95,7 +76,7 @@ export const getESQLKnowledgeBaseTool = ({ ); }; - const esqlKBTool = tool(async (input) => { + const esqlTool = tool(async (input) => { const generateEvent = await callNaturalLanguageToEsql(input); const answer = generateEvent.content ?? 'An error occurred in the tool'; @@ -103,5 +84,5 @@ export const getESQLKnowledgeBaseTool = ({ return answer; }, toolParams); - return esqlKBTool; + return esqlTool; }; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts index 4c83598b95b0f..3221f2c5ae4b3 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts @@ -5,15 +5,15 @@ * 2.0. */ +import type { StructuredTool } from '@langchain/core/tools'; import type { ChatModel } from '../../../../actions_client_chat'; -import type { ESQLKnowledgeBaseTool } from './tools/esql_knowledge_base_tool'; import type { translateRuleState } from './state'; export type TranslateRuleState = typeof translateRuleState.State; export interface TranslateRuleGraphParams { model: ChatModel; - esqlKnowledgeBaseTool: ESQLKnowledgeBaseTool; + tools: StructuredTool[]; } export interface TranslateRuleNodeParams { diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts index 3346164851a8e..64e22943707cd 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts @@ -17,7 +17,7 @@ import { SPLUNK_TRANSLATE_RULE_PATH } from '../../../../../../../common/api/siem import { getTranslateRuleGraph } from '../agent/graph'; import type { TranslateRuleState } from '../agent/types'; import { ActionsClientChat } from '../../../../actions_client_chat'; -import { getESQLKnowledgeBaseTool } from '../agent/tools/esql_knowledge_base_tool'; +import { getEsqlTranslatorTool } from '../agent/tools/esql_translator_tool'; type SplunkTranslateRuleRouteResponse = IKibanaResponse; @@ -48,7 +48,7 @@ export const registerSplunkTranslateRuleRoute = ( const ctx = await context.resolve(['core', 'actions', 'securitySolution']); const inferenceClient = ctx.securitySolution.getInferenceClient({ request: req }); - const esqlKnowledgeBaseTool = getESQLKnowledgeBaseTool({ + const esqlTranslatorTool = getEsqlTranslatorTool({ inferenceClient, connectorId, logger, @@ -73,7 +73,7 @@ export const registerSplunkTranslateRuleRoute = ( ...getLangSmithTracer({ ...langSmithOptions, logger }), ], }; - const graph = getTranslateRuleGraph({ model, esqlKnowledgeBaseTool }); + const graph = getTranslateRuleGraph({ model, tools: [esqlTranslatorTool] }); const translateRuleState = await graph.invoke(parameters, options); // const { response, messages } = translateRuleState as TranslateRuleState; @@ -96,4 +96,4 @@ export const registerSplunkTranslateRuleRoute = ( } } ); -}; \ No newline at end of file +}; From 93a4b2855654c275beb30897d6afd3af26b475ea Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Wed, 16 Oct 2024 09:44:34 +0200 Subject: [PATCH 15/17] fix bindTools --- .../splunk/rules/translate_rule/agent/graph.ts | 6 ++++-- .../rules/translate_rule/agent/prompts.ts | 2 +- .../agent/tools/esql_translator_tool.ts | 18 +++++++++++------- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts index 868dabbbf6b7e..b80ae23dbd6c7 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts @@ -41,10 +41,12 @@ export function getTranslateRuleGraph({ model, tools }: TranslateRuleGraphParams if (model.bindTools === undefined) { throw new Error(`The ${model.name} model does not support tools`); } - model.bindTools(tools); + const modelWithTools = model.bindTools(tools); const translateRuleGraph = new StateGraph(translateRuleState) - .addNode('callModel', (state: TranslateRuleState) => callModel({ state, model })) + .addNode('callModel', (state: TranslateRuleState) => + callModel({ state, model: modelWithTools }) + ) .addNode('tools', new ToolNode(tools)) .addEdge(START, 'callModel') .addConditionalEdges('callModel', toolConditionalEdge) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts index b8bd1d8fc43e4..c7431cd93cdd9 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -14,7 +14,7 @@ export const TRANSLATE_RULE_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ You will be provided with a Splunk rule information: the title, description and the SPL (Search Processing Language) query. Your goal is to translate the SPL query into an equivalent Elastic Security Query Language (ES|QL) query. -IMPORTANT: Use the tools provided to translate the ES|QL query and summarize the translation, rather than making assumptions about how the ES|QL language works. +IMPORTANT: Always use the tools provided to translate the ES|QL query and summarize the translation, rather than making assumptions about how the ES|QL language works. The output should contain: - First, the ES|QL query inside an \`\`\`esql code block. diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts index 9b022f93d280d..9554a9d3f020c 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts @@ -8,7 +8,7 @@ import type { Logger } from '@kbn/core/server'; import { naturalLanguageToEsql, type InferenceClient } from '@kbn/inference-plugin/server'; import type { StructuredTool } from '@langchain/core/tools'; -import { tool } from '@langchain/core/tools'; +import { DynamicStructuredTool } from '@langchain/core/tools'; import { lastValueFrom } from 'rxjs'; import { z } from '@kbn/zod'; @@ -76,13 +76,17 @@ export const getEsqlTranslatorTool = ({ ); }; - const esqlTool = tool(async (input) => { - const generateEvent = await callNaturalLanguageToEsql(input); - const answer = generateEvent.content ?? 'An error occurred in the tool'; + const esqlTool = new DynamicStructuredTool({ + ...toolParams, + responseFormat: 'markdown', + func: async (input) => { + const generateEvent = await callNaturalLanguageToEsql(input); + const answer = generateEvent.content ?? 'An error occurred in the tool'; - logger.debug(`Received response from NL to ESQL tool: ${answer}`); - return answer; - }, toolParams); + logger.debug(`Received response from NL to ESQL tool: ${answer}`); + return answer; + }, + }); return esqlTool; }; From 51544875d2e6c4043f299167a1f59af0b71c38a7 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Thu, 17 Oct 2024 11:46:23 +0200 Subject: [PATCH 16/17] extract esql KB and call it directly as a graph node --- .../rules/translate_rule/agent/graph.ts | 116 ++++++++++++------ .../rules/translate_rule/agent/prompts.ts | 31 +++++ .../rules/translate_rule/agent/state.ts | 8 +- .../agent/tools/esql_translator_tool.ts | 23 ++++ .../rules/translate_rule/agent/types.ts | 10 +- .../translate_rule/api/translate_rule.ts | 19 ++- 6 files changed, 155 insertions(+), 52 deletions(-) diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts index b80ae23dbd6c7..86fddb701e75c 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/graph.ts @@ -8,61 +8,97 @@ import { END, START, StateGraph } from '@langchain/langgraph'; import { ToolNode } from '@langchain/langgraph/prebuilt'; // import { StringOutputParser } from '@langchain/core/output_parsers'; -import type { AIMessage } from '@langchain/core/messages'; +import { AIMessage } from '@langchain/core/messages'; +import type { Runnable } from '@langchain/core/runnables'; import { translateRuleState } from './state'; -import type { - TranslateRuleGraphParams, - TranslateRuleNodeParams, - TranslateRuleState, -} from './types'; -import { TRANSLATE_RULE_MAIN_PROMPT } from './prompts'; +import type { TranslateRuleGraphParams, TranslateRuleState } from './types'; +import { TRANSLATE_RULE_MAIN_PROMPT, getEsqlTranslationPrompt } from './prompts'; +import { getEsqlKnowledgeBase, type EsqlKnowledgeBaseCaller } from './esql_knowledge_base_caller'; -const callModel = async ({ - state, - model, -}: TranslateRuleNodeParams): Promise> => { - const response = await model.invoke(state.messages); - return { messages: [response] }; +type GraphNode = (state: TranslateRuleState) => Promise>; + +const initMessages: GraphNode = async (state: TranslateRuleState) => { + return { + messages: await TRANSLATE_RULE_MAIN_PROMPT.formatMessages({ + splunkRuleTitle: state.splunkRuleTitle, + splunkRuleDescription: state.splunkRuleDescription, + splunkRuleQuery: state.splunkRuleQuery, + }), + }; +}; + +const createCallModelNode = + (model: Runnable): GraphNode => + async (state) => { + const response = await model.invoke(state.messages); + return { messages: [response] }; + }; + +const createTranslationNode = (esqlKnowledgeBaseCaller: EsqlKnowledgeBaseCaller): GraphNode => { + return async (state) => { + const input = getEsqlTranslationPrompt({ + splunkRuleTitle: state.splunkRuleTitle, + splunkRuleDescription: state.splunkRuleDescription, + splunkRuleQuery: state.splunkRuleQuery, + }); + const response = await esqlKnowledgeBaseCaller(input); + return { messages: [new AIMessage(response)] }; + }; +}; + +const responseNode: GraphNode = async (state) => { + const messages = state.messages; + const lastMessage = messages[messages.length - 1] as AIMessage; + return { response: lastMessage.content as string }; }; -// Define the function that determines whether to continue or not -// We can extract the state typing via `StateAnnotation.State` function toolConditionalEdge(state: TranslateRuleState) { const messages = state.messages; const lastMessage = messages[messages.length - 1] as AIMessage; - // If the LLM makes a tool call, then we route to the "tools" node if (lastMessage.tool_calls?.length) { return 'tools'; } - return END; + return 'processResponse'; } -export function getTranslateRuleGraph({ model, tools }: TranslateRuleGraphParams) { - if (model.bindTools === undefined) { - throw new Error(`The ${model.name} model does not support tools`); - } - const modelWithTools = model.bindTools(tools); +// export function getTranslateRuleGraph({ model, tools }: TranslateRuleGraphParams) { +// if (model.bindTools === undefined) { +// throw new Error(`The ${model.name} model does not support tools`); +// } +// const modelWithTools: Runnable = model.bindTools(tools); +// const callModel = createCallModelNode(modelWithTools); +// const toolsNode = new ToolNode(tools); - const translateRuleGraph = new StateGraph(translateRuleState) - .addNode('callModel', (state: TranslateRuleState) => - callModel({ state, model: modelWithTools }) - ) - .addNode('tools', new ToolNode(tools)) - .addEdge(START, 'callModel') - .addConditionalEdges('callModel', toolConditionalEdge) - .addEdge('tools', 'callModel'); +// const translateRuleGraph = new StateGraph(translateRuleState) +// .addNode('initMessages', initMessages) +// .addNode('callModel', callModel) +// .addNode('tools', toolsNode) +// .addNode('processResponse', responseNode) - const graph = translateRuleGraph.compile(); +// .addEdge(START, 'initMessages') +// .addEdge('initMessages', 'callModel') +// .addConditionalEdges('callModel', toolConditionalEdge) +// .addEdge('tools', 'callModel') +// .addEdge('processResponse', END); - const invoke: typeof graph.invoke = async (state, options) => { - const mainPrompts = await TRANSLATE_RULE_MAIN_PROMPT.formatMessages({ - splunkRuleTitle: state.splunkRuleTitle, - splunkRuleDescription: state.splunkRuleDescription, - splunkRuleQuery: state.splunkRuleQuery, - }); +// return translateRuleGraph.compile(); +// } - return graph.invoke({ ...state, messages: mainPrompts }, options); - }; +export function getTranslateRuleGraph({ + inferenceClient, + connectorId, + logger, +}: TranslateRuleGraphParams) { + const esqlKnowledgeBaseCaller = getEsqlKnowledgeBase({ inferenceClient, connectorId, logger }); + const translationNode = createTranslationNode(esqlKnowledgeBaseCaller); + + const translateRuleGraph = new StateGraph(translateRuleState) + .addNode('translation', translationNode) + .addNode('processResponse', responseNode) + + .addEdge(START, 'translation') + .addEdge('translation', 'processResponse') + .addEdge('processResponse', END); - return { invoke }; + return translateRuleGraph.compile(); } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts index c7431cd93cdd9..c5e7895ee5f59 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/prompts.ts @@ -6,6 +6,7 @@ */ import { ChatPromptTemplate } from '@langchain/core/prompts'; +import type { TranslateRuleState } from './types'; export const TRANSLATE_RULE_MAIN_PROMPT = ChatPromptTemplate.fromMessages([ [ @@ -39,3 +40,33 @@ The output should contain: `, ], ]); + +export const getEsqlTranslationPrompt = ({ + splunkRuleTitle, + splunkRuleDescription, + splunkRuleQuery, +}: Pick< + TranslateRuleState, + 'splunkRuleTitle' | 'splunkRuleDescription' | 'splunkRuleQuery' +>): string => `You are a helpful cybersecurity (SIEM) expert agent. Your task is to migrate "detection rules" from Splunk to Elastic Security. +Below you will find Splunk rule information: the title, description and the SPL (Search Processing Language) query. +Your goal is to translate the SPL query into an equivalent Elastic Security Query Language (ES|QL) query. + +The output will be parsed and should contain: +- First, the ES|QL query inside an \`\`\`esql code block. +- At the end, the summary of the translation process followed in markdown, starting with "## Migration Summary". + +This is the Splunk rule information: + +<> +${splunkRuleTitle} +<> + +<> +${splunkRuleDescription} +<> + +<> +${splunkRuleQuery} +<> +`; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts index 6f6023c9893cb..29ba9a1433d2b 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/state.ts @@ -13,8 +13,8 @@ export const translateRuleState = Annotation.Root({ reducer: messagesStateReducer, default: () => [], }), - splunkRuleTitle: Annotation, - splunkRuleDescription: Annotation, - splunkRuleQuery: Annotation, - response: Annotation, + splunkRuleTitle: Annotation(), + splunkRuleDescription: Annotation(), + splunkRuleQuery: Annotation(), + response: Annotation(), }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts index 9554a9d3f020c..588437e7d2a4a 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/tools/esql_translator_tool.ts @@ -90,3 +90,26 @@ export const getEsqlTranslatorTool = ({ return esqlTool; }; + +export type EsqlKnowledgeBaseCaller = (input: string) => Promise; +export const getEsqlKnowledgeBase = + ({ + inferenceClient: client, + connectorId, + logger, + }: GetEsqlTranslatorToolParams): EsqlKnowledgeBaseCaller => + async (input: string) => { + const { content } = await lastValueFrom( + naturalLanguageToEsql({ + client, + connectorId, + input, + logger: { + debug: (source) => { + logger.debug(typeof source === 'function' ? source() : source); + }, + }, + }) + ); + return content; + }; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts index 3221f2c5ae4b3..16eec347d57fc 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/types.ts @@ -5,7 +5,9 @@ * 2.0. */ +import type { Logger } from '@kbn/core/server'; import type { StructuredTool } from '@langchain/core/tools'; +import type { InferenceClient } from '@kbn/inference-plugin/server'; import type { ChatModel } from '../../../../actions_client_chat'; import type { translateRuleState } from './state'; @@ -14,9 +16,7 @@ export type TranslateRuleState = typeof translateRuleState.State; export interface TranslateRuleGraphParams { model: ChatModel; tools: StructuredTool[]; -} - -export interface TranslateRuleNodeParams { - model: ChatModel; - state: TranslateRuleState; + inferenceClient: InferenceClient; + connectorId: string; + logger: Logger; } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts index 64e22943707cd..d595e906cf2ad 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/api/translate_rule.ts @@ -14,7 +14,7 @@ import type { SplunkRuleMigrationTranslateRuleResponse } from '../../../../../.. import { SplunkRuleMigrationTranslateRuleRequestBody } from '../../../../../../../common/api/siem_migrations/splunk/rules/translate_rule.gen'; import type { SecuritySolutionPluginRouter } from '../../../../../../types'; import { SPLUNK_TRANSLATE_RULE_PATH } from '../../../../../../../common/api/siem_migrations/splunk/rules/constants'; -import { getTranslateRuleGraph } from '../agent/graph'; +import { getTranslateRuleGraph, getTranslateRuleGraphNew } from '../agent/graph'; import type { TranslateRuleState } from '../agent/types'; import { ActionsClientChat } from '../../../../actions_client_chat'; import { getEsqlTranslatorTool } from '../agent/tools/esql_translator_tool'; @@ -73,7 +73,20 @@ export const registerSplunkTranslateRuleRoute = ( ...getLangSmithTracer({ ...langSmithOptions, logger }), ], }; - const graph = getTranslateRuleGraph({ model, tools: [esqlTranslatorTool] }); + // const graph = getTranslateRuleGraph({ + // model, + // tools: [esqlTranslatorTool], + // inferenceClient, + // connectorId, + // logger, + // }); + const graph = getTranslateRuleGraph({ + model, + tools: [esqlTranslatorTool], + inferenceClient, + connectorId, + logger, + }); const translateRuleState = await graph.invoke(parameters, options); // const { response, messages } = translateRuleState as TranslateRuleState; @@ -88,7 +101,7 @@ export const registerSplunkTranslateRuleRoute = ( // messages, // }; - return res.ok({ body: { messages: translateRuleState.messages } }); + return res.ok({ body: translateRuleState }); } catch (err) { return res.badRequest({ body: err.message, From 91b658380b42c04be028053681aebc65a4992278 Mon Sep 17 00:00:00 2001 From: Sergi Massaneda Date: Thu, 17 Oct 2024 11:46:36 +0200 Subject: [PATCH 17/17] extract esql KB and call it directly as a graph node --- .../agent/esql_knowledge_base_caller.ts | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/esql_knowledge_base_caller.ts diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/esql_knowledge_base_caller.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/esql_knowledge_base_caller.ts new file mode 100644 index 0000000000000..963eec6ace830 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/splunk/rules/translate_rule/agent/esql_knowledge_base_caller.ts @@ -0,0 +1,36 @@ +/* + * 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 { naturalLanguageToEsql, type InferenceClient } from '@kbn/inference-plugin/server'; +import type { Logger } from '@kbn/core/server'; +import { lastValueFrom } from 'rxjs'; + +export type EsqlKnowledgeBaseCaller = (input: string) => Promise; + +type GetEsqlTranslatorToolParams = (params: { + inferenceClient: InferenceClient; + connectorId: string; + logger: Logger; +}) => EsqlKnowledgeBaseCaller; + +export const getEsqlKnowledgeBase: GetEsqlTranslatorToolParams = + ({ inferenceClient: client, connectorId, logger }) => + async (input: string) => { + const { content } = await lastValueFrom( + naturalLanguageToEsql({ + client, + connectorId, + input, + logger: { + debug: (source) => { + logger.debug(typeof source === 'function' ? source() : source); + }, + }, + }) + ); + return content; + };