Skip to content

Commit

Permalink
Adapt the commands event generator to job-scheduler changes (#577)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
QU3B1M authored Dec 19, 2024
1 parent a6476b2 commit 08f935a
Showing 1 changed file with 34 additions and 22 deletions.
56 changes: 34 additions & 22 deletions ecs/command/event-generator/event_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand All @@ -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:
Expand All @@ -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.')

Expand All @@ -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))


Expand Down

0 comments on commit 08f935a

Please sign in to comment.