Added extensions to drakcore and drakpush. (#205)

* Added extensions to drakcore and drakpush.

* fix

* fix prev. fix

Co-authored-by: Konstanty Cieslinski <konstanty.cieslinski@cert.pl>

drakcore/drakcore/app.py

@@ -58,24 +58,32 @@ def upload():
             sample = Resource("sample", fr.read())
 
     task = Task({"type": "sample", "stage": "recognized", "platform": "win32"})
-    task.add_resource("override_uid", task.uid)
+    task.add_payload("override_uid", task.uid)
 
     # Add analysis timeout to task
     timeout = request.form.get("timeout")
     if timeout:
-        task.add_resource("timeout", int(timeout))
+        task.add_payload("timeout", int(timeout))
 
     # Add filename override to task
-    filename = request.form.get("file_name")
-    if filename:
-        task.add_resource("file_name", filename)
+    if request.form.get("file_name"):
+        filename = request.form.get("file_name")
+    else:
+        filename = request.files['file'].filename
+    task.add_payload("file_name", os.path.splitext(filename)[0])
+
+    # Extract and add extension
+    extension = os.path.splitext(filename)[1][1:]
+    if extension:
+        task.headers['extension'] = extension
 
     # Add startup command to task
     start_command = request.form.get("start_command")
     if start_command:
-        task.add_resource("start_command", start_command)
+        task.add_payload("start_command", start_command)
 
     task.add_resource("sample", sample)
+
     producer.send_task(task)
 
     return jsonify({"task_uid": task.uid})

drakrun/drakrun/drakpush.py

@@ -8,17 +8,30 @@
 def main():
     parser = argparse.ArgumentParser(description='Push sample to the karton')
     parser.add_argument('sample', help='Path to the sample')
+    parser.add_argument('--start_command', help='e.g. start %f, %f will be replaced by file name', required=False)
     args = parser.parse_args()
 
     conf = Config(os.path.join(ETC_DIR, 'config.ini'))
     producer = Producer(conf)
 
+    task = Task({"type": "sample", "stage": "recognized", "platform": "win32"})
+
     with open(args.sample, "rb") as f:
         sample = Resource("sample", f.read())
-
-    task = Task({"type": "sample", "stage": "recognized", "platform": "win32"})
     task.add_resource("sample", sample)
 
+    # Add filename
+    filename = os.path.basename(args.sample)
+    task.add_payload("file_name", os.path.splitext(filename)[0])
+
+    # Extract and add extension
+    extension = os.path.splitext(filename)[1][1:]
+    if extension:
+        task.headers['extension'] = extension
+
+    if args.start_command is not None:
+        task.add_payload("start_command", args.start_command)
+
     producer.send_task(task)
 
 

drakrun/drakrun/main.py

@@ -297,7 +297,7 @@ def process(self):
             extension = 'dll'
         self.log.info("Running file as %s", extension)
 
-        file_name = self.current_task.payload.get("file_name", f"malwar.{extension}")
+        file_name = self.current_task.payload.get("file_name", "malwar") + f".{extension}"
         # Alphanumeric, dot, underscore, dash
         if not re.match(r"^[a-zA-Z0-9\._\-]+$", file_name):
             self.log.error("Filename contains invalid characters")