From a40076b65856e75edd222e3786b8e898a1f2581f Mon Sep 17 00:00:00 2001
From: Marshall Main <55718608+marshallmain@users.noreply.github.com>
Date: Wed, 3 Jun 2020 19:23:47 -0400
Subject: [PATCH] Add generator function that creates multiple alerts (#67713)

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
---
 .../siem/common/endpoint/generate_data.ts     | 35 ++++++++++++
 .../scripts/endpoint/resolver_generator.ts    | 56 ++++++++++---------
 2 files changed, 64 insertions(+), 27 deletions(-)

diff --git a/x-pack/plugins/siem/common/endpoint/generate_data.ts b/x-pack/plugins/siem/common/endpoint/generate_data.ts
index 597ad4df64dfeb..57d41b65549073 100644
--- a/x-pack/plugins/siem/common/endpoint/generate_data.ts
+++ b/x-pack/plugins/siem/common/endpoint/generate_data.ts
@@ -445,6 +445,41 @@ export class EndpointDocGenerator {
     };
   }
 
+  /**
+   * Wrapper generator for fullResolverTreeGenerator to make it easier to quickly stream
+   * many resolver trees to Elasticsearch.
+   * @param numAlerts - number of alerts to generate
+   * @param alertAncestors - number of ancestor generations to create relative to the alert
+   * @param childGenerations - number of child generations to create relative to the alert
+   * @param maxChildrenPerNode - maximum number of children for any given node in the tree
+   * @param relatedEventsPerNode - number of related events (file, registry, etc) to create for each process event in the tree
+   * @param percentNodesWithRelated - percent of nodes which should have related events
+   * @param percentTerminated - percent of nodes which will have process termination events
+   * @param alwaysGenMaxChildrenPerNode - flag to always return the max children per node instead of it being a random number of children
+   */
+  public *alertsGenerator(
+    numAlerts: number,
+    alertAncestors?: number,
+    childGenerations?: number,
+    maxChildrenPerNode?: number,
+    relatedEventsPerNode?: number,
+    percentNodesWithRelated?: number,
+    percentTerminated?: number,
+    alwaysGenMaxChildrenPerNode?: boolean
+  ) {
+    for (let i = 0; i < numAlerts; i++) {
+      yield* this.fullResolverTreeGenerator(
+        alertAncestors,
+        childGenerations,
+        maxChildrenPerNode,
+        relatedEventsPerNode,
+        percentNodesWithRelated,
+        percentTerminated,
+        alwaysGenMaxChildrenPerNode
+      );
+    }
+  }
+
   /**
    * Generator function that creates the full set of events needed to render resolver.
    * The number of nodes grows exponentially with the number of generations and children per node.
diff --git a/x-pack/plugins/siem/scripts/endpoint/resolver_generator.ts b/x-pack/plugins/siem/scripts/endpoint/resolver_generator.ts
index 26c6e5ccc28a88..63c15077181373 100644
--- a/x-pack/plugins/siem/scripts/endpoint/resolver_generator.ts
+++ b/x-pack/plugins/siem/scripts/endpoint/resolver_generator.ts
@@ -221,6 +221,7 @@ async function main() {
     console.log(`No seed supplied, using random seed: ${seed}`);
   }
   const random = seedrandom(seed);
+  const startTime = new Date().getTime();
   for (let i = 0; i < argv.numHosts; i++) {
     const generator = new EndpointDocGenerator(random);
     const timeBetweenDocs = 6 * 3600 * 1000; // 6 hours between metadata documents
@@ -241,36 +242,37 @@ async function main() {
       });
     }
 
-    for (let j = 0; j < argv.alertsPerHost; j++) {
-      const resolverDocGenerator = generator.fullResolverTreeGenerator(
-        argv.ancestors,
-        argv.generations,
-        argv.children,
-        argv.relatedEvents,
-        argv.percentWithRelated,
-        argv.percentTerminated,
-        argv.maxChildrenPerNode
-      );
-      let result = resolverDocGenerator.next();
-      while (!result.done) {
-        let k = 0;
-        const resolverDocs: Event[] = [];
-        while (k < 1000 && !result.done) {
-          resolverDocs.push(result.value);
-          result = resolverDocGenerator.next();
-          k++;
-        }
-        const body = resolverDocs.reduce(
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
-          (array: Array<Record<string, any>>, doc) => (
-            array.push({ index: { _index: argv.eventIndex } }, doc), array
-          ),
-          []
-        );
-        await client.bulk({ body });
+    const alertGenerator = generator.alertsGenerator(
+      argv.alertsPerHost,
+      argv.ancestors,
+      argv.generations,
+      argv.children,
+      argv.relatedEvents,
+      argv.percentWithRelated,
+      argv.percentTerminated,
+      argv.maxChildrenPerNode
+    );
+    let result = alertGenerator.next();
+    while (!result.done) {
+      let k = 0;
+      const resolverDocs: Event[] = [];
+      while (k < 1000 && !result.done) {
+        resolverDocs.push(result.value);
+        result = alertGenerator.next();
+        k++;
       }
+      const body = resolverDocs.reduce(
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (array: Array<Record<string, any>>, doc) => (
+          array.push({ index: { _index: argv.eventIndex } }, doc), array
+        ),
+        []
+      );
+      await client.bulk({ body });
     }
   }
+  // eslint-disable-next-line no-console
+  console.log(`Creating and indexing documents took: ${new Date().getTime() - startTime}ms`);
 }
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any