diff --git a/x-pack/plugins/security_solution/docs/siem_migration/img/agent_graph.png b/x-pack/plugins/security_solution/docs/siem_migration/img/agent_graph.png index ecccb602e2016..a1f9cca5a34a7 100644 Binary files a/x-pack/plugins/security_solution/docs/siem_migration/img/agent_graph.png and b/x-pack/plugins/security_solution/docs/siem_migration/img/agent_graph.png differ diff --git a/x-pack/plugins/security_solution/scripts/siem_migration/draw_graphs_script.ts b/x-pack/plugins/security_solution/scripts/siem_migration/draw_graphs_script.ts index 4cb2b5c740474..4d4c8ea025c80 100644 --- a/x-pack/plugins/security_solution/scripts/siem_migration/draw_graphs_script.ts +++ b/x-pack/plugins/security_solution/scripts/siem_migration/draw_graphs_script.ts @@ -16,9 +16,7 @@ import { FakeLLM } from '@langchain/core/utils/testing'; import fs from 'fs/promises'; import path from 'path'; import { getRuleMigrationAgent } from '../../server/lib/siem_migrations/rules/task/agent'; -import type { IntegrationRetriever } from '../../server/lib/siem_migrations/rules/task/util/integration_retriever'; -import type { PrebuiltRulesMapByName } from '../../server/lib/siem_migrations/rules/task/util/prebuilt_rules'; -import type { RuleResourceRetriever } from '../../server/lib/siem_migrations/rules/task/util/rule_resource_retriever'; +import type { RuleMigrationsRetriever } from '../../server/lib/siem_migrations/rules/task/retrievers'; interface Drawable { drawMermaidPng: () => Promise; @@ -30,9 +28,7 @@ const mockLlm = new FakeLLM({ const inferenceClient = {} as InferenceClient; const connectorId = 'draw_graphs'; -const prebuiltRulesMap = {} as PrebuiltRulesMapByName; -const resourceRetriever = {} as RuleResourceRetriever; -const integrationRetriever = {} as IntegrationRetriever; +const ruleMigrationsRetriever = {} as RuleMigrationsRetriever; const createLlmInstance = () => { return mockLlm; @@ -43,9 +39,7 @@ async function getAgentGraph(logger: Logger): Promise { const graph = getRuleMigrationAgent({ model, inferenceClient, - prebuiltRulesMap, - resourceRetriever, - integrationRetriever, + ruleMigrationsRetriever, connectorId, logger, }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/create.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/create.ts index 7c94631be6b65..19669fa75cd3d 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/create.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/create.ts @@ -50,7 +50,7 @@ export const registerSiemRuleMigrationsCreateRoute = ( migration_id: migrationId, original_rule: originalRule, })); - await ruleMigrationsClient.data.integrations.create(); + await ruleMigrationsClient.data.rules.create(ruleMigrations); return res.ok({ body: { migration_id: migrationId } }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/start.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/start.ts index e9369c0e8d19d..4e50d3d583c65 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/start.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/api/start.ts @@ -6,15 +6,15 @@ */ import type { IKibanaResponse, Logger } from '@kbn/core/server'; -import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; import { APMTracer } from '@kbn/langchain/server/tracers/apm'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { buildRouteValidationWithZod } from '@kbn/zod-helpers'; +import { SIEM_RULE_MIGRATION_START_PATH } from '../../../../../common/siem_migrations/constants'; import { StartRuleMigrationRequestBody, StartRuleMigrationRequestParams, type StartRuleMigrationResponse, } from '../../../../../common/siem_migrations/model/api/rules/rule_migration.gen'; -import { SIEM_RULE_MIGRATION_START_PATH } from '../../../../../common/siem_migrations/constants'; import type { SecuritySolutionPluginRouter } from '../../../../types'; import { withLicense } from './util/with_license'; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_client.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_client.ts index 8960edd0cce21..c06c889482360 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_client.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_client.ts @@ -7,6 +7,7 @@ import type { ElasticsearchClient, Logger } from '@kbn/core/server'; import { RuleMigrationsDataIntegrationsClient } from './rule_migrations_data_integrations_client'; +import { RuleMigrationsDataPrebuiltRulesClient } from './rule_migrations_data_prebuilt_rules_client'; import { RuleMigrationsDataResourcesClient } from './rule_migrations_data_resources_client'; import { RuleMigrationsDataRulesClient } from './rule_migrations_data_rules_client'; import type { AdapterId } from './rule_migrations_data_service'; @@ -18,6 +19,7 @@ export class RuleMigrationsDataClient { public readonly rules: RuleMigrationsDataRulesClient; public readonly resources: RuleMigrationsDataResourcesClient; public readonly integrations: RuleMigrationsDataIntegrationsClient; + public readonly prebuiltRules: RuleMigrationsDataPrebuiltRulesClient; constructor( indexNameProviders: IndexNameProviders, @@ -43,5 +45,11 @@ export class RuleMigrationsDataClient { esClient, logger ); + this.prebuiltRules = new RuleMigrationsDataPrebuiltRulesClient( + indexNameProviders.prebuiltrules, + username, + esClient, + logger + ); } } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_integrations_client.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_integrations_client.ts index 3fdf1d11de36c..fdb063836f9e4 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_integrations_client.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_integrations_client.ts @@ -26,24 +26,27 @@ export class RuleMigrationsDataIntegrationsClient extends RuleMigrationsDataBase async create(): Promise { const index = await this.getIndexName(); await this.esClient - .bulk({ - refresh: 'wait_for', - operations: INTEGRATIONS.flatMap((integration) => [ - { update: { _index: index, _id: integration.id } }, - { - doc: { - title: integration.title, - description: integration.description, - data_streams: integration.data_streams, - elser_embedding: integration.elser_embedding, - '@timestamp': new Date().toISOString(), + .bulk( + { + refresh: 'wait_for', + operations: INTEGRATIONS.flatMap((integration) => [ + { update: { _index: index, _id: integration.id } }, + { + doc: { + title: integration.title, + description: integration.description, + data_streams: integration.data_streams, + elser_embedding: integration.elser_embedding, + '@timestamp': new Date().toISOString(), + }, + doc_as_upsert: true, }, - doc_as_upsert: true, - }, - ]), - }) + ]), + }, + { requestTimeout: 10 * 60 * 1000 } + ) .catch((error) => { - this.logger.error(`Error indexing integration details for ELSER: ${error.message}`); + this.logger.error(`Error preparing integrations for SIEM migration ${error.message}`); throw error; }); } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_prebuilt_rules_client.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_prebuilt_rules_client.ts new file mode 100644 index 0000000000000..ccd158c347c77 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_prebuilt_rules_client.ts @@ -0,0 +1,137 @@ +/* + * 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 { createPrebuiltRuleAssetsClient } from '../../../detection_engine/prebuilt_rules/logic/rule_assets/prebuilt_rule_assets_client'; +import { createPrebuiltRuleObjectsClient } from '../../../detection_engine/prebuilt_rules/logic/rule_objects/prebuilt_rule_objects_client'; +import { fetchRuleVersionsTriad } from '../../../detection_engine/prebuilt_rules/logic/rule_versions/fetch_rule_versions_triad'; +import type { RuleMigrationPrebuiltRule } from '../types'; +import { RuleMigrationsDataBaseClient } from './rule_migrations_data_base_client'; + +interface RetrievePrebuiltRulesParams { + soClient: SavedObjectsClientContract; + rulesClient: RulesClient; +} + +/* The minimum score required for a integration to be considered correct, might need to change this later */ +const MIN_SCORE = 40 as const; +/* The number of integrations the RAG will return, sorted by score */ +const RETURNED_RULES = 5 as const; + +/* BULK_MAX_SIZE defines the number to break down the bulk operations by. + * The 500 number was chosen as a reasonable number to avoid large payloads. It can be adjusted if needed. + */ +const BULK_MAX_SIZE = 500 as const; + +export class RuleMigrationsDataPrebuiltRulesClient extends RuleMigrationsDataBaseClient { + /** Indexes an array of integrations to be used with ELSER semantic search queries */ + async create({ soClient, rulesClient }: RetrievePrebuiltRulesParams): Promise { + const ruleAssetsClient = createPrebuiltRuleAssetsClient(soClient); + const ruleObjectsClient = createPrebuiltRuleObjectsClient(rulesClient); + + const ruleVersionsMap = await fetchRuleVersionsTriad({ + ruleAssetsClient, + ruleObjectsClient, + }); + + const filteredRules: RuleMigrationPrebuiltRule[] = []; + ruleVersionsMap.forEach((ruleVersions) => { + const rule = ruleVersions.target || ruleVersions.current; + if (rule) { + const mitreAttackIds = rule?.threat?.flatMap( + ({ technique }) => technique?.map(({ id }) => id) ?? [] + ); + + filteredRules.push({ + rule_id: rule.rule_id, + name: rule.name, + installedRuleId: ruleVersions.current?.id, + description: rule.description, + elser_embedding: `${rule.name} - ${rule.description}`, + ...(mitreAttackIds?.length && { mitre_attack_ids: mitreAttackIds }), + }); + } + }); + + const index = await this.getIndexName(); + const createdAt = new Date().toISOString(); + let prebuiltRuleSlice: RuleMigrationPrebuiltRule[]; + while ((prebuiltRuleSlice = filteredRules.splice(0, BULK_MAX_SIZE)).length) { + await this.esClient + .bulk( + { + refresh: 'wait_for', + operations: prebuiltRuleSlice.flatMap((prebuiltRule) => [ + { update: { _index: index, _id: prebuiltRule.rule_id } }, + { + doc: { + ...prebuiltRule, + '@timestamp': createdAt, + }, + doc_as_upsert: true, + }, + ]), + }, + { requestTimeout: 10 * 60 * 1000 } + ) + .catch((error) => { + this.logger.error(`Error preparing prebuilt rules for SIEM migration: ${error.message}`); + throw error; + }); + } + } + + /** Based on a LLM generated semantic string, returns the 5 best results with a score above 40 */ + async retrieveRules( + semanticString: string, + techniqueIds: string + ): Promise { + const index = await this.getIndexName(); + const query = { + bool: { + should: [ + { + semantic: { + query: semanticString, + field: 'elser_embedding', + boost: 1.5, + }, + }, + { + multi_match: { + query: semanticString, + fields: ['name^2', 'description'], + boost: 3, + }, + }, + { + multi_match: { + query: techniqueIds, + fields: ['mitre_attack_ids'], + boost: 2, + }, + }, + ], + }, + }; + const results = await this.esClient + .search({ + index, + query, + size: RETURNED_RULES, + min_score: MIN_SCORE, + }) + .then(this.processResponseHits.bind(this)) + .catch((error) => { + this.logger.error(`Error querying prebuilt rule details for ELSER: ${error.message}`); + throw error; + }); + + return results; + } +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.test.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.test.ts index f8cc0c3f1c076..e991ce2684f3e 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.test.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.test.ts @@ -42,7 +42,7 @@ describe('SiemRuleMigrationsDataService', () => { describe('constructor', () => { it('should create IndexPatternAdapters', () => { new RuleMigrationsDataService(logger, kibanaVersion); - expect(MockedIndexPatternAdapter).toHaveBeenCalledTimes(3); + expect(MockedIndexPatternAdapter).toHaveBeenCalledTimes(4); }); it('should create component templates', () => { @@ -57,6 +57,9 @@ describe('SiemRuleMigrationsDataService', () => { expect(indexPatternAdapter.setComponentTemplate).toHaveBeenCalledWith( expect.objectContaining({ name: `${INDEX_PATTERN}-integrations` }) ); + expect(indexPatternAdapter.setComponentTemplate).toHaveBeenCalledWith( + expect.objectContaining({ name: `${INDEX_PATTERN}-prebuiltrules` }) + ); }); it('should create index templates', () => { @@ -71,6 +74,9 @@ describe('SiemRuleMigrationsDataService', () => { expect(indexPatternAdapter.setIndexTemplate).toHaveBeenCalledWith( expect.objectContaining({ name: `${INDEX_PATTERN}-integrations` }) ); + expect(indexPatternAdapter.setIndexTemplate).toHaveBeenCalledWith( + expect.objectContaining({ name: `${INDEX_PATTERN}-prebuiltrules` }) + ); }); }); @@ -102,6 +108,7 @@ describe('SiemRuleMigrationsDataService', () => { rulesIndexPatternAdapter, resourcesIndexPatternAdapter, integrationsIndexPatternAdapter, + prebuiltrulesIndexPatternAdapter, ] = MockedIndexPatternAdapter.mock.instances; (rulesIndexPatternAdapter.install as jest.Mock).mockResolvedValueOnce(undefined); @@ -111,6 +118,7 @@ describe('SiemRuleMigrationsDataService', () => { await mockIndexNameProviders.rules(); await mockIndexNameProviders.resources(); await mockIndexNameProviders.integrations(); + await mockIndexNameProviders.prebuiltrules(); expect(rulesIndexPatternAdapter.createIndex).toHaveBeenCalledWith('space1'); expect(rulesIndexPatternAdapter.getIndexName).toHaveBeenCalledWith('space1'); @@ -120,6 +128,9 @@ describe('SiemRuleMigrationsDataService', () => { expect(integrationsIndexPatternAdapter.createIndex).toHaveBeenCalledWith('space1'); expect(integrationsIndexPatternAdapter.getIndexName).toHaveBeenCalledWith('space1'); + + expect(prebuiltrulesIndexPatternAdapter.createIndex).toHaveBeenCalledWith('space1'); + expect(prebuiltrulesIndexPatternAdapter.getIndexName).toHaveBeenCalledWith('space1'); }); }); }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.ts index ceff8e05f9f2f..5799e5ab84c07 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_data_service.ts @@ -10,6 +10,7 @@ import type { IndexNameProvider, IndexNameProviders } from './rule_migrations_da import { RuleMigrationsDataClient } from './rule_migrations_data_client'; import { integrationsFieldMap, + prebuiltRulesFieldMap, ruleMigrationResourcesFieldMap, ruleMigrationsFieldMap, } from './rule_migrations_field_maps'; @@ -17,7 +18,7 @@ import { const TOTAL_FIELDS_LIMIT = 2500; export const INDEX_PATTERN = '.kibana-siem-rule-migrations'; -export type AdapterId = 'rules' | 'resources' | 'integrations'; +export type AdapterId = 'rules' | 'resources' | 'integrations' | 'prebuiltrules'; interface CreateClientParams { spaceId: string; @@ -33,6 +34,7 @@ export class RuleMigrationsDataService { rules: this.createAdapter({ id: 'rules', fieldMap: ruleMigrationsFieldMap }), resources: this.createAdapter({ id: 'resources', fieldMap: ruleMigrationResourcesFieldMap }), integrations: this.createAdapter({ id: 'integrations', fieldMap: integrationsFieldMap }), + prebuiltrules: this.createAdapter({ id: 'prebuiltrules', fieldMap: prebuiltRulesFieldMap }), }; } @@ -52,6 +54,7 @@ export class RuleMigrationsDataService { this.adapters.rules.install({ ...params, logger: this.logger }), this.adapters.resources.install({ ...params, logger: this.logger }), this.adapters.integrations.install({ ...params, logger: this.logger }), + this.adapters.prebuiltrules.install({ ...params, logger: this.logger }), ]); } @@ -60,6 +63,7 @@ export class RuleMigrationsDataService { rules: this.createIndexNameProvider('rules', spaceId), resources: this.createIndexNameProvider('resources', spaceId), integrations: this.createIndexNameProvider('integrations', spaceId), + prebuiltrules: this.createIndexNameProvider('prebuiltrules', spaceId), }; return new RuleMigrationsDataClient( diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_field_maps.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_field_maps.ts index f63953192844b..7aca804c12890 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_field_maps.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/data/rule_migrations_field_maps.ts @@ -62,3 +62,12 @@ export const integrationsFieldMap: FieldMap = { 'data_streams.index_pattern': { type: 'keyword', required: true }, elser_embeddings: { type: 'semantic_text', required: true }, }; + +export const prebuiltRulesFieldMap: FieldMap = { + '@timestamp': { type: 'date', required: true }, + name: { type: 'text', required: true }, + description: { type: 'text', required: true }, + elser_embedding: { type: 'semantic_text', required: true }, + rule_id: { type: 'keyword', required: true }, + mitre_attack_ids: { type: 'keyword', array: true, required: false }, +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.test.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.test.ts index eece827726a33..ba43c369c6321 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.test.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.test.ts @@ -12,9 +12,7 @@ import type { } from '@kbn/langchain/server/language_models'; import { loggerMock } from '@kbn/logging-mocks'; import { FakeLLM } from '@langchain/core/utils/testing'; -import type { IntegrationRetriever } from '../util/integration_retriever'; -import type { PrebuiltRulesMapByName } from '../util/prebuilt_rules'; -import type { RuleResourceRetriever } from '../util/rule_resource_retriever'; +import type { RuleMigrationsRetriever } from '../retrievers'; import { getRuleMigrationAgent } from './graph'; describe('getRuleMigrationAgent', () => { @@ -24,9 +22,7 @@ describe('getRuleMigrationAgent', () => { const inferenceClient = {} as InferenceClient; const connectorId = 'draw_graphs'; - const prebuiltRulesMap = {} as PrebuiltRulesMapByName; - const resourceRetriever = {} as RuleResourceRetriever; - const integrationRetriever = {} as IntegrationRetriever; + const ruleMigrationsRetriever = {} as RuleMigrationsRetriever; const logger = loggerMock.create(); it('Ensures that the graph compiles', async () => { @@ -34,9 +30,7 @@ describe('getRuleMigrationAgent', () => { await getRuleMigrationAgent({ model, inferenceClient, - prebuiltRulesMap, - resourceRetriever, - integrationRetriever, + ruleMigrationsRetriever, connectorId, logger, }); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.ts index 4ce5e2d87a3fe..20b7cc6f1361d 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.ts @@ -6,7 +6,10 @@ */ import { END, START, StateGraph } from '@langchain/langgraph'; +import { getCreateSemanticQueryNode } from './nodes/create_semantic_query'; import { getMatchPrebuiltRuleNode } from './nodes/match_prebuilt_rule'; +import { getProcessQueryNode } from './nodes/process_query'; + import { migrateRuleState } from './state'; import { getTranslateRuleGraph } from './sub_graphs/translate_rule'; import type { MigrateRuleGraphParams, MigrateRuleState } from './types'; @@ -14,29 +17,37 @@ import type { MigrateRuleGraphParams, MigrateRuleState } from './types'; export function getRuleMigrationAgent({ model, inferenceClient, - prebuiltRulesMap, - resourceRetriever, - integrationRetriever, + ruleMigrationsRetriever, connectorId, logger, }: MigrateRuleGraphParams) { - const matchPrebuiltRuleNode = getMatchPrebuiltRuleNode({ model, prebuiltRulesMap }); - const translationSubGraph = getTranslateRuleGraph({ + const matchPrebuiltRuleNode = getMatchPrebuiltRuleNode({ model, + ruleMigrationsRetriever, + }); + const translationSubGraph = getTranslateRuleGraph({ inferenceClient, - resourceRetriever, - integrationRetriever, + ruleMigrationsRetriever, connectorId, logger, }); + const createSemanticQueryNode = getCreateSemanticQueryNode({ model }); + const processQueryNode = getProcessQueryNode({ model, ruleMigrationsRetriever }); const siemMigrationAgentGraph = new StateGraph(migrateRuleState) // Nodes + .addNode('processQuery', processQueryNode) + .addNode('createSemanticQuery', createSemanticQueryNode) .addNode('matchPrebuiltRule', matchPrebuiltRuleNode) .addNode('translationSubGraph', translationSubGraph) // Edges - .addEdge(START, 'matchPrebuiltRule') - .addConditionalEdges('matchPrebuiltRule', matchedPrebuiltRuleConditional) + .addEdge(START, 'processQuery') + .addEdge('processQuery', 'createSemanticQuery') + .addEdge('createSemanticQuery', 'matchPrebuiltRule') + .addConditionalEdges('matchPrebuiltRule', matchedPrebuiltRuleConditional, [ + 'translationSubGraph', + END, + ]) .addEdge('translationSubGraph', END); const graph = siemMigrationAgentGraph.compile(); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/create_semantic_query.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/create_semantic_query.ts new file mode 100644 index 0000000000000..446b96234711a --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/create_semantic_query.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. + */ + +import { JsonOutputParser } from '@langchain/core/output_parsers'; +import type { ChatModel } from '../../../util/actions_client_chat'; +import type { GraphNode } from '../../types'; +import { CREATE_SEMANTIC_QUERY_PROMPT } from './prompts'; + +interface GetCreateSemanticQueryNodeParams { + model: ChatModel; +} + +interface GetSemanticQueryResponse { + semantic_query: string; +} + +export const getCreateSemanticQueryNode = ({ + model, +}: GetCreateSemanticQueryNodeParams): GraphNode => { + const jsonParser = new JsonOutputParser(); + const semanticQueryChain = CREATE_SEMANTIC_QUERY_PROMPT.pipe(model).pipe(jsonParser); + return async (state) => { + const query = state.original_rule.query; + const integrationQuery = (await semanticQueryChain.invoke({ + title: state.original_rule.title, + description: state.original_rule.description, + query, + })) as GetSemanticQueryResponse; + if (!integrationQuery.semantic_query) { + return {}; + } + + return { semantic_query: integrationQuery.semantic_query }; + }; +}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/index.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/index.ts new file mode 100644 index 0000000000000..84bf247965a26 --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/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 { getCreateSemanticQueryNode } from './create_semantic_query'; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/prompts.ts similarity index 83% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/prompts.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/prompts.ts index 962de190acd02..54be39eb193f7 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/create_semantic_query/prompts.ts @@ -6,7 +6,6 @@ */ import { ChatPromptTemplate } from '@langchain/core/prompts'; - export const CREATE_SEMANTIC_QUERY_PROMPT = ChatPromptTemplate.fromMessages([ [ 'system', @@ -29,7 +28,7 @@ Go through the relevant title, description and data sources from the above query - The query should be short and concise. - Include keywords that are relevant to the use case. - Add related keywords you detected from the above query, like one or more vendor, product, cloud provider, OS platform etc. -- Always reply with a JSON object with the key "query" and the value as the semantic search query inside three backticks as shown in the below example. +- Always reply with a JSON object with the key "semantic_query" and the value as the semantic search query inside three backticks as shown in the below example. @@ -38,11 +37,11 @@ Title: Processes created by netsh Description: This search looks for processes launching netsh.exe to execute various commands via the netsh command-line utility. Netsh.exe is a command-line scripting utility that allows you to, either locally or remotely, display or modify the network configuration of a computer that is currently running. Netsh can be used as a persistence proxy technique to execute a helper .dll when netsh.exe is executed. In this search, we are looking for processes spawned by netsh.exe that are executing commands via the command line. Deprecated because we have another detection of the same type. Data Sources: -A: Please find the query keywords JSON object below: +A: Please find the semantic_query keywords JSON object below: \`\`\`json -{{"query": "windows host endpoint netsh.exe process creation command-line utility network configuration persistence proxy dll execution sysmon event id 1"}} +{{"semantic_query": "windows host endpoint netsh.exe process creation command-line utility network configuration persistence proxy dll execution sysmon event id 1"}} \`\`\` `, ], - ['ai', 'Please find the query keywords JSON object below:'], + ['ai', 'Please find the semantic_query keywords JSON object below:'], ]); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/match_prebuilt_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/match_prebuilt_rule.ts index 5900f45912599..ea403c5c4ffa7 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/match_prebuilt_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/match_prebuilt_rule.ts @@ -7,14 +7,14 @@ import { JsonOutputParser } from '@langchain/core/output_parsers'; import { SiemMigrationRuleTranslationResult } from '../../../../../../../../common/siem_migrations/constants'; +import type { RuleMigrationsRetriever } from '../../../retrievers'; import type { ChatModel } from '../../../util/actions_client_chat'; -import { filterPrebuiltRules, type PrebuiltRulesMapByName } from '../../../util/prebuilt_rules'; import type { GraphNode } from '../../types'; import { MATCH_PREBUILT_RULE_PROMPT } from './prompts'; interface GetMatchPrebuiltRuleNodeParams { model: ChatModel; - prebuiltRulesMap: PrebuiltRulesMapByName; + ruleMigrationsRetriever: RuleMigrationsRetriever; } interface GetMatchedRuleResponse { @@ -22,40 +22,42 @@ interface GetMatchedRuleResponse { } export const getMatchPrebuiltRuleNode = - ({ model, prebuiltRulesMap }: GetMatchPrebuiltRuleNodeParams): GraphNode => + ({ model, ruleMigrationsRetriever }: GetMatchPrebuiltRuleNodeParams): GraphNode => async (state) => { - const mitreAttackIds = state.original_rule.annotations?.mitre_attack; - if (!mitreAttackIds?.length) { - return {}; - } - - const filteredPrebuiltRulesMap = filterPrebuiltRules(prebuiltRulesMap, mitreAttackIds); - if (filteredPrebuiltRulesMap.size === 0) { - return {}; - } + const query = state.semantic_query; + const techniqueIds = state.original_rule.annotations?.mitre_attack || []; + const prebuiltRules = await ruleMigrationsRetriever.prebuiltRules.getRules( + query, + techniqueIds.join(',') + ); const outputParser = new JsonOutputParser(); const matchPrebuiltRule = MATCH_PREBUILT_RULE_PROMPT.pipe(model).pipe(outputParser); - const elasticSecurityRules = [...filteredPrebuiltRulesMap.keys()].join('\n'); + const elasticSecurityRules = prebuiltRules.map((rule) => { + return { + name: rule.name, + description: rule.description, + }; + }); + const response = (await matchPrebuiltRule.invoke({ - elasticSecurityRules, + rules: JSON.stringify(elasticSecurityRules, null, 2), ruleTitle: state.original_rule.title, })) as GetMatchedRuleResponse; if (response.match) { - const result = filteredPrebuiltRulesMap.get(response.match); - if (result != null) { + const matchedRule = prebuiltRules.find((r) => r.name === response.match); + if (matchedRule) { return { elastic_rule: { - title: result.rule.name, - description: result.rule.description, - prebuilt_rule_id: result.rule.rule_id, - id: result.installedRuleId, + title: matchedRule.name, + description: matchedRule.description, + id: matchedRule.installedRuleId, + prebuilt_rule_id: matchedRule.rule_id, }, translation_result: SiemMigrationRuleTranslationResult.FULL, }; } } - return {}; }; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/prompts.ts index ab5d7383e27d4..60fea54250bb3 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/prompts.ts @@ -15,7 +15,7 @@ Here are some context for you to reference for your task, read it carefully as y -{elasticSecurityRules} +{rules} `, diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/index.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/index.ts similarity index 100% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/index.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/index.ts diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/process_query.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/process_query.ts similarity index 81% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/process_query.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/process_query.ts index ae0e93ee0c4bb..27a9bca16390d 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/process_query.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/process_query.ts @@ -7,23 +7,23 @@ import { StringOutputParser } from '@langchain/core/output_parsers'; import { isEmpty } from 'lodash/fp'; -import type { ChatModel } from '../../../../../util/actions_client_chat'; -import type { RuleResourceRetriever } from '../../../../../util/rule_resource_retriever'; +import type { RuleMigrationsRetriever } from '../../../retrievers'; +import type { ChatModel } from '../../../util/actions_client_chat'; import type { GraphNode } from '../../types'; import { REPLACE_QUERY_RESOURCE_PROMPT, getResourcesContext } from './prompts'; interface GetProcessQueryNodeParams { model: ChatModel; - resourceRetriever: RuleResourceRetriever; + ruleMigrationsRetriever: RuleMigrationsRetriever; } export const getProcessQueryNode = ({ model, - resourceRetriever, + ruleMigrationsRetriever, }: GetProcessQueryNodeParams): GraphNode => { return async (state) => { let query = state.original_rule.query; - const resources = await resourceRetriever.getResources(state.original_rule); + const resources = await ruleMigrationsRetriever.resources.getResources(state.original_rule); if (!isEmpty(resources)) { const replaceQueryParser = new StringOutputParser(); const replaceQueryResourcePrompt = diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/prompts.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/prompts.ts similarity index 98% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/prompts.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/prompts.ts index b4c6b0e74aaa9..be19ca8b0bf10 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/prompts.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/process_query/prompts.ts @@ -6,7 +6,7 @@ */ import { ChatPromptTemplate } from '@langchain/core/prompts'; -import type { RuleMigrationResources } from '../../../../../util/rule_resource_retriever'; +import type { RuleMigrationResources } from '../../../retrievers/rule_resource_retriever'; interface ResourceContext { macros: string; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/state.ts index 512406d6577de..edd33e2ec69b6 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/state.ts @@ -13,6 +13,7 @@ import type { OriginalRule, RuleMigration, } from '../../../../../../common/siem_migrations/model/rule_migration.gen'; +import type { Integration } from '../../types'; export const migrateRuleState = Annotation.Root({ messages: Annotation({ @@ -23,6 +24,18 @@ export const migrateRuleState = Annotation.Root({ elastic_rule: Annotation({ reducer: (state, action) => ({ ...state, ...action }), }), + semantic_query: Annotation({ + reducer: (current, value) => value ?? current, + default: () => '', + }), + inline_query: Annotation({ + reducer: (current, value) => value ?? current, + default: () => '', + }), + integrations: Annotation({ + reducer: (current, value) => value ?? current, + default: () => [], + }), translation_result: Annotation(), comments: Annotation({ reducer: (current, value) => (value ? (current ?? []).concat(value) : current), diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/graph.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/graph.ts index 32f41e54619be..267a5bb0dd520 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/graph.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/graph.ts @@ -9,7 +9,6 @@ import { END, START, StateGraph } from '@langchain/langgraph'; import { isEmpty } from 'lodash/fp'; import { SiemMigrationRuleTranslationResult } from '../../../../../../../../common/siem_migrations/constants'; import { getFixQueryErrorsNode } from './nodes/fix_query_errors'; -import { getProcessQueryNode } from './nodes/process_query'; import { getRetrieveIntegrationsNode } from './nodes/retrieve_integrations'; import { getTranslateRuleNode } from './nodes/translate_rule'; import { getValidationNode } from './nodes/validation'; @@ -20,45 +19,32 @@ import type { TranslateRuleGraphParams, TranslateRuleState } from './types'; const MAX_VALIDATION_ITERATIONS = 3; export function getTranslateRuleGraph({ - model, inferenceClient, - resourceRetriever, - integrationRetriever, connectorId, + ruleMigrationsRetriever, logger, }: TranslateRuleGraphParams) { const translateRuleNode = getTranslateRuleNode({ - model, inferenceClient, - resourceRetriever, connectorId, logger, }); - const processQueryNode = getProcessQueryNode({ - model, - resourceRetriever, - }); - const retrieveIntegrationsNode = getRetrieveIntegrationsNode({ - model, - integrationRetriever, - }); const validationNode = getValidationNode({ logger }); const fixQueryErrorsNode = getFixQueryErrorsNode({ inferenceClient, connectorId, logger }); + const retrieveIntegrationsNode = getRetrieveIntegrationsNode({ ruleMigrationsRetriever }); const translateRuleGraph = new StateGraph(translateRuleState) // Nodes - .addNode('processQuery', processQueryNode) - .addNode('retrieveIntegrations', retrieveIntegrationsNode) .addNode('translateRule', translateRuleNode) .addNode('validation', validationNode) .addNode('fixQueryErrors', fixQueryErrorsNode) + .addNode('retrieveIntegrations', retrieveIntegrationsNode) // Edges - .addEdge(START, 'processQuery') - .addEdge('processQuery', 'retrieveIntegrations') + .addEdge(START, 'retrieveIntegrations') .addEdge('retrieveIntegrations', 'translateRule') .addEdge('translateRule', 'validation') .addEdge('fixQueryErrors', 'validation') - .addConditionalEdges('validation', validationRouter); + .addConditionalEdges('validation', validationRouter, ['fixQueryErrors', END]); const graph = translateRuleGraph.compile(); graph.name = 'Translate Rule Graph'; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/retrieve_integrations.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/retrieve_integrations.ts index 18577532fdf66..fa5b761806b5d 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/retrieve_integrations.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/retrieve_integrations/retrieve_integrations.ts @@ -5,38 +5,20 @@ * 2.0. */ -import { JsonOutputParser } from '@langchain/core/output_parsers'; -import type { ChatModel } from '../../../../../util/actions_client_chat'; -import type { IntegrationRetriever } from '../../../../../util/integration_retriever'; +import type { RuleMigrationsRetriever } from '../../../../../retrievers'; import type { GraphNode } from '../../types'; -import { CREATE_SEMANTIC_QUERY_PROMPT } from './prompts'; interface GetRetrieveIntegrationsNodeParams { - model: ChatModel; - integrationRetriever: IntegrationRetriever; -} - -interface GetSemanticQueryResponse { - query: string; + ruleMigrationsRetriever: RuleMigrationsRetriever; } export const getRetrieveIntegrationsNode = ({ - model, - integrationRetriever, + ruleMigrationsRetriever, }: GetRetrieveIntegrationsNodeParams): GraphNode => { - const jsonParser = new JsonOutputParser(); - const semanticQueryChain = CREATE_SEMANTIC_QUERY_PROMPT.pipe(model).pipe(jsonParser); - return async (state) => { - const query = state.inline_query; - - const integrationQuery = (await semanticQueryChain.invoke({ - title: state.original_rule.title, - description: state.original_rule.description, - query, - })) as GetSemanticQueryResponse; + const query = state.semantic_query; - const integrations = await integrationRetriever.getIntegrations(integrationQuery.query); + const integrations = await ruleMigrationsRetriever.integrations.getIntegrations(query); return { integrations, }; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/translate_rule/translate_rule.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/translate_rule/translate_rule.ts index 6ba5edee11b22..85f5e7279d2b9 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/translate_rule/translate_rule.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/translate_rule/translate_rule.ts @@ -8,17 +8,13 @@ import type { Logger } from '@kbn/core/server'; import type { InferenceClient } from '@kbn/inference-plugin/server'; import { SiemMigrationRuleTranslationResult } from '../../../../../../../../../../common/siem_migrations/constants'; -import type { ChatModel } from '../../../../../util/actions_client_chat'; import { getEsqlKnowledgeBase } from '../../../../../util/esql_knowledge_base_caller'; -import type { RuleResourceRetriever } from '../../../../../util/rule_resource_retriever'; import type { GraphNode } from '../../types'; import { SIEM_RULE_MIGRATION_CIM_ECS_MAP } from './cim_ecs_map'; import { ESQL_TRANSLATION_PROMPT } from './prompts'; interface GetTranslateRuleNodeParams { - model: ChatModel; inferenceClient: InferenceClient; - resourceRetriever: RuleResourceRetriever; connectorId: string; logger: Logger; } diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/state.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/state.ts index 391d7a54f9ea8..ac8799cb09d74 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/state.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/state.ts @@ -30,6 +30,10 @@ export const translateRuleState = Annotation.Root({ reducer: (current, value) => value ?? current, default: () => '', }), + semantic_query: Annotation({ + reducer: (current, value) => value ?? current, + default: () => '', + }), elastic_rule: Annotation({ reducer: (state, action) => ({ ...state, ...action }), default: () => ({} as ElasticRule), diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/types.ts index 44a5750812be0..eddc415f23392 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/types.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/types.ts @@ -7,9 +7,7 @@ import type { Logger } from '@kbn/core/server'; import type { InferenceClient } from '@kbn/inference-plugin/server'; -import type { ChatModel } from '../../../util/actions_client_chat'; -import type { IntegrationRetriever } from '../../../util/integration_retriever'; -import type { RuleResourceRetriever } from '../../../util/rule_resource_retriever'; +import type { RuleMigrationsRetriever } from '../../../retrievers'; import type { translateRuleState } from './state'; export type TranslateRuleState = typeof translateRuleState.State; @@ -17,10 +15,8 @@ export type GraphNode = (state: TranslateRuleState) => Promise { + return this.prebuiltRulesRetriever(semanticString, techniqueIds); + } + + private prebuiltRulesRetriever = async ( + semanticString: string, + techniqueIds: string + ): Promise => { + const rules = await this.dataClient.prebuiltRules.retrieveRules(semanticString, techniqueIds); + + return rules; + }; +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_migrations_retriever.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_migrations_retriever.ts new file mode 100644 index 0000000000000..22c884fa4043b --- /dev/null +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_migrations_retriever.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 { RuleMigrationsDataClient } from '../../data/rule_migrations_data_client'; +import { IntegrationRetriever } from './integration_retriever'; +import { PrebuiltRulesRetriever } from './prebuilt_rules_retriever'; +import { RuleResourceRetriever } from './rule_resource_retriever'; + +export class RuleMigrationsRetriever { + public readonly resources: RuleResourceRetriever; + public readonly integrations: IntegrationRetriever; + public readonly prebuiltRules: PrebuiltRulesRetriever; + + constructor(dataClient: RuleMigrationsDataClient, migrationId: string) { + this.resources = new RuleResourceRetriever(migrationId, dataClient); + this.integrations = new IntegrationRetriever(dataClient); + this.prebuiltRules = new PrebuiltRulesRetriever(dataClient); + } +} diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/rule_resource_retriever.test.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_resource_retriever.test.ts similarity index 100% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/rule_resource_retriever.test.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_resource_retriever.test.ts diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/rule_resource_retriever.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_resource_retriever.ts similarity index 100% rename from x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/rule_resource_retriever.ts rename to x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/retrievers/rule_resource_retriever.ts diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/rule_migrations_task_client.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/rule_migrations_task_client.ts index a6ea5c9040e16..6dbee5c64ee47 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/rule_migrations_task_client.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/rule_migrations_task_client.ts @@ -9,14 +9,15 @@ import type { AuthenticatedUser, Logger } from '@kbn/core/server'; import { AbortError, abortSignalToPromise } from '@kbn/kibana-utils-plugin/server'; import type { RunnableConfig } from '@langchain/core/runnables'; import { - SiemMigrationTaskStatus, SiemMigrationStatus, + SiemMigrationTaskStatus, } from '../../../../../common/siem_migrations/constants'; import type { RuleMigrationTaskStats } from '../../../../../common/siem_migrations/model/rule_migration.gen'; import type { RuleMigrationsDataClient } from '../data/rule_migrations_data_client'; import type { RuleMigrationDataStats } from '../data/rule_migrations_data_rules_client'; import { getRuleMigrationAgent } from './agent'; import type { MigrateRuleState } from './agent/types'; +import { RuleMigrationsRetriever } from './retrievers'; import type { MigrationAgent, RuleMigrationTaskPrepareParams, @@ -26,9 +27,6 @@ import type { RuleMigrationTaskStopResult, } from './types'; import { ActionsClientChat } from './util/actions_client_chat'; -import { IntegrationRetriever } from './util/integration_retriever'; -import { retrievePrebuiltRulesMap } from './util/prebuilt_rules'; -import { RuleResourceRetriever } from './util/rule_resource_retriever'; const ITERATION_BATCH_SIZE = 50 as const; const ITERATION_SLEEP_SECONDS = 10 as const; @@ -67,14 +65,12 @@ export class RuleMigrationsTaskClient { const abortController = new AbortController(); - // Await the preparation to make sure the agent is created properly so the task can run - const agent = await this.prepare({ ...params, abortController }); - - // not awaiting the `run` promise to execute the task in the background - this.run({ ...params, agent, abortController }).catch((err) => { - // All errors in the `run` method are already catch, this should never happen, but just in case - this.logger.error(`Unexpected error running the migration ID:${migrationId}`, err); - }); + // Retrieve agent from prepare and pass it to run right after without awaiting but using .then + this.prepare({ ...params, abortController }) + .then((agent) => this.run({ ...params, agent, abortController })) + .catch((error) => { + this.logger.error(`Error starting migration ID:${migrationId} with error:${error}`, error); + }); return { exists: true, started: true }; } @@ -88,9 +84,17 @@ export class RuleMigrationsTaskClient { soClient, abortController, }: RuleMigrationTaskPrepareParams): Promise { - const prebuiltRulesMap = await retrievePrebuiltRulesMap({ soClient, rulesClient }); - const resourceRetriever = new RuleResourceRetriever(migrationId, this.data); - const integrationRetriever = new IntegrationRetriever(this.data); + await Promise.all([ + // Populates the indices used for RAG searches on prebuilt rules and integrations. + await this.data.prebuiltRules.create({ rulesClient, soClient }), + // Will use Fleet API client for integration retrieval as an argument once feature is available + await this.data.integrations.create(), + ]).catch((error) => { + this.logger.error(`Error preparing RAG indices for migration ID:${migrationId}`, error); + throw error; + }); + + const ruleMigrationsRetriever = new RuleMigrationsRetriever(this.data, migrationId); const actionsClientChat = new ActionsClientChat(connectorId, actionsClient, this.logger); const model = await actionsClientChat.createModel({ @@ -102,9 +106,7 @@ export class RuleMigrationsTaskClient { connectorId, model, inferenceClient, - prebuiltRulesMap, - resourceRetriever, - integrationRetriever, + ruleMigrationsRetriever, logger: this.logger, }); return agent; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/prebuilt_rules.test.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/prebuilt_rules.test.ts deleted file mode 100644 index 55256d0ad0fdd..0000000000000 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/prebuilt_rules.test.ts +++ /dev/null @@ -1,105 +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 { savedObjectsClientMock } from '@kbn/core-saved-objects-api-server-mocks'; -import type { PrebuiltRulesMapByName } from './prebuilt_rules'; -import { filterPrebuiltRules, retrievePrebuiltRulesMap } from './prebuilt_rules'; -import { rulesClientMock } from '@kbn/alerting-plugin/server/mocks'; - -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 mitreAttackIds = 'T1234'; -const rule1 = { - name: 'rule one', - id: 'rule1', - threat: [ - { - framework: 'MITRE ATT&CK', - technique: [{ id: mitreAttackIds, 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 defaultParams = { - soClient: savedObjectsClientMock.create(), - rulesClient: rulesClientMock.create(), -}; - -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({ installedRuleId: undefined }) - ); - expect(prebuiltRulesMap.get('rule two')).toEqual( - expect.objectContaining({ installedRuleId: rule2.id }) - ); - }); - }); -}); - -describe('filterPrebuiltRules', () => { - let prebuiltRulesMap: PrebuiltRulesMapByName; - - beforeEach(async () => { - prebuiltRulesMap = await retrievePrebuiltRulesMap(defaultParams); - jest.clearAllMocks(); - }); - - describe('when splunk rule contains empty mitreAttackIds', () => { - it('should return empty rules map', async () => { - const filteredPrebuiltRules = filterPrebuiltRules(prebuiltRulesMap, []); - expect(filteredPrebuiltRules.size).toBe(0); - }); - }); - - describe('when splunk rule does not match mitreAttackIds', () => { - it('should return empty rules map', async () => { - const filteredPrebuiltRules = filterPrebuiltRules(prebuiltRulesMap, [`${mitreAttackIds}_2`]); - expect(filteredPrebuiltRules.size).toBe(0); - }); - }); - - describe('when splunk rule contains matching mitreAttackIds', () => { - it('should return the filtered rules map', async () => { - const filteredPrebuiltRules = filterPrebuiltRules(prebuiltRulesMap, [mitreAttackIds]); - expect(filteredPrebuiltRules.size).toBe(1); - expect(filteredPrebuiltRules.get('rule one')).toEqual( - expect.objectContaining({ rule: rule1 }) - ); - }); - }); -}); diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/prebuilt_rules.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/prebuilt_rules.ts deleted file mode 100644 index ade6632aaa5b5..0000000000000 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/util/prebuilt_rules.ts +++ /dev/null @@ -1,77 +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 { RulesClient } from '@kbn/alerting-plugin/server'; -import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; -import type { PrebuiltRuleAsset } from '../../../../detection_engine/prebuilt_rules'; -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'; - -export interface PrebuiltRuleMapped { - rule: PrebuiltRuleAsset; - installedRuleId?: string; -} - -export type PrebuiltRulesMapByName = Map; - -interface RetrievePrebuiltRulesParams { - soClient: SavedObjectsClientContract; - rulesClient: RulesClient; -} - -export const retrievePrebuiltRulesMap = async ({ - soClient, - rulesClient, -}: 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, { - rule, - installedRuleId: ruleVersions.current?.id, - }); - } - }); - return prebuiltRulesByName; -}; - -export const filterPrebuiltRules = ( - prebuiltRulesByName: PrebuiltRulesMapByName, - mitreAttackIds: string[] -) => { - const filteredPrebuiltRulesByName = new Map(); - if (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 skip it - return; - } - - const sameTechnique = mitreAttackThreat.find((threat) => - threat.technique?.some(({ id }) => mitreAttackIds?.includes(id)) - ); - - if (sameTechnique) { - filteredPrebuiltRulesByName.set(ruleName, prebuiltRulesByName.get(ruleName)); - } - }); - } - return filteredPrebuiltRulesByName; -}; diff --git a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/types.ts b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/types.ts index f8a0f0b3b25a7..f13a407ee2500 100644 --- a/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/types.ts +++ b/x-pack/plugins/security_solution/server/lib/siem_migrations/rules/types.ts @@ -22,3 +22,12 @@ export interface Integration { data_streams: Array<{ dataset: string; title: string; index_pattern: string }>; elser_embedding: string; } + +export interface RuleMigrationPrebuiltRule { + rule_id: string; + installedRuleId?: string; + name: string; + description: string; + elser_embedding: string; + mitre_attack_ids?: string[]; +}