Update event_generator.py for vulnerability detector

wazuh · Dec 7, 2023 · 0293469 · 0293469
1 parent d90dbf5
commit 0293469
Showing 1 changed file with 24 additions and 11 deletions.
diff --git a/ecs/vulnerability-detector/event-generator/event_generator.py b/ecs/vulnerability-detector/event-generator/event_generator.py
@@ -64,7 +64,7 @@ def generate_random_event():
         'created': generate_random_date(),
         'dataset': random.choice(['process', 'file', 'registry', 'socket', 'dns', 'http', 'tls', 'alert',
                                  'authentication', 'authorization', 'configuration', 'communication', 'file',
-                                 'network', 'process', 'registry', 'storage', 'system', 'web']),
+                                  'network', 'process', 'registry', 'storage', 'system', 'web']),
         'duration': random.randint(0, 99999),
         'end': generate_random_date(),
         'hash': str(hash(f'hash{random.randint(0, 99999)}')),
@@ -74,12 +74,12 @@ def generate_random_event():
                               'state', 'pipeline_error', 'signal']),
         'module': random.choice(['process', 'file', 'registry', 'socket', 'dns', 'http', 'tls', 'alert',
                                 'authentication', 'authorization', 'configuration', 'communication', 'file',
-                                'network', 'process', 'registry', 'storage', 'system', 'web']),
+                                 'network', 'process', 'registry', 'storage', 'system', 'web']),
         'original': f'original{random.randint(0, 99999)}',
         'outcome': random.choice(['success', 'failure', 'unknown']),
         'provider': random.choice(['process', 'file', 'registry', 'socket', 'dns', 'http', 'tls', 'alert',
                                   'authentication', 'authorization', 'configuration', 'communication', 'file',
-                                  'network', 'process', 'registry', 'storage', 'system', 'web']),
+                                   'network', 'process', 'registry', 'storage', 'system', 'web']),
         'reason': f'This event happened due to reason{random.randint(0, 99999)}',
         'reference': f'https://system.example.com/event/#{random.randint(0, 99999)}',
         'risk_score': round(random.uniform(0, 10), 1),
@@ -89,15 +89,16 @@ def generate_random_event():
         'start': generate_random_date(),
         'timezone': random.choice(['UTC', 'GMT', 'PST', 'EST', 'CST', 'MST', 'PDT', 'EDT', 'CDT', 'MDT']),
         'type': random.choice(['access', 'admin', 'allowed', 'change', 'connection', 'creation', 'deletion',
-                             'denied', 'end', 'error', 'group', 'indicator', 'info', 'installation', 'protocol',
-                             'start', 'user']),
+                               'denied', 'end', 'error', 'group', 'indicator', 'info', 'installation', 'protocol',
+                               'start', 'user']),
         'url': f'http://mysystem.example.com/alert/{random.randint(0, 99999)}'
     }
     return event
 
 
 def generate_random_host():
-    family = random.choice(['debian', 'ubuntu', 'macos', 'ios', 'android', 'RHEL'])
+    family = random.choice(
+        ['debian', 'ubuntu', 'macos', 'ios', 'android', 'RHEL'])
     version = f'{random.randint(0, 99)}.{random.randint(0, 99)}'
     host = {
         'os': {
@@ -114,7 +115,8 @@ def generate_random_host():
 
 
 def generate_random_labels():
-    labels = {'label1': f'label{random.randint(0, 99)}', 'label2': f'label{random.randint(0, 99)}'}
+    labels = {
+        'label1': f'label{random.randint(0, 99)}', 'label2': f'label{random.randint(0, 99)}'}
     return labels
 
 
@@ -133,7 +135,7 @@ def generate_random_package():
         'size': random.randint(0, 99999),
         'type': random.choice(['deb', 'rpm', 'msi', 'pkg', 'app', 'apk', 'exe', 'zip', 'tar', 'gz', '7z',
                               'rar', 'cab', 'iso', 'dmg', 'tar.gz', 'tar.bz2', 'tar.xz', 'tar.Z', 'tar.lz4',
-                              'tar.sz', 'tar.zst']),
+                               'tar.sz', 'tar.zst']),
         'version': f'v{random.randint(0, 9)}-stable'
     }
     return package
@@ -166,20 +168,30 @@ def generate_random_vulnerability():
     return vulnerability
 
 
+def generate_random_wazuh():
+    wazuh = {
+        'cluster': {
+            'name': f'wazuh-cluster-{random.randint(0,10)}'
+        }
+    }
+    return wazuh
+
+
 def generate_random_data(number):
     data = []
     for _ in range(number):
         event_data = {
             '@timestamp': generate_random_date(),
             'agent': generate_random_agent(),
             'ecs': {'version': '1.7.0'},
-            'event': generate_random_event(),
+            # 'event': generate_random_event(),
             'host': generate_random_host(),
             'labels': generate_random_labels(),
             'message': f'message{random.randint(0, 99999)}',
             'package': generate_random_package(),
             'tags': generate_random_tags(),
-            'vulnerability': generate_random_vulnerability()
+            'vulnerability': generate_random_vulnerability(),
+            'wazuh': generate_random_wazuh()
         }
         data.append(event_data)
     return data
@@ -221,7 +233,8 @@ def main():
 
     logging.info('Data generation completed.')
 
-    inject = input("Do you want to inject the generated data into your indexer? (y/n) ").strip().lower()
+    inject = input(
+        "Do you want to inject the generated data into your indexer? (y/n) ").strip().lower()
     if inject == 'y':
         ip = input("Enter the IP of your Indexer: ")
         port = input("Enter the port of your Indexer: ")