diff --git a/examples/feed-generator-from-redis/generator.py b/examples/feed-generator-from-redis/generator.py index d0291f8..db383f9 100755 --- a/examples/feed-generator-from-redis/generator.py +++ b/examples/feed-generator-from-redis/generator.py @@ -6,9 +6,8 @@ import json import os import sys import time -import uuid -from pymisp import MISPEvent +from pymisp import MISPEvent, MISPOrganisation import settings @@ -35,11 +34,6 @@ def get_system_templates(): return templates -def gen_uuid(): - """Generate a random UUID and returns its string representation""" - return str(uuid.uuid4()) - - class FeedGenerator: """Helper object to create MISP feed. @@ -164,7 +158,7 @@ class FeedGenerator: self.create_daily_event() def flush_event(self, new_event=None): - print('Writting event on disk'+' '*50) + print('Writting event on disk' + ' ' * 50) if new_event is not None: event_uuid = new_event['uuid'] event = new_event @@ -172,9 +166,8 @@ class FeedGenerator: event_uuid = self.current_event_uuid event = self.current_event - eventFile = open(os.path.join(settings.outputdir, event_uuid+'.json'), 'w') - eventFile.write(event.to_json()) - eventFile.close() + with open(os.path.join(settings.outputdir, event_uuid + '.json'), 'w') as eventFile: + json.dump(event.to_feed(), eventFile) self.save_hashes() @@ -197,27 +190,11 @@ class FeedGenerator: hashFile.write('{},{}\n'.format(element[0], element[1])) hashFile.close() self.attributeHashes = [] - print('Hash saved' + ' '*30) + print('Hash saved' + ' ' * 30) except Exception as e: print(e) sys.exit('Could not create the quick hash lookup file.') - def _addEventToManifest(self, event): - event_dict = event.to_dict()['Event'] - tags = [] - for eventTag in event_dict.get('EventTag', []): - tags.append({'name': eventTag['Tag']['name'], - 'colour': eventTag['Tag']['colour']}) - return { - 'Orgc': event_dict.get('Orgc', []), - 'Tag': tags, - 'info': event_dict['info'], - 'date': event_dict['date'], - 'analysis': event_dict['analysis'], - 'threat_level_id': event_dict['threat_level_id'], - 'timestamp': event_dict.get('timestamp', int(time.time())) - } - def get_last_event_from_manifest(self): """Retreive last event from the manifest. @@ -240,7 +217,7 @@ class FeedGenerator: # Sort by date then by event name dated_events.sort(key=lambda k: (k[0], k[2]), reverse=True) return dated_events[0] - except FileNotFoundError as e: + except FileNotFoundError: print('Manifest not found, generating a fresh one') self._init_manifest() return self.get_last_event_from_manifest() @@ -263,11 +240,9 @@ class FeedGenerator: return event def create_daily_event(self): - new_uuid = gen_uuid() today = str(datetime.date.today()) event_dict = { - 'uuid': new_uuid, - 'id': len(self.manifest)+1, + 'id': len(self.manifest) + 1, 'Tag': settings.Tag, 'info': self.daily_event_name.format(today), 'analysis': settings.analysis, # [0-2] @@ -279,14 +254,14 @@ class FeedGenerator: event.from_dict(**event_dict) # reference org - org_dict = {} - org_dict['name'] = settings.org_name - org_dict['uuid'] = settings.org_uuid - event['Orgc'] = org_dict + org = MISPOrganisation() + org.name = settings.org_name + org.uuid = settings.org_uuid + event.Orgc = org # save event on disk self.flush_event(new_event=event) # add event to manifest - self.manifest[event['uuid']] = self._addEventToManifest(event) + self.manifest.update(event.manifest) self.save_manifest() return event