diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.test.tsx new file mode 100644 index 0000000000000..1aaf7879b1c0b --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.test.tsx @@ -0,0 +1,86 @@ +/* + * 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 { fireEvent, render, screen } from '@testing-library/react'; +import React from 'react'; + +import { AlertsRange } from './alerts_range'; +import { + MAX_LATEST_ALERTS, + MIN_LATEST_ALERTS, +} from '../assistant/settings/alerts_settings/alerts_settings'; +import { KnowledgeBaseConfig } from '../assistant/types'; + +const nonDefaultMin = MIN_LATEST_ALERTS + 5000; +const nonDefaultMax = nonDefaultMin + 5000; + +describe('AlertsRange', () => { + beforeEach(() => jest.clearAllMocks()); + + it('renders the expected default min alerts', () => { + render(); + + expect(screen.getByText(`${MIN_LATEST_ALERTS}`)).toBeInTheDocument(); + }); + + it('renders the expected NON-default min alerts', () => { + render( + + ); + + expect(screen.getByText(`${nonDefaultMin}`)).toBeInTheDocument(); + }); + + it('renders the expected default max alerts', () => { + render(); + + expect(screen.getByText(`${MAX_LATEST_ALERTS}`)).toBeInTheDocument(); + }); + + it('renders the expected NON-default max alerts', () => { + render( + + ); + + expect(screen.getByText(`${nonDefaultMax}`)).toBeInTheDocument(); + }); + + it('calls onChange when the range value changes', () => { + const mockOnChange = jest.fn(); + render(); + + fireEvent.click(screen.getByText(`${MAX_LATEST_ALERTS}`)); + + expect(mockOnChange).toHaveBeenCalled(); + }); + + it('calls setUpdatedKnowledgeBaseSettings with the expected arguments', () => { + const mockSetUpdatedKnowledgeBaseSettings = jest.fn(); + const knowledgeBase: KnowledgeBaseConfig = { latestAlerts: 150 }; + + render( + + ); + + fireEvent.click(screen.getByText(`${MAX_LATEST_ALERTS}`)); + + expect(mockSetUpdatedKnowledgeBaseSettings).toHaveBeenCalledWith({ + ...knowledgeBase, + latestAlerts: MAX_LATEST_ALERTS, + }); + }); + + it('renders with the correct initial value', () => { + render(); + + expect(screen.getByTestId('alertsRange')).toHaveValue('250'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/__mocks__/raw_attack_discoveries.ts b/x-pack/plugins/elastic_assistant/server/__mocks__/raw_attack_discoveries.ts new file mode 100644 index 0000000000000..3c2df8aa1ab45 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/__mocks__/raw_attack_discoveries.ts @@ -0,0 +1,106 @@ +/* + * 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'; + +/** + * A mock response from invoking the `attack-discovery` tool. + * This is a JSON string that represents the response from the tool + */ +export const getRawAttackDiscoveriesMock = () => + '{\n "insights": [\n {\n "alertIds": [\n "cced5cec88026ccb68fc0c01c096d6330873ee80838fa367a24c5cd04b679df1",\n "40a4242b163d2552ad24c208dc7ab754f3b2c9cd76fb961ea72391cb5f654580",\n "42ac2ecf60173edff8ef10b32c3b706b866845e75e5107870d7f43f681c819dc",\n "bd8204c37db970bf86c2713325652710d8e5ac2cd43a0f0f2234a65e8e5a0157",\n "b7a073c94cccde9fc4164a1f5aba5169b3ef5e349797326f8b166314c8cdb60d"\n ],\n "detailsMarkdown": "- {{ user.name 1ee7566b-9b26-4f3e-8d2f-0eaafc40cd5d }} executed a suspicious process {{ process.name unix1 }} on {{ host.name 3abc855f-65b6-49b0-ac2f-123e34355b83 }}. The process was located at {{ file.path /Users/james/unix1 }} and had a hash of {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\\n- The process {{ process.name unix1 }} attempted to access sensitive files such as {{ process.args /Users/james/library/Keychains/login.keychain-db }}.\\n- The process {{ process.name unix1 }} was executed with the command line {{ process.command_line /Users/james/unix1 /Users/james/library/Keychains/login.keychain-db TempTemp1234!! }}.\\n- The process {{ process.name unix1 }} was not trusted as indicated by the code signature status {{ process.code_signature.status code failed to satisfy specified code requirement(s) }}.\\n- Another process {{ process.name My Go Application.app }} was also detected on the same host, indicating potential lateral movement or additional malicious activity.",\n "entitySummaryMarkdown": "{{ host.name 3abc855f-65b6-49b0-ac2f-123e34355b83 }} and {{ user.name 1ee7566b-9b26-4f3e-8d2f-0eaafc40cd5d }} were involved in the suspicious activity.",\n "mitreAttackTactics": [\n "Execution",\n "Credential Access",\n "Persistence"\n ],\n "summaryMarkdown": "A critical malware detection alert was triggered on {{ host.name 3abc855f-65b6-49b0-ac2f-123e34355b83 }} involving the user {{ user.name 1ee7566b-9b26-4f3e-8d2f-0eaafc40cd5d }}. The process {{ process.name unix1 }} was executed with suspicious arguments and attempted to access sensitive files.",\n "title": "Critical Malware Detection on macOS Host"\n },\n {\n "alertIds": [\n "1b9c52673b184e6b9bd29b3378f90ec5e7b917c17018ce2d40188a065f145087",\n "881c8cd24296c3efc066f894b2f60e28c86b6398e8d81fcdb0a21e2d4e6f37fb",\n "6ae56534e1246b42afbb0658586bfe03717ee9853cc80d462b9f0aceb44194d3",\n "94dda5ac846d122cf2e582ade68123f036b1b78c63752a30bcf8acdbbbba83ce",\n "250f7967181328c67d1de251c606fd4a791fd81964f431e3d7d76149f531be00"\n ],\n "detailsMarkdown": "- {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} executed a suspicious PowerShell script via Microsoft Office on {{ host.name 5e15d911-50a1-486a-a520-baa449451358 }}. The script was located at {{ file.path C:\\\\ProgramData\\\\WindowsAppPool\\\\AppPool.ps1 }}.\\n- The process {{ process.name powershell.exe }} was executed with the command line {{ process.command_line \\"C:\\\\Windows\\\\System32\\\\WindowsPowerShell\\\\v1.0\\\\powershell.exe\\" -exec bypass -file C:\\\\ProgramData\\\\WindowsAppPool\\\\AppPool.ps1 }}.\\n- The parent process {{ process.parent.name wscript.exe }} was executed by {{ process.parent.command_line wscript C:\\\\ProgramData\\\\WindowsAppPool\\\\AppPool.vbs }}.\\n- The process {{ process.name powershell.exe }} was not trusted as indicated by the code signature status {{ process.code_signature.status trusted }}.\\n- The process {{ process.name powershell.exe }} was detected with a high integrity level, indicating potential privilege escalation.",\n "entitySummaryMarkdown": "{{ host.name 5e15d911-50a1-486a-a520-baa449451358 }} and {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} were involved in the suspicious activity.",\n "mitreAttackTactics": [\n "Initial Access",\n "Execution",\n "Persistence"\n ],\n "summaryMarkdown": "A critical malicious behavior detection alert was triggered on {{ host.name 5e15d911-50a1-486a-a520-baa449451358 }} involving the user {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }}. The process {{ process.name powershell.exe }} was executed with suspicious arguments and attempted to execute a PowerShell script.",\n "title": "Suspicious PowerShell Execution via Microsoft Office"\n },\n {\n "alertIds": [\n "6cbbf7fb63ffed6e091ae21866043df699c839603ec573d3173b36e2d0e66ea3",\n "e7b6f978336961522b0753ffe79cc4a2aa6e2c08c491657ade3eccdb58033852",\n "d3ef244bda90960c091f516874a87b9cf01d206844c2e6ba324e3034472787f5"\n ],\n "detailsMarkdown": "- {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} executed a suspicious PowerShell script via MsiExec on {{ host.name 2068fbbd-341a-477a-b06c-7097ddecd024 }}. The script was located at {{ file.path C:\\\\Users\\\\ADMINI~1\\\\AppData\\\\Local\\\\Temp\\\\2\\\\Package Installation Dir\\\\chch.ps1 }}.\\n- The process {{ process.name powershell.exe }} was executed with the command line {{ process.command_line \\"C:\\\\Windows\\\\System32\\\\WindowsPowerShell\\\\v1.0\\\\powershell.exe\\" -ep bypass -file \\"C:\\\\Users\\\\ADMINI~1\\\\AppData\\\\Local\\\\Temp\\\\2\\\\Package Installation Dir\\\\chch.ps1\\" }}.\\n- The parent process {{ process.parent.name msiexec.exe }} was executed by {{ process.parent.command_line C:\\\\Windows\\\\system32\\\\msiexec.exe /V }}.\\n- The process {{ process.name powershell.exe }} was not trusted as indicated by the code signature status {{ process.code_signature.status trusted }}.\\n- The process {{ process.name powershell.exe }} was detected with a high integrity level, indicating potential privilege escalation.",\n "entitySummaryMarkdown": "{{ host.name 2068fbbd-341a-477a-b06c-7097ddecd024 }} and {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} were involved in the suspicious activity.",\n "mitreAttackTactics": [\n "Defense Evasion",\n "Execution"\n ],\n "summaryMarkdown": "A critical malicious behavior detection alert was triggered on {{ host.name 2068fbbd-341a-477a-b06c-7097ddecd024 }} involving the user {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }}. The process {{ process.name powershell.exe }} was executed with suspicious arguments and attempted to execute a PowerShell script.",\n "title": "Suspicious PowerShell Execution via MsiExec"\n },\n {\n "alertIds": [\n "8b1ccd0bfb927caeb5f9818098eebde9a091b99334c84bfffd36aa83db8b36ee",\n "0ae1370d0c08d651a05421009ed8358d9037f3d6af0cf5f3417979489ca80f12",\n "bed4a026232fb8e67f248771a99af722116556ace7ef9aaddefc082da4209c61",\n "d28f2c32ae8e6bc33edfe51ace4621c0e7b826c087386c46ce9138be92baf3f9"\n ],\n "detailsMarkdown": "- {{ user.name 00468e82-e37f-4224-80c1-c62e594c74b1 }} executed a suspicious process {{ process.name unzip }} on {{ host.name b557bb12-8206-44b6-b2a5-dbcce5b1e65e }}. The process was located at {{ file.path /home/ubuntu/74ef6cc38f5a1a80148752b63c117e6846984debd2af806c65887195a8eccc56 }} and had a hash of {{ file.hash.sha256 74ef6cc38f5a1a80148752b63c117e6846984debd2af806c65887195a8eccc56 }}.\\n- The process {{ process.name unzip }} was executed with the command line {{ process.command_line unzip 9415656314.zip }}.\\n- The process {{ process.name unzip }} was detected with a high integrity level, indicating potential privilege escalation.\\n- Another process {{ process.name kdmtmpflush }} was also detected on the same host, indicating potential lateral movement or additional malicious activity.",\n "entitySummaryMarkdown": "{{ host.name b557bb12-8206-44b6-b2a5-dbcce5b1e65e }} and {{ user.name 00468e82-e37f-4224-80c1-c62e594c74b1 }} were involved in the suspicious activity.",\n "mitreAttackTactics": [\n "Execution",\n "Persistence"\n ],\n "summaryMarkdown": "A critical malware detection alert was triggered on {{ host.name b557bb12-8206-44b6-b2a5-dbcce5b1e65e }} involving the user {{ user.name 00468e82-e37f-4224-80c1-c62e594c74b1 }}. The process {{ process.name unzip }} was executed with suspicious arguments and attempted to extract a potentially malicious file.",\n "title": "Suspicious File Extraction on Linux Host"\n },\n {\n "alertIds": [\n "15c3053659b3bccbcc2c75eb90963596bbba707496e6b8c4927b5dc3995e0e11",\n "461fedbfddd0d8d42c11630d5cdb9a103fac05327dff5bcdbf51505f01ec39da",\n "03ef2d6a825993d08f545cfa25e8dab765dd1f4688124e7d12d8d81a2f324464",\n "bfd4f9a71c9ca6a8dc68a41ea96b5ca14380da9669fb62ccae06769ad931eef2"\n ],\n "detailsMarkdown": "- {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} executed a suspicious process {{ process.name MsMpEng.exe }} on {{ host.name b808feb3-7ab3-4006-9c67-3cf7aeffe572 }}. The process was located at {{ file.path C:\\\\Windows\\\\MsMpEng.exe }} and had a hash of {{ file.hash.sha256 33bc14d231a4afaa18f06513766d5f69d8b88f1e697cd127d24fb4b72ad44c7a }}.\\n- The process {{ process.name MsMpEng.exe }} was executed with the command line {{ process.command_line \\"C:\\\\Windows\\\\MsMpEng.exe\\" }}.\\n- The parent process {{ process.parent.name d55f983c994caa160ec63a59f6b4250fe67fb3e8c43a388aec60a4a6978e9f1e.exe }} was executed by {{ process.parent.command_line \\"C:\\\\Users\\\\Administrator\\\\Desktop\\\\8813719803\\\\d55f983c994caa160ec63a59f6b4250fe67fb3e8c43a388aec60a4a6978e9f1e.exe\\" }}.\\n- The process {{ process.name MsMpEng.exe }} was not trusted as indicated by the code signature status {{ process.code_signature.status trusted }}.\\n- The process {{ process.name MsMpEng.exe }} was detected with a high integrity level, indicating potential privilege escalation.",\n "entitySummaryMarkdown": "{{ host.name b808feb3-7ab3-4006-9c67-3cf7aeffe572 }} and {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} were involved in the suspicious activity.",\n "mitreAttackTactics": [\n "Execution",\n "Persistence"\n ],\n "summaryMarkdown": "A critical malware detection alert was triggered on {{ host.name b808feb3-7ab3-4006-9c67-3cf7aeffe572 }} involving the user {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }}. The process {{ process.name MsMpEng.exe }} was executed with suspicious arguments and attempted to execute a potentially malicious file.",\n "title": "Suspicious Process Execution on Windows Host"\n }\n ]\n}'; + +export const getParsedAttackDiscoveriesMock = ( + attackDiscoveryTimestamp: string +): AttackDiscovery[] => [ + { + alertIds: [ + 'cced5cec88026ccb68fc0c01c096d6330873ee80838fa367a24c5cd04b679df1', + '40a4242b163d2552ad24c208dc7ab754f3b2c9cd76fb961ea72391cb5f654580', + '42ac2ecf60173edff8ef10b32c3b706b866845e75e5107870d7f43f681c819dc', + 'bd8204c37db970bf86c2713325652710d8e5ac2cd43a0f0f2234a65e8e5a0157', + 'b7a073c94cccde9fc4164a1f5aba5169b3ef5e349797326f8b166314c8cdb60d', + ], + detailsMarkdown: + '- {{ user.name 1ee7566b-9b26-4f3e-8d2f-0eaafc40cd5d }} executed a suspicious process {{ process.name unix1 }} on {{ host.name 3abc855f-65b6-49b0-ac2f-123e34355b83 }}. The process was located at {{ file.path /Users/james/unix1 }} and had a hash of {{ file.hash.sha256 0b18d6880dc9670ab2b955914598c96fc3d0097dc40ea61157b8c79e75edf231 }}.\n- The process {{ process.name unix1 }} attempted to access sensitive files such as {{ process.args /Users/james/library/Keychains/login.keychain-db }}.\n- The process {{ process.name unix1 }} was executed with the command line {{ process.command_line /Users/james/unix1 /Users/james/library/Keychains/login.keychain-db TempTemp1234!! }}.\n- The process {{ process.name unix1 }} was not trusted as indicated by the code signature status {{ process.code_signature.status code failed to satisfy specified code requirement(s) }}.\n- Another process {{ process.name My Go Application.app }} was also detected on the same host, indicating potential lateral movement or additional malicious activity.', + entitySummaryMarkdown: + '{{ host.name 3abc855f-65b6-49b0-ac2f-123e34355b83 }} and {{ user.name 1ee7566b-9b26-4f3e-8d2f-0eaafc40cd5d }} were involved in the suspicious activity.', + mitreAttackTactics: ['Execution', 'Credential Access', 'Persistence'], + summaryMarkdown: + 'A critical malware detection alert was triggered on {{ host.name 3abc855f-65b6-49b0-ac2f-123e34355b83 }} involving the user {{ user.name 1ee7566b-9b26-4f3e-8d2f-0eaafc40cd5d }}. The process {{ process.name unix1 }} was executed with suspicious arguments and attempted to access sensitive files.', + title: 'Critical Malware Detection on macOS Host', + timestamp: attackDiscoveryTimestamp, + }, + { + alertIds: [ + '1b9c52673b184e6b9bd29b3378f90ec5e7b917c17018ce2d40188a065f145087', + '881c8cd24296c3efc066f894b2f60e28c86b6398e8d81fcdb0a21e2d4e6f37fb', + '6ae56534e1246b42afbb0658586bfe03717ee9853cc80d462b9f0aceb44194d3', + '94dda5ac846d122cf2e582ade68123f036b1b78c63752a30bcf8acdbbbba83ce', + '250f7967181328c67d1de251c606fd4a791fd81964f431e3d7d76149f531be00', + ], + detailsMarkdown: + '- {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} executed a suspicious PowerShell script via Microsoft Office on {{ host.name 5e15d911-50a1-486a-a520-baa449451358 }}. The script was located at {{ file.path C:\\ProgramData\\WindowsAppPool\\AppPool.ps1 }}.\n- The process {{ process.name powershell.exe }} was executed with the command line {{ process.command_line "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe" -exec bypass -file C:\\ProgramData\\WindowsAppPool\\AppPool.ps1 }}.\n- The parent process {{ process.parent.name wscript.exe }} was executed by {{ process.parent.command_line wscript C:\\ProgramData\\WindowsAppPool\\AppPool.vbs }}.\n- The process {{ process.name powershell.exe }} was not trusted as indicated by the code signature status {{ process.code_signature.status trusted }}.\n- The process {{ process.name powershell.exe }} was detected with a high integrity level, indicating potential privilege escalation.', + entitySummaryMarkdown: + '{{ host.name 5e15d911-50a1-486a-a520-baa449451358 }} and {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} were involved in the suspicious activity.', + mitreAttackTactics: ['Initial Access', 'Execution', 'Persistence'], + summaryMarkdown: + 'A critical malicious behavior detection alert was triggered on {{ host.name 5e15d911-50a1-486a-a520-baa449451358 }} involving the user {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }}. The process {{ process.name powershell.exe }} was executed with suspicious arguments and attempted to execute a PowerShell script.', + title: 'Suspicious PowerShell Execution via Microsoft Office', + timestamp: attackDiscoveryTimestamp, + }, + { + alertIds: [ + '6cbbf7fb63ffed6e091ae21866043df699c839603ec573d3173b36e2d0e66ea3', + 'e7b6f978336961522b0753ffe79cc4a2aa6e2c08c491657ade3eccdb58033852', + 'd3ef244bda90960c091f516874a87b9cf01d206844c2e6ba324e3034472787f5', + ], + detailsMarkdown: + '- {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} executed a suspicious PowerShell script via MsiExec on {{ host.name 2068fbbd-341a-477a-b06c-7097ddecd024 }}. The script was located at {{ file.path C:\\Users\\ADMINI~1\\AppData\\Local\\Temp\\2\\Package Installation Dir\\chch.ps1 }}.\n- The process {{ process.name powershell.exe }} was executed with the command line {{ process.command_line "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe" -ep bypass -file "C:\\Users\\ADMINI~1\\AppData\\Local\\Temp\\2\\Package Installation Dir\\chch.ps1" }}.\n- The parent process {{ process.parent.name msiexec.exe }} was executed by {{ process.parent.command_line C:\\Windows\\system32\\msiexec.exe /V }}.\n- The process {{ process.name powershell.exe }} was not trusted as indicated by the code signature status {{ process.code_signature.status trusted }}.\n- The process {{ process.name powershell.exe }} was detected with a high integrity level, indicating potential privilege escalation.', + entitySummaryMarkdown: + '{{ host.name 2068fbbd-341a-477a-b06c-7097ddecd024 }} and {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} were involved in the suspicious activity.', + mitreAttackTactics: ['Defense Evasion', 'Execution'], + summaryMarkdown: + 'A critical malicious behavior detection alert was triggered on {{ host.name 2068fbbd-341a-477a-b06c-7097ddecd024 }} involving the user {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }}. The process {{ process.name powershell.exe }} was executed with suspicious arguments and attempted to execute a PowerShell script.', + title: 'Suspicious PowerShell Execution via MsiExec', + timestamp: attackDiscoveryTimestamp, + }, + { + alertIds: [ + '8b1ccd0bfb927caeb5f9818098eebde9a091b99334c84bfffd36aa83db8b36ee', + '0ae1370d0c08d651a05421009ed8358d9037f3d6af0cf5f3417979489ca80f12', + 'bed4a026232fb8e67f248771a99af722116556ace7ef9aaddefc082da4209c61', + 'd28f2c32ae8e6bc33edfe51ace4621c0e7b826c087386c46ce9138be92baf3f9', + ], + detailsMarkdown: + '- {{ user.name 00468e82-e37f-4224-80c1-c62e594c74b1 }} executed a suspicious process {{ process.name unzip }} on {{ host.name b557bb12-8206-44b6-b2a5-dbcce5b1e65e }}. The process was located at {{ file.path /home/ubuntu/74ef6cc38f5a1a80148752b63c117e6846984debd2af806c65887195a8eccc56 }} and had a hash of {{ file.hash.sha256 74ef6cc38f5a1a80148752b63c117e6846984debd2af806c65887195a8eccc56 }}.\n- The process {{ process.name unzip }} was executed with the command line {{ process.command_line unzip 9415656314.zip }}.\n- The process {{ process.name unzip }} was detected with a high integrity level, indicating potential privilege escalation.\n- Another process {{ process.name kdmtmpflush }} was also detected on the same host, indicating potential lateral movement or additional malicious activity.', + entitySummaryMarkdown: + '{{ host.name b557bb12-8206-44b6-b2a5-dbcce5b1e65e }} and {{ user.name 00468e82-e37f-4224-80c1-c62e594c74b1 }} were involved in the suspicious activity.', + mitreAttackTactics: ['Execution', 'Persistence'], + summaryMarkdown: + 'A critical malware detection alert was triggered on {{ host.name b557bb12-8206-44b6-b2a5-dbcce5b1e65e }} involving the user {{ user.name 00468e82-e37f-4224-80c1-c62e594c74b1 }}. The process {{ process.name unzip }} was executed with suspicious arguments and attempted to extract a potentially malicious file.', + title: 'Suspicious File Extraction on Linux Host', + timestamp: attackDiscoveryTimestamp, + }, + { + alertIds: [ + '15c3053659b3bccbcc2c75eb90963596bbba707496e6b8c4927b5dc3995e0e11', + '461fedbfddd0d8d42c11630d5cdb9a103fac05327dff5bcdbf51505f01ec39da', + '03ef2d6a825993d08f545cfa25e8dab765dd1f4688124e7d12d8d81a2f324464', + 'bfd4f9a71c9ca6a8dc68a41ea96b5ca14380da9669fb62ccae06769ad931eef2', + ], + detailsMarkdown: + '- {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} executed a suspicious process {{ process.name MsMpEng.exe }} on {{ host.name b808feb3-7ab3-4006-9c67-3cf7aeffe572 }}. The process was located at {{ file.path C:\\Windows\\MsMpEng.exe }} and had a hash of {{ file.hash.sha256 33bc14d231a4afaa18f06513766d5f69d8b88f1e697cd127d24fb4b72ad44c7a }}.\n- The process {{ process.name MsMpEng.exe }} was executed with the command line {{ process.command_line "C:\\Windows\\MsMpEng.exe" }}.\n- The parent process {{ process.parent.name d55f983c994caa160ec63a59f6b4250fe67fb3e8c43a388aec60a4a6978e9f1e.exe }} was executed by {{ process.parent.command_line "C:\\Users\\Administrator\\Desktop\\8813719803\\d55f983c994caa160ec63a59f6b4250fe67fb3e8c43a388aec60a4a6978e9f1e.exe" }}.\n- The process {{ process.name MsMpEng.exe }} was not trusted as indicated by the code signature status {{ process.code_signature.status trusted }}.\n- The process {{ process.name MsMpEng.exe }} was detected with a high integrity level, indicating potential privilege escalation.', + entitySummaryMarkdown: + '{{ host.name b808feb3-7ab3-4006-9c67-3cf7aeffe572 }} and {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }} were involved in the suspicious activity.', + mitreAttackTactics: ['Execution', 'Persistence'], + summaryMarkdown: + 'A critical malware detection alert was triggered on {{ host.name b808feb3-7ab3-4006-9c67-3cf7aeffe572 }} involving the user {{ user.name 37764a98-eeb5-459f-ab04-f8b70e8239cb }}. The process {{ process.name MsMpEng.exe }} was executed with suspicious arguments and attempted to execute a potentially malicious file.', + title: 'Suspicious Process Execution on Windows Host', + timestamp: attackDiscoveryTimestamp, + }, +]; diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymization_fields.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymization_fields.ts new file mode 100644 index 0000000000000..ed487e4705c27 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/__mocks__/mock_anonymization_fields.ts @@ -0,0 +1,65 @@ +/* + * 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 { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; + +export const getMockAnonymizationFieldResponse = (): AnonymizationFieldResponse[] => [ + { + id: '6UDO45IBoEQSo_rIK1EW', + timestamp: '2024-10-31T18:19:52.468Z', + field: '_id', + allowed: true, + anonymized: false, + createdAt: '2024-10-31T18:19:52.468Z', + namespace: 'default', + }, + { + id: '6kDO45IBoEQSo_rIK1EW', + timestamp: '2024-10-31T18:19:52.468Z', + field: '@timestamp', + allowed: true, + anonymized: false, + createdAt: '2024-10-31T18:19:52.468Z', + namespace: 'default', + }, + { + id: '60DO45IBoEQSo_rIK1EW', + timestamp: '2024-10-31T18:19:52.468Z', + field: 'cloud.availability_zone', + allowed: true, + anonymized: false, + createdAt: '2024-10-31T18:19:52.468Z', + namespace: 'default', + }, + { + id: '_EDO45IBoEQSo_rIK1EW', + timestamp: '2024-10-31T18:19:52.468Z', + field: 'host.name', + allowed: true, + anonymized: true, + createdAt: '2024-10-31T18:19:52.468Z', + namespace: 'default', + }, + { + id: 'SkDO45IBoEQSo_rIK1IW', + timestamp: '2024-10-31T18:19:52.468Z', + field: 'user.name', + allowed: true, + anonymized: true, + createdAt: '2024-10-31T18:19:52.468Z', + namespace: 'default', + }, + { + id: 'TUDO45IBoEQSo_rIK1IW', + timestamp: '2024-10-31T18:19:52.468Z', + field: 'user.target.name', + allowed: true, + anonymized: true, + createdAt: '2024-10-31T18:19:52.468Z', + namespace: 'default', + }, +]; diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_alerts_context_prompt/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_alerts_context_prompt/index.test.ts index 287f5e6b2130a..098d2b81f4914 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_alerts_context_prompt/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/helpers/get_alerts_context_prompt/index.test.ts @@ -12,7 +12,7 @@ describe('getAlertsContextPrompt', () => { it('generates the correct prompt', () => { const anonymizedAlerts = ['Alert 1', 'Alert 2', 'Alert 3']; - const expected = `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). + const expected = `${getDefaultAttackDiscoveryPrompt()} Use context from the following alerts to provide insights: 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 index da815aad9795a..07c3e0007f851 100644 --- 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 @@ -5,8 +5,8 @@ * 2.0. */ +import type { Logger } from '@kbn/core/server'; import type { ActionsClientLlm } from '@kbn/langchain/server'; -import { loggerMock } from '@kbn/logging-mocks'; import { FakeLLM } from '@langchain/core/utils/testing'; import { getGenerateNode } from '.'; @@ -16,7 +16,15 @@ import { } 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 { getDefaultAttackDiscoveryPrompt } from '../helpers/get_default_attack_discovery_prompt'; +import { getDefaultRefinePrompt } from '../refine/helpers/get_default_refine_prompt'; import { GraphState } from '../../types'; +import { + getParsedAttackDiscoveriesMock, + getRawAttackDiscoveriesMock, +} from '../../../../../../__mocks__/raw_attack_discoveries'; + +const attackDiscoveryTimestamp = '2024-10-11T17:55:59.702Z'; jest.mock('../helpers/get_chain_with_format_instructions', () => { const mockInvoke = jest.fn().mockResolvedValue(''); @@ -27,19 +35,21 @@ jest.mock('../helpers/get_chain_with_format_instructions', () => { invoke: mockInvoke, }, formatInstructions: ['mock format instructions'], - llmType: 'fake', + llmType: 'openai', mockInvoke, // <-- added for testing }), }; }); -const mockLogger = loggerMock.create(); +const mockLogger = { + debug: (x: Function) => x(), +} as unknown as Logger; + 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).", + attackDiscoveryPrompt: getDefaultAttackDiscoveryPrompt(), anonymizedAlerts: [...mockAnonymizedAlerts], combinedGenerations: '', combinedRefinements: '', @@ -51,8 +61,7 @@ const initialGraphState: GraphState = { 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:', + refinePrompt: getDefaultRefinePrompt(), replacements: { ...mockAnonymizedAlertsReplacements, }, @@ -63,11 +72,18 @@ describe('getGenerateNode', () => { beforeEach(() => { jest.clearAllMocks(); + jest.useFakeTimers(); + jest.setSystemTime(new Date(attackDiscoveryTimestamp)); + mockLlm = new FakeLLM({ - response: JSON.stringify({}, null, 2), + response: '', }) as unknown as ActionsClientLlm; }); + afterEach(() => { + jest.useRealTimers(); + }); + it('returns a function', () => { const generateNode = getGenerateNode({ llm: mockLlm, @@ -77,9 +93,8 @@ describe('getGenerateNode', () => { 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); + it('invokes the chain with the expected alerts from state and formatting instructions', async () => { + const mockInvoke = getChainWithFormatInstructions(mockLlm).chain.invoke as jest.Mock; const generateNode = getGenerateNode({ llm: mockLlm, @@ -100,4 +115,214 @@ ${getAnonymizedAlertsFromState(initialGraphState).join('\n\n')} `, }); }); + + it('removes the surrounding json from the response', async () => { + const response = + 'You asked for some JSON, here it is:\n```json\n{"key": "value"}\n```\nI hope that works for you.'; + + const mockLlmWithResponse = new FakeLLM({ response }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const generateNode = getGenerateNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const state = await generateNode(initialGraphState); + + expect(state).toEqual({ + ...initialGraphState, + combinedGenerations: '{"key": "value"}', + errors: [ + 'generate node is unable to parse (fake) response from attempt 0; (this may be an incomplete response from the model): [\n {\n "code": "invalid_type",\n "expected": "array",\n "received": "undefined",\n "path": [\n "insights"\n ],\n "message": "Required"\n }\n]', + ], + generationAttempts: 1, + generations: ['{"key": "value"}'], + }); + }); + + it('handles hallucinations', async () => { + const hallucinatedResponse = + 'tactics like **Credential Access**, **Command and Control**, and **Persistence**.",\n "entitySummaryMarkdown": "Malware detected on host **{{ host.name hostNameValue }}**'; + + const mockLlmWithHallucination = new FakeLLM({ + response: hallucinatedResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithHallucination).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(hallucinatedResponse); + + const generateNode = getGenerateNode({ + llm: mockLlmWithHallucination, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: '{"key": "value"}', + generationAttempts: 1, + generations: ['{"key": "value"}'], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedGenerations: '', // <-- reset + generationAttempts: 2, // <-- incremented + generations: [], // <-- reset + hallucinationFailures: 1, // <-- incremented + }); + }); + + it('discards previous generations and starts over when the maxRepeatedGenerations limit is reached', async () => { + const repeatedResponse = 'gen1'; + + const mockLlmWithRepeatedGenerations = new FakeLLM({ + response: repeatedResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithRepeatedGenerations).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(repeatedResponse); + + const generateNode = getGenerateNode({ + llm: mockLlmWithRepeatedGenerations, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: 'gen1gen1', + generationAttempts: 2, + generations: ['gen1', 'gen1'], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedGenerations: '', + generationAttempts: 3, // <-- incremented + generations: [], + }); + }); + + it('combines the response with the previous generations', async () => { + const response = 'gen1'; + + const mockLlmWithResponse = new FakeLLM({ + response, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const generateNode = getGenerateNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: 'gen0', + generationAttempts: 1, + generations: ['gen0'], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedGenerations: 'gen0gen1', + errors: [ + 'generate node is unable to parse (fake) response from attempt 1; (this may be an incomplete response from the model): SyntaxError: Unexpected token \'g\', "gen0gen1" is not valid JSON', + ], + generationAttempts: 2, + generations: ['gen0', 'gen1'], + }); + }); + + it('returns unrefined results when combined responses pass validation', async () => { + // split the response into two parts to simulate a valid response + const splitIndex = 100; // arbitrary index + const firstResponse = getRawAttackDiscoveriesMock().slice(0, splitIndex); + const secondResponse = getRawAttackDiscoveriesMock().slice(splitIndex); + + const mockLlmWithResponse = new FakeLLM({ + response: secondResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(secondResponse); + + const generateNode = getGenerateNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: firstResponse, + generationAttempts: 1, + generations: [firstResponse], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + attackDiscoveries: null, + combinedGenerations: firstResponse.concat(secondResponse), + errors: [], + generationAttempts: 2, + generations: [firstResponse, secondResponse], + unrefinedResults: getParsedAttackDiscoveriesMock(attackDiscoveryTimestamp), // <-- generated from the combined response + }); + }); + + it('skips the refinements step if the max number of retries has already been reached', async () => { + // split the response into two parts to simulate a valid response + const splitIndex = 100; // arbitrary index + const firstResponse = getRawAttackDiscoveriesMock().slice(0, splitIndex); + const secondResponse = getRawAttackDiscoveriesMock().slice(splitIndex); + + const mockLlmWithResponse = new FakeLLM({ + response: secondResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(secondResponse); + + const generateNode = getGenerateNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: firstResponse, + generationAttempts: 9, + generations: [firstResponse], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + attackDiscoveries: getParsedAttackDiscoveriesMock(attackDiscoveryTimestamp), // <-- skip the refinement step + combinedGenerations: firstResponse.concat(secondResponse), + errors: [], + generationAttempts: 10, + generations: [firstResponse, secondResponse], + unrefinedResults: getParsedAttackDiscoveriesMock(attackDiscoveryTimestamp), // <-- generated from the combined response + }); + }); }); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.ts index 1fcd81622f0fe..0dfe1b0629f58 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/generate/index.ts @@ -58,10 +58,10 @@ export const getGenerateNode = ({ () => `generate node is invoking the chain (${llmType}), attempt ${generationAttempts}` ); - const rawResponse = (await chain.invoke({ + const rawResponse = await chain.invoke({ format_instructions: formatInstructions, query, - })) as unknown as string; + }); // LOCAL MUTATION: partialResponse = extractJson(rawResponse); // remove the surrounding ```json``` @@ -86,7 +86,7 @@ export const getGenerateNode = ({ generationsAreRepeating({ currentGeneration: partialResponse, previousGenerations: generations, - sampleLastNGenerations: maxRepeatedGenerations, + sampleLastNGenerations: maxRepeatedGenerations - 1, }) ) { logger?.debug( diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts new file mode 100644 index 0000000000000..4c95cb05faae0 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts @@ -0,0 +1,46 @@ +/* + * 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 { addTrailingBackticksIfNecessary } from '.'; + +describe('addTrailingBackticksIfNecessary', () => { + it('adds trailing backticks when necessary', () => { + const input = '```json\n{\n "key": "value"\n}'; + const expected = '```json\n{\n "key": "value"\n}\n```'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(expected); + }); + + it('does NOT add trailing backticks when they are already present', () => { + const input = '```json\n{\n "key": "value"\n}\n```'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); + + it("does NOT add trailing backticks when there's no leading JSON wrapper", () => { + const input = '{\n "key": "value"\n}'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); + + it('handles empty string input', () => { + const input = ''; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); + + it('handles input without a JSON wrapper, but with trailing backticks', () => { + const input = '{\n "key": "value"\n}\n```'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.test.ts index 5e13ec9f0dafe..7a2ced64163c5 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.test.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.test.ts @@ -8,6 +8,24 @@ import { extractJson } from '.'; describe('extractJson', () => { + it('returns an empty string if input is undefined', () => { + const input = undefined; + + expect(extractJson(input)).toBe(''); + }); + + it('returns an empty string if input an array', () => { + const input = ['some', 'array']; + + expect(extractJson(input)).toBe(''); + }); + + it('returns an empty string if input is an object', () => { + const input = {}; + + expect(extractJson(input)).toBe(''); + }); + it('returns the JSON text surrounded by ```json and ``` with no whitespace or additional text', () => { const input = '```json{"key": "value"}```'; diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.ts index 79d3f9c0d0599..089756840e568 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/extract_json/index.ts @@ -5,7 +5,11 @@ * 2.0. */ -export const extractJson = (input: string): string => { +export const extractJson = (input: unknown): string => { + if (typeof input !== 'string') { + return ''; + } + const regex = /```json\s*([\s\S]*?)(?:\s*```|$)/; const match = input.match(regex); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined/index.test.ts new file mode 100644 index 0000000000000..75d7d83db3e92 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_combined/index.test.ts @@ -0,0 +1,37 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { getCombined } from '.'; + +describe('getCombined', () => { + it('combines two strings correctly', () => { + const combinedGenerations = 'generation1'; + const partialResponse = 'response1'; + const expected = 'generation1response1'; + const result = getCombined({ combinedGenerations, partialResponse }); + + expect(result).toEqual(expected); + }); + + it('handles empty combinedGenerations', () => { + const combinedGenerations = ''; + const partialResponse = 'response1'; + const expected = 'response1'; + const result = getCombined({ combinedGenerations, partialResponse }); + + expect(result).toEqual(expected); + }); + + it('handles an empty partialResponse', () => { + const combinedGenerations = 'generation1'; + const partialResponse = ''; + const expected = 'generation1'; + const result = getCombined({ combinedGenerations, partialResponse }); + + expect(result).toEqual(expected); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_continue_prompt/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_continue_prompt/index.test.ts new file mode 100644 index 0000000000000..35dae31a3ae6a --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_continue_prompt/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 { getContinuePrompt } from '.'; + +describe('getContinuePrompt', () => { + it('returns the expected prompt string', () => { + const expectedPrompt = `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: +`; + + expect(getContinuePrompt()).toBe(expectedPrompt); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_default_attack_discovery_prompt/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_default_attack_discovery_prompt/index.test.ts new file mode 100644 index 0000000000000..6fce86bfceb6f --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/get_default_attack_discovery_prompt/index.test.ts @@ -0,0 +1,16 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { getDefaultAttackDiscoveryPrompt } from '.'; + +describe('getDefaultAttackDiscoveryPrompt', () => { + it('returns the default attack discovery prompt', () => { + expect(getDefaultAttackDiscoveryPrompt()).toEqual( + "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)." + ); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/parse_combined_or_throw/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/parse_combined_or_throw/index.test.ts new file mode 100644 index 0000000000000..cfbc837a83f66 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/parse_combined_or_throw/index.test.ts @@ -0,0 +1,118 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { Logger } from '@kbn/core/server'; + +import { parseCombinedOrThrow } from '.'; +import { getRawAttackDiscoveriesMock } from '../../../../../../../__mocks__/raw_attack_discoveries'; + +describe('parseCombinedOrThrow', () => { + const mockLogger: Logger = { + debug: jest.fn(), + } as unknown as Logger; + + const nodeName = 'testNodeName'; + const llmType = 'testLlmType'; + + const validCombinedResponse = getRawAttackDiscoveriesMock(); + + const invalidCombinedResponse = 'invalid json'; + + const defaultArgs = { + combinedResponse: validCombinedResponse, + generationAttempts: 0, + nodeName, + llmType, + logger: mockLogger, + }; + + it('returns an Attack discovery for each insight in a valid combined response', () => { + const discoveries = parseCombinedOrThrow({ + ...defaultArgs, + }); + + expect(discoveries).toHaveLength(5); + }); + + it('adds a timestamp to all discoveries in a valid response', () => { + const discoveries = parseCombinedOrThrow({ + ...defaultArgs, + }); + + expect(discoveries.every((discovery) => discovery.timestamp != null)).toBe(true); + }); + + it('adds trailing backticks to the combined response if necessary', () => { + const withLeadingJson = '```json\n'.concat(validCombinedResponse); + + const discoveries = parseCombinedOrThrow({ + ...defaultArgs, + combinedResponse: withLeadingJson, + }); + + expect(discoveries).toHaveLength(5); + }); + + it('logs the parsing step', () => { + const generationAttempts = 0; + + parseCombinedOrThrow({ + ...defaultArgs, + generationAttempts, + }); + + expect((mockLogger.debug as jest.Mock).mock.calls[0][0]()).toBe( + `${nodeName} node is parsing extractedJson (${llmType}) from attempt ${generationAttempts}` + ); + }); + + it('logs the validation step', () => { + const generationAttempts = 0; + + parseCombinedOrThrow({ + ...defaultArgs, + generationAttempts, + }); + + expect((mockLogger.debug as jest.Mock).mock.calls[1][0]()).toBe( + `${nodeName} node is validating combined response (${llmType}) from attempt ${generationAttempts}` + ); + }); + + it('logs the successful validation step', () => { + const generationAttempts = 0; + + parseCombinedOrThrow({ + ...defaultArgs, + generationAttempts, + }); + + expect((mockLogger.debug as jest.Mock).mock.calls[2][0]()).toBe( + `${nodeName} node successfully validated Attack discoveries response (${llmType}) from attempt ${generationAttempts}` + ); + }); + + it('throws the expected error when JSON parsing fails', () => { + expect(() => + parseCombinedOrThrow({ + ...defaultArgs, + combinedResponse: invalidCombinedResponse, + }) + ).toThrowError('Unexpected token \'i\', "invalid json" is not valid JSON'); + }); + + it('throws the expected error when JSON validation fails', () => { + const invalidJson = '{ "insights": "not an array" }'; + + expect(() => + parseCombinedOrThrow({ + ...defaultArgs, + combinedResponse: invalidJson, + }) + ).toThrowError('Expected array, received string'); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts new file mode 100644 index 0000000000000..1409b3d47473c --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts @@ -0,0 +1,82 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { discardPreviousRefinements } from '.'; +import { mockAttackDiscoveries } from '../../../../../../evaluation/__mocks__/mock_attack_discoveries'; +import { GraphState } from '../../../../types'; + +const initialState: GraphState = { + anonymizedAlerts: [], + attackDiscoveries: null, + attackDiscoveryPrompt: 'attackDiscoveryPrompt', + combinedGenerations: 'generation1generation2', + combinedRefinements: 'refinement1', // <-- existing refinements + errors: [], + generationAttempts: 3, + generations: ['generation1', 'generation2'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: ['refinement1'], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: [...mockAttackDiscoveries], +}; + +describe('discardPreviousRefinements', () => { + describe('common state updates', () => { + let result: GraphState; + + beforeEach(() => { + result = discardPreviousRefinements({ + generationAttempts: initialState.generationAttempts, + hallucinationFailures: initialState.hallucinationFailures, + isHallucinationDetected: true, + state: initialState, + }); + }); + + it('resets the combined refinements', () => { + expect(result.combinedRefinements).toBe(''); + }); + + it('increments the generation attempts', () => { + expect(result.generationAttempts).toBe(initialState.generationAttempts + 1); + }); + + it('resets the refinements', () => { + expect(result.refinements).toEqual([]); + }); + + it('increments the hallucination failures when hallucinations are detected', () => { + expect(result.hallucinationFailures).toBe(initialState.hallucinationFailures + 1); + }); + }); + + it('increments the hallucination failures when hallucinations are detected', () => { + const result = discardPreviousRefinements({ + generationAttempts: initialState.generationAttempts, + hallucinationFailures: initialState.hallucinationFailures, + isHallucinationDetected: true, // <-- hallucinations detected + state: initialState, + }); + + expect(result.hallucinationFailures).toBe(initialState.hallucinationFailures + 1); + }); + + it('does NOT increment the hallucination failures when hallucinations are NOT detected', () => { + const result = discardPreviousRefinements({ + generationAttempts: initialState.generationAttempts, + hallucinationFailures: initialState.hallucinationFailures, + isHallucinationDetected: false, // <-- no hallucinations detected + state: initialState, + }); + + expect(result.hallucinationFailures).toBe(initialState.hallucinationFailures); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts new file mode 100644 index 0000000000000..f955f1b175b5b --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts @@ -0,0 +1,77 @@ +/* + * 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 { getCombinedRefinePrompt } from '.'; +import { mockAttackDiscoveries } from '../../../../../../evaluation/__mocks__/mock_attack_discoveries'; +import { getContinuePrompt } from '../../../helpers/get_continue_prompt'; + +describe('getCombinedRefinePrompt', () => { + it('returns the base query when combinedRefinements is empty', () => { + const result = getCombinedRefinePrompt({ + attackDiscoveryPrompt: 'Initial query', + combinedRefinements: '', + refinePrompt: 'Refine prompt', + unrefinedResults: [...mockAttackDiscoveries], + }); + + expect(result).toEqual(`Initial query + +Refine prompt + +""" +${JSON.stringify(mockAttackDiscoveries, null, 2)} +""" + +`); + }); + + it('returns the combined prompt when combinedRefinements is not empty', () => { + const result = getCombinedRefinePrompt({ + attackDiscoveryPrompt: 'Initial query', + combinedRefinements: 'Combined refinements', + refinePrompt: 'Refine prompt', + unrefinedResults: [...mockAttackDiscoveries], + }); + + expect(result).toEqual(`Initial query + +Refine prompt + +""" +${JSON.stringify(mockAttackDiscoveries, null, 2)} +""" + + + +${getContinuePrompt()} + +""" +Combined refinements +""" + +`); + }); + + it('handles null unrefinedResults', () => { + const result = getCombinedRefinePrompt({ + attackDiscoveryPrompt: 'Initial query', + combinedRefinements: '', + refinePrompt: 'Refine prompt', + unrefinedResults: null, + }); + + expect(result).toEqual(`Initial query + +Refine prompt + +""" +null +""" + +`); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_default_refine_prompt/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_default_refine_prompt/index.test.ts new file mode 100644 index 0000000000000..95a68524ca31e --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_default_refine_prompt/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 { getDefaultRefinePrompt } from '.'; + +describe('getDefaultRefinePrompt', () => { + it('returns the default refine prompt string', () => { + const result = getDefaultRefinePrompt(); + + expect(result) + .toEqual(`You previously generated the following insights, but sometimes they represent the same attack. + +Combine the insights below, when they represent the same attack; leave any insights that are not combined unchanged:`); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts new file mode 100644 index 0000000000000..3b9aa160b4918 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts @@ -0,0 +1,46 @@ +/* + * 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 '.'; + +describe('getUseUnrefinedResults', () => { + it('returns true if both maxHallucinationFailuresReached and maxRetriesReached are true', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: true, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); + + it('returns true if maxHallucinationFailuresReached is true and maxRetriesReached is false', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true if maxHallucinationFailuresReached is false and maxRetriesReached is true', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); + + it('returns false if both maxHallucinationFailuresReached and maxRetriesReached are false', () => { + const result = getUseUnrefinedResults({ + 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/nodes/refine/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/index.test.ts new file mode 100644 index 0000000000000..d5b5a333f48f2 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/index.test.ts @@ -0,0 +1,342 @@ +/* + * 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'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; +import { loggerMock } from '@kbn/logging-mocks'; +import { FakeLLM } from '@langchain/core/utils/testing'; + +import { getRefineNode } from '.'; +import { + mockAnonymizedAlerts, + mockAnonymizedAlertsReplacements, +} from '../../../../evaluation/__mocks__/mock_anonymized_alerts'; +import { getChainWithFormatInstructions } from '../helpers/get_chain_with_format_instructions'; +import { getDefaultAttackDiscoveryPrompt } from '../helpers/get_default_attack_discovery_prompt'; +import { getDefaultRefinePrompt } from './helpers/get_default_refine_prompt'; +import { GraphState } from '../../types'; +import { + getParsedAttackDiscoveriesMock, + getRawAttackDiscoveriesMock, +} from '../../../../../../__mocks__/raw_attack_discoveries'; + +const attackDiscoveryTimestamp = '2024-10-11T17:55:59.702Z'; + +export const mockUnrefinedAttackDiscoveries: AttackDiscovery[] = [ + { + title: 'unrefinedTitle1', + alertIds: ['unrefinedAlertId1', 'unrefinedAlertId2', 'unrefinedAlertId3'], + timestamp: '2024-10-10T22:59:52.749Z', + detailsMarkdown: 'unrefinedDetailsMarkdown1', + summaryMarkdown: 'unrefinedSummaryMarkdown1 - entity A', + mitreAttackTactics: ['Input Capture'], + entitySummaryMarkdown: 'entitySummaryMarkdown1', + }, + { + title: 'unrefinedTitle2', + alertIds: ['unrefinedAlertId3', 'unrefinedAlertId4', 'unrefinedAlertId5'], + timestamp: '2024-10-10T22:59:52.749Z', + detailsMarkdown: 'unrefinedDetailsMarkdown2', + summaryMarkdown: 'unrefinedSummaryMarkdown2 - also entity A', + mitreAttackTactics: ['Credential Access'], + entitySummaryMarkdown: 'entitySummaryMarkdown2', + }, +]; + +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: 'openai', + mockInvoke, // <-- added for testing + }), + }; +}); + +const mockLogger = loggerMock.create(); +let mockLlm: ActionsClientLlm; + +const initialGraphState: GraphState = { + attackDiscoveries: null, + attackDiscoveryPrompt: getDefaultAttackDiscoveryPrompt(), + anonymizedAlerts: [...mockAnonymizedAlerts], + combinedGenerations: 'gen1', + combinedRefinements: '', + errors: [], + generationAttempts: 1, + generations: ['gen1'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: getDefaultRefinePrompt(), + replacements: { + ...mockAnonymizedAlertsReplacements, + }, + unrefinedResults: [...mockUnrefinedAttackDiscoveries], +}; + +describe('getRefineNode', () => { + beforeEach(() => { + jest.clearAllMocks(); + + jest.useFakeTimers(); + jest.setSystemTime(new Date(attackDiscoveryTimestamp)); + + mockLlm = new FakeLLM({ + response: '', + }) as unknown as ActionsClientLlm; + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('returns a function', () => { + const refineNode = getRefineNode({ + llm: mockLlm, + logger: mockLogger, + }); + + expect(typeof refineNode).toBe('function'); + }); + + it('invokes the chain with the unrefinedResults from state and format instructions', async () => { + const mockInvoke = getChainWithFormatInstructions(mockLlm).chain.invoke as jest.Mock; + + const refineNode = getRefineNode({ + llm: mockLlm, + logger: mockLogger, + }); + + await refineNode(initialGraphState); + + expect(mockInvoke).toHaveBeenCalledWith({ + format_instructions: ['mock format instructions'], + query: `${initialGraphState.attackDiscoveryPrompt} + +${getDefaultRefinePrompt()} + +\"\"\" +${JSON.stringify(initialGraphState.unrefinedResults, null, 2)} +\"\"\" + +`, + }); + }); + + it('removes the surrounding json from the response', async () => { + const response = + 'You asked for some JSON, here it is:\n```json\n{"key": "value"}\n```\nI hope that works for you.'; + + const mockLlmWithResponse = new FakeLLM({ response }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const refineNode = getRefineNode({ + llm: mockLlm, + logger: mockLogger, + }); + + const state = await refineNode(initialGraphState); + + expect(state).toEqual({ + ...initialGraphState, + combinedRefinements: '{"key": "value"}', + errors: [ + 'refine node is unable to parse (fake) response from attempt 1; (this may be an incomplete response from the model): [\n {\n "code": "invalid_type",\n "expected": "array",\n "received": "undefined",\n "path": [\n "insights"\n ],\n "message": "Required"\n }\n]', + ], + generationAttempts: 2, + refinements: ['{"key": "value"}'], + }); + }); + + it('handles hallucinations', async () => { + const hallucinatedResponse = + 'tactics like **Credential Access**, **Command and Control**, and **Persistence**.",\n "entitySummaryMarkdown": "Malware detected on host **{{ host.name hostNameValue }}**'; + + const mockLlmWithHallucination = new FakeLLM({ + response: hallucinatedResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithHallucination).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(hallucinatedResponse); + + const refineNode = getRefineNode({ + llm: mockLlmWithHallucination, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedRefinements: '{"key": "value"}', + refinements: ['{"key": "value"}'], + }; + + const state = await refineNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedRefinements: '', // <-- reset + generationAttempts: 2, // <-- incremented + refinements: [], // <-- reset + hallucinationFailures: 1, // <-- incremented + }); + }); + + it('discards previous refinements and starts over when the maxRepeatedGenerations limit is reached', async () => { + const repeatedResponse = '{"key": "value"}'; + + const mockLlmWithRepeatedGenerations = new FakeLLM({ + response: repeatedResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithRepeatedGenerations).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(repeatedResponse); + + const refineNode = getRefineNode({ + llm: mockLlmWithRepeatedGenerations, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedRefinements: '{"key": "value"}{"key": "value"}', + generationAttempts: 3, + refinements: ['{"key": "value"}', '{"key": "value"}'], + }; + + const state = await refineNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedRefinements: '', + generationAttempts: 4, // <-- incremented + refinements: [], + }); + }); + + it('combines the response with the previous refinements', async () => { + const response = 'refine1'; + + const mockLlmWithResponse = new FakeLLM({ + response, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const refineNode = getRefineNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedRefinements: 'refine0', + generationAttempts: 2, + refinements: ['refine0'], + }; + + const state = await refineNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedRefinements: 'refine0refine1', + errors: [ + 'refine node is unable to parse (fake) response from attempt 2; (this may be an incomplete response from the model): SyntaxError: Unexpected token \'r\', "refine0refine1" is not valid JSON', + ], + generationAttempts: 3, + refinements: ['refine0', 'refine1'], + }); + }); + + it('returns refined results when combined responses pass validation', async () => { + // split the response into two parts to simulate a valid response + const splitIndex = 100; // arbitrary index + const firstResponse = getRawAttackDiscoveriesMock().slice(0, splitIndex); + const secondResponse = getRawAttackDiscoveriesMock().slice(splitIndex); + + const mockLlmWithResponse = new FakeLLM({ + response: secondResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(secondResponse); + + const refineNode = getRefineNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedRefinements: firstResponse, + generationAttempts: 2, + refinements: [firstResponse], + }; + + const state = await refineNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + attackDiscoveries: getParsedAttackDiscoveriesMock(attackDiscoveryTimestamp), + combinedRefinements: firstResponse.concat(secondResponse), + generationAttempts: 3, + refinements: [firstResponse, secondResponse], + }); + }); + + it('uses the unrefined results when the max number of retries has already been reached', async () => { + const response = 'this will not pass JSON parsing'; + + const mockLlmWithResponse = new FakeLLM({ + response, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions(mockLlmWithResponse).chain + .invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const refineNode = getRefineNode({ + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedRefinements: 'refine1', + generationAttempts: 9, + refinements: ['refine1'], + }; + + const state = await refineNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + attackDiscoveries: state.unrefinedResults, // <-- the unrefined results are returned + combinedRefinements: 'refine1this will not pass JSON parsing', + errors: [ + 'refine node is unable to parse (fake) response from attempt 9; (this may be an incomplete response from the model): SyntaxError: Unexpected token \'r\', "refine1thi"... is not valid JSON', + ], + generationAttempts: 10, + refinements: ['refine1', response], + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/index.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/index.ts index 0c7987eef92bc..d1bed136f6a1c 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/refine/index.ts @@ -89,7 +89,7 @@ export const getRefineNode = ({ generationsAreRepeating({ currentGeneration: partialResponse, previousGenerations: refinements, - sampleLastNGenerations: maxRepeatedGenerations, + sampleLastNGenerations: maxRepeatedGenerations - 1, }) ) { logger?.debug( diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/anonymized_alerts_retriever/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/anonymized_alerts_retriever/index.test.ts new file mode 100644 index 0000000000000..dab2d57b20edc --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/anonymized_alerts_retriever/index.test.ts @@ -0,0 +1,104 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import type { ElasticsearchClient } from '@kbn/core/server'; +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; + +import { AnonymizedAlertsRetriever } from '.'; +import { getMockAnonymizationFieldResponse } from '../../../../../evaluation/__mocks__/mock_anonymization_fields'; +import { getAnonymizedAlerts } from '../helpers/get_anonymized_alerts'; + +const anonymizationFields = getMockAnonymizationFieldResponse(); + +const rawAlerts = [ + '@timestamp,2024-11-05T15:42:48.034Z\n_id,07d86d116ff754f4aa57c00e23a5273c2efbc9450416823ebd1d7b343b42d11a\nevent.category,malware,intrusion_detection,process\nevent.dataset,endpoint.alerts\nevent.module,endpoint\nevent.outcome,success\nfile.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nfile.name,My Go Application.app\nfile.path,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app\nhost.name,d26e9abd-6cbb-4620-a802-a22b97845d5c\nhost.os.name,macOS\nhost.os.version,13.4\nkibana.alert.original_time,2023-06-19T00:28:06.888Z\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,Malware Detection Alert\nkibana.alert.severity,critical\nkibana.alert.workflow_status,open\nmessage,Malware Detection Alert\nprocess.args,xpcproxy,application.Appify by Machine Box.My Go Application.20.23\nprocess.code_signature.exists,true\nprocess.code_signature.signing_id,a.out\nprocess.code_signature.status,code failed to satisfy specified code requirement(s)\nprocess.code_signature.subject_name,\nprocess.code_signature.trusted,false\nprocess.command_line,xpcproxy application.Appify by Machine Box.My Go Application.20.23\nprocess.executable,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app\nprocess.hash.md5,e62bdd3eaf2be436fca2e67b7eede603\nprocess.hash.sha1,58a3bddbc7c45193ecbefa22ad0496b60a29dff2\nprocess.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nprocess.name,My Go Application.app\nprocess.parent.args,/sbin/launchd\nprocess.parent.args_count,1\nprocess.parent.code_signature.exists,true\nprocess.parent.code_signature.status,No error.\nprocess.parent.code_signature.subject_name,Software Signing\nprocess.parent.code_signature.trusted,true\nprocess.parent.command_line,/sbin/launchd\nprocess.parent.executable,/sbin/launchd\nprocess.parent.name,launchd\nprocess.pid,1200\nuser.name,81c3db40-f3da-4c6a-b3c8-48c776148102', + '@timestamp,2024-11-05T15:42:48.033Z\n_id,f2d2d8bd15402e8efff81d48b70ef8cb890d5502576fb92365ee2328f5fcb123\nevent.category,malware,intrusion_detection,process\nevent.dataset,endpoint.alerts\nevent.module,endpoint\nevent.outcome,success\nfile.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nfile.name,My Go Application.app\nfile.path,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/3C4D44B9-4838-4613-BACC-BD00A9CE4025/d/Setup.app/Contents/MacOS/My Go Application.app\nhost.name,d26e9abd-6cbb-4620-a802-a22b97845d5c\nhost.os.name,macOS\nhost.os.version,13.4\nkibana.alert.original_time,2023-06-19T00:27:47.362Z\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,Malware Detection Alert\nkibana.alert.severity,critical\nkibana.alert.workflow_status,open\nmessage,Malware Detection Alert\nprocess.args,xpcproxy,application.Appify by Machine Box.My Go Application.20.23\nprocess.code_signature.exists,true\nprocess.code_signature.signing_id,a.out\nprocess.code_signature.status,code failed to satisfy specified code requirement(s)\nprocess.code_signature.subject_name,\nprocess.code_signature.trusted,false\nprocess.command_line,xpcproxy application.Appify by Machine Box.My Go Application.20.23\nprocess.executable,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/3C4D44B9-4838-4613-BACC-BD00A9CE4025/d/Setup.app/Contents/MacOS/My Go Application.app\nprocess.hash.md5,e62bdd3eaf2be436fca2e67b7eede603\nprocess.hash.sha1,58a3bddbc7c45193ecbefa22ad0496b60a29dff2\nprocess.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nprocess.name,My Go Application.app\nprocess.parent.args,/sbin/launchd\nprocess.parent.args_count,1\nprocess.parent.code_signature.exists,true\nprocess.parent.code_signature.status,No error.\nprocess.parent.code_signature.subject_name,Software Signing\nprocess.parent.code_signature.trusted,true\nprocess.parent.command_line,/sbin/launchd\nprocess.parent.executable,/sbin/launchd\nprocess.parent.name,launchd\nprocess.pid,1169\nuser.name,81c3db40-f3da-4c6a-b3c8-48c776148102', +]; + +jest.mock('../helpers/get_anonymized_alerts', () => ({ + getAnonymizedAlerts: jest.fn(), +})); + +describe('AnonymizedAlertsRetriever', () => { + let esClient: ElasticsearchClient; + + beforeEach(() => { + jest.clearAllMocks(); + + esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + + (getAnonymizedAlerts as jest.Mock).mockResolvedValue([...rawAlerts]); + }); + + it('returns the expected pageContent and metadata', async () => { + const retriever = new AnonymizedAlertsRetriever({ + alertsIndexPattern: 'test-pattern', + anonymizationFields, + esClient, + size: 10, + }); + + const documents = await retriever._getRelevantDocuments('test-query'); + + expect(documents).toEqual([ + { + pageContent: + '@timestamp,2024-11-05T15:42:48.034Z\n_id,07d86d116ff754f4aa57c00e23a5273c2efbc9450416823ebd1d7b343b42d11a\nevent.category,malware,intrusion_detection,process\nevent.dataset,endpoint.alerts\nevent.module,endpoint\nevent.outcome,success\nfile.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nfile.name,My Go Application.app\nfile.path,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app\nhost.name,d26e9abd-6cbb-4620-a802-a22b97845d5c\nhost.os.name,macOS\nhost.os.version,13.4\nkibana.alert.original_time,2023-06-19T00:28:06.888Z\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,Malware Detection Alert\nkibana.alert.severity,critical\nkibana.alert.workflow_status,open\nmessage,Malware Detection Alert\nprocess.args,xpcproxy,application.Appify by Machine Box.My Go Application.20.23\nprocess.code_signature.exists,true\nprocess.code_signature.signing_id,a.out\nprocess.code_signature.status,code failed to satisfy specified code requirement(s)\nprocess.code_signature.subject_name,\nprocess.code_signature.trusted,false\nprocess.command_line,xpcproxy application.Appify by Machine Box.My Go Application.20.23\nprocess.executable,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/6D63F08A-011C-4511-8556-EAEF9AFD6340/d/Setup.app/Contents/MacOS/My Go Application.app\nprocess.hash.md5,e62bdd3eaf2be436fca2e67b7eede603\nprocess.hash.sha1,58a3bddbc7c45193ecbefa22ad0496b60a29dff2\nprocess.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nprocess.name,My Go Application.app\nprocess.parent.args,/sbin/launchd\nprocess.parent.args_count,1\nprocess.parent.code_signature.exists,true\nprocess.parent.code_signature.status,No error.\nprocess.parent.code_signature.subject_name,Software Signing\nprocess.parent.code_signature.trusted,true\nprocess.parent.command_line,/sbin/launchd\nprocess.parent.executable,/sbin/launchd\nprocess.parent.name,launchd\nprocess.pid,1200\nuser.name,81c3db40-f3da-4c6a-b3c8-48c776148102', + metadata: {}, + }, + { + pageContent: + '@timestamp,2024-11-05T15:42:48.033Z\n_id,f2d2d8bd15402e8efff81d48b70ef8cb890d5502576fb92365ee2328f5fcb123\nevent.category,malware,intrusion_detection,process\nevent.dataset,endpoint.alerts\nevent.module,endpoint\nevent.outcome,success\nfile.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nfile.name,My Go Application.app\nfile.path,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/3C4D44B9-4838-4613-BACC-BD00A9CE4025/d/Setup.app/Contents/MacOS/My Go Application.app\nhost.name,d26e9abd-6cbb-4620-a802-a22b97845d5c\nhost.os.name,macOS\nhost.os.version,13.4\nkibana.alert.original_time,2023-06-19T00:27:47.362Z\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,Malware Detection Alert\nkibana.alert.severity,critical\nkibana.alert.workflow_status,open\nmessage,Malware Detection Alert\nprocess.args,xpcproxy,application.Appify by Machine Box.My Go Application.20.23\nprocess.code_signature.exists,true\nprocess.code_signature.signing_id,a.out\nprocess.code_signature.status,code failed to satisfy specified code requirement(s)\nprocess.code_signature.subject_name,\nprocess.code_signature.trusted,false\nprocess.command_line,xpcproxy application.Appify by Machine Box.My Go Application.20.23\nprocess.executable,/private/var/folders/_b/rmcpc65j6nv11ygrs50ctcjr0000gn/T/AppTranslocation/3C4D44B9-4838-4613-BACC-BD00A9CE4025/d/Setup.app/Contents/MacOS/My Go Application.app\nprocess.hash.md5,e62bdd3eaf2be436fca2e67b7eede603\nprocess.hash.sha1,58a3bddbc7c45193ecbefa22ad0496b60a29dff2\nprocess.hash.sha256,2c63ba2b1a5131b80e567b7a1a93997a2de07ea20d0a8f5149701c67b832c097\nprocess.name,My Go Application.app\nprocess.parent.args,/sbin/launchd\nprocess.parent.args_count,1\nprocess.parent.code_signature.exists,true\nprocess.parent.code_signature.status,No error.\nprocess.parent.code_signature.subject_name,Software Signing\nprocess.parent.code_signature.trusted,true\nprocess.parent.command_line,/sbin/launchd\nprocess.parent.executable,/sbin/launchd\nprocess.parent.name,launchd\nprocess.pid,1169\nuser.name,81c3db40-f3da-4c6a-b3c8-48c776148102', + metadata: {}, + }, + ]); + }); + + it('calls getAnonymizedAlerts with the expected parameters', async () => { + const onNewReplacements = jest.fn(); + const mockReplacements = { + replacement1: 'SRVMAC08', + replacement2: 'SRVWIN01', + replacement3: 'SRVWIN02', + }; + + const retriever = new AnonymizedAlertsRetriever({ + alertsIndexPattern: 'test-pattern', + anonymizationFields, + esClient, + onNewReplacements, + replacements: mockReplacements, + size: 10, + }); + + await retriever._getRelevantDocuments('test-query'); + + expect(getAnonymizedAlerts as jest.Mock).toHaveBeenCalledWith({ + alertsIndexPattern: 'test-pattern', + anonymizationFields, + esClient, + onNewReplacements, + replacements: mockReplacements, + size: 10, + }); + }); + + it('handles empty anonymized alerts', async () => { + (getAnonymizedAlerts as jest.Mock).mockResolvedValue([]); + + const retriever = new AnonymizedAlertsRetriever({ + esClient, + alertsIndexPattern: 'test-pattern', + anonymizationFields, + size: 10, + }); + + const documents = await retriever._getRelevantDocuments('test-query'); + + expect(documents).toHaveLength(0); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.test.ts new file mode 100644 index 0000000000000..bfd8bf2ce6953 --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.test.ts @@ -0,0 +1,111 @@ +/* + * 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 { ElasticsearchClient, Logger } from '@kbn/core/server'; +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { Replacements } from '@kbn/elastic-assistant-common'; + +import { getRetrieveAnonymizedAlertsNode } from '.'; +import { mockAnonymizedAlerts } from '../../../../evaluation/__mocks__/mock_anonymized_alerts'; +import { getDefaultAttackDiscoveryPrompt } from '../helpers/get_default_attack_discovery_prompt'; +import { getDefaultRefinePrompt } from '../refine/helpers/get_default_refine_prompt'; +import type { GraphState } from '../../types'; + +const initialGraphState: GraphState = { + attackDiscoveries: null, + attackDiscoveryPrompt: getDefaultAttackDiscoveryPrompt(), + anonymizedAlerts: [], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: getDefaultRefinePrompt(), + replacements: {}, + unrefinedResults: null, +}; + +jest.mock('./anonymized_alerts_retriever', () => ({ + AnonymizedAlertsRetriever: jest + .fn() + .mockImplementation( + ({ + onNewReplacements, + replacements, + }: { + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; + }) => ({ + withConfig: jest.fn().mockReturnValue({ + invoke: jest.fn(async () => { + if (onNewReplacements != null && replacements != null) { + onNewReplacements(replacements); + } + + return mockAnonymizedAlerts; + }), + }), + }) + ), +})); + +describe('getRetrieveAnonymizedAlertsNode', () => { + const logger = { + debug: jest.fn(), + } as unknown as Logger; + + let esClient: ElasticsearchClient; + + beforeEach(() => { + esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + }); + + it('returns a function', () => { + const result = getRetrieveAnonymizedAlertsNode({ + esClient, + logger, + }); + expect(typeof result).toBe('function'); + }); + + it('updates state with anonymized alerts', async () => { + const state: GraphState = { ...initialGraphState }; + + const retrieveAnonymizedAlerts = getRetrieveAnonymizedAlertsNode({ + esClient, + logger, + }); + + const result = await retrieveAnonymizedAlerts(state); + + expect(result).toHaveProperty('anonymizedAlerts', mockAnonymizedAlerts); + }); + + it('calls onNewReplacements with updated replacements', async () => { + const state: GraphState = { ...initialGraphState }; + const onNewReplacements = jest.fn(); + const replacements = { key: 'value' }; + + const retrieveAnonymizedAlerts = getRetrieveAnonymizedAlertsNode({ + esClient, + logger, + onNewReplacements, + replacements, + }); + + await retrieveAnonymizedAlerts(state); + + expect(onNewReplacements).toHaveBeenCalledWith({ + ...replacements, + }); + }); +}); diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.ts index 951ae3bca8854..a5d31fa14770a 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/retriever/index.ts @@ -60,11 +60,3 @@ export const getRetrieveAnonymizedAlertsNode = ({ return retrieveAnonymizedAlerts; }; - -/** - * Retrieve documents - * - * @param {GraphState} state The current state of the graph. - * @param {RunnableConfig | undefined} config The configuration object for tracing. - * @returns {Promise} The new state object. - */ diff --git a/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/state/index.test.ts b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/state/index.test.ts new file mode 100644 index 0000000000000..dcc8ee1e4292d --- /dev/null +++ b/x-pack/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/state/index.test.ts @@ -0,0 +1,115 @@ +/* + * 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 { getDefaultGraphState } from '.'; +import { + DEFAULT_MAX_GENERATION_ATTEMPTS, + DEFAULT_MAX_HALLUCINATION_FAILURES, + DEFAULT_MAX_REPEATED_GENERATIONS, +} from '../constants'; +import { getDefaultAttackDiscoveryPrompt } from '../nodes/helpers/get_default_attack_discovery_prompt'; +import { getDefaultRefinePrompt } from '../nodes/refine/helpers/get_default_refine_prompt'; + +const defaultAttackDiscoveryPrompt = getDefaultAttackDiscoveryPrompt(); +const defaultRefinePrompt = getDefaultRefinePrompt(); + +describe('getDefaultGraphState', () => { + it('returns the expected default attackDiscoveries', () => { + const state = getDefaultGraphState(); + + expect(state.attackDiscoveries?.default?.()).toBeNull(); + }); + + it('returns the expected default attackDiscoveryPrompt', () => { + const state = getDefaultGraphState(); + + expect(state.attackDiscoveryPrompt?.default?.()).toEqual(defaultAttackDiscoveryPrompt); + }); + + it('returns the expected default empty collection of anonymizedAlerts', () => { + const state = getDefaultGraphState(); + + expect(state.anonymizedAlerts?.default?.()).toHaveLength(0); + }); + + it('returns the expected default combinedGenerations state', () => { + const state = getDefaultGraphState(); + + expect(state.combinedGenerations?.default?.()).toBe(''); + }); + + it('returns the expected default combinedRefinements state', () => { + const state = getDefaultGraphState(); + + expect(state.combinedRefinements?.default?.()).toBe(''); + }); + + it('returns the expected default errors state', () => { + const state = getDefaultGraphState(); + + expect(state.errors?.default?.()).toHaveLength(0); + }); + + it('return the expected default generationAttempts state', () => { + const state = getDefaultGraphState(); + + expect(state.generationAttempts?.default?.()).toBe(0); + }); + + it('returns the expected default generations state', () => { + const state = getDefaultGraphState(); + + expect(state.generations?.default?.()).toHaveLength(0); + }); + + it('returns the expected default hallucinationFailures state', () => { + const state = getDefaultGraphState(); + + expect(state.hallucinationFailures?.default?.()).toBe(0); + }); + + it('returns the expected default refinePrompt state', () => { + const state = getDefaultGraphState(); + + expect(state.refinePrompt?.default?.()).toEqual(defaultRefinePrompt); + }); + + it('returns the expected default maxGenerationAttempts state', () => { + const state = getDefaultGraphState(); + + expect(state.maxGenerationAttempts?.default?.()).toBe(DEFAULT_MAX_GENERATION_ATTEMPTS); + }); + + it('returns the expected default maxHallucinationFailures state', () => { + const state = getDefaultGraphState(); + expect(state.maxHallucinationFailures?.default?.()).toBe(DEFAULT_MAX_HALLUCINATION_FAILURES); + }); + + it('returns the expected default maxRepeatedGenerations state', () => { + const state = getDefaultGraphState(); + + expect(state.maxRepeatedGenerations?.default?.()).toBe(DEFAULT_MAX_REPEATED_GENERATIONS); + }); + + it('returns the expected default refinements state', () => { + const state = getDefaultGraphState(); + + expect(state.refinements?.default?.()).toHaveLength(0); + }); + + it('returns the expected default replacements state', () => { + const state = getDefaultGraphState(); + + expect(state.replacements?.default?.()).toEqual({}); + }); + + it('returns the expected default unrefinedResults state', () => { + const state = getDefaultGraphState(); + + expect(state.unrefinedResults?.default?.()).toBeNull(); + }); +}); diff --git a/x-pack/plugins/security_solution/public/attack_discovery/pages/empty_states/helpers/show_empty_states/index.test.ts b/x-pack/plugins/security_solution/public/attack_discovery/pages/empty_states/helpers/show_empty_states/index.test.ts new file mode 100644 index 0000000000000..0211dc8d51eba --- /dev/null +++ b/x-pack/plugins/security_solution/public/attack_discovery/pages/empty_states/helpers/show_empty_states/index.test.ts @@ -0,0 +1,87 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +import { showEmptyStates } from '.'; +import { + showEmptyPrompt, + showFailurePrompt, + showNoAlertsPrompt, + showWelcomePrompt, +} from '../../../helpers'; + +jest.mock('../../../helpers', () => ({ + showEmptyPrompt: jest.fn().mockReturnValue(false), + showFailurePrompt: jest.fn().mockReturnValue(false), + showNoAlertsPrompt: jest.fn().mockReturnValue(false), + showWelcomePrompt: jest.fn().mockReturnValue(false), +})); + +const defaultArgs = { + aiConnectorsCount: 0, + alertsContextCount: 0, + attackDiscoveriesCount: 0, + connectorId: undefined, + failureReason: null, + isLoading: false, +}; + +describe('showEmptyStates', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('returns true if showWelcomePrompt returns true', () => { + (showWelcomePrompt as jest.Mock).mockReturnValue(true); + + const result = showEmptyStates({ + ...defaultArgs, + }); + expect(result).toBe(true); + }); + + it('returns true if showFailurePrompt returns true', () => { + (showFailurePrompt as jest.Mock).mockReturnValue(true); + + const result = showEmptyStates({ + ...defaultArgs, + connectorId: 'test', + failureReason: 'error', + }); + expect(result).toBe(true); + }); + + it('returns true if showNoAlertsPrompt returns true', () => { + (showNoAlertsPrompt as jest.Mock).mockReturnValue(true); + + const result = showEmptyStates({ + ...defaultArgs, + connectorId: 'test', + }); + expect(result).toBe(true); + }); + + it('returns true if showEmptyPrompt returns true', () => { + (showEmptyPrompt as jest.Mock).mockReturnValue(true); + + const result = showEmptyStates({ + ...defaultArgs, + }); + expect(result).toBe(true); + }); + + it('returns false if all prompts return false', () => { + (showWelcomePrompt as jest.Mock).mockReturnValue(false); + (showFailurePrompt as jest.Mock).mockReturnValue(false); + (showNoAlertsPrompt as jest.Mock).mockReturnValue(false); + (showEmptyPrompt as jest.Mock).mockReturnValue(false); + + const result = showEmptyStates({ + ...defaultArgs, + }); + expect(result).toBe(false); + }); +}); diff --git a/x-pack/plugins/security_solution/public/attack_discovery/pages/generate/index.test.tsx b/x-pack/plugins/security_solution/public/attack_discovery/pages/generate/index.test.tsx new file mode 100644 index 0000000000000..e818d1c4140f6 --- /dev/null +++ b/x-pack/plugins/security_solution/public/attack_discovery/pages/generate/index.test.tsx @@ -0,0 +1,46 @@ +/* + * 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 { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import React from 'react'; + +import { Generate } from '.'; +import * as i18n from '../empty_prompt/translations'; + +describe('Generate Component', () => { + it('calls onGenerate when the button is clicked', () => { + const onGenerate = jest.fn(); + + render(); + + fireEvent.click(screen.getByTestId('generate')); + + expect(onGenerate).toHaveBeenCalled(); + }); + + it('disables the generate button when isLoading is true', () => { + render(); + + expect(screen.getByTestId('generate')).toBeDisabled(); + }); + + it('disables the generate button when isDisabled is true', () => { + render(); + + expect(screen.getByTestId('generate')).toBeDisabled(); + }); + + it('shows tooltip content when the button is disabled', async () => { + render(); + + fireEvent.mouseOver(screen.getByTestId('generate')); + + await waitFor(() => { + expect(screen.getByText(i18n.SELECT_A_CONNECTOR)).toBeInTheDocument(); + }); + }); +}); diff --git a/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.test.tsx b/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.test.tsx new file mode 100644 index 0000000000000..958c9094fabf3 --- /dev/null +++ b/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.test.tsx @@ -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 { render, screen, fireEvent } from '@testing-library/react'; +import React from 'react'; + +import { AlertsSettings, MAX_ALERTS } from '.'; + +const maxAlerts = '150'; + +const setMaxAlerts = jest.fn(); + +describe('AlertsSettings', () => { + it('calls setMaxAlerts when the alerts range changes', () => { + render(); + + fireEvent.click(screen.getByText(`${MAX_ALERTS}`)); + + expect(setMaxAlerts).toHaveBeenCalledWith(`${MAX_ALERTS}`); + }); + + it('displays the correct maxAlerts value', () => { + render(); + + expect(screen.getByTestId('alertsRange')).toHaveValue(maxAlerts); + }); + + it('displays the expected text for anonymization settings', () => { + render(); + + expect(screen.getByTestId('latestAndRiskiest')).toHaveTextContent( + 'Send Attack discovery information about your 150 newest and riskiest open or acknowledged alerts.' + ); + }); +}); diff --git a/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.tsx b/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.tsx index b51a1fc3f85c8..336da549f55ea 100644 --- a/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.tsx +++ b/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/alerts_settings/index.tsx @@ -51,7 +51,9 @@ const AlertsSettingsComponent: React.FC = ({ maxAlerts, setMaxAlerts }) = - {i18n.LATEST_AND_RISKIEST_OPEN_ALERTS(Number(maxAlerts))} + + {i18n.LATEST_AND_RISKIEST_OPEN_ALERTS(Number(maxAlerts))} + diff --git a/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/footer/index.test.tsx b/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/footer/index.test.tsx new file mode 100644 index 0000000000000..e487304c41350 --- /dev/null +++ b/x-pack/plugins/security_solution/public/attack_discovery/pages/header/settings_modal/footer/index.test.tsx @@ -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 React from 'react'; +import { fireEvent, render, screen } from '@testing-library/react'; + +import { Footer } from '.'; + +describe('Footer', () => { + const closeModal = jest.fn(); + const onReset = jest.fn(); + const onSave = jest.fn(); + + beforeEach(() => jest.clearAllMocks()); + + it('calls onReset when the reset button is clicked', () => { + render(