From 08f935a1751471b0ae050b2ebdba4dab94a94fcd Mon Sep 17 00:00:00 2001
From: Kevin Ledesma <kevinledesmam95@gmail.com>
Date: Thu, 19 Dec 2024 13:20:31 -0300
Subject: [PATCH] Adapt the commands event generator to job-scheduler changes
 (#577)

* Adapt commands event_generator to job-scheduler changes

Use new commands array on the body for API bulks

Add tool argument to select http or https protocols

* Update post request body to use JSON dumps

* Fix generated data dump to log file
---
 .../event-generator/event_generator.py        | 56 +++++++++++--------
 1 file changed, 34 insertions(+), 22 deletions(-)

diff --git a/ecs/command/event-generator/event_generator.py b/ecs/command/event-generator/event_generator.py
index 2ae50a82bced6..7b7ac4c1566ff 100644
--- a/ecs/command/event-generator/event_generator.py
+++ b/ecs/command/event-generator/event_generator.py
@@ -75,34 +75,41 @@ def generate_random_data(number, include_all_fields=False):
     data = []
     for _ in range(number):
         data.append(generate_random_command(include_all_fields))
+    if not include_all_fields:
+        return {"commands": data}
     return data
 
 
-def inject_events(ip, port, index, username, password, data, use_index=False):
-    session = requests.Session()
-    session.auth = (username, password)
-    session.verify = False
-    headers = {'Content-Type': 'application/json'}
-
+def inject_events(protocol, ip, port, index, username, password, data, use_index=False):
     try:
+        if not use_index:
+            # Use the command-manager API
+            url = f'{protocol}://{ip}:{port}/_plugins/_command_manager/commands'
+            send_post_request(username, password, url, data)
+            return
         for event_data in data:
-            if use_index:
-                # Generate UUIDs for the document id
-                doc_id = str(uuid.uuid4())
-                url = f'http://{ip}:{port}/{index}/_doc/{doc_id}'
-            else:
-                # Default URL for command manager API without the index
-                url = f'http://{ip}:{port}/_plugins/_command_manager/commands'
-            response = session.post(url, json=event_data, headers=headers)
-            if response.status_code != 201:
-                logging.error(f'Error: {response.status_code}')
-                logging.error(response.text)
-                break
+            # Generate UUIDs for the document id
+            doc_id = str(uuid.uuid4())
+            url = f'{protocol}://{ip}:{port}/{index}/_doc/{doc_id}'
+            send_post_request(username, password, url, event_data)
         logging.info('Data injection completed successfully.')
     except Exception as e:
         logging.error(f'Error: {str(e)}')
 
 
+def send_post_request(username, password, url, event_data):
+    session = requests.Session()
+    session.auth = (username, password)
+    session.verify = False
+    headers = {'Content-Type': 'application/json'}
+    # Send request
+    response = session.post(url, data=json.dumps(event_data), headers=headers)
+    if response.status_code not in [201, 200]:
+        logging.error(f'Error: {response.status_code}')
+        logging.error(response.text)
+    return response
+
+
 def main():
     parser = argparse.ArgumentParser(
         description="Generate and optionally inject events into an OpenSearch index or Command Manager."
@@ -112,6 +119,12 @@ def main():
         action="store_true",
         help="Generate additional fields for indexing and inject into a specific index."
     )
+    parser.add_argument(
+        "--protocol",
+        choices=['http', 'https'],
+        default='https',
+        help="Specify the protocol to use: http or https."
+    )
     args = parser.parse_args()
 
     try:
@@ -124,9 +137,8 @@ def main():
     data = generate_random_data(number, include_all_fields=args.index)
 
     with open(GENERATED_DATA_FILE, 'a') as outfile:
-        for event_data in data:
-            json.dump(event_data, outfile)
-            outfile.write('\n')
+        json.dump(data, outfile)
+        outfile.write('\n')
 
     logging.info('Data generation completed.')
 
@@ -145,7 +157,7 @@ def main():
         username = input(f"Username (default: '{USERNAME}'): ") or USERNAME
         password = input(f"Password (default: '{PASSWORD}'): ") or PASSWORD
 
-        inject_events(ip, port, index, username, password,
+        inject_events(args.protocol, ip, port, index, username, password,
                       data, use_index=bool(args.index))