From 9df636cd3725473fce307739726d5ae3f663fead Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Wed, 4 Sep 2019 13:59:20 +0200 Subject: [PATCH] chg: Update upload malware/attachment example script Fix #447 Make data at attibute level more generic with getter/setter methods --- examples/upload.py | 69 ++++++++++++++++++++++++++++----------------- pymisp/mispevent.py | 23 +++++++++++---- 2 files changed, 61 insertions(+), 31 deletions(-) diff --git a/examples/upload.py b/examples/upload.py index 4c6708d..447a9ea 100755 --- a/examples/upload.py +++ b/examples/upload.py @@ -1,43 +1,60 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 # -*- coding: utf-8 -*- -from pymisp import PyMISP +from pymisp import ExpandedPyMISP, MISPEvent, MISPAttribute from keys import misp_url, misp_key, misp_verifycert import argparse -import os -import glob - - -def init(url, key): - return PyMISP(url, key, misp_verifycert, 'json') - - -def upload_files(m, eid, paths, distrib, ids, categ, comment, info, analysis, threat): - out = m.upload_samplelist(paths, eid, distrib, ids, categ, comment, info, analysis, threat) - print(out) +from pathlib import Path if __name__ == '__main__': parser = argparse.ArgumentParser(description='Send malware sample to MISP.') parser.add_argument("-u", "--upload", type=str, required=True, help="File or directory of files to upload.") - parser.add_argument("-e", "--event", type=int, help="Not supplying an event ID will cause MISP to create a single new event for all of the POSTed malware samples.") parser.add_argument("-d", "--distrib", type=int, help="The distribution setting used for the attributes and for the newly created event, if relevant. [0-3].") - parser.add_argument("-ids", action='store_true', help="You can flag all attributes created during the transaction to be marked as \"to_ids\" or not.") - parser.add_argument("-c", "--categ", help="The category that will be assigned to the uploaded samples. Valid options are: Payload delivery, Artifacts dropped, Payload Installation, External Analysis.") + parser.add_argument("-c", "--comment", type=str, help="Comment for the uploaded file(s).") + parser.add_argument('-m', '--is-malware', action='store_true', help='The file(s) to upload are malwares') + parser.add_argument('--expand', action='store_true', help='(Only if the file is a malware) Run lief expansion (creates objects)') + parser.add_argument("-e", "--event", type=int, default=None, help="Not supplying an event ID will cause MISP to create a single new event for all of the POSTed malware samples.") parser.add_argument("-i", "--info", help="Used to populate the event info field if no event ID supplied.") - parser.add_argument("-a", "--analysis", type=int, help="The analysis level of the newly created event, if applicatble. [0-2]") - parser.add_argument("-t", "--threat", type=int, help="The threat level ID of the newly created event, if applicatble. [1-4]") - parser.add_argument("-co", "--comment", type=str, help="Comment for the uploaded file(s).") args = parser.parse_args() - misp = init(misp_url, misp_key) + misp = ExpandedPyMISP(misp_url, misp_key, misp_verifycert) files = [] - if os.path.isfile(args.upload): - files = [args.upload] - elif os.path.isdir(args.upload): - files = [f for f in glob.iglob(os.path.join(args.upload + '*'))] + p = Path(args.upload) + if p.is_file(): + files = [p] + elif p.is_dir(): + files = [f for f in p.glob('**/*') if f.is_file()] else: - print('invalid file') + print('invalid upload path (must be file or dir)') exit(0) - upload_files(misp, args.event, files, args.distrib, args.ids, args.categ, args.comment, args.info, args.analysis, args.threat) + if args.is_malware: + arg_type = 'malware-sample' + else: + arg_type = 'attachment' + + # Create attributes + attributes = [] + for f in files: + a = MISPAttribute() + a.type = arg_type + a.value = f.name + a.data = f + a.comment = args.comment + a.distribution = args.distrib + if args.expand and arg_type == 'malware-sample': + a.expand = 'binary' + attributes.append(a) + + if args.event: + for a in attributes: + misp.add_attribute(args.event, a) + else: + m = MISPEvent() + m.info = args.info + m.distribution = args.distrib + m.attributes = attributes + if args.expand and arg_type == 'malware-sample': + m.run_expansions() + misp.add_event(m) diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py index 15a3550..c9fba61 100644 --- a/pymisp/mispevent.py +++ b/pymisp/mispevent.py @@ -41,6 +41,8 @@ if (3, 0) <= sys.version_info < (3, 6): else: OLD_PY3 = False +if sys.version_info >= (3, 6): + from pathlib import Path try: from dateutil.parser import parse @@ -120,6 +122,7 @@ class MISPAttribute(AbstractMISP): self.__category_type_mapping = describe_types['category_type_mappings'] self.__sane_default = describe_types['sane_defaults'] self.__strict = strict + self._data = None self.uuid = str(uuid.uuid4()) self.ShadowAttribute = [] self.Sighting = [] @@ -251,7 +254,6 @@ class MISPAttribute(AbstractMISP): # other possible values if kwargs.get('data'): self.data = kwargs.pop('data') - self._load_data() if kwargs.get('id'): self.id = int(kwargs.pop('id')) if kwargs.get('event_id'): @@ -294,7 +296,7 @@ class MISPAttribute(AbstractMISP): def to_dict(self): to_return = super(MISPAttribute, self).to_dict() - if to_return.get('data'): + if self.data: to_return['data'] = base64.b64encode(self.data.getvalue()).decode() return to_return @@ -328,9 +330,20 @@ class MISPAttribute(AbstractMISP): return False return True - def _load_data(self): - if not isinstance(self.data, BytesIO): - self.data = BytesIO(base64.b64decode(self.data)) + @property + def data(self): + return self._data if self._data else None + + @data.setter + def data(self, data): + if sys.version_info >= (3, 6): + if isinstance(data, Path): + with data.open('rb') as f: + self._data = BytesIO(f.read()) + if isinstance(data, (str, bytes)): + self._data = BytesIO(base64.b64decode(data)) + elif isinstance(data, BytesIO): + self._data = data if self.type == 'malware-sample': try: with ZipFile(self.data) as f: