Alert generator improvements

README.md

@@ -57,7 +57,7 @@ Examles of config:
 ### Alerts
 `yarn start help` - To see the commands list
 
-`yarn start generate-alerts <n>` - Generate *n* alerts
+`yarn start generate-alerts -n <number of alerts> -h <number of hosts within the alerts> -u <number of users within the alerts>
 
 `yarn start delete-alerts` - Delete all alerts
 
@@ -76,7 +76,7 @@ To modify alert document, you can change `createAlert.ts` file.
 Example list of command for testing Risk Score API woth 10.000 alerts.
 ```
 yarn start delete-alerts
-yarn start generate-alerts 10000
+yarn start generate-alerts -n 10000 -h 100 -u 100
 yarn start test-risk-score
 ```
 

package.json

@@ -16,6 +16,7 @@
     "@faker-js/faker": "^8.0.0",
     "@types/lodash-es": "^4.17.12",
     "chalk": "^5.2.0",
+    "cli-progress": "^3.12.0",
     "commander": "^10.0.0",
     "conf": "^11.0.1",
     "fp-ts": "^2.16.5",
@@ -24,17 +25,19 @@
     "lodash-es": "^4.17.21",
     "moment": "^2.29.4",
     "node-fetch": "^3.3.1",
+    "p-map": "^7.0.2",
     "tsx": "^4.7.1",
     "url-join": "^5.0.0",
     "uuid": "^9.0.1"
   },
   "devDependencies": {
-    "eslint": "<9.0.0",
+    "@types/cli-progress": "^3.11.5",
     "@types/inquirer": "^9.0.7",
     "@types/uuid": "^9.0.8",
     "@typescript-eslint/eslint-plugin": "^7.7.1",
     "@typescript-eslint/parser": "^7.7.1",
     "esbuild": "^0.20.2",
+    "eslint": "<9.0.0",
     "eslint-config-airbnb": "^19.0.4",
     "eslint-config-airbnb-typescript": "^18.0.0",
     "typescript": "^5.3.3"

src/commands/documents.ts

@@ -1,17 +1,19 @@
 
 import createAlerts from '../createAlerts';
 import createEvents from '../createEvents';
-import alertMappings from '../mappings/alertMappings.json' assert { type: 'json' };
 import eventMappings from '../mappings/eventMappings.json' assert { type: 'json' };
 import { getEsClient, indexCheck } from './utils/index';
 import { getConfig } from '../get_config';
 import { MappingTypeMapping, BulkOperationContainer } from '@elastic/elasticsearch/lib/api/types';
+import pMap from 'p-map';
+import _ from 'lodash';
+import cliProgress from 'cli-progress';
+import { faker } from '@faker-js/faker';
 
 const config = getConfig();
 const client = getEsClient(); 
 
 const ALERT_INDEX = '.alerts-security.alerts-default';
-const EVENT_INDEX = config.eventIndex;
 
 const generateDocs = async ({ createDocs, amount, index }: {createDocs: DocumentCreator; amount: number; index: string}) => {
   if (!client) {
@@ -28,21 +30,39 @@ const generateDocs = async ({ createDocs, amount, index }: {createDocs: Document
       index
     );
     try {
-      const result = await client.bulk({ body: docs, refresh: true });
+      const result = await bulkUpsert(docs);
       generated += result.items.length / 2;
-      console.log(
-        `${result.items.length} documents created, ${amount - generated} left`
-      );
     } catch (err) {
       console.log('Error: ', err);
+      process.exit(1);
     }
   }
 };
 
+const bulkUpsert = async (docs: unknown[]) => {
+  if (!client) {
+    throw new Error('failed to create ES client');
+  }
+  try {
+    return client.bulk({ body: docs, refresh: true });
+  } catch (err) {
+    console.log('Error: ', err);
+    process.exit(1);
+  }
+};
+
 interface DocumentCreator {
 	(descriptor: { id_field: string, id_value: string }): object;
 }
 
+const alertToBatchOps = (alert: object, index: string): unknown[] => {
+  return [
+    { index: { _index: index } },
+    { ...alert },
+  ];
+
+}
+
 const createDocuments = (n: number, generated: number, createDoc: DocumentCreator, index: string): unknown[] => {
   return Array(n)
     .fill(null)
@@ -64,29 +84,71 @@ const createDocuments = (n: number, generated: number, createDoc: DocumentCreato
 };
 
 
-export const generateAlerts = async (n: number) => {
-  await indexCheck(ALERT_INDEX, alertMappings as MappingTypeMapping);
+export const generateAlerts = async (alertCount: number, hostCount: number, userCount: number) => {
+
+  if (userCount > alertCount) {
+    console.log('User count should be less than alert count');
+    process.exit(1);
+  }
 
-  console.log('Generating alerts...');
+  if (hostCount > alertCount) {
+    console.log('Host count should be less than alert count');
+    process.exit(1);
+  }
 
-  await generateDocs({
-    createDocs: createAlerts,
-    amount: n,
-    index: ALERT_INDEX,
-  });
+  console.log(`Generating ${alertCount} alerts containing ${hostCount} hosts and ${userCount} users.`);
+  const concurrency = 10; // how many batches to send in parallel
+  const batchSize = 2500; // number of alerts in a batch
+  const no_overrides = {};
+
+  const batchOpForIndex = ({ userName, hostName } : { userName: string, hostName: string }) => alertToBatchOps(createAlerts(no_overrides, { userName, hostName }), ALERT_INDEX);
+
+
+  console.log('Generating entity names...');
+  const userNames = Array.from({ length: userCount }, () => faker.internet.userName());
+  const hostNames = Array.from({ length: hostCount }, () => faker.internet.domainName());
+
+  console.log('Assigning entity names...')
+  const alertEntityNames = Array.from({ length: alertCount }, (_, i) => ({
+    userName: userNames[i % userCount],
+    hostName: hostNames[i % hostCount],
+  }));
+
+  console.log('Entity names assigned. Batching...');
+  const operationBatches = _.chunk(alertEntityNames, batchSize).map((batch) => 
+    batch.flatMap(batchOpForIndex)
+  );
+
+  console.log('Batching complete. Sending to ES...');
+
+  console.log(`Sending in ${operationBatches.length} batches of ${batchSize} alerts, with up to ${concurrency} batches in parallel\n\n`);
+  const progress = new cliProgress.SingleBar({}, cliProgress.Presets.shades_classic);
+
+  progress.start(operationBatches.length, 0);
+
+  await pMap(
+    operationBatches,
+    async (operations) => {
+      await bulkUpsert(operations);
+      progress.increment();
+    },  
+    { concurrency }
+  );
 
-  console.log('Finished gerating alerts');
+  progress.stop();
 };
 
+// this creates asset criticality not events? 
 export const generateEvents = async (n: number) => {
-  await indexCheck(EVENT_INDEX, eventMappings as MappingTypeMapping);
+  if(!config.eventIndex) { throw new Error('eventIndex not defined in config'); }
+  await indexCheck(config.eventIndex, eventMappings as MappingTypeMapping);
 
   console.log('Generating events...');
 
   await generateDocs({
     createDocs: createEvents,
     amount: n,
-    index: EVENT_INDEX,
+    index: config.eventIndex,
   });
 
   console.log('Finished generating events');
@@ -112,10 +174,10 @@ export const generateGraph = async ({ users = 100, maxHosts = 3 }) => {
     for (let j = 0; j < maxHosts; j++) {
       const alert = createAlerts({
         host: {
-          name: `Host ${i}${j}`,
+          name: 'Host mark',
         },
         user: {
-          name: `User ${i}`,
+          name: 'User pablo',
         },
       });
       userCluster.push(alert);
@@ -179,11 +241,12 @@ export const deleteAllAlerts = async () => {
 
 export const deleteAllEvents = async () => {
   console.log('Deleting all events...');
+  if (!config.eventIndex) { throw new Error('eventIndex not defined in config'); }
   try {
     console.log('Deleted all events');
     if (!client) throw new Error;
     await client.deleteByQuery({
-      index: EVENT_INDEX,
+      index: config.eventIndex,
       refresh: true,
       body: {
         query: {

src/createAlerts.ts

@@ -1,7 +1,16 @@
 import { faker } from '@faker-js/faker';
 
-function baseCreateAlerts() {
+function baseCreateAlerts({
+  userName = 'user-1',
+  hostName = 'host-1',
+} : {
+  userName?: string,
+  hostName?: string,
+} = {
+}) {
   return {
+    'host.name': hostName,
+    'user.name': userName,
     'kibana.alert.start': '2023-04-11T20:18:15.816Z',
     'kibana.alert.last_detected': '2023-04-11T20:18:15.816Z',
     'kibana.version': '8.7.0',
@@ -59,7 +68,7 @@ function baseCreateAlerts() {
     'kibana.alert.status': 'active',
     'kibana.alert.workflow_status': 'open',
     'kibana.alert.depth': 1,
-    'kibana.alert.reason': 'event on Host 4 created low alert 1.',
+    'kibana.alert.reason': 'event on ' + hostName +  'created low alert 1.',
     'kibana.alert.severity': 'low',
     'kibana.alert.risk_score': 21,
     'kibana.alert.rule.actions': [],
@@ -96,6 +105,13 @@ function baseCreateAlerts() {
   }
 }
 
-export default function createAlerts<O extends object>(override: O): O & ReturnType<typeof baseCreateAlerts> {
-  return { ...baseCreateAlerts(), ...override };
+export default function createAlerts<O extends object>(override: O, {
+  userName,
+  hostName,
+} : {
+  userName?: string,
+  hostName?: string,
+} = {
+}): O & ReturnType<typeof baseCreateAlerts> {
+  return { ...baseCreateAlerts({ userName, hostName}), ...override };
 }

src/get_config.ts

@@ -18,8 +18,8 @@ const Node = t.union([NodeWithCredentials, NodeWithAPIKey]);
 const Config = t.type({
   elastic: Node,
   kibana: Node,
-  eventIndex: t.string,
-  eventDateOffsetHours: t.number,
+  eventIndex: t.union([t.string, t.undefined]),
+  eventDateOffsetHours: t.union([t.number, t.undefined]),
 });
 
 type ConfigType = t.TypeOf<typeof Config>;

src/index.ts

@@ -17,9 +17,16 @@ import { ENTITY_STORE_OPTIONS, generateNewSeed } from './constants';
 
 program
   .command('generate-alerts')
-  .argument('<n>', 'integer argument', parseInt)
+  .option('-n <n>', 'number of alerts')
+  .option('-h <h>', 'number of hosts')
+  .option('-u <h>', 'number of users')
   .description('Generate fake alerts')
-  .action(generateAlerts);
+  .action((options) => {
+    const alertsCount = parseInt(options.n || 1);
+    const hostCount = parseInt(options.h || 1);
+    const userCount = parseInt(options.u || 1);
+    generateAlerts(alertsCount, userCount, hostCount);
+  });
 
 program
   .command('generate-events')

yarn.lock

@@ -346,6 +346,13 @@
     "@nodelib/fs.scandir" "2.1.5"
     fastq "^1.6.0"
 
+"@types/cli-progress@^3.11.5":
+  version "3.11.5"
+  resolved "https://registry.yarnpkg.com/@types/cli-progress/-/cli-progress-3.11.5.tgz#9518c745e78557efda057e3f96a5990c717268c3"
+  integrity sha512-D4PbNRbviKyppS5ivBGyFO29POlySLmA2HyUFE4p5QGazAMM3CwkKWcvTl8gvElSuxRh6FPKL8XmidX873ou4g==
+  dependencies:
+    "@types/node" "*"
+
 "@types/inquirer@^9.0.7":
   version "9.0.7"
   resolved "https://registry.yarnpkg.com/@types/inquirer/-/inquirer-9.0.7.tgz#61bb8d0e42f038b9a1738b08fba7fa98ad9b4b24"
@@ -650,6 +657,13 @@ cli-cursor@^3.1.0:
   dependencies:
     restore-cursor "^3.1.0"
 
+cli-progress@^3.12.0:
+  version "3.12.0"
+  resolved "https://registry.yarnpkg.com/cli-progress/-/cli-progress-3.12.0.tgz#807ee14b66bcc086258e444ad0f19e7d42577942"
+  integrity sha512-tRkV3HJ1ASwm19THiiLIXLO7Im7wlTuKnvkYaTkyoAPefqjNg7W7DHKUlGRxy9vxDvbyCYQkQozvptuMkGCg8A==
+  dependencies:
+    string-width "^4.2.3"
+
 cli-spinners@^2.5.0:
   version "2.9.2"
   resolved "https://registry.yarnpkg.com/cli-spinners/-/cli-spinners-2.9.2.tgz#1773a8f4b9c4d6ac31563df53b3fc1d79462fe41"
@@ -1570,6 +1584,11 @@ p-locate@^5.0.0:
   dependencies:
     p-limit "^3.0.2"
 
+p-map@^7.0.2:
+  version "7.0.2"
+  resolved "https://registry.yarnpkg.com/p-map/-/p-map-7.0.2.tgz#7c5119fada4755660f70199a66aa3fe2f85a1fe8"
+  integrity sha512-z4cYYMMdKHzw4O5UkWJImbZynVIo0lSGTXc7bzB1e/rrDqkgGUNysK/o4bTr+0+xKvvLoTyGqYC4Fgljy9qe1Q==
+
 parent-module@^1.0.0:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/parent-module/-/parent-module-1.0.1.tgz#691d2709e78c79fae3a156622452d00762caaaa2"