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))