Work on passing execute field to karton

mwdb/core/karton.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Optional, Dict, Any
 
 from flask import g
 from karton.core import Config as KartonConfig
@@ -9,6 +9,10 @@
 
 from ..version import app_version
 from .config import app_config
+from ..model.file import File
+from ..model.config import Config
+from ..model.blob import TextBlob
+from ..model.object import Object
 
 logger = logging.getLogger("mwdb.karton")
 
@@ -35,27 +39,45 @@ def get_karton_producer() -> Optional[Producer]:
     return karton_producer
 
 
-def send_file_to_karton(file) -> str:
+def prepare_headers(obj: Object, arguments: Dict[str, Any]) -> Dict[str, Any]:
+    """Prepare headers to use when submitting this object to Karton.
+    Takes into account object arguments to this analysis, some attributes,
+    and share_3rd_party field (in this order of precedence).
+    """
+    headers={
+        "share_3rd_party": obj.share_3rd_party,
+    }
+
+    passthrough_attributes = ["execute"]
+    for attribute in obj.attributes:
+        if attribute.key in passthrough_attributes:
+            headers[attribute.key] = attribute.value
+
+    for attribute, value in arguments.items():
+        headers[attribute] = value
+
+    return headers
+
+
+def send_file_to_karton(file: File, arguments: Dict[str, Any]) -> str:
     producer = get_karton_producer()
 
     if producer is None:
         raise RuntimeError("Karton is not enabled or failed to load properly")
 
+    feed_quality = g.auth_user.feed_quality
+    headers_persistent=prepare_headers(file, arguments)
+    headers_persistent["quality"] = feed_quality
+    task_priority = TaskPriority.NORMAL if feed_quality == "high" else TaskPriority.LOW
+
     file_stream = file.open()
     try:
-        feed_quality = g.auth_user.feed_quality
-        task_priority = (
-            TaskPriority.NORMAL if feed_quality == "high" else TaskPriority.LOW
-        )
         task = Task(
             headers={
                 "type": "sample",
                 "kind": "raw",
             },
-            headers_persistent={
-                "quality": feed_quality,
-                "share_3rd_party": file.share_3rd_party,
-            },
+            headers_persistent=headers_persistent,
             payload={
                 "sample": Resource(file.file_name, fd=file_stream, sha256=file.sha256),
                 "attributes": file.get_attributes(
@@ -72,7 +94,7 @@ def send_file_to_karton(file) -> str:
     return task.root_uid
 
 
-def send_config_to_karton(config) -> str:
+def send_config_to_karton(config: Config, arguments: Dict[str, Any]) -> str:
     producer = get_karton_producer()
 
     if producer is None:
@@ -84,9 +106,7 @@ def send_config_to_karton(config) -> str:
             "kind": config.config_type,
             "family": config.family,
         },
-        headers_persistent={
-            "share_3rd_party": config.share_3rd_party,
-        },
+        headers_persistent=prepare_headers(config, arguments),
         payload={
             "config": config.cfg,
             "dhash": config.dhash,
@@ -99,7 +119,7 @@ def send_config_to_karton(config) -> str:
     return task.root_uid
 
 
-def send_blob_to_karton(blob) -> str:
+def send_blob_to_karton(blob: TextBlob, arguments: Dict[str, Any]) -> str:
     producer = get_karton_producer()
 
     if producer is None:
@@ -110,9 +130,7 @@ def send_blob_to_karton(blob) -> str:
             "type": "blob",
             "kind": blob.blob_type,
         },
-        headers_persistent={
-            "share_3rd_party": blob.share_3rd_party,
-        },
+        headers_persistent=prepare_headers(blob, arguments),
         payload={
             "content": blob.content,
             "dhash": blob.dhash,

mwdb/model/blob.py

@@ -1,5 +1,7 @@
 import datetime
 import hashlib
+from typing import Dict, Any
+from mwdb.model import attribute
 
 from sqlalchemy.ext.hybrid import hybrid_property
 
@@ -63,5 +65,5 @@ def get_or_create(
 
         return blob_obj, is_new
 
-    def _send_to_karton(self):
-        return send_blob_to_karton(self)
+    def _send_to_karton(self, arguments: Dict[str, Any]):
+        return send_blob_to_karton(self, arguments)

mwdb/model/config.py

@@ -1,5 +1,6 @@
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.ext.hybrid import hybrid_property
+from typing import Dict, Any
 
 from mwdb.core.karton import send_config_to_karton
 from mwdb.core.util import config_decode, config_dhash, config_encode
@@ -57,8 +58,8 @@ def get_or_create(
             tags=tags,
         )
 
-    def _send_to_karton(self):
-        return send_config_to_karton(self)
+    def _send_to_karton(self, arguments: Dict[str, Any]):
+        return send_config_to_karton(self, arguments)
 
 
 # Compatibility reasons

mwdb/model/file.py

@@ -3,6 +3,7 @@
 import os
 import shutil
 import tempfile
+from typing import Dict, Any
 
 import pyzipper
 from Cryptodome.Util.strxor import strxor_c
@@ -356,5 +357,5 @@ def get_by_download_token(download_token):
             return None
         return File.get(download_req["identifier"]).first()
 
-    def _send_to_karton(self):
-        return send_file_to_karton(self)
+    def _send_to_karton(self, arguments: Dict[str, Any]):
+        return send_file_to_karton(self, arguments)

mwdb/model/object.py

@@ -1,6 +1,6 @@
 import datetime
 from collections import namedtuple
-from typing import Optional
+from typing import Optional, Dict, Any
 from uuid import UUID
 
 from flask import g
@@ -988,14 +988,14 @@ def get_shares(self):
             .order_by(ObjectPermission.access_time.asc())
         ).all()
 
-    def _send_to_karton(self):
+    def _send_to_karton(self, arguments: Dict[str, Any]):
         raise NotImplementedError
 
-    def spawn_analysis(self, arguments, commit=True):
+    def spawn_analysis(self, arguments: Dict[str, Any], commit=True):
         """
         Spawns new KartonAnalysis for this object
         """
-        analysis_id = self._send_to_karton()
+        analysis_id = self._send_to_karton(arguments)
         analysis = KartonAnalysis.create(
             analysis_id=UUID(analysis_id),
             initial_object=self,