From 59f9b37b5e89ef34ccf8f71c89b7c1b84055c631 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Tue, 18 Jul 2017 11:26:09 +0200 Subject: [PATCH] Allow to pass a bytestream to upload_sample. Fix #101 --- pymisp/api.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/pymisp/api.py b/pymisp/api.py index 5ddd387..87656a7 100644 --- a/pymisp/api.py +++ b/pymisp/api.py @@ -685,16 +685,20 @@ class PyMISP(object): to_post['request']['comment'] = comment return to_post - def _encode_file_to_upload(self, path): - with open(path, 'rb') as f: - return base64.b64encode(f.read()).decode() + def _encode_file_to_upload(self, filepath_or_bytes): + if isinstance(filepath_or_bytes, basestring) and os.path.isfile(filepath_or_bytes): + with open(filepath_or_bytes, 'rb') as f: + binblob = f.read() + else: + binblob = filepath_or_bytes + return base64.b64encode(binblob).decode() - def upload_sample(self, filename, filepath, event_id, distribution=None, + def upload_sample(self, filename, filepath_or_bytes, event_id, distribution=None, to_ids=True, category=None, comment=None, info=None, analysis=None, threat_level_id=None): to_post = self._prepare_upload(event_id, distribution, to_ids, category, 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_or_bytes)}] return self._upload_sample(to_post) def upload_samplelist(self, filepaths, event_id, distribution=None,