[Asset Management] Osquery agent picker tests/fixes. (#97580)

* general refactoring, tests, and fixes around host data munging

* fix kql, pull and display offline agents in search

x-pack/plugins/osquery/public/agents/agent_grouper.test.ts

@@ -0,0 +1,140 @@
+/*
+ * 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 { AgentGrouper } from './agent_grouper';
+import { AGENT_GROUP_KEY, Group, GroupedAgent, GroupOptionValue } from './types';
+import uuid from 'uuid';
+import { ALL_AGENTS_LABEL } from './translations';
+
+type GroupData = {
+  [key in Exclude<AGENT_GROUP_KEY, AGENT_GROUP_KEY.All | AGENT_GROUP_KEY.Agent>]: Group[];
+};
+export function genGroup(name: string) {
+  return {
+    name,
+    id: uuid.v4(),
+    size: 5,
+  };
+}
+
+export function genAgent(policyId: string, hostname: string, id: string): GroupedAgent {
+  return {
+    status: 'online',
+    policy_id: policyId,
+    local_metadata: {
+      elastic: {
+        agent: {
+          id,
+        },
+      },
+      os: {
+        platform: 'test platform',
+      },
+      host: {
+        hostname,
+      },
+    },
+  };
+}
+export const groupData: GroupData = {
+  [AGENT_GROUP_KEY.Platform]: new Array(3).fill('test platform ').map((el, i) => genGroup(el + i)),
+  [AGENT_GROUP_KEY.Policy]: new Array(3).fill('test policy ').map((el, i) => genGroup(el + i)),
+};
+
+describe('AgentGrouper', () => {
+  describe('All agents', () => {
+    it('should handle empty groups properly', () => {
+      const agentGrouper = new AgentGrouper();
+      expect(agentGrouper.generateOptions()).toEqual([]);
+    });
+
+    it('should ignore calls to add things to the "all" group', () => {
+      const agentGrouper = new AgentGrouper();
+      agentGrouper.updateGroup(AGENT_GROUP_KEY.All, [{}]);
+      expect(agentGrouper.generateOptions()).toEqual([]);
+    });
+
+    it('should omit the "all agents" option when total is set to <= 0', () => {
+      const agentGrouper = new AgentGrouper();
+      agentGrouper.setTotalAgents(0);
+      expect(agentGrouper.generateOptions()).toEqual([]);
+      agentGrouper.setTotalAgents(-1);
+      expect(agentGrouper.generateOptions()).toEqual([]);
+    });
+
+    it('should add the "all agents" option when the total is set to > 0', () => {
+      const agentGrouper = new AgentGrouper();
+      agentGrouper.setTotalAgents(100);
+      const groups = agentGrouper.generateOptions();
+
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      const allGroup = groups[AGENT_GROUP_KEY.All].options![0];
+      expect(allGroup.label).toEqual(ALL_AGENTS_LABEL);
+      const size: number = (allGroup.value as GroupOptionValue).size;
+
+      expect(size).toEqual(100);
+      agentGrouper.setTotalAgents(0);
+      expect(agentGrouper.generateOptions()).toEqual([]);
+    });
+  });
+
+  describe('Policies and platforms', () => {
+    function genGroupTest(
+      key: AGENT_GROUP_KEY.Platform | AGENT_GROUP_KEY.Policy,
+      dataName: string
+    ) {
+      return () => {
+        const agentGrouper = new AgentGrouper();
+        const data = groupData[key];
+        agentGrouper.updateGroup(key, data);
+
+        const groups = agentGrouper.generateOptions();
+        const options = groups[0].options;
+        expect(options).toBeTruthy();
+
+        data.forEach((datum, i) => {
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          const opt = options![i];
+          expect(opt.label).toEqual(`test ${dataName} ${i} (${datum.id})`);
+          expect(opt.key).toEqual(datum.id);
+          expect(opt.value).toEqual({
+            groupType: key,
+            id: datum.id,
+            size: 5,
+          });
+        });
+      };
+    }
+    it('should generate policy options', genGroupTest(AGENT_GROUP_KEY.Policy, 'policy'));
+    it('should generate platform options', genGroupTest(AGENT_GROUP_KEY.Platform, 'platform'));
+  });
+
+  describe('agents', () => {
+    it('should generate agent options', () => {
+      const agentGrouper = new AgentGrouper();
+      const policyId = uuid.v4();
+      const agentData: GroupedAgent[] = [
+        genAgent(policyId, `agent host 1`, uuid.v4()),
+        genAgent(policyId, `agent host 2`, uuid.v4()),
+      ];
+      agentGrouper.updateGroup(AGENT_GROUP_KEY.Agent, agentData);
+
+      const groups = agentGrouper.generateOptions();
+      const options = groups[0].options;
+      expect(options).toBeTruthy();
+      agentData.forEach((ag, i) => {
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        const opt = options![i];
+        expect(opt.label).toEqual(
+          `${ag.local_metadata.host.hostname} (${ag.local_metadata.elastic.agent.id})`
+        );
+        expect(opt.key).toEqual(ag.local_metadata.elastic.agent.id);
+        expect(opt.value?.id).toEqual(ag.local_metadata.elastic.agent.id);
+      });
+    });
+  });
+});

x-pack/plugins/osquery/public/agents/agent_grouper.ts

@@ -5,7 +5,6 @@
  * 2.0.
  */
 
-import { Agent } from '../../common/shared_imports';
 import { generateColorPicker } from './helpers';
 import {
   ALL_AGENTS_LABEL,
@@ -13,7 +12,7 @@ import {
   AGENT_POLICY_LABEL,
   AGENT_SELECTION_LABEL,
 } from './translations';
-import { AGENT_GROUP_KEY, Group, GroupOption } from './types';
+import { AGENT_GROUP_KEY, Group, GroupOption, GroupedAgent } from './types';
 
 const getColor = generateColorPicker();
 
@@ -27,6 +26,38 @@ const generateGroup = <T = Group>(label: string, groupType: AGENT_GROUP_KEY) =>
   };
 };
 
+export const generateAgentOption = (
+  label: string,
+  groupType: AGENT_GROUP_KEY,
+  data: GroupedAgent[]
+) => ({
+  label,
+  options: data.map((agent) => ({
+    label: `${agent.local_metadata.host.hostname} (${agent.local_metadata.elastic.agent.id})`,
+    key: agent.local_metadata.elastic.agent.id,
+    color: getColor(groupType),
+    value: {
+      groupType,
+      groups: {
+        policy: agent.policy_id ?? '',
+        platform: agent.local_metadata.os.platform,
+      },
+      id: agent.local_metadata.elastic.agent.id,
+      status: agent.status ?? 'unknown',
+    },
+  })),
+});
+
+export const generateGroupOption = (label: string, groupType: AGENT_GROUP_KEY, data: Group[]) => ({
+  label,
+  options: (data as Group[]).map(({ name, id, size }) => ({
+    label: name !== id ? `${name} (${id})` : name,
+    key: id,
+    color: getColor(groupType),
+    value: { groupType, id, size },
+  })),
+});
+
 export class AgentGrouper {
   groupOrder = [
     AGENT_GROUP_KEY.All,
@@ -38,12 +69,15 @@ export class AgentGrouper {
     [AGENT_GROUP_KEY.All]: generateGroup(ALL_AGENTS_LABEL, AGENT_GROUP_KEY.All),
     [AGENT_GROUP_KEY.Platform]: generateGroup(AGENT_PLATFORMS_LABEL, AGENT_GROUP_KEY.Platform),
     [AGENT_GROUP_KEY.Policy]: generateGroup(AGENT_POLICY_LABEL, AGENT_GROUP_KEY.Policy),
-    [AGENT_GROUP_KEY.Agent]: generateGroup<Agent>(AGENT_SELECTION_LABEL, AGENT_GROUP_KEY.Agent),
+    [AGENT_GROUP_KEY.Agent]: generateGroup<GroupedAgent>(
+      AGENT_SELECTION_LABEL,
+      AGENT_GROUP_KEY.Agent
+    ),
   };
 
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   updateGroup(key: AGENT_GROUP_KEY, data: any[], append = false) {
-    if (!data?.length) {
+    if (!data?.length || key === AGENT_GROUP_KEY.All) {
       return;
     }
     const group = this.groups[key];
@@ -56,6 +90,9 @@ export class AgentGrouper {
   }
 
   setTotalAgents(total: number): void {
+    if (total < 0) {
+      return;
+    }
     this.groups[AGENT_GROUP_KEY.All].size = total;
   }
 
@@ -82,34 +119,10 @@ export class AgentGrouper {
           break;
         case AGENT_GROUP_KEY.Platform:
         case AGENT_GROUP_KEY.Policy:
-          opts.push({
-            label,
-            options: (data as Group[]).map(({ name, id, size: groupSize }) => ({
-              label: name !== id ? `${name} (${id})` : name,
-              key: id,
-              color: getColor(groupType),
-              value: { groupType, id, size: groupSize },
-            })),
-          });
+          opts.push(generateGroupOption(label, key, data as Group[]));
           break;
         case AGENT_GROUP_KEY.Agent:
-          opts.push({
-            label,
-            options: (data as Agent[]).map((agent: Agent) => ({
-              label: `${agent.local_metadata.host.hostname} (${agent.local_metadata.elastic.agent.id})`,
-              key: agent.local_metadata.elastic.agent.id,
-              color,
-              value: {
-                groupType,
-                groups: {
-                  policy: agent.policy_id ?? '',
-                  platform: agent.local_metadata.os.platform,
-                },
-                id: agent.local_metadata.elastic.agent.id,
-                online: agent.active,
-              },
-            })),
-          });
+          opts.push(generateAgentOption(label, key, data as GroupedAgent[]));
           break;
       }
     }

x-pack/plugins/osquery/public/agents/agents_table.tsx

@@ -134,7 +134,7 @@ const AgentsTableComponent: React.FC<AgentsTableProps> = ({ agentSelection, onCh
   const renderOption = useCallback((option, searchVal, contentClassName) => {
     const { label, value } = option;
     return value?.groupType === AGENT_GROUP_KEY.Agent ? (
-      <EuiHealth color={value?.online ? 'success' : 'danger'}>
+      <EuiHealth color={value?.status === 'online' ? 'success' : 'danger'}>
         <span className={contentClassName}>
           <EuiHighlight search={searchVal}>{label}</EuiHighlight>
         </span>

x-pack/plugins/osquery/public/agents/helpers.test.ts

@@ -5,8 +5,59 @@
  * 2.0.
  */
 
-import { getNumOverlapped, getNumAgentsInGrouping, processAggregations } from './helpers';
-import { Overlap, SelectedGroups } from './types';
+import uuid from 'uuid';
+import { generateGroupOption } from './agent_grouper';
+import {
+  getNumOverlapped,
+  getNumAgentsInGrouping,
+  processAggregations,
+  generateAgentSelection,
+} from './helpers';
+import { AGENT_GROUP_KEY, GroupOption, Overlap, SelectedGroups } from './types';
+
+describe('generateAgentSelection', () => {
+  it('should handle empty input', () => {
+    const options: GroupOption[] = [];
+    const { newAgentSelection, selectedGroups, selectedAgents } = generateAgentSelection(options);
+    expect(newAgentSelection).toEqual({
+      agents: [],
+      allAgentsSelected: false,
+      platformsSelected: [],
+      policiesSelected: [],
+    });
+    expect(selectedAgents).toEqual([]);
+    expect(selectedGroups).toEqual({
+      policy: {},
+      platform: {},
+    });
+  });
+
+  it('should properly pull out group ids', () => {
+    const options: GroupOption[] = [];
+    const policyOptions = generateGroupOption('policy', AGENT_GROUP_KEY.Policy, [
+      { name: 'policy 1', id: 'policy 1', size: 5 },
+      { name: 'policy 2', id: uuid.v4(), size: 5 },
+    ]).options;
+    options.push(...policyOptions);
+
+    const platformOptions = generateGroupOption('platform', AGENT_GROUP_KEY.Platform, [
+      { name: 'platform 1', id: 'platform 1', size: 5 },
+      { name: 'platform 2', id: uuid.v4(), size: 5 },
+    ]).options;
+    options.push(...platformOptions);
+
+    const { newAgentSelection, selectedGroups, selectedAgents } = generateAgentSelection(options);
+    expect(newAgentSelection).toEqual({
+      agents: [],
+      allAgentsSelected: false,
+      platformsSelected: platformOptions.map(({ value: { id } }) => id),
+      policiesSelected: policyOptions.map(({ value: { id } }) => id),
+    });
+    expect(selectedAgents).toEqual([]);
+    expect(Object.keys(selectedGroups.platform).length).toEqual(2);
+    expect(Object.keys(selectedGroups.policy).length).toEqual(2);
+  });
+});
 
 describe('processAggregations', () => {
   it('should handle empty inputs properly', () => {

x-pack/plugins/osquery/public/agents/helpers.ts

@@ -114,9 +114,10 @@ export const generateAgentSelection = (selection: GroupOption[]) => {
     platform: {},
   };
 
-  // TODO: clean this up, make it less awkward
   for (const opt of selection) {
     const groupType = opt.value?.groupType;
+    // best effort to get the proper identity
+    const key = opt.key ?? opt.value?.id ?? opt.label;
     let value;
     switch (groupType) {
       case AGENT_GROUP_KEY.All:
@@ -126,17 +127,17 @@ export const generateAgentSelection = (selection: GroupOption[]) => {
         value = opt.value as GroupOptionValue;
         if (!newAgentSelection.allAgentsSelected) {
           // we don't need to calculate diffs when all agents are selected
-          selectedGroups.platform[opt.value?.id ?? opt.label] = value.size;
+          selectedGroups.platform[key] = value.size;
         }
-        newAgentSelection.platformsSelected.push(opt.label);
+        newAgentSelection.platformsSelected.push(key);
         break;
       case AGENT_GROUP_KEY.Policy:
         value = opt.value as GroupOptionValue;
         if (!newAgentSelection.allAgentsSelected) {
           // we don't need to calculate diffs when all agents are selected
-          selectedGroups.policy[opt.value?.id ?? opt.label] = value.size;
+          selectedGroups.policy[key] = value.size;
         }
-        newAgentSelection.policiesSelected.push(opt.label);
+        newAgentSelection.policiesSelected.push(key);
         break;
       case AGENT_GROUP_KEY.Agent:
         value = opt.value as AgentOptionValue;
@@ -144,9 +145,7 @@ export const generateAgentSelection = (selection: GroupOption[]) => {
           // we don't need to count how many agents are selected if they are all selected
           selectedAgents.push(value);
         }
-        if (value?.id) {
-          newAgentSelection.agents.push(value.id);
-        }
+        newAgentSelection.agents.push(key);
         break;
       default:
         // this should never happen!

x-pack/plugins/osquery/public/agents/types.ts

@@ -7,6 +7,7 @@
 
 import { TermsAggregate } from '@elastic/elasticsearch/api/types';
 import { EuiComboBoxOptionOption } from '@elastic/eui';
+import { Agent } from '../../common/shared_imports';
 
 interface BaseDataPoint {
   key: string;
@@ -30,6 +31,8 @@ export interface SelectedGroups {
   [groupType: string]: { [groupName: string]: number };
 }
 
+export type GroupedAgent = Pick<Agent, 'local_metadata' | 'policy_id' | 'status'>;
+
 export type GroupOption = EuiComboBoxOptionOption<AgentOptionValue | GroupOptionValue>;
 
 export interface AgentSelection {
@@ -46,7 +49,7 @@ interface BaseGroupOption {
 
 export type AgentOptionValue = BaseGroupOption & {
   groups: { [groupType: string]: string };
-  online: boolean;
+  status: string;
 };
 
 export type GroupOptionValue = BaseGroupOption & {
@@ -57,5 +60,6 @@ export enum AGENT_GROUP_KEY {
   All,
   Platform,
   Policy,
+  // eslint-disable-next-line @typescript-eslint/no-shadow
   Agent,
 }

x-pack/plugins/osquery/public/agents/use_all_agents.ts

@@ -31,14 +31,22 @@ export const useAllAgents = (
   const { isLoading: agentsLoading, data: agentData } = useQuery<GetAgentsResponse>(
     ['agents', osqueryPolicies, searchValue, perPage],
     () => {
-      let kuery = `(${osqueryPolicies.map((p) => `policy_id:${p}`).join(' or ')})`;
-      if (searchValue) {
-        kuery += ` and (local_metadata.host.hostname:/${searchValue}/ or local_metadata.elastic.agent.id:/${searchValue}/)`;
+      const kueryFragments: string[] = [];
+      if (osqueryPolicies.length) {
+        kueryFragments.push(`${osqueryPolicies.map((p) => `policy_id:${p}`).join(' or ')}`);
       }
+
+      if (searchValue) {
+        kueryFragments.push(
+          `local_metadata.host.hostname:*${searchValue}* or local_metadata.elastic.agent.id:*${searchValue}*`
+        );
+      }
+
       return http.get(agentRouteService.getListPath(), {
         query: {
-          kuery,
+          kuery: kueryFragments.map((frag) => `(${frag})`).join(' and '),
           perPage,
+          showInactive: true,
         },
       });
     },