Merge pull request #1243 from dragsu/feat-attribute-by-id

feat: Adds methods to get attribute by id/uuid

pymisp/exceptions.py

@@ -67,18 +67,26 @@ class NoKey(PyMISPError):
     pass
 
 
-class MISPObjectException(PyMISPError):
-    pass
+class MISPAttributeException(PyMISPError):
+    """A base class for attribute specific exceptions"""
 
+class MISPObjectException(PyMISPError):
+    """A base class for object specific exceptions"""
+
+
+class InvalidMISPAttribute(MISPAttributeException):
+    """Exception raised when an attribute doesn't respect the constraints in the definition"""
+
+class InvalidMISPObjectAttribute(MISPAttributeException):
+    """Exception raised when an object attribute doesn't respect the constraints in the definition"""
 
 class InvalidMISPObject(MISPObjectException):
-    """Exception raised when an object doesn't respect the contrains in the definition"""
-    pass
+    """Exception raised when an object doesn't respect the constraints in the definition"""
 
 
 class UnknownMISPObjectTemplate(MISPObjectException):
     """Exception raised when the template is unknown"""
-    pass
+
 
 
 class InvalidMISPGalaxy(PyMISPError):

pymisp/mispevent.py

@@ -23,8 +23,8 @@ except ImportError:
     import json
 
 from .abstract import AbstractMISP, MISPTag
-from .exceptions import (NewNoteError, NewOpinionError, NewRelationshipError, UnknownMISPObjectTemplate, InvalidMISPGalaxy, InvalidMISPObject,
-                         PyMISPError, NewEventError, NewAttributeError, NewEventReportError,
+from .exceptions import (NewNoteError, NewOpinionError, NewRelationshipError, UnknownMISPObjectTemplate, InvalidMISPGalaxy, InvalidMISPAttribute,
+                         InvalidMISPObject, InvalidMISPObjectAttribute, PyMISPError, NewEventError, NewAttributeError, NewEventReportError,
                          NewGalaxyClusterError, NewGalaxyClusterRelationError, NewAnalystDataError)
 
 logger = logging.getLogger('pymisp')
@@ -1027,6 +1027,26 @@ class MISPObject(AnalystDataBehaviorMixin):
         self.edited = True
         return reference
 
+    def get_attribute_by_id(self, attribute_id: str | int) -> MISPObjectAttribute:
+        """Get an object attribute by ID
+
+        :param attribute_id: The ID of the seeking object attribute"""
+        for attribute in self.attributes:
+            if hasattr(attribute, 'id') and attribute.id == attribute_id:
+                return attribute
+
+        raise InvalidMISPObjectAttribute(f'Object attribute with {attribute_id} does not exist in this event')
+
+    def get_attribute_by_uuid(self, attribute_uuid: str) -> MISPObjectAttribute:
+        """Get an object attribute by UUID
+
+        :param attribute_uuid: The UUID of the seeking object attribute"""
+        for attribute in self.attributes:
+            if hasattr(attribute, 'uuid') and attribute.uuid == attribute_uuid:
+                return attribute
+
+        raise InvalidMISPObjectAttribute(f'Object attribute with {attribute_uuid} does not exist in this event')
+
     def get_attributes_by_relation(self, object_relation: str) -> list[MISPAttribute]:
         '''Returns the list of attributes with the given object relation in the object'''
         return self._fast_attribute_access.get(object_relation, [])
@@ -2039,6 +2059,25 @@ class MISPEvent(AnalystDataBehaviorMixin):
         self.galaxies.append(misp_galaxy)
         return misp_galaxy
 
+    def get_attribute_by_id(self, attribute_id: str | int) -> MISPAttribute:
+        """Get an attribute by ID
+
+        :param attribute_id: The ID of the seeking attribute"""
+        for attribute in self.attributes:
+            if hasattr(attribute, 'id') and int(attribute.id) == int(attribute_id):
+                return attribute
+        raise InvalidMISPAttribute(f'Attribute with {attribute_id} does not exist in this event')
+
+    def get_attribute_by_uuid(self, attribute_uuid: str) -> MISPAttribute:
+        """Get an attribute by UUID
+
+        :param attribute_uuid: The UUID of the seeking attribute"""
+        for attribute in self.attributes:
+            if hasattr(attribute, 'uuid') and attribute.uuid == attribute_uuid:
+                return attribute
+
+        raise InvalidMISPAttribute(f'Attribute with {attribute_uuid} does not exist in this event')
+
     def get_object_by_id(self, object_id: str | int) -> MISPObject:
         """Get an object by ID