From 341a569de54c56ddb927bc80a67f0319824db6f4 Mon Sep 17 00:00:00 2001 From: Jesse Hedden Date: Sun, 21 Jun 2020 19:52:17 -0700 Subject: [PATCH] ready for code review --- .../modules/expansion/trustar_enrich.py | 113 ++++++++++++------ 1 file changed, 76 insertions(+), 37 deletions(-) diff --git a/misp_modules/modules/expansion/trustar_enrich.py b/misp_modules/modules/expansion/trustar_enrich.py index e786ff3..38f5d16 100644 --- a/misp_modules/modules/expansion/trustar_enrich.py +++ b/misp_modules/modules/expansion/trustar_enrich.py @@ -3,61 +3,100 @@ from pymisp import MISPAttribute, MISPEvent, MISPObject from trustar import TruStar misperrors = {'error': "Error"} -mispattributes = {'input': ["btc", "domain","email-src", "filename", "hostname", "ip-src", "ip-dst", "malware-type", "md5", "sha1", "sha256", "url"], 'format': 'misp_standard'} +mispattributes = { + 'input': ["btc", "domain", "email-src", "filename", "hostname", "ip-src", "ip-dst", "malware-type", "md5", "sha1", + "sha256", "url"], 'format': 'misp_standard'} moduleinfo = {'version': "0.1", 'author': "Jesse Hedden", 'description': "Enrich data with TruSTAR", 'module-type': ["hover", "expansion"]} -moduleconfig = ["api_key", "api_secret", "enclave_ids"] +moduleconfig = ["user_api_key", "user_api_secret", "enclave_ids"] -def get_results(misp_event): - event = json.loads(misp_event.to_json()) - results = {key: event[key] for key in ('Attribute', 'Object')} - return {'results': results} +class TruSTARParser: + ENTITY_TYPE_MAPPINGS = { + 'BITCOIN_ADDRESS': "btc", + 'CIDR_BLOCK': "ip-src", + 'CVE': "vulnerability", + 'URL': "url", + 'EMAIL_ADDRESS': "email-src", + 'SOFTWARE': "filename", + 'IP': "ip-src", + 'MALWARE': "malware-type", + 'MD5': "md5", + 'REGISTRY_KEY': "regkey", + 'SHA1': "sha1", + 'SHA256': "sha256" + } -def parse_indicator_summary(attribute, summary): - misp_event = MISPEvent() - misp_attribute = MISPAttribute().from_dict(**attribute) - misp_event.add_attribute(**misp_attribute) + REPORT_BASE_URL = "https://station.trustar.co/constellation/reports/{}" - mapping = {'value': 'text', 'reportId': 'text', 'enclaveId': 'text', 'description': 'text'} + def __init__(self, attribute, config): + config['enclave_ids'] = config.get('enclave_ids', "").split(',') + self.ts_client = TruStar(config=config) - for item in summary.get('items'): - trustar_obj = MISPObject(attribute.value) - for key, attribute_type in mapping.items(): - trustar_obj.add_attribute(key, attribute_type=attribute_type, value=item[key]) - trustar_obj.add_reference(misp_attribute.uuid, 'associated-to') - misp_event.add_object(**trustar_obj) + self.misp_event = MISPEvent() + self.misp_attribute = MISPAttribute() + self.misp_attribute.from_dict(**attribute) + self.misp_event.add_attribute(**self.misp_attribute) - return misp_event + def get_results(self): + event = json.loads(self.misp_event.to_json()) + results = {key: event[key] for key in ('Attribute', 'Object') if (key in event and event[key])} + return {'results': results} + def generate_trustar_links(self, entity_value): + """ + Generates links to TruSTAR reports if they exist. -def handler(q=False): + :param entity_value: Value of entity. + """ + report_links = list() + trustar_reports = self.ts_client.search_reports(entity_value) + for report in trustar_reports: + report_links.append(self.REPORT_BASE_URL.format(report.id)) - if q is False: - return False + return report_links - request = json.loads(q) - config = request.get('config', {}) - if not config.get('api_key') or not config.get('api_secret'): - misperrors['error'] = "Your TruSTAR API key and secret are required for indicator enrichment." - return misperrors + def parse_indicator_summary(self, attribute, summaries): - enclave_ids = [enclave_id for enclave_id in config.get('enclave_ids', "").split(',')] - ts_client = TruStar(config={'user_api_key': config.get('api_key'), 'user_api_secret': config.get('api_secret'), 'enclave_ids': enclave_ids}) - attribute = request.get('attribute') + for summary in summaries: + trustar_obj = MISPObject('trustar_report') + summary_dict = summary.to_dict() + summary_type = summary_dict.get('type') + summary_value = summary_dict.get('value') + if summary_type in self.ENTITY_TYPE_MAPPINGS: + trustar_obj.add_attribute(summary_type, attribute_type=self.ENTITY_TYPE_MAPPINGS[summary_type], + value=summary_value) + trustar_obj.add_attribute("INDICATOR_SUMMARY", attribute_type="text", + value=json.dumps(summary_dict, sort_keys=True, indent=4)) + report_links = self.generate_trustar_links(summary_value) + for link in report_links: + trustar_obj.add_attribute("REPORT_LINK", attribute_type="link", value=link) + self.misp_event.add_object(**trustar_obj) - summary = ts_client.get_indicator_summaries(attribute) + def handler(q=False): - misp_event = parse_indicator_summary(attribute, summary) - return get_results(misp_event) + if q is False: + return False -def introspection(): - return mispattributes + request = json.loads(q) -def version(): - moduleinfo['config'] = moduleconfig - return moduleinfo + config = request.get('config', {}) + if not config.get('user_api_key') or not config.get('user_api_secret'): + misperrors['error'] = "Your TruSTAR API key and secret are required for indicator enrichment." + return misperrors + attribute = request['attribute'] + trustar_parser = TruSTARParser(attribute, config) + summaries = trustar_parser.ts_client.get_indicator_summaries([attribute['value']]) + trustar_parser.parse_indicator_summary(attribute, summaries) + return trustar_parser.get_results() + + def introspection(): + return mispattributes + + def version(): + moduleinfo['config'] = moduleconfig + return moduleinfo