More cleanup

pull/34/head
Raphaël Vinot 2016-09-28 18:20:37 +02:00
parent e035922949
commit 6482a21834
3 changed files with 97 additions and 102 deletions

View File

@ -23,9 +23,10 @@ except ImportError:
HAVE_REQUESTS = False HAVE_REQUESTS = False
from . import __version__ from . import __version__
from .exceptions import PyMISPError, NewAttributeError, SearchError, MissingDependency, NoURL, NoKey from .exceptions import PyMISPError, SearchError, MissingDependency, NoURL, NoKey
from .mispevent import MISPEvent, MISPAttribute, EncodeUpdate from .mispevent import MISPEvent, MISPAttribute, EncodeUpdate
# Least dirty way to support python 2 and 3 # Least dirty way to support python 2 and 3
try: try:
basestring basestring
@ -297,10 +298,10 @@ class PyMISP(object):
return json.dumps(misp_event, cls=EncodeUpdate) return json.dumps(misp_event, cls=EncodeUpdate)
def _prepare_full_attribute(self, category, type_value, value, to_ids, comment=None, distribution=5): def _prepare_full_attribute(self, category, type_value, value, to_ids, comment=None, distribution=5):
misp_attribute = MISPAttribute(self.categories, self.types, self.category_type_mapping) misp_attribute = MISPAttribute(self.describe_types['result'])
misp_attribute.set_all_values(type=type_value, value=value, category=category, misp_attribute.set_all_values(type=type_value, value=value, category=category,
to_ids=to_ids, comment=comment, distribution=distribution) to_ids=to_ids, comment=comment, distribution=distribution)
return json.dumps(misp_attribute, cls=EncodeUpdate) return misp_attribute
def _one_or_more(self, value): def _one_or_more(self, value):
"""Returns a list/tuple of one or more items, regardless of input.""" """Returns a list/tuple of one or more items, regardless of input."""
@ -356,19 +357,14 @@ class PyMISP(object):
else: else:
e = MISPEvent(self.describe_types['result']) e = MISPEvent(self.describe_types['result'])
e.load(event) e.load(event)
for a in e.attributes: e.attributes += attributes
if a.get('distribution') is None:
a['distribution'] = 5
response = self.update_event(event['Event']['id'], json.dumps(e, cls=EncodeUpdate)) response = self.update_event(event['Event']['id'], json.dumps(e, cls=EncodeUpdate))
return response return response
def add_named_attribute(self, event, category, type_value, value, to_ids=False, comment=None, distribution=None, proposal=False): def add_named_attribute(self, event, type_value, value, category=None, to_ids=False, comment=None, distribution=None, proposal=False):
attributes = [] attributes = []
if value and category and type: for value in self._one_or_more(value):
try: attributes.append(self._prepare_full_attribute(category, type_value, value, to_ids, comment, distribution))
attributes.append(self._prepare_full_attribute(category, type_value, value, to_ids, comment, distribution))
except NewAttributeError as e:
return e
return self._send_attributes(event, attributes, proposal) return self._send_attributes(event, attributes, proposal)
def add_hashes(self, event, category='Artifacts dropped', filename=None, md5=None, sha1=None, sha256=None, ssdeep=None, comment=None, to_ids=True, distribution=None, proposal=False): def add_hashes(self, event, category='Artifacts dropped', filename=None, md5=None, sha1=None, sha256=None, ssdeep=None, comment=None, to_ids=True, distribution=None, proposal=False):
@ -632,8 +628,8 @@ class PyMISP(object):
# ######### Upload samples through the API ######### # ######### Upload samples through the API #########
# ################################################## # ##################################################
def prepare_attribute(self, event_id, distribution, to_ids, category, comment, info, def _prepare_upload(self, event_id, distribution, to_ids, category, comment, info,
analysis, threat_level_id): analysis, threat_level_id):
to_post = {'request': {}} to_post = {'request': {}}
if event_id is not None: if event_id is not None:
@ -647,12 +643,13 @@ class PyMISP(object):
else: else:
to_post['request']['event_id'] = int(event_id) to_post['request']['event_id'] = int(event_id)
if to_ids not in [True, False]: default_values = self.sane_default['malware-sample']
raise NewAttributeError('{} is invalid, to_ids has to be True or False'.format(to_ids)) if to_ids is None or not isinstance(to_ids, bool):
to_ids = bool(int(default_values['to_ids']))
to_post['request']['to_ids'] = to_ids to_post['request']['to_ids'] = to_ids
if category not in self.categories: if category is None or category not in self.categories:
raise NewAttributeError('{} is invalid, category has to be in {}'.format(category, (', '.join(self.categories)))) category = default_values['default_category']
to_post['request']['category'] = category to_post['request']['category'] = category
to_post['request']['comment'] = comment to_post['request']['comment'] = comment
@ -665,16 +662,16 @@ class PyMISP(object):
def upload_sample(self, filename, filepath, event_id, distribution=None, def upload_sample(self, filename, filepath, event_id, distribution=None,
to_ids=True, category=None, comment=None, info=None, to_ids=True, category=None, comment=None, info=None,
analysis=None, threat_level_id=None): analysis=None, threat_level_id=None):
to_post = self.prepare_attribute(event_id, distribution, to_ids, category, to_post = self._prepare_upload(event_id, distribution, to_ids, category,
comment, info, analysis, threat_level_id) comment, info, analysis, threat_level_id)
to_post['request']['files'] = [{'filename': filename, 'data': self._encode_file_to_upload(filepath)}] to_post['request']['files'] = [{'filename': filename, 'data': self._encode_file_to_upload(filepath)}]
return self._upload_sample(to_post) return self._upload_sample(to_post)
def upload_samplelist(self, filepaths, event_id, distribution=None, def upload_samplelist(self, filepaths, event_id, distribution=None,
to_ids=True, category=None, info=None, to_ids=True, category=None, info=None,
analysis=None, threat_level_id=None): analysis=None, threat_level_id=None):
to_post = self.prepare_attribute(event_id, distribution, to_ids, category, to_post = self._prepare_upload(event_id, distribution, to_ids, category,
info, analysis, threat_level_id) info, analysis, threat_level_id)
files = [] files = []
for path in filepaths: for path in filepaths:
if not os.path.isfile(path): if not os.path.isfile(path):
@ -694,18 +691,14 @@ class PyMISP(object):
# ############################ # ############################
def __query_proposal(self, session, path, id, attribute=None): def __query_proposal(self, session, path, id, attribute=None):
path = path.strip('/')
url = urljoin(self.root_url, 'shadow_attributes/{}/{}'.format(path, id)) url = urljoin(self.root_url, 'shadow_attributes/{}/{}'.format(path, id))
query = None
if path in ['add', 'edit']: if path in ['add', 'edit']:
query = {'request': {'ShadowAttribute': attribute}} query = {'request': {'ShadowAttribute': attribute}}
if path == 'view': response = session.post(url, data=json.dumps(query))
elif path == 'view':
response = session.get(url) response = session.get(url)
else: else: # accept or discard
if query is not None: response = session.post(url)
response = session.post(url, data=json.dumps(query))
else:
response = session.post(url)
return self._check_response(response) return self._check_response(response)
def proposal_view(self, event_id=None, proposal_id=None): def proposal_view(self, event_id=None, proposal_id=None):

View File

@ -21,10 +21,11 @@ from .exceptions import PyMISPError, NewEventError, NewAttributeError
class MISPAttribute(object): class MISPAttribute(object):
def __init__(self, categories, types, category_type_mapping): def __init__(self, describe_types):
self.categories = categories self.categories = describe_types['categories']
self.types = types self.types = describe_types['types']
self.category_type_mapping = category_type_mapping self.category_type_mapping = describe_types['category_type_mappings']
self.sane_default = describe_types['sane_defaults']
self._reinitialize_attribute() self._reinitialize_attribute()
def _reinitialize_attribute(self): def _reinitialize_attribute(self):
@ -46,44 +47,59 @@ class MISPAttribute(object):
self.ShadowAttribute = [] self.ShadowAttribute = []
def set_all_values(self, **kwargs): def set_all_values(self, **kwargs):
if kwargs.get('type') and kwargs.get('category'):
if kwargs['type'] not in self.category_type_mapping[kwargs['category']]:
raise NewAttributeError('{} and {} is an invalid combinaison, type for this category has to be in {}'.capitalizeformat(self.type, self.category, (', '.join(self.category_type_mapping[self.category]))))
# Required
if kwargs.get('type'):
self.type = kwargs['type']
if self.type not in self.types:
raise NewAttributeError('{} is invalid, type has to be in {}'.format(self.type, (', '.join(self.types))))
else:
raise NewAttributeError('The type of the attribute is required.')
type_defaults = self.sane_default[self.type]
if kwargs.get('value'):
self.value = kwargs['value']
else:
raise NewAttributeError('The value of the attribute is required.')
# Default values # Default values
if kwargs.get('category', None): if kwargs.get('category'):
self.category = kwargs['category'] self.category = kwargs['category']
if self.category not in self.categories: if self.category not in self.categories:
raise NewAttributeError('{} is invalid, category has to be in {}'.format(self.category, (', '.join(self.categories)))) raise NewAttributeError('{} is invalid, category has to be in {}'.format(self.category, (', '.join(self.categories))))
if kwargs.get('type', None): else:
self.type = kwargs['type'] self.category = type_defaults['default_category']
if self.type not in self.types:
raise NewAttributeError('{} is invalid, type_value has to be in {}'.format(self.type, (', '.join(self.types)))) if kwargs.get('to_ids'):
if self.type not in self.category_type_mapping[self.category]:
raise NewAttributeError('{} and {} is an invalid combinaison, type_value for this category has to be in {}'.capitalizeformat(self.type, self.category, (', '.join(self.category_type_mapping[self.category]))))
if kwargs.get('value', None):
self.value = kwargs['value']
if kwargs.get('to_ids', None):
self.to_ids = kwargs['to_ids'] self.to_ids = kwargs['to_ids']
if not isinstance(self.to_ids, bool): if not isinstance(self.to_ids, bool):
raise NewAttributeError('{} is invalid, to_ids has to be True or False'.format(self.to_ids)) raise NewAttributeError('{} is invalid, to_ids has to be True or False'.format(self.to_ids))
if kwargs.get('comment', None): else:
self.to_ids = bool(int(type_defaults['to_ids']))
if kwargs.get('comment'):
self.comment = kwargs['comment'] self.comment = kwargs['comment']
if kwargs.get('distribution', None): if kwargs.get('distribution'):
self.distribution = int(kwargs['distribution']) self.distribution = int(kwargs['distribution'])
if self.distribution not in [0, 1, 2, 3, 5]: if self.distribution not in [0, 1, 2, 3, 5]:
raise NewAttributeError('{} is invalid, the distribution has to be in 0, 1, 2, 3, 5'.format(self.distribution)) raise NewAttributeError('{} is invalid, the distribution has to be in 0, 1, 2, 3, 5'.format(self.distribution))
# other possible values # other possible values
if kwargs.get('id', None): if kwargs.get('id'):
self.id = int(kwargs['id']) self.id = int(kwargs['id'])
if kwargs.get('uuid', None): if kwargs.get('uuid'):
self.uuid = kwargs['uuid'] self.uuid = kwargs['uuid']
if kwargs.get('timestamp', None): if kwargs.get('timestamp'):
self.timestamp = datetime.datetime.fromtimestamp(int(kwargs['timestamp'])) self.timestamp = datetime.datetime.fromtimestamp(int(kwargs['timestamp']))
if kwargs.get('sharing_group_id', None): if kwargs.get('sharing_group_id'):
self.sharing_group_id = int(kwargs['sharing_group_id']) self.sharing_group_id = int(kwargs['sharing_group_id'])
if kwargs.get('deleted', None): if kwargs.get('deleted'):
self.deleted = kwargs['deleted'] self.deleted = kwargs['deleted']
if kwargs.get('SharingGroup', None): if kwargs.get('SharingGroup'):
self.SharingGroup = kwargs['SharingGroup'] self.SharingGroup = kwargs['SharingGroup']
if kwargs.get('ShadowAttribute', None): if kwargs.get('ShadowAttribute'):
self.ShadowAttribute = kwargs['ShadowAttribute'] self.ShadowAttribute = kwargs['ShadowAttribute']
def _json(self): def _json(self):
@ -136,6 +152,7 @@ class MISPEvent(object):
if not describe_types: if not describe_types:
t = json.load(open(os.path.join(self.ressources_path, 'describeTypes.json'), 'r')) t = json.load(open(os.path.join(self.ressources_path, 'describeTypes.json'), 'r'))
describe_types = t['result'] describe_types = t['result']
self.describe_types = describe_types
self.categories = describe_types['categories'] self.categories = describe_types['categories']
self.types = describe_types['types'] self.types = describe_types['types']
self.category_type_mapping = describe_types['category_type_mappings'] self.category_type_mapping = describe_types['category_type_mappings']
@ -150,7 +167,7 @@ class MISPEvent(object):
self.distribution = 3 self.distribution = 3
self.threat_level_id = 2 self.threat_level_id = 2
self.analysis = 0 self.analysis = 0
self.info = '' self.info = None
self.published = False self.published = False
self.date = datetime.date.today() self.date = datetime.date.today()
self.attributes = [] self.attributes = []
@ -191,24 +208,28 @@ class MISPEvent(object):
self.set_all_values(**e) self.set_all_values(**e)
def set_all_values(self, **kwargs): def set_all_values(self, **kwargs):
# Required value
if kwargs.get('info'):
self.info = kwargs['info']
else:
raise NewAttributeError('The info field of the new event is required.')
# Default values for a valid event to send to a MISP instance # Default values for a valid event to send to a MISP instance
if kwargs.get('distribution', None) is not None: if kwargs.get('distribution') is not None:
self.distribution = int(kwargs['distribution']) self.distribution = int(kwargs['distribution'])
if self.distribution not in [0, 1, 2, 3]: if self.distribution not in [0, 1, 2, 3]:
raise NewEventError('{} is invalid, the distribution has to be in 0, 1, 2, 3'.format(self.distribution)) raise NewEventError('{} is invalid, the distribution has to be in 0, 1, 2, 3'.format(self.distribution))
if kwargs.get('threat_level_id', None) is not None: if kwargs.get('threat_level_id') is not None:
self.threat_level_id = int(kwargs['threat_level_id']) self.threat_level_id = int(kwargs['threat_level_id'])
if self.threat_level_id not in [1, 2, 3, 4]: if self.threat_level_id not in [1, 2, 3, 4]:
raise NewEventError('{} is invalid, the threat_level has to be in 1, 2, 3, 4'.format(self.threat_level_id)) raise NewEventError('{} is invalid, the threat_level has to be in 1, 2, 3, 4'.format(self.threat_level_id))
if kwargs.get('analysis', None) is not None: if kwargs.get('analysis') is not None:
self.analysis = int(kwargs['analysis']) self.analysis = int(kwargs['analysis'])
if self.analysis not in [0, 1, 2]: if self.analysis not in [0, 1, 2]:
raise NewEventError('{} is invalid, the analysis has to be in 0, 1, 2'.format(self.analysis)) raise NewEventError('{} is invalid, the analysis has to be in 0, 1, 2'.format(self.analysis))
if kwargs.get('info', None): if kwargs.get('published') is not None:
self.info = kwargs['info']
if kwargs.get('published', None) is not None:
self.publish() self.publish()
if kwargs.get('date', None): if kwargs.get('date'):
if isinstance(kwargs['date'], str): if isinstance(kwargs['date'], str):
self.date = parse(kwargs['date']) self.date = parse(kwargs['date'])
elif isinstance(kwargs['date'], datetime.datetime): elif isinstance(kwargs['date'], datetime.datetime):
@ -217,42 +238,42 @@ class MISPEvent(object):
self.date = kwargs['date'] self.date = kwargs['date']
else: else:
raise NewEventError('Invalid format for the date: {} - {}'.format(kwargs['date'], type(kwargs['date']))) raise NewEventError('Invalid format for the date: {} - {}'.format(kwargs['date'], type(kwargs['date'])))
if kwargs.get('Attribute', None): if kwargs.get('Attribute'):
for a in kwargs['Attribute']: for a in kwargs['Attribute']:
attribute = MISPAttribute(self.categories, self.types, self.category_type_mapping) attribute = MISPAttribute(self.describe_types)
attribute.set_all_values(**a) attribute.set_all_values(**a)
self.attributes.append(attribute) self.attributes.append(attribute)
# All other keys # All other keys
if kwargs.get('id', None): if kwargs.get('id'):
self.id = int(kwargs['id']) self.id = int(kwargs['id'])
if kwargs.get('orgc_id', None): if kwargs.get('orgc_id'):
self.orgc_id = int(kwargs['orgc_id']) self.orgc_id = int(kwargs['orgc_id'])
if kwargs.get('org_id', None): if kwargs.get('org_id'):
self.org_id = int(kwargs['org_id']) self.org_id = int(kwargs['org_id'])
if kwargs.get('uuid', None): if kwargs.get('uuid'):
self.uuid = kwargs['uuid'] self.uuid = kwargs['uuid']
if kwargs.get('attribute_count', None): if kwargs.get('attribute_count'):
self.attribute_count = int(kwargs['attribute_count']) self.attribute_count = int(kwargs['attribute_count'])
if kwargs.get('timestamp', None): if kwargs.get('timestamp'):
self.timestamp = datetime.datetime.fromtimestamp(int(kwargs['timestamp'])) self.timestamp = datetime.datetime.fromtimestamp(int(kwargs['timestamp']))
if kwargs.get('proposal_email_lock', None): if kwargs.get('proposal_email_lock'):
self.proposal_email_lock = kwargs['proposal_email_lock'] self.proposal_email_lock = kwargs['proposal_email_lock']
if kwargs.get('locked', None): if kwargs.get('locked'):
self.locked = kwargs['locked'] self.locked = kwargs['locked']
if kwargs.get('publish_timestamp', None): if kwargs.get('publish_timestamp'):
self.publish_timestamp = datetime.datetime.fromtimestamp(int(kwargs['publish_timestamp'])) self.publish_timestamp = datetime.datetime.fromtimestamp(int(kwargs['publish_timestamp']))
if kwargs.get('sharing_group_id', None): if kwargs.get('sharing_group_id'):
self.sharing_group_id = int(kwargs['sharing_group_id']) self.sharing_group_id = int(kwargs['sharing_group_id'])
if kwargs.get('Org', None): if kwargs.get('Org'):
self.Org = kwargs['Org'] self.Org = kwargs['Org']
if kwargs.get('Orgc', None): if kwargs.get('Orgc'):
self.Orgc = kwargs['Orgc'] self.Orgc = kwargs['Orgc']
if kwargs.get('ShadowAttribute', None): if kwargs.get('ShadowAttribute'):
self.ShadowAttribute = kwargs['ShadowAttribute'] self.ShadowAttribute = kwargs['ShadowAttribute']
if kwargs.get('RelatedEvent', None): if kwargs.get('RelatedEvent'):
self.RelatedEvent = kwargs['RelatedEvent'] self.RelatedEvent = kwargs['RelatedEvent']
if kwargs.get('Tag', None): if kwargs.get('Tag'):
self.Tag = kwargs['Tag'] self.Tag = kwargs['Tag']
def _json(self): def _json(self):
@ -309,27 +330,7 @@ class MISPEvent(object):
def unpublish(self): def unpublish(self):
self.published = False self.published = False
def add_attribute(self, type_value, value, **kwargs): def add_attribute(self, type, value, **kwargs):
if not self.sane_default.get(type_value): attribute = MISPAttribute(self.describe_types)
raise NewAttributeError("{} is an invalid type. Can only be one of the following: {}".format(type_value, ', '.join(self.types))) attribute.set_all_values(type=type, value=value, **kwargs)
defaults = self.sane_default[type_value]
if kwargs.get('category'):
category = kwargs.get('category')
else:
category = defaults['default_category']
if kwargs.get('to_ids'):
to_ids = bool(int(kwargs.get('to_ids')))
else:
to_ids = bool(int(defaults['to_ids']))
if kwargs.get('comment'):
comment = kwargs.get('comment')
else:
comment = None
if kwargs.get('distribution'):
distribution = int(kwargs.get('distribution'))
else:
distribution = 5
attribute = MISPAttribute(self.categories, self.types, self.category_type_mapping)
attribute.set_all_values(type=type_value, value=value, category=category,
to_ids=to_ids, comment=comment, distribution=distribution)
self.attributes.append(attribute) self.attributes.append(attribute)

View File

@ -111,7 +111,8 @@ class TestOffline(unittest.TestCase):
self.initURI(m) self.initURI(m)
pymisp = PyMISP(self.domain, self.key) pymisp = PyMISP(self.domain, self.key)
m.register_uri('POST', self.domain + 'events', json=error_empty_info) m.register_uri('POST', self.domain + 'events', json=error_empty_info)
response = pymisp.new_event(0, 1, 0) # TODO Add test exception if info field isn't set
response = pymisp.new_event(0, 1, 0, 'Foo')
self.assertEqual(response, error_empty_info_flatten) self.assertEqual(response, error_empty_info_flatten)
m.register_uri('POST', self.domain + 'events', json=self.new_misp_event) m.register_uri('POST', self.domain + 'events', json=self.new_misp_event)
response = pymisp.new_event(0, 1, 0, "This is a test.", '2016-08-26', False) response = pymisp.new_event(0, 1, 0, "This is a test.", '2016-08-26', False)