chg: Update upload malware/attachment example script

Fix #447

Make data at attibute level more generic with getter/setter methods

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)

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: