From 197dfe2d3f0b97cf3ff855d86e2acb3c51f3ab9a Mon Sep 17 00:00:00 2001
From: Alexander Sennhauser <as@conitrade.com>
Date: Thu, 18 Apr 2024 11:58:19 +0200
Subject: [PATCH] allow payloads to be propagated to new tasks

---
 README.md                                     |  6 ++++++
 karton/archive_extractor/archive_extractor.py | 17 +++++++++++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index d586d52..df62785 100644
--- a/README.md
+++ b/README.md
@@ -69,6 +69,12 @@ max_depth = 5
 max_size = 26214400
 # Maximum number of children files for further analysis
 max_children = 1000
+
+# Specify which payloads are to be propagated to new tasks,
+# takes the form payload_name = payload_persistent
+[archive-extractor-payload-propagation]
+ext_origin_id = True
+ext_source_id = False
 ```
 
 To learn more about configuring your karton services, take a look at [karton configuration docs](https://karton-core.readthedocs.io/en/latest/service_configuration.html)
diff --git a/karton/archive_extractor/archive_extractor.py b/karton/archive_extractor/archive_extractor.py
index e605e7e..f4ae243 100644
--- a/karton/archive_extractor/archive_extractor.py
+++ b/karton/archive_extractor/archive_extractor.py
@@ -52,6 +52,12 @@ def __init__(
             "archive-extractor", "max_children", fallback=1000
         )
 
+        # Payloads to propagate to new tasks
+        self.payloads_to_propagate = {}
+        if self.config.has_section("archive-extractor-payload-propagation"):
+            for k, v in self.config["archive-extractor-payload-propagation"].items():
+                self.payloads_to_propagate[k] = v
+
     def debloat_pe(
         self, filename: str, child_contents: bytes
     ) -> Optional[Tuple[str, bytes]]:
@@ -193,7 +199,7 @@ def process(self, task: Task) -> None:
                 )
                 continue
 
-            task = Task(
+            new_task = Task(
                 headers={
                     "type": "sample",
                     "kind": "raw",
@@ -205,4 +211,11 @@ def process(self, task: Task) -> None:
                     "extraction_level": extraction_level + 1,
                 },
             )
-            self.send_task(task)
+
+            for name, persistent in self.payloads_to_propagate.items():
+                payload = task.get_payload(name)
+                if payload is not None:
+                    self.log.info(f"Propagating payload {name} to new task")
+                    new_task.add_payload(name, payload, persistent=persistent)
+
+            self.send_task(new_task)