Skip to content

Commit

Permalink
allow payloads to be propagated to new tasks
Browse files Browse the repository at this point in the history
  • Loading branch information
conitrade-as committed Apr 18, 2024
1 parent a67bc4b commit dcc17bc
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 2 deletions.
6 changes: 6 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
19 changes: 17 additions & 2 deletions karton/archive_extractor/archive_extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand Down Expand Up @@ -193,7 +199,7 @@ def process(self, task: Task) -> None:
)
continue

task = Task(
new_task = Task(
headers={
"type": "sample",
"kind": "raw",
Expand All @@ -205,4 +211,13 @@ 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)

0 comments on commit dcc17bc

Please sign in to comment.