Merge branch 'release/2.0.0'

Author: nadouani

cortexutils/__init__.py

@@ -0,0 +1,3 @@
+def runner(worker_cls):
+    worker = worker_cls()
+    worker.run()

cortexutils/analyzer.py

@@ -2,14 +2,19 @@
 # encoding: utf-8
 
 import json
-from cortexutils.worker import Worker
+import os
+
 from cortexutils.extractor import Extractor
+from cortexutils.worker import Worker
+from shutil import copyfileobj
+import tempfile
+import ntpath
 
 
 class Analyzer(Worker):
 
-    def __init__(self):
-        Worker.__init__(self)
+    def __init__(self, job_directory=None):
+        Worker.__init__(self, job_directory)
 
         # Not breaking compatibility
         self.artifact = self._input
@@ -23,7 +28,17 @@ def get_data(self):
         :return: Data (observable value) given through Cortex"""
         if self.data_type == 'file':
             return self.get_param('filename', None, 'Missing filename.')
-        return self.get_param('data', None, 'Missing data field')
+        else:
+            return self.get_param('data', None, 'Missing data field')
+
+    def get_param(self, name, default=None, message=None):
+        data = super(Analyzer, self).get_param(name, default, message)
+        if name == 'file' and self.data_type == 'file' and self.job_directory is not None:
+            path = '%s/input/%s' % (self.job_directory, data)
+            if os.path.isfile(path):
+                return path
+        else:
+            return data
 
     def build_taxonomy(self, level, namespace, predicate, value):
         """
@@ -37,11 +52,11 @@ def build_taxonomy(self, level, namespace, predicate, value):
         if level not in ['info', 'safe', 'suspicious', 'malicious']:
             level = 'info'
         return {
-                'level': level,
-                'namespace': namespace,
-                'predicate': predicate,
-                'value': value
-                }    
+            'level': level,
+            'namespace': namespace,
+            'predicate': predicate,
+            'value': value
+        }
 
     def summary(self, raw):
         """Returns a summary, needed for 'short.html' template. Overwrite it for your needs!
@@ -58,6 +73,19 @@ def artifacts(self, raw):
         # Return empty list
         return []
 
+    def build_artifact(self, data_type, data, **kwargs):
+        if data_type == 'file':
+            if os.path.isfile(data):
+                (dst, filename) = tempfile.mkstemp(dir=os.path.join(self.job_directory, "output"))
+                with open(data, 'r') as src:
+                    copyfileobj(src, os.fdopen(dst, 'w'))
+                    kwargs.update({'dataType': data_type, 'file': ntpath.basename(filename),
+                                   'filename': ntpath.basename(data)})
+                    return kwargs
+        else:
+            kwargs.update({'dataType': data_type, 'data': data})
+            return kwargs
+
     def report(self, full_report, ensure_ascii=False):
         """Returns a json dict via stdout.
 
@@ -70,13 +98,12 @@ def report(self, full_report, ensure_ascii=False):
         except Exception:
             pass
 
-        report = {
+        super(Analyzer, self).report({
             'success': True,
             'summary': summary,
             'artifacts': self.artifacts(full_report),
             'full': full_report
-        }
-        json.dump(report, self.fpoutput, ensure_ascii=ensure_ascii)
+        }, ensure_ascii)
 
     def run(self):
         """Overwritten by analyzers"""
@@ -103,4 +130,4 @@ def getParam(self, name, default=None, message=None):
     # Not breaking compatibility
     def checkTlp(self, message):
         if not (self.__check_tlp()):
-            self.error(message)
+            self.error(message)

cortexutils/extractor.py

@@ -164,8 +164,8 @@ def check_iterable(self, iterable):
             dt = self.__checktype(iterable)
             if len(dt) > 0:
                 results.append({
-                    'type': dt,
-                    'value': iterable
+                    'dataType': dt,
+                    'data': iterable
                 })
         elif isinstance(iterable, list):
             for item in iterable:
@@ -175,8 +175,8 @@ def check_iterable(self, iterable):
                     dt = self.__checktype(item)
                     if len(dt) > 0:
                         results.append({
-                            'type': dt,
-                            'value': item
+                            'dataType': dt,
+                            'data': item
                         })
         elif isinstance(iterable, dict):
             for _, item in iterable.items():
@@ -186,10 +186,22 @@ def check_iterable(self, iterable):
                     dt = self.__checktype(item)
                     if len(dt) > 0:
                         results.append({
-                            'type': dt,
-                            'value': item
+                            'dataType': dt,
+                            'data': item
                         })
         else:
             raise TypeError('Not supported type.')
 
-        return results
+        return self.deduplicate(results)
+
+    @staticmethod
+    def deduplicate(list_of_objects):
+        dedup_list = []
+        for obj in list_of_objects:
+            present = False
+            for new_object in dedup_list:
+                if obj['dataType'] == new_object['dataType'] and obj['data'] == new_object['data']:
+                    present = True
+            if not present:
+                dedup_list.append(obj)
+        return dedup_list

cortexutils/responder.py

@@ -2,13 +2,14 @@
 # encoding: utf-8
 
 import json
+import os
 from cortexutils.worker import Worker
 
 
 class Responder(Worker):
 
-    def __init__(self):
-        Worker.__init__(self)
+    def __init__(self, job_directory=None):
+        Worker.__init__(self, job_directory)
 
         # Not breaking compatibility
         self.artifact = self._input
@@ -50,13 +51,11 @@ def report(self, full_report, ensure_ascii=False):
             operation_list = self.operations(full_report)
         except Exception:
             pass
-
-        report = {
+        super(Responder, self).report({
             'success': True,
             'full': full_report,
             'operations': operation_list
-        }
-        json.dump(report, self.fpoutput, ensure_ascii=ensure_ascii)
+        }, ensure_ascii)
 
     def run(self):
         """Overwritten by responders"""

cortexutils/worker.py

@@ -1,23 +1,36 @@
 #!/usr/bin/env python
 # encoding: utf-8
+
 import os
 import sys
 import codecs
 import json
+import select
 
 
-class Worker:
-
-    def __init__(self):
-        self.__set_encoding()
-
-        # Prepare in/out/err streams
-        self.fperror = sys.stderr
-        self.fpinput = sys.stdin
-        self.fpoutput = sys.stdout
+class Worker(object):
+    READ_TIMEOUT = 3  # seconds
 
+    def __init__(self, job_directory):
+        if job_directory is None:
+            if len(sys.argv) > 1:
+                job_directory = sys.argv[1]
+            else:
+                job_directory = '/job'
+        self.job_directory = job_directory
         # Load input
-        self._input = json.load(self.fpinput)
+        self._input = {}
+        if os.path.isfile('%s/input/input.json' % self.job_directory):
+            with open('%s/input/input.json' % self.job_directory) as f_input:
+                self._input = json.load(f_input)
+        else:  # If input file doesn't exist, fallback to old behavior and read input from stdin
+            self.job_directory = None
+            self.__set_encoding()
+            r, w, e = select.select([sys.stdin], [], [], self.READ_TIMEOUT)
+            if sys.stdin in r:
+                self._input = json.load(sys.stdin)
+            else:
+                self.error('Input file doesn''t exist')
 
         # Set parameters
         self.data_type = self.get_param('dataType', None, 'Missing dataType field')
@@ -98,10 +111,21 @@ def __check_pap(self):
 
         return not (self.enable_check_pap and self.pap > self.max_pap)
 
+    def __write_output(self, data, ensure_ascii=False):
+        if self.job_directory is None:
+            json.dump(data, sys.stdout, ensure_ascii=ensure_ascii)
+        else:
+            try:
+                os.makedirs('%s/output' % self.job_directory)
+            except:
+                pass
+            with open('%s/output/output.json' % self.job_directory, mode='w') as f_output:
+                json.dump(data, f_output, ensure_ascii=ensure_ascii)
+
     def get_data(self):
         """Wrapper for getting data from input dict.
 
-        :return: Data (observable value) given through Cortex"""        
+        :return: Data (observable value) given through Cortex"""
         return self.get_param('data', None, 'Missing data field')
 
     def get_param(self, name, default=None, message=None):
@@ -128,34 +152,30 @@ def error(self, message, ensure_ascii=False):
         if 'api_key' in analyzer_input.get('config', {}):
             analyzer_input['config']['api_key'] = 'REMOVED'
 
-        json.dump({'success': False,
-                   'input': analyzer_input,
-                   'errorMessage': message},
-                  self.fpoutput,
-                  ensure_ascii=ensure_ascii)
+        self.__write_output({'success': False,
+                             'input': analyzer_input,
+                             'errorMessage': message},
+                            ensure_ascii=ensure_ascii)
 
         # Force exit after error
         sys.exit(1)
 
-    def report(self, full_report, ensure_ascii=False):
+    def summary(self, raw):
+        """Returns a summary, needed for 'short.html' template. Overwrite it for your needs!
+
+        :returns: by default return an empty dict"""
+        return {}
+
+    def artifacts(self, raw):
+        return []
+
+    def report(self, output, ensure_ascii=False):
         """Returns a json dict via stdout.
 
-        :param full_report: Analyzer results as dict.
+        :param output: worker output.
         :param ensure_ascii: Force ascii output. Default: False"""
 
-        summary = {}
-        try:
-            summary = self.summary(full_report)
-        except Exception:
-            pass
-
-        report = {
-            'success': True,
-            'summary': summary,
-            'artifacts': self.artifacts(full_report),
-            'full': full_report
-        }
-        json.dump(report, self.fpoutput, ensure_ascii=ensure_ascii)
+        self.__write_output(output, ensure_ascii=ensure_ascii)
 
     def run(self):
         """Overwritten by analyzers"""

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='cortexutils',
-    version='1.3.0',
+    version='2.0.0',
     description='A Python library for including utility classes for Cortex analyzers and responders',
     long_description=open('README').read(),
     author='TheHive-Project',