Move space handling to kibanaFetch and add space support to entity-store and entity-resolution-demo commands

src/commands/asset_criticality.ts

@@ -13,10 +13,12 @@ export const generateAssetCriticality = async ({
   users,
   hosts,
   seed = generateNewSeed(),
+  space = 'default',
 }: {
   users: number;
   hosts: number;
   seed?: number;
+  space: string;
 }) => {
   faker.seed(seed);
 
@@ -29,9 +31,17 @@ export const generateAssetCriticality = async ({
       count: hosts,
     });
 
-    await assignAssetCriticalityToEntities(generatedUsers, 'user.name');
+    await assignAssetCriticalityToEntities({
+      entities: generatedUsers,
+      field: 'user.name',
+      space,
+    });
     console.log(`Assigned asset criticality to ${generatedUsers.length} users`);
-    await assignAssetCriticalityToEntities(generatedHosts, 'host.name');
+    await assignAssetCriticalityToEntities({
+      entities: generatedHosts,
+      field: 'host.name',
+      space,
+    });
     console.log(`Assigned asset criticality to ${generatedHosts.length} hosts`);
 
     console.log('Finished generating asset criticality');

src/commands/documents.ts

@@ -196,7 +196,6 @@ export const generateEvents = async (n: number) => {
 };
 
 export const generateGraph = async ({ users = 100, maxHosts = 3 }) => {
-  //await alertIndexCheck(); TODO
   console.log('Generating alerts graph...');
 
   type AlertOverride = { host: { name: string }; user: { name: string } };

src/commands/entity-store.ts

@@ -15,6 +15,7 @@ import {
   MappingTypeMapping,
 } from '@elastic/elasticsearch/lib/api/types';
 import { getConfig } from '../get_config';
+import { initializeSpace } from '../utils';
 
 const config = getConfig();
 const client = getEsClient();
@@ -263,10 +264,12 @@ export const generateEvents = <E extends BaseEntity, EV = BaseEvent>(
   }, acc);
 };
 
-export const assignAssetCriticalityToEntities = async (
-  entities: BaseEntity[],
-  field: string,
-) => {
+export const assignAssetCriticalityToEntities = async (opts: {
+  entities: BaseEntity[];
+  field: string;
+  space?: string;
+}) => {
+  const { entities, field, space } = opts;
   const chunks = chunk(entities, 10000);
   for (const chunk of chunks) {
     const records = chunk
@@ -278,7 +281,7 @@ export const assignAssetCriticalityToEntities = async (
       }));
 
     if (records.length > 0) {
-      await assignAssetCriticality(records);
+      await assignAssetCriticality(records, space);
     }
   }
 };
@@ -292,12 +295,14 @@ export const generateEntityStore = async ({
   hosts = 10,
   services = 10,
   seed = generateNewSeed(),
+  space,
   options,
 }: {
   users: number;
   hosts: number;
   services: number;
   seed: number;
+  space?: string;
   options: string[];
 }) => {
   if (options.includes(ENTITY_STORE_OPTIONS.seed)) {
@@ -342,20 +347,32 @@ export const generateEntityStore = async ({
     await ingestEvents(eventsForServices);
     console.log('Services events ingested');
 
+    if (space && space !== 'default') {
+      await initializeSpace(space);
+    }
+
     if (options.includes(ENTITY_STORE_OPTIONS.criticality)) {
-      await assignAssetCriticalityToEntities(generatedUsers, 'user.name');
+      await assignAssetCriticalityToEntities({
+        entities: generatedUsers,
+        field: 'user.name',
+        space,
+      });
       console.log('Assigned asset criticality to users');
-      await assignAssetCriticalityToEntities(generatedHosts, 'host.name');
+      await assignAssetCriticalityToEntities({
+        entities: generatedHosts,
+        field: 'host.name',
+        space,
+      });
       console.log('Assigned asset criticality to hosts');
     }
 
     if (options.includes(ENTITY_STORE_OPTIONS.riskEngine)) {
-      await enableRiskScore();
+      await enableRiskScore(space);
       console.log('Risk score enabled');
     }
 
     if (options.includes(ENTITY_STORE_OPTIONS.rule)) {
-      await createRule();
+      await createRule({ space });
       console.log('Rule created');
     }
 

src/commands/entity_resolution.ts

@@ -4,7 +4,7 @@ import {
   createRule,
   getRule,
   createComponentTemplate,
-  appendPathToKibanaNode,
+  buildKibanaUrl,
 } from '../utils/kibana_api';
 import pMap from 'p-map';
 import cliProgress from 'cli-progress';
@@ -213,7 +213,7 @@ const PACKAGES_TO_INSTALL = [
   'entityanalytics_entra_id',
 ];
 
-const installPackages = async () => {
+const installPackages = async (space: string) => {
   console.log('Installing packages...');
   const progress = new cliProgress.SingleBar(
     {
@@ -225,7 +225,7 @@ const installPackages = async () => {
   await pMap(
     PACKAGES_TO_INSTALL,
     async (packageName) => {
-      await installPackage({ packageName });
+      await installPackage({ packageName, space });
       progress.increment();
     },
     { concurrency: 1 },
@@ -415,8 +415,8 @@ const importFile = async (
   await batchIndexDocsWithProgress(batchGenerator, lineCountInFile);
 };
 
-const createMatchAllRule = async () => {
-  const rule = await getRule(RULE_ID);
+const createMatchAllRule = async (space: string) => {
+  const rule = await getRule(RULE_ID, space);
 
   if (rule) {
     console.log('Match all rule already exists.');
@@ -425,6 +425,7 @@ const createMatchAllRule = async () => {
 
   await createRule({
     id: RULE_ID,
+    space,
   });
   console.log('Match all rule created.');
 };
@@ -462,10 +463,12 @@ export const setupEntityResolutionDemo = async ({
   mini = false,
   deleteData = false,
   keepEmails = false,
+  space,
 }: {
   mini: boolean;
   deleteData: boolean;
   keepEmails: boolean;
+  space: string;
 }) => {
   if (deleteData) {
     console.log('Deleting existing demo data first...');
@@ -474,9 +477,9 @@ export const setupEntityResolutionDemo = async ({
 
   console.log(`Setting up${mini ? ' mini' : ''} entity resolution demo...`);
   // create a rule which matches everything, handy for exploring all the different entity views
-  await createMatchAllRule();
+  await createMatchAllRule(space);
   // install the packages to get the mappings in place
-  await installPackages();
+  await installPackages(space);
   // create @custom component templates to get user.name and user.email field mappings
   // which the inttegrations don't provide
   // we will eventually have to release a new version of the integrations to include these mappings
@@ -493,7 +496,7 @@ Entity resolution demo setup complete.
 
 Now go and install the model!
 
-    CLICK HERE ---->> ${appendPathToKibanaNode('/app/security/entity_analytics_management')} <<---- CLICK HERE
+    CLICK HERE ---->> ${buildKibanaUrl({ path: '/app/security/entity_analytics_management', space })} <<---- CLICK HERE
 
 Once installed, ${mini ? 'Mark Hopkin should have matches' : 'See here:\n\n https://github.com/elastic/security-ml/blob/gus/entity_resoluton_data_generation/projects/entity_resolution_poc_2024/test_data_generation/seed_data_with_name_variations_and_user_agent_gen_and_groups.json \n\nfor all the seed data names'}
   `);

src/constants.ts

@@ -10,33 +10,42 @@ export const generateNewSeed = () => {
   return Math.round(Math.random() * 100000);
 };
 
+export const API_VERSIONS = {
+  public: {
+    v1: '2023-10-31',
+  },
+  internal: {
+    v1: '1',
+  },
+};
+
 // API Endpoint URL's for Kibana
 export const RISK_SCORE_URL = '/internal/risk_score';
 export const RISK_SCORE_DASHBOARD_URL = (entityType: 'host' | 'user') =>
   `/internal/risk_score/prebuilt_content/saved_objects/_bulk_create/${entityType}RiskScoreDashboards`;
 export const RISK_SCORE_SCORES_URL = '/internal/risk_score/scores';
 export const RISK_SCORE_ENGINE_INIT_URL = '/internal/risk_score/engine/init';
 export const ASSET_CRITICALITY_URL = '/api/asset_criticality';
-export const DETECTION_ENGINE_RULES_URL = (space?: string) =>
-  space
-    ? `/s/${space}/api/detection_engine/rules`
-    : '/api/detection_engine/rules';
-export const COMPONENT_TEMPLATES_URL = (space?: string) =>
-  space
-    ? `/s/${space}/api/index_management/component_templates`
-    : '/api/index_management/component_templates';
+export const ASSET_CRITICALITY_BULK_URL = '/api/asset_criticality/bulk';
+export const DETECTION_ENGINE_RULES_URL = '/api/detection_engine/rules';
+export const DETECTION_ENGINE_RULES_BULK_ACTION_URL = `${DETECTION_ENGINE_RULES_URL}/_bulk_action`;
+export const COMPONENT_TEMPLATES_URL =
+  '/api/index_management/component_templates';
 export const FLEET_EPM_PACKAGES_URL = (
   packageName: string,
   version: string = 'latest',
-  space?: string,
 ) => {
-  let url = space
-    ? `/s/${space}/api/fleet/epm/packages/${packageName}`
-    : `/api/fleet/epm/packages/${packageName}`;
+  let url = `/api/fleet/epm/packages/${packageName}`;
   if (version !== 'latest') {
     url = `${url}/${version}`;
   }
   return url;
 };
 export const SPACES_URL = '/api/spaces/space';
 export const SPACE_URL = (space: string) => `/api/spaces/space/${space}`;
+
+export const ENTITY_ENGINES_URL = '/api/entity_store/engines';
+export const ENTITY_ENGINE_URL = (engineType: string) =>
+  `${ENTITY_ENGINES_URL}/${engineType}`;
+export const INIT_ENTITY_ENGINE_URL = (engineType: string) =>
+  `${ENTITY_ENGINE_URL(engineType)}/init`;

src/index.ts

@@ -156,15 +156,17 @@ program
   .option('--mini', 'Only load the mini dataset', false)
   .option('--delete', 'Delete old data', false)
   .option('--keep-emails', 'No Email variants', false)
+  .option('--space', 'space to use', 'default')
   .description('Load entity resolution demo data')
-  .action(({ mini, deleteData, keepEmails }) => {
-    setupEntityResolutionDemo({ mini, deleteData, keepEmails });
+  .action(({ mini, deleteData, keepEmails, space }) => {
+    setupEntityResolutionDemo({ mini, deleteData, keepEmails, space });
   });
 
 program
   .command('entity-store')
   .description('Generate entity store')
-  .action(async () => {
+  .option('--space <space>', 'Space to create entity store in')
+  .action(async (options) => {
     const entityStoreAnswers = await checkbox<
       keyof typeof ENTITY_STORE_OPTIONS
     >({
@@ -226,6 +228,7 @@ program
     }
 
     generateEntityStore({
+      space: options.space,
       users: parseIntBase10(userCount),
       hosts: parseIntBase10(hostCount),
       services: parseIntBase10(serviceCount),
@@ -243,12 +246,14 @@ program
   .command('generate-asset-criticality')
   .option('-h <h>', 'number of hosts')
   .option('-u <u>', 'number of users')
+  .option('-s <s>', 'space')
   .description('Generate asset criticality for entities')
   .action(async (options) => {
     const users = parseInt(options.u || 10);
     const hosts = parseInt(options.h || 10);
+    const space = options.s || 'default';
 
-    generateAssetCriticality({ users, hosts });
+    generateAssetCriticality({ users, hosts, space });
   });
 
 program

src/utils/initialize_space.ts

@@ -86,7 +86,7 @@ const waitForAlertIndexMapping = async (
 
 const ensureSpaceExists = async (space: string) => {
   console.log(`Checking if space ${space} exists`);
-  if (await kibanaApi.getSpace(space)) {
+  if (await kibanaApi.doesSpaceExist(space)) {
     console.log(`Space ${space} exists`);
     return;
   }