Skip to content

Commit

Permalink
Rewrite part of thehivealerter
Browse files Browse the repository at this point in the history
  • Loading branch information
Florian GAULTIER committed Oct 21, 2019
1 parent 325f1df commit 70b8a4e
Showing 1 changed file with 69 additions and 49 deletions.
118 changes: 69 additions & 49 deletions elastalert/alerts.py
Original file line number Diff line number Diff line change
Expand Up @@ -2118,8 +2118,57 @@ class HiveAlerter(Alerter):

required_options = set(['hive_connection', 'hive_alert_config'])

def alert(self, matches):
def get_aggregation_summary_text(self, matches):
text = super(HiveAlerter, self).get_aggregation_summary_text(matches)
if text:
text = u'```\n{0}```\n'.format(text)
return text

def create_artifacts(self, match):
artifacts = []
for mapping in self.rule.get('hive_observable_data_mapping', []):
for observable_type, match_data_key in mapping.iteritems():
try:
artifacts.append(AlertArtifact(dataType=observable_type, data=match_data_key.format(**{'rule': self.rule, 'match': match})))
except KeyError:
raise KeyError('\nformat string\n{}\nmatch data\n{}'.format(match_data_key, {'rule': self.rule, 'match': match}))
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[index]}_{rule[name]}'.format(**context)
}

alert_config.update(self.rule.get('hive_alert_config', {}))

for alert_config_field, alert_config_value in alert_config.iteritems():
if alert_config_field == 'customFields':
custom_fields = CustomFieldHelper()
for cf_key, cf_value in alert_config_value.iteritems():
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, basestring):
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(
Expand All @@ -2128,56 +2177,27 @@ def alert(self, matches):
proxies=connection_details.get('hive_proxies', {'http': '', 'https': ''}),
cert=connection_details.get('hive_verify', False))

for match in matches:
context = {'rule': self.rule, 'match': match}
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 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 observable_type, match_data_key in mapping.items():
try:
match_data_keys = re.findall(r'\{match\[([^\]]*)\]', match_data_key)
rule_data_keys = re.findall(r'\{rule\[([^\]]*)\]', match_data_key)
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 = {
'artifacts': artifacts,
'sourceRef': str(uuid.uuid4())[0:6],
'title': '{rule[index]}_{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

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))
for match in matches:
artifacts += self.create_artifacts(match)

alert_config['artifacts'] = artifacts
alert_config['title'] = self.create_title(matches)
alert_config['description'] = self.create_alert_body(matches)
self.send_to_thehive(alert_config)

def get_info(self):

Expand Down

0 comments on commit 70b8a4e

Please sign in to comment.