chg: Add typing markup

pull/511/head
Raphaël Vinot 2020-01-02 15:55:00 +01:00
parent 6427ce3c84
commit 2e064563c3
2 changed files with 118 additions and 118 deletions

View File

@ -1,27 +1,27 @@
#!/usr/bin/env python #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import datetime import datetime
from deprecated import deprecated from deprecated import deprecated # type: ignore
from json import JSONEncoder from json import JSONEncoder
from uuid import UUID from uuid import UUID
from abc import ABCMeta
try: try:
from rapidjson import load from rapidjson import load # type: ignore
from rapidjson import loads from rapidjson import loads # type: ignore
from rapidjson import dumps from rapidjson import dumps # type: ignore
import rapidjson
HAS_RAPIDJSON = True HAS_RAPIDJSON = True
except ImportError: except ImportError:
from json import load from json import load
from json import loads from json import loads
from json import dumps from json import dumps
import json
HAS_RAPIDJSON = False HAS_RAPIDJSON = False
import logging import logging
from enum import Enum from enum import Enum
from typing import Union, Optional
from .exceptions import PyMISPInvalidFormat, PyMISPError from .exceptions import PyMISPInvalidFormat, PyMISPError
@ -43,7 +43,7 @@ class MISPFileCache(object):
@staticmethod @staticmethod
@lru_cache(maxsize=150) @lru_cache(maxsize=150)
def _load_json(path): def _load_json(path: Path) -> Union[dict, None]:
if not path.exists(): if not path.exists():
return None return None
with path.open('r') as f: with path.open('r') as f:
@ -73,7 +73,7 @@ class Analysis(Enum):
completed = 2 completed = 2
def _int_to_str(d): def _int_to_str(d: dict) -> dict:
# transform all integer back to string # transform all integer back to string
for k, v in d.items(): for k, v in d.items():
if isinstance(v, (int, float)) and not isinstance(v, bool): if isinstance(v, (int, float)) and not isinstance(v, bool):
@ -95,31 +95,7 @@ class MISPEncode(JSONEncoder):
return JSONEncoder.default(self, obj) return JSONEncoder.default(self, obj)
if HAS_RAPIDJSON: class AbstractMISP(MutableMapping, MISPFileCache, metaclass=ABCMeta):
def pymisp_json_default(obj):
if isinstance(obj, AbstractMISP):
return obj.jsonable()
elif isinstance(obj, (datetime.datetime, datetime.date)):
return obj.isoformat()
elif isinstance(obj, Enum):
return obj.value
elif isinstance(obj, UUID):
return str(obj)
return rapidjson.default(obj)
else:
def pymisp_json_default(obj):
if isinstance(obj, AbstractMISP):
return obj.jsonable()
elif isinstance(obj, (datetime.datetime, datetime.date)):
return obj.isoformat()
elif isinstance(obj, Enum):
return obj.value
elif isinstance(obj, UUID):
return str(obj)
return json.default(obj)
class AbstractMISP(MutableMapping, MISPFileCache):
__resources_path = resources_path __resources_path = resources_path
__misp_objects_path = misp_objects_path __misp_objects_path = misp_objects_path
__describe_types = describe_types __describe_types = describe_types
@ -132,15 +108,16 @@ class AbstractMISP(MutableMapping, MISPFileCache):
methods in ExpandedPyMISP/PyMISP. methods in ExpandedPyMISP/PyMISP.
""" """
super().__init__() super().__init__()
self.__edited = True # As we create a new object, we assume it is edited self.__edited: bool = True # As we create a new object, we assume it is edited
self.__not_jsonable = [] self.__not_jsonable: list = []
self.__self_defined_describe_types = None self._fields_for_feed: set = {}
self.__self_defined_describe_types: Union[dict, None] = None
if kwargs.get('force_timestamps') is not None: if kwargs.get('force_timestamps') is not None:
# Ignore the edited objects and keep the timestamps. # Ignore the edited objects and keep the timestamps.
self.__force_timestamps = True self.__force_timestamps: bool = True
else: else:
self.__force_timestamps = False self.__force_timestamps: bool = False
# List of classes having tags # List of classes having tags
from .mispevent import MISPAttribute, MISPEvent from .mispevent import MISPAttribute, MISPEvent
@ -151,30 +128,30 @@ class AbstractMISP(MutableMapping, MISPFileCache):
setattr(AbstractMISP, 'tags', property(AbstractMISP.__get_tags, AbstractMISP.__set_tags)) setattr(AbstractMISP, 'tags', property(AbstractMISP.__get_tags, AbstractMISP.__set_tags))
@property @property
def describe_types(self): def describe_types(self) -> dict:
if self.__self_defined_describe_types: if self.__self_defined_describe_types:
return self.__self_defined_describe_types return self.__self_defined_describe_types
return self.__describe_types return self.__describe_types
@describe_types.setter @describe_types.setter
def describe_types(self, describe_types): def describe_types(self, describe_types: dict):
self.__self_defined_describe_types = describe_types self.__self_defined_describe_types = describe_types
@property @property
def resources_path(self): def resources_path(self) -> Path:
return self.__resources_path return self.__resources_path
@property @property
def misp_objects_path(self): def misp_objects_path(self) -> Path:
return self.__misp_objects_path return self.__misp_objects_path
@misp_objects_path.setter @misp_objects_path.setter
def misp_objects_path(self, misp_objects_path): def misp_objects_path(self, misp_objects_path: Union[str, Path]):
if isinstance(misp_objects_path, str): if isinstance(misp_objects_path, str):
misp_objects_path = Path(misp_objects_path) misp_objects_path = Path(misp_objects_path)
self.__misp_objects_path = misp_objects_path self.__misp_objects_path = misp_objects_path
def from_dict(self, **kwargs): def from_dict(self, **kwargs) -> None:
"""Loading all the parameters as class properties, if they aren't `None`. """Loading all the parameters as class properties, if they aren't `None`.
This method aims to be called when all the properties requiring a special This method aims to be called when all the properties requiring a special
treatment are processed. treatment are processed.
@ -187,19 +164,19 @@ class AbstractMISP(MutableMapping, MISPFileCache):
# We load an existing dictionary, marking it an not-edited # We load an existing dictionary, marking it an not-edited
self.__edited = False self.__edited = False
def update_not_jsonable(self, *args): def update_not_jsonable(self, *args) -> None:
"""Add entries to the __not_jsonable list""" """Add entries to the __not_jsonable list"""
self.__not_jsonable += args self.__not_jsonable += args
def set_not_jsonable(self, *args): def set_not_jsonable(self, args: list) -> None:
"""Set __not_jsonable to a new list""" """Set __not_jsonable to a new list"""
self.__not_jsonable = args self.__not_jsonable = args
def from_json(self, json_string): def from_json(self, json_string: str) -> None:
"""Load a JSON string""" """Load a JSON string"""
self.from_dict(**loads(json_string)) self.from_dict(**loads(json_string))
def to_dict(self): def to_dict(self) -> dict:
"""Dump the class to a dictionary. """Dump the class to a dictionary.
This method automatically removes the timestamp recursively in every object This method automatically removes the timestamp recursively in every object
that has been edited is order to let MISP update the event accordingly.""" that has been edited is order to let MISP update the event accordingly."""
@ -223,15 +200,15 @@ class AbstractMISP(MutableMapping, MISPFileCache):
to_return = _int_to_str(to_return) to_return = _int_to_str(to_return)
return to_return return to_return
def jsonable(self): def jsonable(self) -> dict:
"""This method is used by the JSON encoder""" """This method is used by the JSON encoder"""
return self.to_dict() return self.to_dict()
def _to_feed(self): def _to_feed(self) -> dict:
if not hasattr(self, '_fields_for_feed'): if not hasattr(self, '_fields_for_feed') or not self._fields_for_feed:
raise PyMISPError('Unable to export in the feed format, _fields_for_feed is missing.') raise PyMISPError('Unable to export in the feed format, _fields_for_feed is missing.')
if hasattr(self, '_set_default') and callable(self._set_default): if hasattr(self, '_set_default') and callable(self._set_default): # type: ignore
self._set_default() self._set_default() # type: ignore
to_return = {} to_return = {}
for field in self._fields_for_feed: for field in self._fields_for_feed:
if getattr(self, field, None) is not None: if getattr(self, field, None) is not None:
@ -248,7 +225,7 @@ class AbstractMISP(MutableMapping, MISPFileCache):
raise PyMISPError('The field {} is required in {} when generating a feed.'.format(field, self.__class__.__name__)) raise PyMISPError('The field {} is required in {} when generating a feed.'.format(field, self.__class__.__name__))
return to_return return to_return
def to_json(self, sort_keys=False, indent=None): def to_json(self, sort_keys: bool=False, indent: Optional[int]=None):
"""Dump recursively any class of type MISPAbstract to a json string""" """Dump recursively any class of type MISPAbstract to a json string"""
return dumps(self, default=pymisp_json_default, sort_keys=sort_keys, indent=indent) return dumps(self, default=pymisp_json_default, sort_keys=sort_keys, indent=indent)
@ -274,7 +251,7 @@ class AbstractMISP(MutableMapping, MISPFileCache):
return len([k for k in self.__dict__.keys() if not (k[0] == '_' or k in self.__not_jsonable)]) return len([k for k in self.__dict__.keys() if not (k[0] == '_' or k in self.__not_jsonable)])
@property @property
def edited(self): def edited(self) -> bool:
"""Recursively check if an object has been edited and update the flag accordingly """Recursively check if an object has been edited and update the flag accordingly
to the parent objects""" to the parent objects"""
if self.__edited: if self.__edited:
@ -304,7 +281,7 @@ class AbstractMISP(MutableMapping, MISPFileCache):
self.__edited = True self.__edited = True
super().__setattr__(name, value) super().__setattr__(name, value)
def _datetime_to_timestamp(self, d): def _datetime_to_timestamp(self, d: Union[int, float, str, datetime.datetime]) -> int:
"""Convert a datetime.datetime object to a timestamp (int)""" """Convert a datetime.datetime object to a timestamp (int)"""
if isinstance(d, (int, float, str)): if isinstance(d, (int, float, str)):
# Assume we already have a timestamp # Assume we already have a timestamp
@ -325,10 +302,11 @@ class AbstractMISP(MutableMapping, MISPFileCache):
misp_tag = MISPTag() misp_tag = MISPTag()
misp_tag.from_dict(**kwargs) misp_tag.from_dict(**kwargs)
else: else:
raise PyMISPInvalidFormat("The tag is in an invalid format (can be either string, MISPTag, or an expanded dict): {}".format(tag)) raise PyMISPInvalidFormat(f"The tag is in an invalid format (can be either string, MISPTag, or an expanded dict): {tag}")
if misp_tag not in self.tags: if misp_tag not in self.tags:
self.Tag.append(misp_tag) self.Tag.append(misp_tag)
self.edited = True self.edited = True
return misp_tag
def __get_tags(self): def __get_tags(self):
"""Returns a lost of tags associated to this Attribute""" """Returns a lost of tags associated to this Attribute"""
@ -357,7 +335,7 @@ class AbstractMISP(MutableMapping, MISPFileCache):
class MISPTag(AbstractMISP): class MISPTag(AbstractMISP):
_fields_for_feed = {'name', 'colour'} _fields_for_feed: set = {'name', 'colour'}
def from_dict(self, **kwargs): def from_dict(self, **kwargs):
if kwargs.get('Tag'): if kwargs.get('Tag'):
@ -372,3 +350,25 @@ class MISPTag(AbstractMISP):
if hasattr(self, 'exportable') and not self.exportable: if hasattr(self, 'exportable') and not self.exportable:
return False return False
return super()._to_feed() return super()._to_feed()
if HAS_RAPIDJSON:
def pymisp_json_default(obj: Union[AbstractMISP, datetime.datetime, datetime.date, Enum, UUID]) -> Union[dict, str]:
if isinstance(obj, AbstractMISP):
return obj.jsonable()
elif isinstance(obj, (datetime.datetime, datetime.date)):
return obj.isoformat()
elif isinstance(obj, Enum):
return obj.value
elif isinstance(obj, UUID):
return str(obj)
else:
def pymisp_json_default(obj: Union[AbstractMISP, datetime.datetime, datetime.date, Enum, UUID]) -> Union[dict, str]:
if isinstance(obj, AbstractMISP):
return obj.jsonable()
elif isinstance(obj, (datetime.datetime, datetime.date)):
return obj.isoformat()
elif isinstance(obj, Enum):
return obj.value
elif isinstance(obj, UUID):
return str(obj)

View File

@ -25,26 +25,26 @@ except ImportError:
logger.exception("Cannot import dateutil") logger.exception("Cannot import dateutil")
try: try:
import jsonschema import jsonschema # type: ignore
except ImportError: except ImportError:
logger.exception("Cannot import jsonschema") logger.exception("Cannot import jsonschema")
try: try:
# pyme renamed to gpg the 2016-10-28 # pyme renamed to gpg the 2016-10-28
import gpg import gpg # type: ignore
from gpg.constants.sig import mode from gpg.constants.sig import mode # type: ignore
has_pyme = True has_pyme = True
except ImportError: except ImportError:
try: try:
# pyme renamed to gpg the 2016-10-28 # pyme renamed to gpg the 2016-10-28
import pyme as gpg import pyme as gpg # type: ignore
from pyme.constants.sig import mode from pyme.constants.sig import mode # type: ignore
has_pyme = True has_pyme = True
except ImportError: except ImportError:
has_pyme = False has_pyme = False
def _int_to_str(d): def _int_to_str(d: dict):
# transform all integer back to string # transform all integer back to string
for k, v in d.items(): for k, v in d.items():
if isinstance(v, (int, float)) and not isinstance(v, bool): if isinstance(v, (int, float)) and not isinstance(v, bool):
@ -52,7 +52,7 @@ def _int_to_str(d):
return d return d
def make_bool(value): def make_bool(value: Union[bool, int, str, dict, list, None]) -> bool:
if isinstance(value, bool): if isinstance(value, bool):
return value return value
if isinstance(value, int): if isinstance(value, int):
@ -70,7 +70,7 @@ def make_bool(value):
class MISPOrganisation(AbstractMISP): class MISPOrganisation(AbstractMISP):
_fields_for_feed = {'name', 'uuid'} _fields_for_feed: set = {'name', 'uuid'}
def from_dict(self, **kwargs): def from_dict(self, **kwargs):
if 'Organisation' in kwargs: if 'Organisation' in kwargs:
@ -87,8 +87,8 @@ class MISPShadowAttribute(AbstractMISP):
def __repr__(self) -> str: def __repr__(self) -> str:
if hasattr(self, 'value'): if hasattr(self, 'value'):
return '<{self.__class__.__name__}(type={self.type}, value={self.value})'.format(self=self) return f'<{self.__class__.__name__}(type={self.type}, value={self.value})' # type: ignore
return '<{self.__class__.__name__}(NotInitialized)'.format(self=self) return f'<{self.__class__.__name__}(NotInitialized)'
class MISPSighting(AbstractMISP): class MISPSighting(AbstractMISP):
@ -109,17 +109,17 @@ class MISPSighting(AbstractMISP):
def __repr__(self) -> str: def __repr__(self) -> str:
if hasattr(self, 'value'): if hasattr(self, 'value'):
return '<{self.__class__.__name__}(value={self.value})'.format(self=self) return '<{self.__class__.__name__}(value={self.value})'.format(self=self) # type: ignore
if hasattr(self, 'id'): if hasattr(self, 'id'):
return '<{self.__class__.__name__}(id={self.id})'.format(self=self) return '<{self.__class__.__name__}(id={self.id})'.format(self=self) # type: ignore
if hasattr(self, 'uuid'): if hasattr(self, 'uuid'):
return '<{self.__class__.__name__}(uuid={self.uuid})'.format(self=self) return '<{self.__class__.__name__}(uuid={self.uuid})'.format(self=self) # type: ignore
return '<{self.__class__.__name__}(NotInitialized)'.format(self=self) return '<{self.__class__.__name__}(NotInitialized)'.format(self=self)
class MISPAttribute(AbstractMISP): class MISPAttribute(AbstractMISP):
_fields_for_feed = {'uuid', 'value', 'category', 'type', 'comment', 'data', _fields_for_feed: set = {'uuid', 'value', 'category', 'type', 'comment', 'data',
'timestamp', 'to_ids', 'disable_correlation'} 'timestamp', 'to_ids', 'disable_correlation'}
def __init__(self, describe_types: Optional[dict]=None, strict: bool=False): def __init__(self, describe_types: Optional[dict]=None, strict: bool=False):
"""Represents an Attribute """Represents an Attribute
@ -128,15 +128,15 @@ class MISPAttribute(AbstractMISP):
""" """
super().__init__() super().__init__()
if describe_types: if describe_types:
self.describe_types = describe_types self.describe_types: dict = describe_types
self.__categories = self.describe_types['categories'] self.__categories: List[str] = self.describe_types['categories']
self.__category_type_mapping = self.describe_types['category_type_mappings'] self.__category_type_mapping: dict = self.describe_types['category_type_mappings']
self.__sane_default = self.describe_types['sane_defaults'] self.__sane_default: dict = self.describe_types['sane_defaults']
self.__strict = strict self.__strict: bool = strict
self._data = None self._data = None
self.uuid = str(uuid.uuid4()) self.uuid: str = str(uuid.uuid4())
self.ShadowAttribute = [] self.ShadowAttribute: List[MISPShadowAttribute] = []
self.Sighting = [] self.Sighting: List[MISPSighting] = []
def hash_values(self, algorithm: str='sha512') -> List[str]: def hash_values(self, algorithm: str='sha512') -> List[str]:
"""Compute the hash of every values for fast lookups""" """Compute the hash of every values for fast lookups"""
@ -167,8 +167,8 @@ class MISPAttribute(AbstractMISP):
to_return = super()._to_feed() to_return = super()._to_feed()
if self.data: if self.data:
to_return['data'] = base64.b64encode(self.data.getvalue()).decode() to_return['data'] = base64.b64encode(self.data.getvalue()).decode()
if self.tags: if self.tags: # type: ignore
to_return['Tag'] = list(filter(None, [tag._to_feed() for tag in self.tags])) to_return['Tag'] = list(filter(None, [tag._to_feed() for tag in self.tags])) # type: ignore
return to_return return to_return
@property @property
@ -177,7 +177,7 @@ class MISPAttribute(AbstractMISP):
return self.describe_types['types'] return self.describe_types['types']
@property @property
def malware_binary(self) -> BytesIO: def malware_binary(self) -> Union[BytesIO, None]:
"""Returns a BytesIO of the malware (if the attribute has one, obvs).""" """Returns a BytesIO of the malware (if the attribute has one, obvs)."""
if hasattr(self, '_malware_binary'): if hasattr(self, '_malware_binary'):
return self._malware_binary return self._malware_binary
@ -445,8 +445,8 @@ class MISPAttribute(AbstractMISP):
class MISPObjectReference(AbstractMISP): class MISPObjectReference(AbstractMISP):
_fields_for_feed = {'uuid', 'timestamp', 'relationship_type', 'comment', _fields_for_feed: set = {'uuid', 'timestamp', 'relationship_type', 'comment',
'object_uuid', 'referenced_uuid'} 'object_uuid', 'referenced_uuid'}
def __init__(self): def __init__(self):
super().__init__() super().__init__()
@ -471,9 +471,9 @@ class MISPObjectReference(AbstractMISP):
class MISPObject(AbstractMISP): class MISPObject(AbstractMISP):
_fields_for_feed = {'name', 'meta-category', 'description', 'template_uuid', _fields_for_feed: set = {'name', 'meta-category', 'description', 'template_uuid',
'template_version', 'uuid', 'timestamp', 'distribution', 'template_version', 'uuid', 'timestamp', 'distribution',
'sharing_group_id', 'comment'} 'sharing_group_id', 'comment'}
def __init__(self, name: str, strict: bool=False, standalone: bool=False, default_attributes_parameters: dict={}, **kwargs): def __init__(self, name: str, strict: bool=False, standalone: bool=False, default_attributes_parameters: dict={}, **kwargs):
''' Master class representing a generic MISP object ''' Master class representing a generic MISP object
@ -489,21 +489,21 @@ class MISPObject(AbstractMISP):
:misp_objects_path_custom: Path to custom object templates :misp_objects_path_custom: Path to custom object templates
''' '''
super().__init__(**kwargs) super().__init__(**kwargs)
self._strict = strict self._strict: bool = strict
self.name = name self.name: str = name
self._known_template = False self._known_template: bool = False
self._set_template(kwargs.get('misp_objects_path_custom')) self._set_template(kwargs.get('misp_objects_path_custom'))
self.uuid = str(uuid.uuid4()) self.uuid: str = str(uuid.uuid4())
self.__fast_attribute_access = defaultdict(list) # Hashtable object_relation: [attributes] self.__fast_attribute_access: dict = defaultdict(list) # Hashtable object_relation: [attributes]
self.ObjectReference = [] self.ObjectReference: List[MISPObjectReference] = []
self.Attribute = [] self.Attribute: List[MISPAttribute] = []
if isinstance(default_attributes_parameters, MISPAttribute): if isinstance(default_attributes_parameters, MISPAttribute):
# Just make sure we're not modifying an existing MISPAttribute # Just make sure we're not modifying an existing MISPAttribute
self._default_attributes_parameters = default_attributes_parameters.to_dict() self._default_attributes_parameters: dict = default_attributes_parameters.to_dict()
else: else:
self._default_attributes_parameters = default_attributes_parameters self._default_attributes_parameters: dict = default_attributes_parameters
if self._default_attributes_parameters: if self._default_attributes_parameters:
# Let's clean that up # Let's clean that up
self._default_attributes_parameters.pop('value', None) # duh self._default_attributes_parameters.pop('value', None) # duh
@ -529,7 +529,7 @@ class MISPObject(AbstractMISP):
self.update_not_jsonable('ObjectReference') self.update_not_jsonable('ObjectReference')
def _load_template_path(self, template_path: Union[Path, str]) -> bool: def _load_template_path(self, template_path: Union[Path, str]) -> bool:
self._definition = self._load_json(template_path) self._definition: Union[dict, None] = self._load_json(template_path)
if not self._definition: if not self._definition:
return False return False
setattr(self, 'meta-category', self._definition['meta-category']) setattr(self, 'meta-category', self._definition['meta-category'])
@ -555,7 +555,7 @@ class MISPObject(AbstractMISP):
self.name = object_name self.name = object_name
self._set_template(misp_objects_path_custom) self._set_template(misp_objects_path_custom)
def _set_template(self, misp_objects_path_custom: Union[Path, str]=None): def _set_template(self, misp_objects_path_custom: Optional[Union[Path, str]]=None):
if misp_objects_path_custom: if misp_objects_path_custom:
# If misp_objects_path_custom is given, and an object with the given name exists, use that. # If misp_objects_path_custom is given, and an object with the given name exists, use that.
self.misp_objects_path = misp_objects_path_custom self.misp_objects_path = misp_objects_path_custom
@ -668,7 +668,7 @@ class MISPObject(AbstractMISP):
'''True if all the relations in the list are defined in the object''' '''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) return all(relation in self._fast_attribute_access for relation in list_of_relations)
def add_attribute(self, object_relation: str, simple_value: Union[str, int, float]=None, **value) -> MISPAttribute: def add_attribute(self, object_relation: str, simple_value: Union[str, int, float]=None, **value) -> Union[MISPAttribute, None]:
"""Add an attribute. object_relation is required and the value key is a """Add an attribute. object_relation is required and the value key is a
dictionary with all the keys supported by MISPAttribute""" dictionary with all the keys supported by MISPAttribute"""
if simple_value is not None: # /!\ The value *can* be 0 if simple_value is not None: # /!\ The value *can* be 0
@ -676,7 +676,7 @@ class MISPObject(AbstractMISP):
if value.get('value') in [None, '']: if value.get('value') in [None, '']:
logger.warning("The value of the attribute you're trying to add is None or empty string, skipping it. Object relation: {}".format(object_relation)) logger.warning("The value of the attribute you're trying to add is None or empty string, skipping it. Object relation: {}".format(object_relation))
return None return None
if self._known_template: if self._known_template and self._definition:
if object_relation in self._definition['attributes']: if object_relation in self._definition['attributes']:
attribute = MISPObjectAttribute(self._definition['attributes'][object_relation]) attribute = MISPObjectAttribute(self._definition['attributes'][object_relation])
else: else:
@ -694,7 +694,7 @@ class MISPObject(AbstractMISP):
self.edited = True self.edited = True
return attribute return attribute
def add_attributes(self, object_relation: str, *attributes) -> List[MISPAttribute]: def add_attributes(self, object_relation: str, *attributes) -> List[Optional[MISPAttribute]]:
'''Add multiple attributes with the same object_relation. '''Add multiple attributes with the same object_relation.
Helper for object_relation when multiple is True in the template. Helper for object_relation when multiple is True in the template.
It is the same as calling multiple times add_attribute with the same object_relation. It is the same as calling multiple times add_attribute with the same object_relation.
@ -713,7 +713,7 @@ class MISPObject(AbstractMISP):
self._validate() self._validate()
return super(MISPObject, self).to_dict() return super(MISPObject, self).to_dict()
def to_json(self, strict: bool=False, sort_keys: bool=False, indent: Optional[int]=None): def to_json(self, sort_keys: bool=False, indent: Optional[int]=None, strict: bool=False):
if strict or self._strict and self._known_template: if strict or self._strict and self._known_template:
self._validate() self._validate()
return super(MISPObject, self).to_json(sort_keys=sort_keys, indent=indent) return super(MISPObject, self).to_json(sort_keys=sort_keys, indent=indent)
@ -745,8 +745,8 @@ class MISPObject(AbstractMISP):
class MISPEvent(AbstractMISP): class MISPEvent(AbstractMISP):
_fields_for_feed = {'uuid', 'info', 'threat_level_id', 'analysis', 'timestamp', _fields_for_feed: set = {'uuid', 'info', 'threat_level_id', 'analysis', 'timestamp',
'publish_timestamp', 'published', 'date', 'extends_uuid'} 'publish_timestamp', 'published', 'date', 'extends_uuid'}
def __init__(self, describe_types: dict=None, strict_validation: bool=False, **kwargs): def __init__(self, describe_types: dict=None, strict_validation: bool=False, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
@ -759,10 +759,10 @@ class MISPEvent(AbstractMISP):
# This variable is used in add_attribute in order to avoid duplicating the structure # This variable is used in add_attribute in order to avoid duplicating the structure
self.describe_types = describe_types self.describe_types = describe_types
self.Attribute = [] self.Attribute: List[MISPAttribute] = []
self.Object = [] self.Object: List[MISPObject] = []
self.RelatedEvent = [] self.RelatedEvent: List[MISPEvent] = []
self.ShadowAttribute = [] self.ShadowAttribute: List[MISPShadowAttribute] = []
def _set_default(self): def _set_default(self):
"""There are a few keys that could, or need to be set by default for the feed generator""" """There are a few keys that could, or need to be set by default for the feed generator"""
@ -807,7 +807,7 @@ class MISPEvent(AbstractMISP):
} }
def attributes_hashes(self, algorithm: str='sha512') -> List[str]: def attributes_hashes(self, algorithm: str='sha512') -> List[str]:
to_return = [] to_return: List[str] = []
for attribute in self.attributes: for attribute in self.attributes:
to_return += attribute.hash_values(algorithm) to_return += attribute.hash_values(algorithm)
for obj in self.objects: for obj in self.objects:
@ -828,7 +828,7 @@ class MISPEvent(AbstractMISP):
if (hasattr(self, 'distribution') if (hasattr(self, 'distribution')
and self.distribution is not None and self.distribution is not None
and int(self.distribution) not in valid_distributions): and int(self.distribution) not in valid_distributions):
return return {}
to_return = super()._to_feed() to_return = super()._to_feed()
if with_meta: if with_meta:
@ -1069,7 +1069,7 @@ class MISPEvent(AbstractMISP):
"""Return the tags associated to an attribute or an object attribute. """Return the tags associated to an attribute or an object attribute.
:attribute_identifier: can be an ID, UUID, or the value. :attribute_identifier: can be an ID, UUID, or the value.
""" """
tags = [] tags: List[MISPTag] = []
for a in self.attributes + [attribute for o in self.objects for attribute in o.attributes]: for a in self.attributes + [attribute for o in self.objects for attribute in o.attributes]:
if ((hasattr(a, 'id') and a.id == attribute_identifier) if ((hasattr(a, 'id') and a.id == attribute_identifier)
or (hasattr(a, 'uuid') and a.uuid == attribute_identifier) or (hasattr(a, 'uuid') and a.uuid == attribute_identifier)
@ -1117,10 +1117,10 @@ class MISPEvent(AbstractMISP):
if not found: if not found:
raise PyMISPError('No attribute with UUID/ID {} found.'.format(attribute_id)) raise PyMISPError('No attribute with UUID/ID {} found.'.format(attribute_id))
def add_attribute(self, type: str, value: Union[str, int, float], **kwargs) -> MISPAttribute: def add_attribute(self, type: str, value: Union[str, int, float], **kwargs) -> Union[MISPAttribute, List[MISPAttribute]]:
"""Add an attribute. type and value are required but you can pass all """Add an attribute. type and value are required but you can pass all
other parameters supported by MISPAttribute""" other parameters supported by MISPAttribute"""
attr_list = [] attr_list: List[MISPAttribute] = []
if isinstance(value, list): if isinstance(value, list):
attr_list = [self.add_attribute(type=type, value=a, **kwargs) for a in value] attr_list = [self.add_attribute(type=type, value=a, **kwargs) for a in value]
else: else:
@ -1377,8 +1377,8 @@ class MISPEventDelegation(AbstractMISP):
class MISPObjectAttribute(MISPAttribute): class MISPObjectAttribute(MISPAttribute):
_fields_for_feed = {'uuid', 'object_relation', 'value', 'category', 'type', _fields_for_feed: set = {'uuid', 'object_relation', 'value', 'category', 'type',
'comment', 'data', 'timestamp', 'to_ids', 'disable_correlation'} 'comment', 'data', 'timestamp', 'to_ids', 'disable_correlation'}
def __init__(self, definition): def __init__(self, definition):
super().__init__() super().__init__()