chg: Improve typing

pull/665/head
Raphaël Vinot 2020-11-24 14:40:00 +01:00
parent 35860b49bd
commit 7b2e78246a
1 changed files with 7 additions and 6 deletions

View File

@ -1,7 +1,7 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
from typing import TypeVar, Optional, Tuple, List, Dict, Union, Any, Mapping, Iterator from typing import TypeVar, Optional, Tuple, List, Dict, Union, Any, Mapping, Iterable
from datetime import date, datetime from datetime import date, datetime
import csv import csv
from pathlib import Path from pathlib import Path
@ -563,7 +563,7 @@ class PyMISP:
r = self._prepare_request('HEAD', f'attributes/view/{attribute_id}') r = self._prepare_request('HEAD', f'attributes/view/{attribute_id}')
return self._check_head_response(r) return self._check_head_response(r)
def add_attribute(self, event: Union[MISPEvent, int, str, UUID], attribute: Union[MISPAttribute, list], pythonify: bool = False) -> Union[Dict, MISPAttribute, MISPShadowAttribute]: def add_attribute(self, event: Union[MISPEvent, int, str, UUID], attribute: Union[MISPAttribute, Iterable], pythonify: bool = False) -> Union[Dict, MISPAttribute, MISPShadowAttribute]:
"""Add an attribute to an existing MISP event """Add an attribute to an existing MISP event
:param event: event to extend :param event: event to extend
@ -579,7 +579,7 @@ class PyMISP:
# Multiple attributes were passed at once, the handling is totally different # Multiple attributes were passed at once, the handling is totally different
if not (self.global_pythonify or pythonify): if not (self.global_pythonify or pythonify):
return new_attribute return new_attribute
to_return = {'attributes': []} to_return: Dict[str, List[MISPAttribute]] = {'attributes': []}
if 'errors' in new_attribute: if 'errors' in new_attribute:
to_return['errors'] = new_attribute['errors'] to_return['errors'] = new_attribute['errors']
@ -600,6 +600,7 @@ class PyMISP:
and new_attribute['errors'][1]['message'] == 'You do not have permission to do that.'): and new_attribute['errors'][1]['message'] == 'You do not have permission to do that.'):
# At this point, we assume the user tried to add an attribute on an event they don't own # At this point, we assume the user tried to add an attribute on an event they don't own
# Re-try with a proposal # Re-try with a proposal
if isinstance(attribute, MISPAttribute):
return self.add_attribute_proposal(event_id, attribute, pythonify) return self.add_attribute_proposal(event_id, attribute, pythonify)
if not (self.global_pythonify or pythonify) or 'errors' in new_attribute: if not (self.global_pythonify or pythonify) or 'errors' in new_attribute:
return new_attribute return new_attribute
@ -1094,7 +1095,7 @@ class PyMISP:
warninglist_id = get_uuid_or_id_from_abstract_misp(warninglist) warninglist_id = get_uuid_or_id_from_abstract_misp(warninglist)
return self.toggle_warninglist(warninglist_id=warninglist_id, force_enable=False) return self.toggle_warninglist(warninglist_id=warninglist_id, force_enable=False)
def values_in_warninglist(self, value: Iterator) -> Dict: def values_in_warninglist(self, value: Iterable) -> Dict:
"""Check if IOC values are in warninglist """Check if IOC values are in warninglist
:param value: iterator with values to check :param value: iterator with values to check
@ -3094,7 +3095,7 @@ class PyMISP:
def __repr__(self): def __repr__(self):
return f'<{self.__class__.__name__}(url={self.root_url})' return f'<{self.__class__.__name__}(url={self.root_url})'
def _prepare_request(self, request_type: str, url: str, data: Union[str, Iterator, Mapping, AbstractMISP] = {}, params: Mapping = {}, def _prepare_request(self, request_type: str, url: str, data: Union[str, Iterable, Mapping, AbstractMISP] = {}, params: Mapping = {},
kw_params: Mapping = {}, output_type: str = 'json') -> requests.Response: kw_params: Mapping = {}, output_type: str = 'json') -> requests.Response:
'''Prepare a request for python-requests''' '''Prepare a request for python-requests'''
url = urljoin(self.root_url, url) url = urljoin(self.root_url, url)