Update Cuckoo module to support files and URLs

misp_modules/modules/expansion/cuckoo_submit.py

@@ -1,56 +1,143 @@
+import base64
+import io
 import json
+import logging
 import requests
+import sys
+import urllib.parse
+import zipfile
 
-misperrors = {'error': 'Error'}
+from requests.exceptions import RequestException
+
+log = logging.getLogger('cuckoo_submit')
+log.setLevel(logging.DEBUG)
+sh = logging.StreamHandler(sys.stdout)
+sh.setLevel(logging.DEBUG)
+fmt = logging.Formatter(
+    '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+sh.setFormatter(fmt)
+log.addHandler(sh)
+
+moduleinfo = {
+    "version": "0.1", 'author': "Evert Kors",
+    "description": "Submit files and URLs to Cuckoo Sandbox",
+    "module-type": ["expansion", "hover"]
+}
+misperrors = {"error": "Error"}
+moduleconfig = ["cuckoo_api", "api_key"]
 mispattributes = {
-        'input': ['url'],
-        'output': ['text']
+    "input": ["attachment', 'malware-sample", "url", "domain"],
+    "output": ["text"]
 }
 
-# possible module-types: 'expansion', 'hover' or both
-moduleinfo = {'version': '1', 'author': 'Evert Kors',
-              'description': 'MODULE_DESCRIPTION',
-              'module-type': ['expansion', 'hover']}
 
-# config fields that your code expects from the site admin
-moduleconfig = ['cuckoo_api']
+class APIKeyError(RequestException):
+    """Raised if the Cuckoo API returns a 401. This means no or an invalid
+    bearer token was supplied."""
+    pass
+
+
+class CuckooAPI(object):
+
+    def __init__(self, api_url, api_key=""):
+        self.api_key = api_key
+        if not api_url.startswith("http"):
+            api_url = "https://{}".format(api_url)
+
+        self.api_url = api_url
+
+    def _post_api(self, endpoint, files=None, data={}):
+        data.update({
+            "owner": "MISP"
+        })
+
+        try:
+            response = requests.post(
+                urllib.parse.urljoin(self.api_url, endpoint),
+                files=files, data=data,
+                headers={"Authorization: Bearer {}".format(self.api_key)}
+            )
+        except RequestException as e:
+            log.error("Failed to submit sample to Cuckoo Sandbox. %s", e)
+            return None
+
+        if response.status_code == 401:
+            raise APIKeyError("Invalid or no Cuckoo Sandbox API key provided")
+
+        return response.json()
+
+    def create_task(self, filename, fp):
+        response = self._post_api(
+            "/tasks/create/file", files={"file": (filename, fp)}
+        )
+        if not response:
+            return False
+
+        return response["task_id"]
+
+    def create_url(self, url):
+        response = self._post_api(
+            "/tasks/create/url", data={"url": url}
+        )
+        if not response:
+            return False
+
+        return response["task_id"]
 
 
 def handler(q=False):
     if q is False:
         return False
+
     request = json.loads(q)
-    config = request.get('config')
-    if config is None:
-        misperrors['error'] = 'config is missing'
+
+    # See if the API URL was provided. The API key is optional, as it can
+    # be disabled in the Cuckoo API settings.
+    api_url = request["config"].get("api_url")
+    api_key = request["config"].get("api_key", "")
+    if not api_url:
+        misperrors["error"] = "No Cuckoo API URL provided"
         return misperrors
 
-    cuck = config.get('cuckoo_api')
-    if cuck is None:
-        misperrors['error'] = 'cuckoo api url is missing'
-        return misperrors
+    url = request.get("url") or request.get("domain")
+    data = request.get("data")
+    filename = None
+    if data:
+        data = base64.b64decode(data)
 
-    # The url to submit
-    url = request.get('url')
+        if "malware-sample" in request:
+            filename = request.get("malware-sample").split("|", 1)[0]
+            with zipfile.ZipFile(io.BytesIO(data)) as zipf:
+                data = zipf.read(zipf.namelist()[0], pwd=b"infected")
 
-    HEADERS = {"Authorization": "Bearer S4MPL3"}
-
-    urls = [
-        url
-    ]
+        elif "attachment" in request:
+            filename = request.get("attachment")
 
+    cuckoo_api = CuckooAPI(api_url=api_url, api_key=api_key)
+    task_id = None
     try:
-        r = requests.post(
-            "%s/tasks/create/submit" % (cuck),
-            headers=HEADERS,
-            data={"strings": "\n".join(urls)}
-        )
-    except Exception as e:
-        misperrors['error'] = str(e)
+        if url:
+            log.debug("Submitting URL to Cuckoo Sandbox %s", api_url)
+            task_id = cuckoo_api.create_url(url)
+        elif data and filename:
+            log.debug("Submitting file to Cuckoo Sandbox %s", api_url)
+            task_id = cuckoo_api.create_task(
+                filename=filename, fp=io.BytesIO(data)
+            )
+    except APIKeyError as e:
+        misperrors["error"] = "Failed to submit to Cuckoo: {}".format(e)
         return misperrors
 
-    r = {'results': [{'types': "text", 'values': "cool"}]}
-    return r
+    if not task_id:
+        misperrors["error"] = "File or URL submission failed"
+        return misperrors
+
+    return {
+        "results": [
+            {"types": "text", "values": "Cuckoo task id: {}".format(task_id)}
+        ]
+    }
 
 
 def introspection():