Adds new markdown parser for send to timeline actions for eql, kql an…

…d queryDSL

x-pack/plugins/security_solution/public/security_assistant/security_assistant.tsx

@@ -6,7 +6,6 @@
  */
 
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import { createPortal } from 'react-dom';
 import type { EuiCommentProps } from '@elastic/eui';
 import {
   EuiFlexGroup,
@@ -21,17 +20,18 @@ import {
   EuiAvatar,
   EuiPageHeader,
   EuiMarkdownFormat,
-  EuiIcon,
   EuiToolTip,
+  EuiIcon,
 } from '@elastic/eui';
-import type { DataProvider } from '@kbn/timelines-plugin/common';
 import { CommentType } from '@kbn/cases-plugin/common';
 import styled from 'styled-components';
+import { createPortal } from 'react-dom';
 import * as i18n from './translations';
 
 import { useKibana } from '../common/lib/kibana';
 import { getCombinedMessage, getMessageFromRawResponse, isFileHash } from './helpers';
 import { SendToTimelineButton } from './send_to_timeline_button';
+
 import { SettingsPopover } from './settings_popover';
 import { useSecurityAssistantContext } from './security_assistant_context';
 import { ContextPills } from './context_pills';
@@ -46,7 +46,7 @@ import { getDefaultSystemPrompt, getSuperheroPrompt } from './prompt/helpers';
 import type { Prompt } from './types';
 import { getPromptById } from './prompt_editor/helpers';
 
-const isMac = navigator.platform.toLowerCase().indexOf('mac') >= 0;
+import { analyzeMarkdown } from './use_conversation/helpers';
 
 const CommentsContainer = styled.div`
   max-height: 600px;
@@ -100,6 +100,7 @@ export const SecurityAssistant: React.FC<SecurityAssistantProps> =
 
       const { cases } = useKibana().services;
       const bottomRef = useRef<HTMLDivElement | null>(null);
+      const lastCommentRef = useRef<HTMLDivElement | null>(null);
 
       const [promptTextPreview, setPromptTextPreview] = useState<string>('');
       const [systemPrompts] = useState<Prompt[]>([getDefaultSystemPrompt(), getSuperheroPrompt()]);
@@ -220,49 +221,48 @@ export const SecurityAssistant: React.FC<SecurityAssistantProps> =
       );
 
       // Drill in `Add To Timeline` action
-      // Grab all relevant dom elements
-      const commentBlocks = [...document.getElementsByClassName('euiMarkdownFormat')];
-      // Filter if no code block exists as to not make extra portals
-      commentBlocks.filter((cb) => cb.querySelectorAll('.euiCodeBlock__code').length > 0);
-
-      let commentDetails: Array<{
-        commentBlock: Element;
-        codeBlocks: Element[];
-        codeBlockControls: Element[];
-        dataProviders: DataProvider[];
-      }> =
-        currentConversation.messages.length > 0
-          ? commentBlocks.map((commentBlock) => {
-              return {
-                commentBlock,
-                codeBlocks: [...commentBlock.querySelectorAll('.euiCodeBlock__code')],
-                codeBlockControls: [...commentBlock.querySelectorAll('.euiCodeBlock__controls')],
-                dataProviders: [],
-              };
-            })
-          : [];
-      commentDetails = commentDetails.map((details) => {
-        const dataProviders: DataProvider[] = details.codeBlocks.map((codeBlock, i) => {
-          return {
-            id: 'assistant-data-provider',
-            name: 'Assistant Query',
-            enabled: true,
-            // overriding to use as isEQL
-            excluded: details.commentBlock?.textContent?.includes('EQL') ?? false,
-            kqlQuery: codeBlock.textContent ?? '',
-            queryMatch: {
-              field: 'host.name',
-              operator: ':',
-              value: 'test',
-            },
-            and: [],
-          };
+      // First let's find
+      const messageCodeBlocks = useMemo(() => {
+        const cbd = currentConversation.messages.map(({ content }) => {
+          return analyzeMarkdown(content);
         });
-        return {
-          ...details,
-          dataProviders,
-        };
-      });
+        return cbd.map((codeBlocks, messageIndex) => {
+          return codeBlocks.map((codeBlock, codeBlockIndex) => {
+            return {
+              ...codeBlock,
+              controlContainer: document.querySelectorAll(
+                `.message-${messageIndex} .euiCodeBlock__controls`
+              )[codeBlockIndex],
+              button: (
+                <SendToTimelineButton
+                  asEmptyButton={true}
+                  dataProviders={[
+                    {
+                      id: 'assistant-data-provider',
+                      name: `Assistant Query from conversation ${currentConversation.id}`,
+                      enabled: true,
+                      excluded: false,
+                      queryType: codeBlock.type,
+                      kqlQuery: codeBlock.content ?? '',
+                      queryMatch: {
+                        field: 'host.name',
+                        operator: ':',
+                        value: 'test',
+                      },
+                      and: [],
+                    },
+                  ]}
+                  keepDataView={true}
+                >
+                  <EuiToolTip position="right" content={'Add to timeline'}>
+                    <EuiIcon type="timeline" />
+                  </EuiToolTip>
+                </SendToTimelineButton>
+              ),
+            };
+          });
+        });
+      }, [currentConversation.messages, lastCommentRef.current]);
       ////
 
       // Add min-height to all codeblocks so timeline icon doesn't overflow
@@ -324,31 +324,15 @@ export const SecurityAssistant: React.FC<SecurityAssistantProps> =
           )}
 
           {/* Create portals for each EuiCodeBlock to add the `Investigate in Timeline` action */}
-          {currentConversation.messages.length > 0 &&
-            commentDetails.length > 0 &&
-            // eslint-disable-next-line array-callback-return
-            commentDetails.map((e) => {
-              if (e.dataProviders != null && e.dataProviders.length > 0) {
-                return e.codeBlocks.map((block, i) => {
-                  if (e.codeBlockControls[i] != null) {
-                    return createPortal(
-                      <SendToTimelineButton
-                        asEmptyButton={true}
-                        dataProviders={[e.dataProviders?.[i] ?? []]}
-                        keepDataView={true}
-                      >
-                        <EuiToolTip position="right" content={'Add to timeline'}>
-                          <EuiIcon type="timeline" />
-                        </EuiToolTip>
-                      </SendToTimelineButton>,
-                      e.codeBlockControls[i]
-                    );
-                  } else {
-                    return <></>;
-                  }
-                });
-              }
-            })}
+          {messageCodeBlocks.map((codeBlocks) => {
+            return codeBlocks.map((codeBlock) => {
+              return codeBlock.controlContainer != null ? (
+                createPortal(codeBlock.button, codeBlock.controlContainer)
+              ) : (
+                <></>
+              );
+            });
+          })}
 
           <ContextPills
             promptContexts={promptContexts}
@@ -366,6 +350,14 @@ export const SecurityAssistant: React.FC<SecurityAssistantProps> =
                   username: isUser ? 'You' : 'Assistant',
                   actions: (
                     <>
+                      <EuiToolTip position="top" content={'Add to timeline note'}>
+                        <EuiButtonIcon
+                          onClick={() => handleAddToExistingCaseClick(message.content)}
+                          iconType="editorComment"
+                          color="primary"
+                          aria-label="Add message content as a timeline note"
+                        />
+                      </EuiToolTip>
                       <EuiToolTip position="top" content={'Add to case'}>
                         <EuiButtonIcon
                           onClick={() => handleAddToExistingCaseClick(message.content)}
@@ -389,11 +381,21 @@ export const SecurityAssistant: React.FC<SecurityAssistantProps> =
                     </>
                   ),
                   // event: isUser ? 'Asked a question' : 'Responded with',
-                  children: (
-                    <EuiText>
-                      <EuiMarkdownFormat>{message.content}</EuiMarkdownFormat>
-                    </EuiText>
-                  ),
+                  children:
+                    index !== currentConversation.messages.length - 1 ? (
+                      <EuiText>
+                        <EuiMarkdownFormat className={`message-${index}`}>
+                          {message.content}
+                        </EuiMarkdownFormat>
+                      </EuiText>
+                    ) : (
+                      <EuiText>
+                        <EuiMarkdownFormat className={`message-${index}`}>
+                          {message.content}
+                        </EuiMarkdownFormat>
+                        <span ref={lastCommentRef} />
+                      </EuiText>
+                    ),
                   timelineAvatar: isUser ? (
                     <EuiAvatar name="user" size="l" color="subdued" iconType={'logoSecurity'} />
                   ) : (

x-pack/plugins/security_solution/public/security_assistant/send_to_timeline_button.tsx

@@ -10,22 +10,21 @@ import { EuiButton, EuiButtonEmpty } from '@elastic/eui';
 import type { Filter } from '@kbn/es-query';
 import { useDispatch } from 'react-redux';
 
-import { sourcererSelectors } from '../../public/common/store';
+import { sourcererSelectors } from '../common/store';
 import { InputsModelId } from '../common/store/inputs/constants';
-import { TimeRange } from '../common/store/inputs/model';
-import { inputsActions } from '../../public/common/store/inputs';
+import type { TimeRange } from '../common/store/inputs/model';
+import { inputsActions } from '../common/store/inputs';
 import {
   applyKqlFilterQuery,
   setActiveTabTimeline,
   setFilters,
   updateDataView,
   updateEqlOptions,
 } from '../timelines/store/timeline/actions';
-import { sourcererActions } from '../../public/common/store/actions';
+import { sourcererActions } from '../common/store/actions';
 import { SourcererScopeName } from '../common/store/sourcerer/model';
-import { TimelineTabs } from '../../common/types';
-import { TimelineId, TimelineType } from '../../common/types';
-import { DataProvider } from '../timelines/components/timeline/data_providers/data_provider';
+import { TimelineTabs, TimelineId, TimelineType } from '../../common/types';
+import type { DataProvider } from '../timelines/components/timeline/data_providers/data_provider';
 import { useCreateTimeline } from '../timelines/components/timeline/properties/use_create_timeline';
 import { useDeepEqualSelector } from '../common/hooks/use_selector';
 import { ACTION_INVESTIGATE_IN_TIMELINE } from '../detections/components/alerts_table/translations';
@@ -86,42 +85,66 @@ export const SendToTimelineButton: React.FunctionComponent<SendToTimelineButtonP
           })
         );
 
-        // overriding `excluded` to use as `isEQL`
-        if (dataProviders[0].excluded) {
-          // is EQL
-          dispatch(
-            updateEqlOptions({
-              id: TimelineId.active,
-              field: 'query',
-              value: dataProviders[0].kqlQuery,
-            })
-          );
-          dispatch(
-            setActiveTabTimeline({
-              id: TimelineId.active,
-              activeTab: TimelineTabs.eql,
-            })
-          );
-        } else {
-          // is KQL
-          dispatch(
-            applyKqlFilterQuery({
-              id: TimelineId.active,
-              filterQuery: {
-                kuery: {
-                  kind: 'kuery',
-                  expression: dataProviders[0].kqlQuery,
+        // Added temporary queryType to dataproviders to support EQL/DSL
+        switch (dataProviders[0].queryType) {
+          case 'eql':
+            // is EQL
+            dispatch(
+              updateEqlOptions({
+                id: TimelineId.active,
+                field: 'query',
+                value: dataProviders[0].kqlQuery,
+              })
+            );
+            dispatch(
+              setActiveTabTimeline({
+                id: TimelineId.active,
+                activeTab: TimelineTabs.eql,
+              })
+            );
+            break;
+          case 'kql':
+            // is KQL
+            dispatch(
+              applyKqlFilterQuery({
+                id: TimelineId.active,
+                filterQuery: {
+                  kuery: {
+                    kind: 'kuery',
+                    expression: dataProviders[0].kqlQuery,
+                  },
+                  serializedQuery: dataProviders[0].kqlQuery,
                 },
-                serializedQuery: dataProviders[0].kqlQuery,
+              })
+            );
+            dispatch(
+              setActiveTabTimeline({
+                id: TimelineId.active,
+                activeTab: TimelineTabs.query,
+              })
+            );
+            break;
+          case 'dsl':
+            const filter = {
+              meta: {
+                type: 'custom',
+                disabled: false,
+                negate: false,
+                alias: dataProviders[0].name,
+                key: 'query',
+                value: dataProviders[0].kqlQuery,
               },
-            })
-          );
-          dispatch(
-            setActiveTabTimeline({
-              id: TimelineId.active,
-              activeTab: TimelineTabs.query,
-            })
-          );
+              // TODO: Now you be careful parsing json directly from the robots, mmmkay?
+              query: JSON.parse(dataProviders[0].kqlQuery),
+            };
+            dispatch(setFilters({ id: TimelineId.active, filters: [filter] }));
+            dispatch(
+              setActiveTabTimeline({
+                id: TimelineId.active,
+                activeTab: TimelineTabs.query,
+              })
+            );
+            break;
         }
       }
       // Use filters if more than a certain amount of ids for dom performance.

x-pack/plugins/security_solution/public/security_assistant/use_conversation/helpers.tsx

@@ -0,0 +1,49 @@
+/*
+ * 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 React from 'react';
+import type { QueryType } from '../../timelines/components/timeline/data_providers/data_provider';
+
+export interface CodeBlockDetails {
+  type: QueryType;
+  content: string;
+  start: number;
+  end: number;
+  controlContainer?: HTMLElement;
+  button?: React.FC;
+}
+
+export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => {
+  const codeBlockRegex = /```(\w+)?\s([\s\S]*?)```/g;
+  const matches = [...markdown.matchAll(codeBlockRegex)];
+  const types = {
+    eql: ['Event Query Language', 'EQL sequence query'],
+    kql: ['Kibana Query Language', 'KQL Query'],
+    dsl: ['Elasticsearch QueryDSL', 'Elasticsearch Query DSL', 'Elasticsearch DSL'],
+  };
+
+  const result: CodeBlockDetails[] = matches.map((match) => {
+    let type = match[1] || 'no-type';
+    if (type === 'no-type' || type === 'json') {
+      const start = match.index || 0;
+      const precedingText = markdown.slice(0, start);
+      for (const [typeKey, keywords] of Object.entries(types)) {
+        if (keywords.some((kw) => precedingText.includes(kw))) {
+          type = typeKey;
+          break;
+        }
+      }
+    }
+
+    const content = match[2].trim();
+    const start = match.index || 0;
+    const end = start + match[0].length;
+    return { type: type as QueryType, content, start, end };
+  });
+
+  return result;
+};