From 5e8ff41a15b01e02f2823c51b5779ca7e93040c6 Mon Sep 17 00:00:00 2001 From: Andrew Macri Date: Thu, 17 Oct 2024 14:11:02 -0400 Subject: [PATCH] ### [Security Solution] [Attack discovery] Updates default Attack discovery max alerts for users still using legacy models In consideration of users still using legacy models, (e.g. GPT-4 instead of GPT-4o), this PR updates `DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS` from its previous value `200` in to `100`. This PR also includes additional tests. ## Desk testing 1) Navigate to Security > Attack discovery 2) Click the settings gear 3) Select any value above or below `100` in the Alerts range slider 4) Click `Reset` **Expected result** - The range slider resets to `100` --- .../impl/assistant_context/constants.tsx | 2 +- .../__mocks__/mock_anonymized_alerts.ts | 26 ++++ .../__mocks__/mock_attack_discoveries.ts | 32 ++++ .../__mocks__/mock_experiment_connector.ts | 26 ++++ .../attack_discovery/evaluation/index.test.ts | 143 ++++++++++++++++++ .../evaluation/run_evaluations/index.test.ts | 141 +++++++++++++++++ .../helpers/get_has_results/index.test.ts | 44 ++++++ .../helpers/get_has_zero_alerts/index.test.ts | 19 +++ .../get_refine_or_end_decision/index.test.ts | 116 ++++++++++++++ .../helpers/get_should_end/index.test.ts | 60 ++++++++ .../edges/refine_or_end/index.test.ts | 99 ++++++++++++ .../get_retrieve_or_generate/index.test.ts | 19 +++ .../index.test.ts | 61 ++++++++ .../index.test.ts | 28 ++++ .../get_max_retries_reached/index.test.ts | 22 +++ .../index.test.ts | 94 ++++++++++++ .../index.test.ts | 41 +++++ .../get_use_unrefined_results/index.test.ts | 51 +++++++ .../nodes/generate/index.test.ts | 103 +++++++++++++ .../index.test.ts | 42 +++++ .../index.test.ts | 62 ++++++++ .../response_is_hallucinated/index.test.ts | 26 ++++ 22 files changed, 1256 insertions(+), 1 deletion(-) create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymized_alerts.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_attack_discoveries.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_experiment_connector.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/run_evaluations/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_results/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_zero_alerts/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/get_retrieve_or_generate/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_hallucination_failures_reached/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_retries_reached/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_anonymized_alerts_from_state/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined_attack_discovery_prompt/index.test.ts create mode 100644 x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/response_is_hallucinated/index.test.ts diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx index 6e4a114c14256..c2ec745cc5c64 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx @@ -24,7 +24,7 @@ export const ANONYMIZATION_TABLE_SESSION_STORAGE_KEY = 'anonymizationTable'; export const DEFAULT_LATEST_ALERTS = 100; /** The default maximum number of alerts to be sent as context when generating Attack discoveries */ -export const DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS = 200; +export const DEFAULT_ATTACK_DISCOVERY_MAX_ALERTS = 100; export const DEFAULT_KNOWLEDGE_BASE_SETTINGS: KnowledgeBaseConfig = { latestAlerts: DEFAULT_LATEST_ALERTS, diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymized_alerts.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymized_alerts.ts new file mode 100644 index 0000000000000..e6fb24e3831c7 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymized_alerts.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. + */ + +import type { Document } from '@langchain/core/documents'; + +export const mockAnonymizedAlerts: Document[] = [ + { + pageContent: + '@timestamp,2024-10-16T02:40:08.837Z\n_id,87c42d26897490ee02ba42ec4e872910b29f3c69bda357b8faf197b533b8528a\nevent.category,malware,intrusion_detection\nevent.dataset,endpoint.alerts\nevent.module,endpoint\nevent.outcome,success\nhost.name,f5b69281-3e7e-4b52-9225-e5c30dc29c78\nhost.os.name,Windows\nhost.os.version,21H2 (10.0.20348.1607)\nkibana.alert.original_time,2023-04-01T22:03:26.909Z\nkibana.alert.risk_score,99\nkibana.alert.rule.description,Generates a detection alert each time an Elastic Endpoint Security alert is received. Enabling this rule allows you to immediately begin investigating your Endpoint alerts.\nkibana.alert.rule.name,Malicious Behavior Detection Alert: Execution of a Windows Script File Written by a Suspicious Process\nkibana.alert.severity,critical\nkibana.alert.workflow_status,open\nmessage,Malicious Behavior Detection Alert: Execution of a Windows Script File Written by a Suspicious Process\nprocess.Ext.token.integrity_level_name,high\nprocess.args,wscript,C:\\ProgramData\\WindowsAppPool\\AppPool.vbs\nprocess.code_signature.exists,true\nprocess.code_signature.status,trusted\nprocess.code_signature.subject_name,Microsoft Windows\nprocess.code_signature.trusted,true\nprocess.command_line,wscript C:\\ProgramData\\WindowsAppPool\\AppPool.vbs\nprocess.executable,C:\\Windows\\System32\\wscript.exe\nprocess.hash.md5,3412340ca1bf2f4118cbfe98961ceeda\nprocess.hash.sha1,bcb0568cbf0af0c09b53829ce9ee8ba30db77c56\nprocess.hash.sha256,02c731754bcc8f063a8c7aa53c7b7d5773f389e17582ffaa6eaaa692da183fd7\nprocess.name,wscript.exe\nprocess.parent.args,C:\\Program Files\\Microsoft Office\\Root\\Office16\\WINWORD.EXE,/n,C:\\Users\\Administrator\\Desktop\\9828375091\\7cbad6b3f505a199d6766a86b41ed23786bbb99dab9cae6c18936afdc2512f00.doc,/o,\nprocess.parent.args_count,5\nprocess.parent.command_line,"C:\\Program Files\\Microsoft Office\\Root\\Office16\\WINWORD.EXE" /n "C:\\Users\\Administrator\\Desktop\\9828375091\\7cbad6b3f505a199d6766a86b41ed23786bbb99dab9cae6c18936afdc2512f00.doc" /o ""\nprocess.parent.executable,C:\\Program Files\\Microsoft Office\\root\\Office16\\WINWORD.EXE\nprocess.parent.name,WINWORD.EXE\nprocess.pe.original_file_name,wscript.exe\nprocess.pid,13024\nprocess.working_directory,C:\\Users\\Administrator\\Desktop\\9828375091\\\nrule.name,Execution of a Windows Script File Written by a Suspicious Process\nthreat.framework,MITRE ATT&CK,MITRE ATT&CK\nthreat.tactic.id,TA0002,TA0005\nthreat.tactic.name,Execution,Defense Evasion\nthreat.tactic.reference,https://attack.mitre.org/tactics/TA0002/,https://attack.mitre.org/tactics/TA0005/\nthreat.technique.id,T1059,T1218\nthreat.technique.name,Command and Scripting Interpreter,System Binary Proxy Execution\nthreat.technique.reference,https://attack.mitre.org/techniques/T1059/,https://attack.mitre.org/techniques/T1218/\nthreat.technique.subtechnique.id,T1059.005,T1059.007,T1059.001,T1218.005\nthreat.technique.subtechnique.name,Visual Basic,JavaScript,PowerShell,Mshta\nthreat.technique.subtechnique.reference,https://attack.mitre.org/techniques/T1059/005/,https://attack.mitre.org/techniques/T1059/007/,https://attack.mitre.org/techniques/T1059/001/,https://attack.mitre.org/techniques/T1218/005/\nuser.domain,OMM-WIN-DETECT\nuser.name,42c4e419-c859-47a5-b1cb-f069d48fa509', + metadata: {}, + }, + { + pageContent: + '@timestamp,2024-10-16T02:40:08.836Z\n_id,be6d293f9a71ba209adbcacc3ba04adfd8e9456260f6af342b7cb0478a7a144a\nevent.category,malware,intrusion_detection\nevent.dataset,endpoint.alerts\nevent.module,endpoint\nevent.outcome,success\nfile.name,AppPool.vbs\nfile.path,C:\\ProgramData\\WindowsAppPool\\AppPool.vbs\nhost.name,f5b69281-3e7e-4b52-9225-e5c30dc29c78\nhost.os.name,Windows\nhost.os.version,21H2 (10.0.20348.1607)\nkibana.alert.original_time,2023-04-01T22:03:26.747Z\nkibana.alert.risk_score,99\nkibana.alert.rule.description,Generates a detection alert each time an Elastic Endpoint Security alert is received. Enabling this rule allows you to immediately begin investigating your Endpoint alerts.\nkibana.alert.rule.name,Malicious Behavior Detection Alert: Suspicious Executable File Creation\nkibana.alert.severity,critical\nkibana.alert.workflow_status,open\nmessage,Malicious Behavior Detection Alert: Suspicious Executable File Creation\nprocess.code_signature.exists,true\nprocess.code_signature.status,trusted\nprocess.code_signature.subject_name,Microsoft Corporation\nprocess.code_signature.trusted,true\nprocess.executable,C:\\Program Files\\Microsoft Office\\root\\Office16\\WINWORD.EXE\nprocess.name,WINWORD.EXE\nprocess.pid,13036\nrule.name,Suspicious Executable File Creation\nthreat.framework,MITRE ATT&CK,MITRE ATT&CK\nthreat.tactic.id,TA0011,TA0002\nthreat.tactic.name,Command and Control,Execution\nthreat.tactic.reference,https://attack.mitre.org/tactics/TA0011/,https://attack.mitre.org/tactics/TA0002/\nthreat.technique.id,T1105,T1059\nthreat.technique.name,Ingress Tool Transfer,Command and Scripting Interpreter\nthreat.technique.reference,https://attack.mitre.org/techniques/T1105/,https://attack.mitre.org/techniques/T1059/\nthreat.technique.subtechnique.id,T1059.005,T1059.007\nthreat.technique.subtechnique.name,Visual Basic,JavaScript\nthreat.technique.subtechnique.reference,https://attack.mitre.org/techniques/T1059/005/,https://attack.mitre.org/techniques/T1059/007/\nuser.domain,OMM-WIN-DETECT\nuser.name,42c4e419-c859-47a5-b1cb-f069d48fa509', + metadata: {}, + }, +]; + +export const mockAnonymizedAlertsReplacements: Record = { + '42c4e419-c859-47a5-b1cb-f069d48fa509': 'Administrator', + 'f5b69281-3e7e-4b52-9225-e5c30dc29c78': 'SRVWIN07', +}; diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_attack_discoveries.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_attack_discoveries.ts new file mode 100644 index 0000000000000..d8d66481571d7 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_attack_discoveries.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 type { AttackDiscovery } from '@kbn/elastic-assistant-common'; + +export const mockAttackDiscoveries: AttackDiscovery[] = [ + { + title: 'Critical Malware and Phishing Alerts on host e1cb3cf0-30f3-4f99-a9c8-518b955c6f90', + alertIds: [ + '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16', + 'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b', + '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c', + '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608', + 'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac', + '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080', + '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6', + '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998', + ], + timestamp: '2024-10-10T22:59:52.749Z', + detailsMarkdown: + '- On `2023-06-19T00:28:38.061Z` a critical malware detection alert was triggered on host {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} running {{ host.os.name macOS }} version {{ host.os.version 13.4 }}.\n- The malware was identified as {{ file.name unix1 }} with SHA256 hash {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\n- The process {{ process.name My Go Application.app }} was executed with command line {{ process.command_line /private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app }}.\n- The process was not trusted as its code signature failed to satisfy specified code requirements.\n- The user involved was {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}.\n- Another critical alert was triggered for potential credentials phishing via {{ process.name osascript }} on the same host.\n- The phishing attempt involved displaying a dialog to capture the user\'s password.\n- The process {{ process.name osascript }} was executed with command line {{ process.command_line osascript -e display dialog "MacOS wants to access System Preferences\\n\\nPlease enter your password." with title "System Preferences" with icon file "System:Library:CoreServices:CoreTypes.bundle:Contents:Resources:ToolbarAdvanced.icns" default answer "" giving up after 30 with hidden answer ¬ }}.\n- The MITRE ATT&CK tactics involved include Credential Access and Input Capture.', + summaryMarkdown: + 'Critical malware and phishing alerts detected on {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} involving user {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}. Malware identified as {{ file.name unix1 }} and phishing attempt via {{ process.name osascript }}.', + mitreAttackTactics: ['Credential Access', 'Input Capture'], + entitySummaryMarkdown: + 'Critical malware and phishing alerts detected on {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} involving user {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}.', + }, +]; diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_experiment_connector.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_experiment_connector.ts new file mode 100644 index 0000000000000..1ee32768ab091 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_experiment_connector.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. + */ + +import type { Connector } from '@kbn/actions-plugin/server/application/connector/types'; + +export const mockExperimentConnector: Connector = { + name: 'Gemini 1.5 Pro 002', + actionTypeId: '.gemini', + config: { + apiUrl: 'https://example.com', + defaultModel: 'gemini-1.5-pro-002', + gcpRegion: 'test-region', + gcpProjectID: 'test-project-id', + }, + secrets: { + credentialsJson: '{}', + }, + id: 'gemini-1-5-pro-002', + isPreconfigured: true, + isSystemAction: false, + isDeprecated: false, +} as Connector; diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.test.ts new file mode 100644 index 0000000000000..8154f0b446566 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.test.ts @@ -0,0 +1,143 @@ +/* + * 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 { ActionsClient } from '@kbn/actions-plugin/server'; +import type { Connector } from '@kbn/actions-plugin/server/application/connector/types'; +import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { loggerMock } from '@kbn/logging-mocks'; +import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; + +import { evaluateAttackDiscovery } from '.'; +import { DefaultAttackDiscoveryGraph } from '../graphs/default_attack_discovery_graph'; +import { AttackDiscoveryGraphMetadata } from '../../langchain/graphs'; +import { mockExperimentConnector } from './__mocks__/mock_experiment_connector'; +import { getLlmType } from '../../../routes/utils'; + +jest.mock('@kbn/langchain/server', () => ({ + ...jest.requireActual('@kbn/langchain/server'), + + ActionsClientLlm: jest.fn(), +})); + +jest.mock('langsmith/evaluation', () => ({ + evaluate: jest.fn(async (predict: Function) => + predict({ + overrides: { + errors: ['test-error'], + }, + }) + ), +})); + +jest.mock('./helpers/get_custom_evaluator', () => ({ + getCustomEvaluator: jest.fn(), +})); + +jest.mock('./helpers/get_evaluator_llm', () => { + const mockLlm = jest.fn() as unknown as ActionsClientLlm; + + return { + getEvaluatorLlm: jest.fn().mockResolvedValue(mockLlm), + }; +}); + +const actionsClient = { + get: jest.fn(), +} as unknown as ActionsClient; +const alertsIndexPattern = 'test-alerts-index-pattern'; +const connectorTimeout = 1000; +const datasetName = 'test-dataset'; +const evaluationId = 'test-evaluation-id'; +const evaluatorConnectorId = 'test-evaluator-connector-id'; +const langSmithApiKey = 'test-api-key'; +const langSmithProject = 'test-lang-smith-project'; +const logger = loggerMock.create(); +const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); +const runName = 'test-run-name'; + +const connectors = [mockExperimentConnector]; + +const projectName = 'test-lang-smith-project'; + +const graphs: Array<{ + connector: Connector; + graph: DefaultAttackDiscoveryGraph; + llmType: string | undefined; + name: string; + traceOptions: { + projectName: string | undefined; + tracers: LangChainTracer[]; + }; +}> = connectors.map((connector) => { + const llmType = getLlmType(connector.actionTypeId); + + const traceOptions = { + projectName, + tracers: [ + ...getLangSmithTracer({ + apiKey: langSmithApiKey, + projectName, + logger, + }), + ], + }; + + const graph = { + invoke: jest.fn().mockResolvedValue({}), + } as unknown as DefaultAttackDiscoveryGraph; + + return { + connector, + graph, + llmType, + name: `${runName} - ${connector.name} - ${evaluationId} - Attack discovery`, + traceOptions, + }; +}); + +const attackDiscoveryGraphs: AttackDiscoveryGraphMetadata[] = [ + { + getDefaultAttackDiscoveryGraph: jest.fn().mockReturnValue(graphs[0].graph), + graphType: 'attack-discovery', + }, +]; + +describe('evaluateAttackDiscovery', () => { + beforeEach(() => jest.clearAllMocks()); + + it('evaluates the attack discovery graphs', async () => { + await evaluateAttackDiscovery({ + actionsClient, + attackDiscoveryGraphs, + alertsIndexPattern, + connectors, + connectorTimeout, + datasetName, + esClient: mockEsClient, + evaluationId, + evaluatorConnectorId, + langSmithApiKey, + langSmithProject, + logger, + runName, + size: 20, + }); + + expect(graphs[0].graph.invoke).toHaveBeenCalledWith( + { + errors: ['test-error'], + }, + { + callbacks: [...graphs[0].traceOptions.tracers], + runName: graphs[0].name, + tags: ['evaluation', graphs[0].llmType ?? ''], + } + ); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/run_evaluations/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/run_evaluations/index.test.ts new file mode 100644 index 0000000000000..909c279218f1c --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/run_evaluations/index.test.ts @@ -0,0 +1,141 @@ +/* + * 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 { ActionsClient } from '@kbn/actions-plugin/server'; +import type { Connector } from '@kbn/actions-plugin/server/application/connector/types'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; +import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; +import { loggerMock } from '@kbn/logging-mocks'; +import type { LangChainTracer } from '@langchain/core/tracers/tracer_langchain'; + +import { runEvaluations } from '.'; +import { type DefaultAttackDiscoveryGraph } from '../../graphs/default_attack_discovery_graph'; +import { mockExperimentConnector } from '../__mocks__/mock_experiment_connector'; +import { getLlmType } from '../../../../routes/utils'; + +jest.mock('@kbn/langchain/server', () => ({ + ...jest.requireActual('@kbn/langchain/server'), + + ActionsClientLlm: jest.fn(), +})); + +jest.mock('langsmith/evaluation', () => ({ + evaluate: jest.fn(async (predict: Function) => + predict({ + overrides: { + errors: ['test-error'], + }, + }) + ), +})); + +jest.mock('../helpers/get_custom_evaluator', () => ({ + getCustomEvaluator: jest.fn(), +})); + +jest.mock('../helpers/get_evaluator_llm', () => { + const mockLlm = jest.fn() as unknown as ActionsClientLlm; + + return { + getEvaluatorLlm: jest.fn().mockResolvedValue(mockLlm), + }; +}); + +const actionsClient = { + get: jest.fn(), +} as unknown as ActionsClient; +const connectorTimeout = 1000; +const datasetName = 'test-dataset'; +const evaluatorConnectorId = 'test-evaluator-connector-id'; +const langSmithApiKey = 'test-api-key'; +const logger = loggerMock.create(); +const connectors = [mockExperimentConnector]; + +const projectName = 'test-lang-smith-project'; + +const graphs: Array<{ + connector: Connector; + graph: DefaultAttackDiscoveryGraph; + llmType: string | undefined; + name: string; + traceOptions: { + projectName: string | undefined; + tracers: LangChainTracer[]; + }; +}> = connectors.map((connector) => { + const llmType = getLlmType(connector.actionTypeId); + + const traceOptions = { + projectName, + tracers: [ + ...getLangSmithTracer({ + apiKey: langSmithApiKey, + projectName, + logger, + }), + ], + }; + + const graph = { + invoke: jest.fn().mockResolvedValue({}), + } as unknown as DefaultAttackDiscoveryGraph; + + return { + connector, + graph, + llmType, + name: `testRunName - ${connector.name} - testEvaluationId - Attack discovery`, + traceOptions, + }; +}); + +describe('runEvaluations', () => { + beforeEach(() => jest.clearAllMocks()); + + it('predict() invokes the graph with the expected overrides', async () => { + await runEvaluations({ + actionsClient, + connectorTimeout, + datasetName, + evaluatorConnectorId, + graphs, + langSmithApiKey, + logger, + }); + + expect(graphs[0].graph.invoke).toHaveBeenCalledWith( + { + errors: ['test-error'], + }, + { + callbacks: [...graphs[0].traceOptions.tracers], + runName: graphs[0].name, + tags: ['evaluation', graphs[0].llmType ?? ''], + } + ); + }); + + it('catches and logs errors that occur during evaluation', async () => { + const error = new Error('Test error'); + + (graphs[0].graph.invoke as jest.Mock).mockRejectedValue(error); + + await runEvaluations({ + actionsClient, + connectorTimeout, + datasetName, + evaluatorConnectorId, + graphs, + langSmithApiKey, + logger, + }); + + expect(logger.error).toHaveBeenCalledWith( + 'Error evaluating connector "Gemini 1.5 Pro 002" (gemini), running experiment "testRunName - Gemini 1.5 Pro 002 - testEvaluationId - Attack discovery": Error: Test error' + ); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_results/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_results/index.test.ts new file mode 100644 index 0000000000000..b589fab8e5797 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_results/index.test.ts @@ -0,0 +1,44 @@ +/* + * 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 { AttackDiscovery } from '@kbn/elastic-assistant-common'; + +import { getHasResults } from '.'; + +const attackDiscoveries: AttackDiscovery[] = [ + { + title: 'Critical Malware and Phishing Alerts on host e1cb3cf0-30f3-4f99-a9c8-518b955c6f90', + alertIds: [ + '4af5689eb58c2420efc0f7fad53c5bf9b8b6797e516d6ea87d6044ce25d54e16', + 'c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b', + '021b27d6bee0650a843be1d511119a3b5c7c8fdaeff922471ce0248ad27bd26c', + '6cc8d5f0e1c2b6c75219b001858f1be64194a97334be7a1e3572f8cfe6bae608', + 'f39a4013ed9609584a8a22dca902e896aa5b24d2da03e0eaab5556608fa682ac', + '909968e926e08a974c7df1613d98ebf1e2422afcb58e4e994beb47b063e85080', + '2c25a4dc31cd1ec254c2b19ea663fd0b09a16e239caa1218b4598801fb330da6', + '3bf907becb3a4f8e39a3b673e0d50fc954a7febef30c12891744c603760e4998', + ], + timestamp: '2024-10-10T22:59:52.749Z', + detailsMarkdown: + '- On `2023-06-19T00:28:38.061Z` a critical malware detection alert was triggered on host {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} running {{ host.os.name macOS }} version {{ host.os.version 13.4 }}.\n- The malware was identified as {{ file.name unix1 }} with SHA256 hash {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\n- The process {{ process.name My Go Application.app }} was executed with command line {{ process.command_line /private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app }}.\n- The process was not trusted as its code signature failed to satisfy specified code requirements.\n- The user involved was {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}.\n- Another critical alert was triggered for potential credentials phishing via {{ process.name osascript }} on the same host.\n- The phishing attempt involved displaying a dialog to capture the user\'s password.\n- The process {{ process.name osascript }} was executed with command line {{ process.command_line osascript -e display dialog "MacOS wants to access System Preferences\\n\\nPlease enter your password." with title "System Preferences" with icon file "System:Library:CoreServices:CoreTypes.bundle:Contents:Resources:ToolbarAdvanced.icns" default answer "" giving up after 30 with hidden answer ¬ }}.\n- The MITRE ATT&CK tactics involved include Credential Access and Input Capture.', + summaryMarkdown: + 'Critical malware and phishing alerts detected on {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} involving user {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}. Malware identified as {{ file.name unix1 }} and phishing attempt via {{ process.name osascript }}.', + mitreAttackTactics: ['Credential Access', 'Input Capture'], + entitySummaryMarkdown: + 'Critical malware and phishing alerts detected on {{ host.name e1cb3cf0-30f3-4f99-a9c8-518b955c6f90 }} involving user {{ user.name 039c15c5-3964-43e7-a891-42fe2ceeb9ff }}.', + }, +]; + +describe('getHasResults', () => { + it('returns true when attackDiscoveries is not null', () => { + expect(getHasResults(attackDiscoveries)).toBe(true); + }); + + it('returns false when attackDiscoveries is null', () => { + expect(getHasResults(null)).toBe(false); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_zero_alerts/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_zero_alerts/index.test.ts new file mode 100644 index 0000000000000..2c500c375db0b --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/helpers/get_has_zero_alerts/index.test.ts @@ -0,0 +1,19 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { getHasZeroAlerts } from '.'; +import { mockAnonymizedAlerts } from '../../../../../evaluation/__mocks__/mock_anonymized_alerts'; + +describe('getHasZeroAlerts', () => { + it('returns true when there are no alerts', () => { + expect(getHasZeroAlerts([])).toBe(true); + }); + + it('returns false when there are alerts', () => { + expect(getHasZeroAlerts(mockAnonymizedAlerts)).toBe(false); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts new file mode 100644 index 0000000000000..87f73402a3a2d --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts @@ -0,0 +1,116 @@ +/* + * 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 { getRefineOrEndDecision } from '.'; + +describe('getRefineOrEndDecision', () => { + it("returns 'end' when the refined results were generated", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: true, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('end'); + }); + + describe('limits shared by both the generate and refine steps', () => { + it("returns 'end' when the (shared) max hallucinations limit was reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when the (shared) max generation attempts limit was reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when multiple limits are reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, // <-- limit reached + maxRetriesReached: true, // <-- another limit reached + }); + + expect(result).toEqual('end'); + }); + }); + + it("returns 'refine' when there are unrefined results, and limits have NOT been reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('refine'); + }); + + describe('getRefineOrEndDecision', () => { + it("returns 'end' when the refined results were generated", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: true, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('end'); + }); + + describe('limits shared by both the generate and refine steps', () => { + it("returns 'end' when the (shared) max hallucinations limit was reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when the (shared) max generation attempts limit was reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when multiple limits are reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, // <-- limit reached + maxRetriesReached: true, // <-- another limit reached + }); + + expect(result).toEqual('end'); + }); + }); + + it("returns 'refine' when there are unrefined results, and limits have NOT been reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('refine'); + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts new file mode 100644 index 0000000000000..8c35773f8bea2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts @@ -0,0 +1,60 @@ +/* + * 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 { getShouldEnd } from '.'; + +describe('getShouldEnd', () => { + it('returns true when hasFinalResults is true', () => { + const result = getShouldEnd({ + hasFinalResults: true, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true when maxHallucinationFailuresReached is true', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true when maxRetriesReached is true', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); + + it('returns true when both maxHallucinationFailuresReached and maxRetriesReached are true', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, // <-- limit reached + maxRetriesReached: true, // <-- another limit reached + }); + + expect(result).toBe(true); + }); + + it('returns false when all conditions are false', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(false); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/index.test.ts new file mode 100644 index 0000000000000..39934255c069c --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/refine_or_end/index.test.ts @@ -0,0 +1,99 @@ +/* + * 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 { loggerMock } from '@kbn/logging-mocks'; + +import { getRefineOrEndEdge } from '.'; +import { mockAttackDiscoveries } from '../../../../evaluation/__mocks__/mock_attack_discoveries'; +import { + mockAnonymizedAlerts, + mockAnonymizedAlertsReplacements, +} from '../../../../evaluation/__mocks__/mock_anonymized_alerts'; +import { GraphState } from '../../types'; + +const logger = loggerMock.create(); + +const initialGraphState: GraphState = { + attackDiscoveries: null, // <-- no refined results + attackDiscoveryPrompt: 'prompt', + anonymizedAlerts: [...mockAnonymizedAlerts], + combinedGenerations: 'generations', + combinedRefinements: '', + errors: [], + generationAttempts: 2, + generations: ['gen', 'erations'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: { + ...mockAnonymizedAlertsReplacements, + }, + unrefinedResults: mockAttackDiscoveries, +}; + +describe('getRefineOrEndEdge', () => { + beforeEach(() => jest.clearAllMocks()); + + it("returns 'end' when the refined results were generated", () => { + const state: GraphState = { + ...initialGraphState, + attackDiscoveries: mockAttackDiscoveries, // <-- attackDiscoveries are NOT null + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'refine' when there are unrefined results, and limits have NOT been reached", () => { + const edge = getRefineOrEndEdge(logger); + const result = edge(initialGraphState); + + expect(result).toEqual('refine'); + }); + + it("returns 'end' when the max generation attempts limit was reached", () => { + const state: GraphState = { + ...initialGraphState, + generationAttempts: initialGraphState.maxGenerationAttempts, + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when the max hallucination failures limit was reached", () => { + const state: GraphState = { + ...initialGraphState, + hallucinationFailures: initialGraphState.maxHallucinationFailures, + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when multiple limits are reached", () => { + const state: GraphState = { + ...initialGraphState, + generationAttempts: initialGraphState.maxGenerationAttempts, + hallucinationFailures: initialGraphState.maxHallucinationFailures, + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/get_retrieve_or_generate/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/get_retrieve_or_generate/index.test.ts new file mode 100644 index 0000000000000..61dba4fb3d479 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/get_retrieve_or_generate/index.test.ts @@ -0,0 +1,19 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { getRetrieveOrGenerate } from '.'; +import { mockAnonymizedAlerts } from '../../../../../evaluation/__mocks__/mock_anonymized_alerts'; + +describe('getRetrieveOrGenerate', () => { + it("returns 'retrieve_anonymized_alerts' when anonymizedAlerts is empty", () => { + expect(getRetrieveOrGenerate([])).toBe('retrieve_anonymized_alerts'); + }); + + it("returns 'generate' when anonymizedAlerts is not empty", () => { + expect(getRetrieveOrGenerate(mockAnonymizedAlerts)).toBe('generate'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/index.test.ts new file mode 100644 index 0000000000000..06377aa565a12 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/retrieve_anonymized_alerts_or_generate/index.test.ts @@ -0,0 +1,61 @@ +/* + * 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 { loggerMock } from '@kbn/logging-mocks'; + +import { getRetrieveAnonymizedAlertsOrGenerateEdge } from '.'; +import { mockAnonymizedAlerts } from '../../../../evaluation/__mocks__/mock_anonymized_alerts'; +import type { GraphState } from '../../types'; + +const logger = loggerMock.create(); + +const initialGraphState: GraphState = { + attackDiscoveries: null, + attackDiscoveryPrompt: 'prompt', + anonymizedAlerts: [], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('getRetrieveAnonymizedAlertsOrGenerateEdge', () => { + beforeEach(() => jest.clearAllMocks()); + + it('returns "generate" when anonymizedAlerts is NOT empty, so there are alerts for the generate step', () => { + const state: GraphState = { + ...initialGraphState, + anonymizedAlerts: mockAnonymizedAlerts, + }; + + const edge = getRetrieveAnonymizedAlertsOrGenerateEdge(logger); + const result = edge(state); + + expect(result).toEqual('generate'); + }); + + it('returns "retrieve_anonymized_alerts" when anonymizedAlerts is empty, so they can be retrieved', () => { + const state: GraphState = { + ...initialGraphState, + anonymizedAlerts: [], // <-- empty + }; + + const edge = getRetrieveAnonymizedAlertsOrGenerateEdge(logger); + const result = edge(state); + + expect(result).toEqual('retrieve_anonymized_alerts'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_hallucination_failures_reached/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_hallucination_failures_reached/index.test.ts new file mode 100644 index 0000000000000..138179109708e --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_hallucination_failures_reached/index.test.ts @@ -0,0 +1,28 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { getMaxHallucinationFailuresReached } from '.'; + +describe('getMaxHallucinationFailuresReached', () => { + it('return true when hallucination failures is equal to the max hallucination failures', () => { + expect( + getMaxHallucinationFailuresReached({ hallucinationFailures: 2, maxHallucinationFailures: 2 }) + ).toBe(true); + }); + + it('returns true when hallucination failures is greater than the max hallucination failures', () => { + expect( + getMaxHallucinationFailuresReached({ hallucinationFailures: 3, maxHallucinationFailures: 2 }) + ).toBe(true); + }); + + it('returns false when hallucination failures is less than the max hallucination failures', () => { + expect( + getMaxHallucinationFailuresReached({ hallucinationFailures: 1, maxHallucinationFailures: 2 }) + ).toBe(false); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_retries_reached/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_retries_reached/index.test.ts new file mode 100644 index 0000000000000..47f49a75415c9 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/helpers/get_max_retries_reached/index.test.ts @@ -0,0 +1,22 @@ +/* + * 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 { getMaxRetriesReached } from '.'; + +describe('getMaxRetriesReached', () => { + it('returns true when generation attempts is equal to the max generation attempts', () => { + expect(getMaxRetriesReached({ generationAttempts: 2, maxGenerationAttempts: 2 })).toBe(true); + }); + + it('returns true when generation attempts is greater than the max generation attempts', () => { + expect(getMaxRetriesReached({ generationAttempts: 3, maxGenerationAttempts: 2 })).toBe(true); + }); + + it('returns false when generation attempts is less than the max generation attempts', () => { + expect(getMaxRetriesReached({ generationAttempts: 1, maxGenerationAttempts: 2 })).toBe(false); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts new file mode 100644 index 0000000000000..6f3b3b1b909a2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts @@ -0,0 +1,94 @@ +/* + * 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 { discardPreviousGenerations } from '.'; +import { GraphState } from '../../../../types'; + +const graphState: GraphState = { + attackDiscoveries: null, + attackDiscoveryPrompt: 'prompt', + anonymizedAlerts: [ + { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', + }, + { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', + }, + ], + combinedGenerations: 'combinedGenerations', + combinedRefinements: '', + errors: [], + generationAttempts: 2, + generations: ['combined', 'Generations'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('discardPreviousGenerations', () => { + describe('common state updates', () => { + let result: GraphState; + + beforeEach(() => { + result = discardPreviousGenerations({ + generationAttempts: graphState.generationAttempts, + hallucinationFailures: graphState.hallucinationFailures, + isHallucinationDetected: false, + state: graphState, + }); + }); + + it('resets the combined generations', () => { + expect(result.combinedGenerations).toBe(''); + }); + + it('increments the generation attempts', () => { + expect(result.generationAttempts).toBe(graphState.generationAttempts + 1); + }); + + it('resets the collection of generations', () => { + expect(result.generations).toEqual([]); + }); + }); + + it('increments hallucinationFailures when a hallucination is detected', () => { + const result = discardPreviousGenerations({ + generationAttempts: graphState.generationAttempts, + hallucinationFailures: graphState.hallucinationFailures, + isHallucinationDetected: true, // <-- hallucination detected + state: graphState, + }); + + expect(result.hallucinationFailures).toBe(graphState.hallucinationFailures + 1); + }); + + it('does NOT increment hallucinationFailures when a hallucination is NOT detected', () => { + const result = discardPreviousGenerations({ + generationAttempts: graphState.generationAttempts, + hallucinationFailures: graphState.hallucinationFailures, + isHallucinationDetected: false, // <-- no hallucination detected + state: graphState, + }); + + expect(result.hallucinationFailures).toBe(graphState.hallucinationFailures); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_anonymized_alerts_from_state/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_anonymized_alerts_from_state/index.test.ts new file mode 100644 index 0000000000000..fb3d541e670df --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_anonymized_alerts_from_state/index.test.ts @@ -0,0 +1,41 @@ +/* + * 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 { getAnonymizedAlertsFromState } from '.'; + +import { mockAnonymizedAlerts } from '../../../../../../evaluation/__mocks__/mock_anonymized_alerts'; +import { GraphState } from '../../../../types'; + +const graphState: GraphState = { + attackDiscoveries: null, + attackDiscoveryPrompt: 'prompt', + anonymizedAlerts: mockAnonymizedAlerts, // <-- mockAnonymizedAlerts is an array of objects with a pageContent property + combinedGenerations: 'combinedGenerations', + combinedRefinements: '', + errors: [], + generationAttempts: 2, + generations: ['combined', 'Generations'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('getAnonymizedAlertsFromState', () => { + it('returns the anonymized alerts from the state', () => { + const result = getAnonymizedAlertsFromState(graphState); + + expect(result).toEqual([ + mockAnonymizedAlerts[0].pageContent, + mockAnonymizedAlerts[1].pageContent, + ]); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts new file mode 100644 index 0000000000000..e9a75d7feb338 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts @@ -0,0 +1,51 @@ +/* + * 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 { getUseUnrefinedResults } from '.'; +import { mockAttackDiscoveries } from '../../../../../../evaluation/__mocks__/mock_attack_discoveries'; + +describe('getUseUnrefinedResults', () => { + it('returns true when the next attempt would exceed the limit, and we have unrefined results', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 2, + maxGenerationAttempts: 3, + unrefinedResults: mockAttackDiscoveries, + }) + ).toBe(true); + }); + + it('returns false when the next attempt would NOT exceed the limit', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 1, + maxGenerationAttempts: 3, + unrefinedResults: mockAttackDiscoveries, + }) + ).toBe(false); + }); + + it('returns false when unrefined results is null', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 2, + maxGenerationAttempts: 3, + unrefinedResults: null, + }) + ).toBe(false); + }); + + it('returns false when unrefined results is empty', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 2, + maxGenerationAttempts: 3, + unrefinedResults: [], + }) + ).toBe(false); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.test.ts new file mode 100644 index 0000000000000..da815aad9795a --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.test.ts @@ -0,0 +1,103 @@ +/* + * 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 { ActionsClientLlm } from '@kbn/langchain/server'; +import { loggerMock } from '@kbn/logging-mocks'; +import { FakeLLM } from '@langchain/core/utils/testing'; + +import { getGenerateNode } from '.'; +import { + mockAnonymizedAlerts, + mockAnonymizedAlertsReplacements, +} from '../../../../evaluation/__mocks__/mock_anonymized_alerts'; +import { getAnonymizedAlertsFromState } from './helpers/get_anonymized_alerts_from_state'; +import { getChainWithFormatInstructions } from '../helpers/get_chain_with_format_instructions'; +import { GraphState } from '../../types'; + +jest.mock('../helpers/get_chain_with_format_instructions', () => { + const mockInvoke = jest.fn().mockResolvedValue(''); + + return { + getChainWithFormatInstructions: jest.fn().mockReturnValue({ + chain: { + invoke: mockInvoke, + }, + formatInstructions: ['mock format instructions'], + llmType: 'fake', + mockInvoke, // <-- added for testing + }), + }; +}); + +const mockLogger = loggerMock.create(); +let mockLlm: ActionsClientLlm; + +const initialGraphState: GraphState = { + attackDiscoveries: null, + attackDiscoveryPrompt: + "You are a cyber security analyst tasked with analyzing security events from Elastic Security to identify and report on potential cyber attacks or progressions. Your report should focus on high-risk incidents that could severely impact the organization, rather than isolated alerts. Present your findings in a way that can be easily understood by anyone, regardless of their technical expertise, as if you were briefing the CISO. Break down your response into sections based on timing, hosts, and users involved. When correlating alerts, use kibana.alert.original_time when it's available, otherwise use @timestamp. Include appropriate context about the affected hosts and users. Describe how the attack progression might have occurred and, if feasible, attribute it to known threat groups. Prioritize high and critical alerts, but include lower-severity alerts if desired. In the description field, provide as much detail as possible, in a bulleted list explaining any attack progressions. Accuracy is of utmost importance. You MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds).", + anonymizedAlerts: [...mockAnonymizedAlerts], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: + 'You previously generated the following insights, but sometimes they represent the same attack.\n\nCombine the insights below, when they represent the same attack; leave any insights that are not combined unchanged:', + replacements: { + ...mockAnonymizedAlertsReplacements, + }, + unrefinedResults: null, +}; + +describe('getGenerateNode', () => { + beforeEach(() => { + jest.clearAllMocks(); + + mockLlm = new FakeLLM({ + response: JSON.stringify({}, null, 2), + }) as unknown as ActionsClientLlm; + }); + + it('returns a function', () => { + const generateNode = getGenerateNode({ + llm: mockLlm, + logger: mockLogger, + }); + + expect(typeof generateNode).toBe('function'); + }); + + it('invokes the chain with the alerts from state and format instructions', async () => { + // @ts-expect-error + const { mockInvoke } = getChainWithFormatInstructions(mockLlm); + + const generateNode = getGenerateNode({ + llm: mockLlm, + logger: mockLogger, + }); + + await generateNode(initialGraphState); + + expect(mockInvoke).toHaveBeenCalledWith({ + format_instructions: ['mock format instructions'], + query: `${initialGraphState.attackDiscoveryPrompt} + +Use context from the following alerts to provide insights: + +\"\"\" +${getAnonymizedAlertsFromState(initialGraphState).join('\n\n')} +\"\"\" +`, + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts new file mode 100644 index 0000000000000..fd98af61150b8 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts @@ -0,0 +1,42 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { FakeLLM } from '@langchain/core/utils/testing'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; + +import { getChainWithFormatInstructions } from '.'; + +describe('getChainWithFormatInstructions', () => { + const mockLlm = new FakeLLM({ + response: JSON.stringify({}, null, 2), + }) as unknown as ActionsClientLlm; + + it('returns the chain with format instructions', () => { + const expectedFormatInstructions = `You must format your output as a JSON value that adheres to a given "JSON Schema" instance. + +"JSON Schema" is a declarative language that allows you to annotate and validate JSON documents. + +For example, the example "JSON Schema" instance {{"properties": {{"foo": {{"description": "a list of test words", "type": "array", "items": {{"type": "string"}}}}}}, "required": ["foo"]}}}} +would match an object with one required property, "foo". The "type" property specifies "foo" must be an "array", and the "description" property semantically describes it as "a list of test words". The items within "foo" must be strings. +Thus, the object {{"foo": ["bar", "baz"]}} is a well-formatted instance of this example "JSON Schema". The object {{"properties": {{"foo": ["bar", "baz"]}}}} is not well-formatted. + +Your output will be parsed and type-checked according to the provided schema instance, so make sure all fields in your output match the schema exactly and there are no trailing commas! + +Here is the JSON Schema instance your output must adhere to. Include the enclosing markdown codeblock: +\`\`\`json +{"type":"object","properties":{"insights":{"type":"array","items":{"type":"object","properties":{"alertIds":{"type":"array","items":{"type":"string"},"description":"The alert IDs that the insight is based on."},"detailsMarkdown":{"type":"string","description":"A detailed insight with markdown, where each markdown bullet contains a description of what happened that reads like a story of the attack as it played out and always uses special {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax for field names and values from the source data. Examples of CORRECT syntax (includes field names and values): {{ host.name hostNameValue }} {{ user.name userNameValue }} {{ source.ip sourceIpValue }} Examples of INCORRECT syntax (bad, because the field names are not included): {{ hostNameValue }} {{ userNameValue }} {{ sourceIpValue }}"},"entitySummaryMarkdown":{"type":"string","description":"A short (no more than a sentence) summary of the insight featuring only the host.name and user.name fields (when they are applicable), using the same {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax"},"mitreAttackTactics":{"type":"array","items":{"type":"string"},"description":"An array of MITRE ATT&CK tactic for the insight, using one of the following values: Reconnaissance,Initial Access,Execution,Persistence,Privilege Escalation,Discovery,Lateral Movement,Command and Control,Exfiltration"},"summaryMarkdown":{"type":"string","description":"A markdown summary of insight, using the same {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax"},"title":{"type":"string","description":"A short, no more than 7 words, title for the insight, NOT formatted with special syntax or markdown. This must be as brief as possible."}},"required":["alertIds","detailsMarkdown","summaryMarkdown","title"],"additionalProperties":false},"description":"Insights with markdown that always uses special {{ field.name fieldValue1 fieldValue2 fieldValueN }} syntax for field names and values from the source data. Examples of CORRECT syntax (includes field names and values): {{ host.name hostNameValue }} {{ user.name userNameValue }} {{ source.ip sourceIpValue }} Examples of INCORRECT syntax (bad, because the field names are not included): {{ hostNameValue }} {{ userNameValue }} {{ sourceIpValue }}"}},"required":["insights"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"} +\`\`\` +`; + + const chainWithFormatInstructions = getChainWithFormatInstructions(mockLlm); + expect(chainWithFormatInstructions).toEqual({ + chain: expect.any(Object), + formatInstructions: expectedFormatInstructions, + llmType: 'fake', + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined_attack_discovery_prompt/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined_attack_discovery_prompt/index.test.ts new file mode 100644 index 0000000000000..069dd77bed874 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined_attack_discovery_prompt/index.test.ts @@ -0,0 +1,62 @@ +/* + * 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 { getCombinedAttackDiscoveryPrompt } from '.'; + +describe('getCombinedAttackDiscoveryPrompt', () => { + it('returns the initial query when there are no partial results', () => { + const result = getCombinedAttackDiscoveryPrompt({ + anonymizedAlerts: ['alert1', 'alert2'], + attackDiscoveryPrompt: 'attackDiscoveryPrompt', + combinedMaybePartialResults: '', + }); + + expect(result).toBe(`attackDiscoveryPrompt + +Use context from the following alerts to provide insights: + +""" +alert1 + +alert2 +""" +`); + }); + + it('returns the initial query combined with a continuation prompt and partial results', () => { + const result = getCombinedAttackDiscoveryPrompt({ + anonymizedAlerts: ['alert1', 'alert2'], + attackDiscoveryPrompt: 'attackDiscoveryPrompt', + combinedMaybePartialResults: 'partialResults', + }); + + expect(result).toBe(`attackDiscoveryPrompt + +Use context from the following alerts to provide insights: + +""" +alert1 + +alert2 +""" + + +Continue exactly where you left off in the JSON output below, generating only the additional JSON output when it's required to complete your work. The additional JSON output MUST ALWAYS follow these rules: +1) it MUST conform to the schema above, because it will be checked against the JSON schema +2) it MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds), because it will be parsed as JSON +3) it MUST NOT repeat any the previous output, because that would prevent partial results from being combined +4) it MUST NOT restart from the beginning, because that would prevent partial results from being combined +5) it MUST NOT be prefixed or suffixed with additional text outside of the JSON, because that would prevent it from being combined and parsed as JSON: + + +""" +partialResults +""" + +`); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/response_is_hallucinated/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/response_is_hallucinated/index.test.ts new file mode 100644 index 0000000000000..3730d6a7c4b96 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/response_is_hallucinated/index.test.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. + */ + +import { responseIsHallucinated } from '.'; + +describe('responseIsHallucinated', () => { + it('returns true when the response is hallucinated', () => { + expect( + responseIsHallucinated( + 'tactics like **Credential Access**, **Command and Control**, and **Persistence**.",\n "entitySummaryMarkdown": "Malware detected on host **{{ host.name hostNameValue }}**' + ) + ).toBe(true); + }); + + it('returns false when the response is not hallucinated', () => { + expect( + responseIsHallucinated( + 'A malicious file {{ file.name WsmpRExIFs.dll }} was detected on {{ host.name 082a86fa-b87d-45ce-813e-eed6b36ef0a9 }}\\n- The file was executed by' + ) + ).toBe(false); + }); +});