fix: Initial round of cleanup on redis feed generator

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