update TheHiveAlerter module

salt/elastalert/files/modules/so/thehive.py

@@ -1,84 +1,107 @@
 # -*- coding: utf-8 -*-
-from __future__ import unicode_literals
+# HiveAlerter modified from original at: https://raw.githubusercontent.com/Nclose-ZA/elastalert_hive_alerter/master/elastalert_hive_alerter/hive_alerter.py
+
 import uuid
-import re
 
 from elastalert.alerts import Alerter
 from thehive4py.api import TheHiveApi
 from thehive4py.models import Alert, AlertArtifact, CustomFieldHelper
 
+
 class TheHiveAlerter(Alerter):
     """
     Use matched data to create alerts containing observables in an instance of TheHive
-    This is a modified version for use with Security Onion
     """
 
     required_options = set(['hive_connection', 'hive_alert_config'])
 
-    def alert(self, matches):
+    def get_aggregation_summary_text(self, matches):
+        text = super(TheHiveAlerter, self).get_aggregation_summary_text(matches)
+        if text:
+            text = '```\n{0}```\n'.format(text)
+        return text
 
+    def create_artifacts(self, match):
+        artifacts = []
+        context = {'rule': self.rule, 'match': match}
+        for mapping in self.rule.get('hive_observable_data_mapping', []):
+            for observable_type, match_data_key in mapping.items():
+                try:
+                    artifacts.append(AlertArtifact(dataType=observable_type, data=match_data_key.format(**context)))
+                except KeyError as e:
+                    print(('format string {} fail cause no key {} in {}'.format(e, match_data_key, context)))
+        return artifacts
+
+    def create_alert_config(self, match):
+        context = {'rule': self.rule, 'match': match}
+        alert_config = {
+            'artifacts': self.create_artifacts(match),
+            'sourceRef': str(uuid.uuid4())[0:6],
+            'title': '{rule[name]}'.format(**context)
+        }
+
+        alert_config.update(self.rule.get('hive_alert_config', {}))
+
+        for alert_config_field, alert_config_value in alert_config.items():
+            if alert_config_field == 'customFields':
+                custom_fields = CustomFieldHelper()
+                for cf_key, cf_value in alert_config_value.items():
+                    try:
+                        func = getattr(custom_fields, 'add_{}'.format(cf_value['type']))
+                    except AttributeError:
+                        raise Exception('unsupported custom field type {}'.format(cf_value['type']))
+                    value = cf_value['value'].format(**context)
+                    func(cf_key, value)
+                alert_config[alert_config_field] = custom_fields.build()
+            elif isinstance(alert_config_value, str):
+                alert_config[alert_config_field] = alert_config_value.format(**context)
+            elif isinstance(alert_config_value, (list, tuple)):
+                formatted_list = []
+                for element in alert_config_value:
+                    try:
+                        formatted_list.append(element.format(**context))
+                    except (AttributeError, KeyError, IndexError):
+                        formatted_list.append(element)
+                alert_config[alert_config_field] = formatted_list
+
+        return alert_config
+
+    def send_to_thehive(self, alert_config):
         connection_details = self.rule['hive_connection']
-
         api = TheHiveApi(
-            connection_details.get('hive_host'),
+            connection_details.get('hive_host', ''),
             connection_details.get('hive_apikey', ''),
             proxies=connection_details.get('hive_proxies', {'http': '', 'https': ''}),
             cert=connection_details.get('hive_verify', False))
 
-        for match in matches:
+        alert = Alert(**alert_config)
-            context = {'rule': self.rule, 'match': match}
+        response = api.create_alert(alert)
 
+        if response.status_code != 201:
+            raise Exception('alert not successfully created in TheHive\n{}'.format(response.text))
+
+    def alert(self, matches):
+        if self.rule.get('hive_alert_config_type', 'custom') != 'classic':
+            for match in matches:
+                alert_config = self.create_alert_config(match)
+                self.send_to_thehive(alert_config)
+        else:
+            alert_config = self.create_alert_config(matches[0])
             artifacts = []
-            for mapping in self.rule.get('hive_observable_data_mapping', []):
+            for match in matches:
-                for observable_type, match_data_key in mapping.items():
+                artifacts += self.create_artifacts(match)
-                    try:
+                if 'related_events' in match:
-                        match_data_keys = re.findall(r'\{match\[([^\]]*)\]', match_data_key)
+                    for related_event in match['related_events']:
-                        rule_data_keys = re.findall(r'\{rule\[([^\]]*)\]', match_data_key)
+                        artifacts += self.create_artifacts(related_event)
-                        data_keys = match_data_keys + rule_data_keys
-                        context_keys = list(context['match'].keys()) + list(context['rule'].keys())
-                        if all([True if k in context_keys else False for k in data_keys]):
-                            artifacts.append(AlertArtifact(dataType=observable_type, data=match_data_key.format(**context)))
-                    except KeyError:
-                        raise KeyError('\nformat string\n{}\nmatch data\n{}'.format(match_data_key, context))
 
-            alert_config = {
+            alert_config['artifacts'] = artifacts
-                'artifacts': artifacts,
+            alert_config['title'] = self.create_title(matches)
-                'sourceRef': str(uuid.uuid4())[0:6],
+            alert_config['description'] = self.create_alert_body(matches)
-                'title': '{rule[index]}_{rule[name]}'.format(**context)
+            self.send_to_thehive(alert_config)
-            }
-            alert_config.update(self.rule.get('hive_alert_config', {}))
-
-            for alert_config_field, alert_config_value in alert_config.items():
-                if alert_config_field == 'customFields':
-                    custom_fields = CustomFieldHelper()
-                    for cf_key, cf_value in alert_config_value.items():
-                        try:
-                            func = getattr(custom_fields, 'add_{}'.format(cf_value['type']))
-                        except AttributeError:
-                            raise Exception('unsupported custom field type {}'.format(cf_value['type']))
-                        value = cf_value['value'].format(**context)
-                        func(cf_key, value)
-                    alert_config[alert_config_field] = custom_fields.build()
-                elif isinstance(alert_config_value, str):
-                    alert_config[alert_config_field] = alert_config_value.format(**context)
-                elif isinstance(alert_config_value, (list, tuple)):
-                    formatted_list = []
-                    for element in alert_config_value:
-                        try:
-                            formatted_list.append(element.format(**context))
-                        except (AttributeError, KeyError, IndexError):
-                            formatted_list.append(element)
-                    alert_config[alert_config_field] = formatted_list
-
-            alert = Alert(**alert_config)
-            response = api.create_alert(alert)
-
-            if response.status_code != 201:
-                raise Exception('alert not successfully created in TheHive\n{}'.format(response.text))
 
     def get_info(self):
 
         return {
             'type': 'hivealerter',
             'hive_host': self.rule.get('hive_connection', {}).get('hive_host', '')
-       } 
+        }