feat: Adds methods to get attribute by id/uuid

pull/1243/head
Sura De Silva 2024-06-24 22:37:05 +10:00
parent 0e21dadd0f
commit 7c8df85a98
2 changed files with 54 additions and 7 deletions

View File

@ -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):

View File

@ -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