Skip to content

Commit

Permalink
Merge pull request #1243 from dragsu/feat-attribute-by-id
Browse files Browse the repository at this point in the history
feat: Adds methods to get attribute by id/uuid
  • Loading branch information
Rafiot authored Jun 25, 2024
2 parents 0e21dad + 7c8df85 commit 8c59415
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 6 deletions.
16 changes: 12 additions & 4 deletions pymisp/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,18 +67,26 @@ class NoKey(PyMISPError):
pass


class MISPAttributeException(PyMISPError):
"""A base class for attribute specific exceptions"""

class MISPObjectException(PyMISPError):
pass
"""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):
Expand Down
43 changes: 41 additions & 2 deletions pymisp/mispevent.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@
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')
Expand Down Expand Up @@ -1027,6 +1027,26 @@ def add_reference(self, referenced_uuid: AbstractMISP | str, relationship_type:
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, [])
Expand Down Expand Up @@ -2039,6 +2059,25 @@ def add_galaxy(self, galaxy: MISPGalaxy | dict[str, Any] | None = None, **kwargs
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
Expand Down

0 comments on commit 8c59415

Please sign in to comment.