Merge pull request #151 from MISP/refactor

chg: Make the library easier to use

pymisp/abstract.py

@@ -27,6 +27,9 @@ class AbstractMISP(collections.MutableMapping):
 
     __not_jsonable = []
 
+    def __init__(self, **kwargs):
+        super(AbstractMISP, self).__init__()
+
     def properties(self):
         to_return = []
         for prop, value in vars(self).items():
@@ -67,7 +70,11 @@ class AbstractMISP(collections.MutableMapping):
         return json.dumps(self, cls=MISPEncode)
 
     def __getitem__(self, key):
-        return getattr(self, key)
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            # Expected by pop and other dict-related methods
+            raise KeyError
 
     def __setitem__(self, key, value):
         setattr(self, key, value)

pymisp/mispevent.py

@@ -304,7 +304,6 @@ class MISPEvent(AbstractMISP):
             describe_types = t['result']
 
         self._types = describe_types['types']
-        self.attributes = []
         self.Tag = []
 
     def _reinitialize_event(self):
@@ -315,7 +314,6 @@ class MISPEvent(AbstractMISP):
         self.info = None
         self.published = False
         self.date = datetime.date.today()
-        self.attributes = []
 
         # All other keys
         self.sig = None
@@ -475,9 +473,9 @@ class MISPEvent(AbstractMISP):
             for a in kwargs.pop('Attribute'):
                 attribute = MISPAttribute()
                 attribute.set_all_values(**a)
-                if not hasattr(self, 'attributes'):
-                    self.attributes = []
-                self.attributes.append(attribute)
+                if not hasattr(self, 'Attribute'):
+                    self.Attribute = []
+                self.Attribute.append(attribute)
 
         # All other keys
         if kwargs.get('id'):
@@ -518,22 +516,24 @@ class MISPEvent(AbstractMISP):
         to_return = super(MISPEvent, self).to_dict()
         if to_return.get('date'):
             to_return['date'] = self.date.isoformat()
-        if to_return.get('attributes'):
-            attributes = to_return.pop('attributes')
-            to_return['Attribute'] = [attribute.to_dict(with_timestamp) for attribute in attributes]
-        if to_return.get('RelatedEvent'):
-            to_return['RelatedEvent'] = [rel_event.to_dict() for rel_event in self.RelatedEvent]
         if with_timestamp and to_return.get('timestamp'):
-            to_return['timestamp'] = int(time.mktime(self.timestamp.timetuple()))
+            if sys.version_info >= (3, 3):
+                to_return['timestamp'] = self.timestamp.timestamp()
+            else:
+                from datetime import timezone  # Only for Python < 3.3
+                to_return['timestamp'] = (self.timestamp - datetime(1970, 1, 1, tzinfo=timezone.utc)).total_seconds()
         else:
             to_return.pop('timestamp', None)
         if with_timestamp and to_return.get('publish_timestamp'):
-            to_return['publish_timestamp'] = int(time.mktime(self.publish_timestamp.timetuple()))
+            if sys.version_info >= (3, 3):
+                to_return['publish_timestamp'] = self.publish_timestamp.timestamp()
+            else:
+                from datetime import timezone  # Only for Python < 3.3
+                to_return['publish_timestamp'] = (self.publish_timestamp - datetime(1970, 1, 1, tzinfo=timezone.utc)).total_seconds()
         else:
             to_return.pop('publish_timestamp', None)
         to_return = _int_to_str(to_return)
         to_return = {'Event': to_return}
-        jsonschema.validate(to_return, self.__json_schema)
         return to_return
 
     def add_tag(self, tag):
@@ -580,6 +580,38 @@ class MISPEvent(AbstractMISP):
                 self.attributes = []
             self.attributes.append(attribute)
 
+    @property
+    def attributes(self):
+        return self.Attribute
+
+    @property
+    def related_events(self):
+        return self.RelatedEvent
+
+    @property
+    def objects(self):
+        return self.Object
+
+    @property
+    def tags(self):
+        return self.Tag
+
+    def get_object_by_id(self, object_id):
+        for obj in self.objects:
+            if hasattr(obj, 'id') and obj.id == object_id:
+                return obj
+        raise InvalidMISPObject('Object with {} does not exists in ths event'.format(object_id))
+
+    def add_object(self, obj):
+        if isinstance(obj, MISPObject):
+            self.Object.append(obj)
+        elif isinstance(obj, dict):
+            tmp_object = MISPObject(obj['name'])
+            tmp_object.from_dict(**obj)
+            self.Object.append(tmp_object)
+        else:
+            raise InvalidMISPObject("An object to add to an existing Event needs to be either a MISPObject, or a plain python dictionary")
+
 
 class MISPObjectReference(AbstractMISP):
 
@@ -644,8 +676,18 @@ class MISPObjectAttribute(MISPAttribute):
 
 class MISPObject(AbstractMISP):
 
-    def __init__(self, name, strict=False):
-        super(MISPObject, self).__init__()
+    def __init__(self, name, strict=False, standalone=False, default_attributes_paramaters={}, **kwargs):
+        ''' Master class representing a generic MISP object
+        :name: Name of the object
+
+        :strict: Enforce validation with the object templates
+
+        :standalone: The object will be pushed as directly on MISP, not as a part of an event.
+            In this case the ObjectReference needs to be pushed manually and cannot be in the JSON dump.
+
+        :default_attributes_paramaters: Used as template for the attributes if they are not overwritten in add_attribute
+        '''
+        super(MISPObject, self).__init__(**kwargs)
         self.__strict = strict
         self.name = name
         self.__misp_objects_path = os.path.join(
@@ -666,11 +708,30 @@ class MISPObject(AbstractMISP):
             self.description = self.__definition['description']
             self.template_version = self.__definition['version']
         else:
-            # FIXME We need to set something for meta-category, template_uuid, description and template_version
+            # Then we have no meta-category, template_uuid, description and template_version
             pass
         self.uuid = str(uuid.uuid4())
-        self.Attribute = []
+        self.__fast_attribute_access = {}  # Hashtable object_relation: [attributes]
+        self._default_attributes_paramaters = default_attributes_paramaters
+        if self._default_attributes_paramaters:
+            # Let's clean that up
+            self._default_attributes_paramaters.pop('value', None)  # duh
+            self._default_attributes_paramaters.pop('uuid', None)  # duh
+            self._default_attributes_paramaters.pop('id', None)  # duh
+            self._default_attributes_paramaters.pop('object_id', None)  # duh
+            self._default_attributes_paramaters.pop('type', None)  # depends on the value
+            self._default_attributes_paramaters.pop('object_relation', None)  # depends on the value
+            self._default_attributes_paramaters.pop('disable_correlation', None)  # depends on the value
+            self._default_attributes_paramaters.pop('to_ids', None)  # depends on the value
+            self._default_attributes_paramaters.pop('category', None)  # depends on the value
+            self._default_attributes_paramaters.pop('deleted', None)  # doesn't make sense to pre-set it
+            self._default_attributes_paramaters.pop('data', None)  # in case the original in a sample or an attachment
+            self.distribution = self._default_attributes_paramaters.distribution
         self.ObjectReference = []
+        self._standalone = standalone
+        if self._standalone:
+            # Mark as non_jsonable because we need to add the references manually after the object(s) have been created
+            self.update_not_jsonable('ObjectReference')
 
     def from_dict(self, **kwargs):
         if self.__known_template:
@@ -696,6 +757,8 @@ class MISPObject(AbstractMISP):
                 setattr(self, key, value)
 
     def to_dict(self, strict=False):
+        # Set the expected key (Attributes)
+        self.Attribute = self.attributes
         if strict or self.__strict and self.__known_template:
             self._validate()
         return super(MISPObject, self).to_dict()
@@ -708,7 +771,7 @@ class MISPObject(AbstractMISP):
     def _validate(self):
         """Make sure the object we're creating has the required fields"""
         all_object_relations = []
-        for a in self.Attribute:
+        for a in self.attributes:
             all_object_relations.append(a.object_relation)
         count_relations = dict(Counter(all_object_relations))
         for key, counter in count_relations.items():
@@ -739,6 +802,22 @@ class MISPObject(AbstractMISP):
                             relationship_type=relationship_type, comment=comment, **kwargs)
         self.ObjectReference.append(reference)
 
+    def get_attributes_by_relation(self, object_relation):
+        '''Returns the list of attributes with the given object relation in the object'''
+        return self.__fast_attribute_access.get(object_relation, [])
+
+    def has_attributes_by_relation(self, list_of_relations):
+        '''True if all the relations in the list are defined in the object'''
+        return all(relation in self.__fast_attribute_access for relation in list_of_relations)
+
+    @property
+    def attributes(self):
+        return [a for sublist in self.__fast_attribute_access.values() for a in sublist]
+
+    @property
+    def references(self):
+        return self.ObjectReference
+
     def add_attribute(self, object_relation, **value):
         if value.get('value') is None:
             return None
@@ -751,6 +830,9 @@ class MISPObject(AbstractMISP):
                 attribute = MISPObjectAttribute({})
         else:
             attribute = MISPObjectAttribute({})
-        attribute.from_dict(object_relation, **value)
-        self.Attribute.append(attribute)
+        # Overwrite the parameters of self._default_attributes_paramaters with the ones of value
+        attribute.from_dict(object_relation=object_relation, **dict(self._default_attributes_paramaters, **value))
+        if not self.__fast_attribute_access.get(object_relation):
+            self.__fast_attribute_access[object_relation] = []
+        self.__fast_attribute_access[object_relation].append(attribute)
         return attribute

pymisp/tools/create_misp_object.py

@@ -22,8 +22,8 @@ class FileTypeNotImplemented(MISPObjectException):
     pass
 
 
-def make_pe_objects(lief_parsed, misp_file):
-    pe_object = PEObject(parsed=lief_parsed)
+def make_pe_objects(lief_parsed, misp_file, standalone=True, default_attributes_paramaters={}):
+    pe_object = PEObject(parsed=lief_parsed, standalone=standalone, default_attributes_paramaters=default_attributes_paramaters)
     misp_file.add_reference(pe_object.uuid, 'included-in', 'PE indicators')
     pe_sections = []
     for s in pe_object.sections:
@@ -31,8 +31,8 @@ def make_pe_objects(lief_parsed, misp_file):
     return misp_file, pe_object, pe_sections
 
 
-def make_elf_objects(lief_parsed, misp_file):
-    elf_object = ELFObject(parsed=lief_parsed)
+def make_elf_objects(lief_parsed, misp_file, standalone=True, default_attributes_paramaters={}):
+    elf_object = ELFObject(parsed=lief_parsed, standalone=standalone, default_attributes_paramaters=default_attributes_paramaters)
     misp_file.add_reference(elf_object.uuid, 'included-in', 'ELF indicators')
     elf_sections = []
     for s in elf_object.sections:
@@ -40,8 +40,8 @@ def make_elf_objects(lief_parsed, misp_file):
     return misp_file, elf_object, elf_sections
 
 
-def make_macho_objects(lief_parsed, misp_file):
-    macho_object = MachOObject(parsed=lief_parsed)
+def make_macho_objects(lief_parsed, misp_file, standalone=True, default_attributes_paramaters={}):
+    macho_object = MachOObject(parsed=lief_parsed, standalone=standalone, default_attributes_paramaters=default_attributes_paramaters)
     misp_file.add_reference(macho_object.uuid, 'included-in', 'MachO indicators')
     macho_sections = []
     for s in macho_object.sections:
@@ -49,8 +49,9 @@ def make_macho_objects(lief_parsed, misp_file):
     return misp_file, macho_object, macho_sections
 
 
-def make_binary_objects(filepath=None, pseudofile=None, filename=None):
-    misp_file = FileObject(filepath=filepath, pseudofile=pseudofile, filename=filename)
+def make_binary_objects(filepath=None, pseudofile=None, filename=None, standalone=True, default_attributes_paramaters={}):
+    misp_file = FileObject(filepath=filepath, pseudofile=pseudofile, filename=filename,
+                           standalone=standalone, default_attributes_paramaters=default_attributes_paramaters)
     if HAS_LIEF and filepath or (pseudofile and filename):
         try:
             if filepath:
@@ -62,11 +63,11 @@ def make_binary_objects(filepath=None, pseudofile=None, filename=None):
                 else:
                     lief_parsed = lief.parse(raw=pseudofile.getvalue(), name=filename)
             if isinstance(lief_parsed, lief.PE.Binary):
-                return make_pe_objects(lief_parsed, misp_file)
+                return make_pe_objects(lief_parsed, misp_file, standalone, default_attributes_paramaters)
             elif isinstance(lief_parsed, lief.ELF.Binary):
-                return make_elf_objects(lief_parsed, misp_file)
+                return make_elf_objects(lief_parsed, misp_file, standalone, default_attributes_paramaters)
             elif isinstance(lief_parsed, lief.MachO.Binary):
-                return make_macho_objects(lief_parsed, misp_file)
+                return make_macho_objects(lief_parsed, misp_file, standalone, default_attributes_paramaters)
         except lief.bad_format as e:
             logger.warning('Bad format: {}'.format(e))
         except lief.bad_file as e:

pymisp/tools/elfobject.py

@@ -24,7 +24,7 @@ except ImportError:
 
 class ELFObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, parsed=None, filepath=None, pseudofile=None):
+    def __init__(self, parsed=None, filepath=None, pseudofile=None, standalone=True, **kwargs):
         if not HAS_PYDEEP:
             logger.warning("Please install pydeep: pip install git+https://github.com/kbandla/pydeep.git")
         if not HAS_LIEF:
@@ -44,10 +44,8 @@ class ELFObject(AbstractMISPObjectGenerator):
                 self.__elf = parsed
             else:
                 raise InvalidMISPObject('Not a lief.ELF.Binary: {}'.format(type(parsed)))
-        super(ELFObject, self).__init__('elf')
+        super(ELFObject, self).__init__('elf', standalone=standalone, **kwargs)
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable('ObjectReference')
 
     def generate_attributes(self):
         # General information
@@ -60,7 +58,7 @@ class ELFObject(AbstractMISPObjectGenerator):
         if self.__elf.sections:
             pos = 0
             for section in self.__elf.sections:
-                s = ELFSectionObject(section)
+                s = ELFSectionObject(section, self._standalone, default_attributes_paramaters=self._default_attributes_paramaters)
                 self.add_reference(s.uuid, 'included-in', 'Section {} of ELF'.format(pos))
                 pos += 1
                 self.sections.append(s)
@@ -69,15 +67,13 @@ class ELFObject(AbstractMISPObjectGenerator):
 
 class ELFSectionObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, section):
+    def __init__(self, section, standalone=True, **kwargs):
         # Python3 way
         # super().__init__('pe-section')
-        super(ELFSectionObject, self).__init__('elf-section')
+        super(ELFSectionObject, self).__init__('elf-section', standalone=standalone, **kwargs)
         self.__section = section
         self.__data = bytes(self.__section.content)
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable('ObjectReference')
 
     def generate_attributes(self):
         self.add_attribute('name', value=self.__section.name)

pymisp/tools/fileobject.py

@@ -28,7 +28,7 @@ except ImportError:
 
 class FileObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, filepath=None, pseudofile=None, filename=None):
+    def __init__(self, filepath=None, pseudofile=None, filename=None, standalone=True, **kwargs):
         if not HAS_PYDEEP:
             logger.warning("Please install pydeep: pip install git+https://github.com/kbandla/pydeep.git")
         if not HAS_MAGIC:
@@ -51,11 +51,9 @@ class FileObject(AbstractMISPObjectGenerator):
             raise InvalidMISPObject('File buffer (BytesIO) or a path is required.')
         # PY3 way:
         # super().__init__('file')
-        super(FileObject, self).__init__('file')
+        super(FileObject, self).__init__('file', standalone=standalone, **kwargs)
         self.__data = self.__pseudofile.getvalue()
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable('ObjectReference')
 
     def generate_attributes(self):
         self.add_attribute('filename', value=self.__filename)

pymisp/tools/machoobject.py

@@ -25,7 +25,7 @@ except ImportError:
 
 class MachOObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, parsed=None, filepath=None, pseudofile=None):
+    def __init__(self, parsed=None, filepath=None, pseudofile=None, standalone=True, **kwargs):
         if not HAS_PYDEEP:
             logger.warning("Please install pydeep: pip install git+https://github.com/kbandla/pydeep.git")
         if not HAS_LIEF:
@@ -47,10 +47,8 @@ class MachOObject(AbstractMISPObjectGenerator):
                 raise InvalidMISPObject('Not a lief.MachO.Binary: {}'.format(type(parsed)))
         # Python3 way
         # super().__init__('elf')
-        super(MachOObject, self).__init__('macho')
+        super(MachOObject, self).__init__('macho', standalone=standalone, **kwargs)
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable(['ObjectReference'])
 
     def generate_attributes(self):
         self.add_attribute('type', value=str(self.__macho.header.file_type).split('.')[1])
@@ -63,7 +61,7 @@ class MachOObject(AbstractMISPObjectGenerator):
         if self.__macho.sections:
             pos = 0
             for section in self.__macho.sections:
-                s = MachOSectionObject(section)
+                s = MachOSectionObject(section, self._standalone, default_attributes_paramaters=self._default_attributes_paramaters)
                 self.add_reference(s.uuid, 'included-in', 'Section {} of MachO'.format(pos))
                 pos += 1
                 self.sections.append(s)
@@ -72,15 +70,13 @@ class MachOObject(AbstractMISPObjectGenerator):
 
 class MachOSectionObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, section):
+    def __init__(self, section, standalone=True, **kwargs):
         # Python3 way
         # super().__init__('pe-section')
-        super(MachOSectionObject, self).__init__('macho-section')
+        super(MachOSectionObject, self).__init__('macho-section', standalone=standalone, **kwargs)
         self.__section = section
         self.__data = bytes(self.__section.content)
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable(['ObjectReference'])
 
     def generate_attributes(self):
         self.add_attribute('name', value=self.__section.name)

pymisp/tools/peobject.py

@@ -25,7 +25,7 @@ except ImportError:
 
 class PEObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, parsed=None, filepath=None, pseudofile=None):
+    def __init__(self, parsed=None, filepath=None, pseudofile=None, standalone=True, **kwargs):
         if not HAS_PYDEEP:
             logger.warning("Please install pydeep: pip install git+https://github.com/kbandla/pydeep.git")
         if not HAS_LIEF:
@@ -47,10 +47,8 @@ class PEObject(AbstractMISPObjectGenerator):
                 raise InvalidMISPObject('Not a lief.PE.Binary: {}'.format(type(parsed)))
         # Python3 way
         # super().__init__('pe')
-        super(PEObject, self).__init__('pe')
+        super(PEObject, self).__init__('pe', standalone=standalone, **kwargs)
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable('ObjectReference')
 
     def _is_exe(self):
         if not self._is_dll() and not self._is_driver():
@@ -106,7 +104,7 @@ class PEObject(AbstractMISPObjectGenerator):
         if self.__pe.sections:
             pos = 0
             for section in self.__pe.sections:
-                s = PESectionObject(section)
+                s = PESectionObject(section, self._standalone, default_attributes_paramaters=self._default_attributes_paramaters)
                 self.add_reference(s.uuid, 'included-in', 'Section {} of PE'.format(pos))
                 if ((self.__pe.entrypoint >= section.virtual_address) and
                         (self.__pe.entrypoint < (section.virtual_address + section.virtual_size))):
@@ -119,15 +117,13 @@ class PEObject(AbstractMISPObjectGenerator):
 
 class PESectionObject(AbstractMISPObjectGenerator):
 
-    def __init__(self, section):
+    def __init__(self, section, standalone=True, **kwargs):
         # Python3 way
         # super().__init__('pe-section')
-        super(PESectionObject, self).__init__('pe-section')
+        super(PESectionObject, self).__init__('pe-section', standalone=standalone, **kwargs)
         self.__section = section
         self.__data = bytes(self.__section.content)
         self.generate_attributes()
-        # Mark as non_jsonable because we need to add them manually
-        self.update_not_jsonable('ObjectReference')
 
     def generate_attributes(self):
         self.add_attribute('name', value=self.__section.name)

pymisp/tools/vtreportobject.py

@@ -23,10 +23,10 @@ class VTReportObject(AbstractMISPObjectGenerator):
 
     :indicator: IOC to search VirusTotal for
     '''
-    def __init__(self, apikey, indicator, vt_proxies=None):
+    def __init__(self, apikey, indicator, vt_proxies=None, standalone=True, **kwargs):
         # PY3 way:
         # super().__init__("virustotal-report")
-        super(VTReportObject, self).__init__("virustotal-report")
+        super(VTReportObject, self).__init__("virustotal-report", standalone=standalone, **kwargs)
         indicator = indicator.strip()
         self._resource_type = self.__validate_resource(indicator)
         if self._resource_type:
@@ -36,8 +36,6 @@ class VTReportObject(AbstractMISPObjectGenerator):
         else:
             error_msg = "A valid indicator is required. (One of type url, md5, sha1, sha256). Received '{}' instead".format(indicator)
             raise InvalidMISPObject(error_msg)
-        # Mark as non_jsonable because we need to add the references manually after the object(s) have been created
-        self.update_not_jsonable('ObjectReference')
 
     def get_report(self):
         return self._report