assemblyline-core 4.6.1.dev229__tar.gz → 4.6.1.dev231__tar.gz

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline-core
-Version: 4.6.1.dev229
+Version: 4.6.1.dev231
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

assemblyline_core-4.6.1.dev231/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.6.1.dev231

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/assemblyline_core/dispatching/dispatcher.py

@@ -73,6 +73,7 @@ if TYPE_CHECKING:
     from redis import Redis
 
     from assemblyline.odm.models.file import File
+    from assemblyline.odm.models.config import Config
 
 
 APM_SPAN_TYPE = 'handle_message'
@@ -218,8 +219,18 @@ class TemporaryFileData:
 class SubmissionTask:
     """Dispatcher internal model for submissions"""
 
-    def __init__(self, submission, completed_queue, scheduler, datastore: AssemblylineDatastore, results=None,
-                 file_infos=None, file_tree=None, errors: Optional[Iterable[str]] = None):
+    def __init__(
+        self,
+        submission,
+        completed_queue,
+        scheduler,
+        datastore: AssemblylineDatastore,
+        config: Config,
+        results=None,
+        file_infos=None,
+        file_tree=None,
+        errors: Optional[Iterable[str]] = None,
+    ):
         self.submission: Submission = Submission(submission)
         submitter: Optional[User] = datastore.user.get_if_exists(self.submission.params.submitter)
         self.service_access_control: Optional[str] = None
@@ -227,6 +238,7 @@ class SubmissionTask:
             self.service_access_control = submitter.classification.value
 
         self.completed_queue = None
+
         if completed_queue:
             self.completed_queue = str(completed_queue)
 
@@ -265,9 +277,31 @@ class SubmissionTask:
                     recurse_tree(file_data['children'], depth + 1)
 
             recurse_tree(file_tree, 0)
+            sorted_file_depth = [(k, v) for k, v in sorted(self.file_depth.items(), key=lambda fd: fd[1])]
+        else:
+            sorted_file_depth = [(self.submission.files[0].sha256, 0)]
+
+        for sha256, depth in sorted_file_depth:
+            # populate temporary data to root level files
+            if depth == 0:
+                # Apply initial data parameter
+                temp_key_config = dict(config.submission.default_temporary_keys)
+                temp_key_config.update(config.submission.temporary_keys)
+                temporary_data = TemporaryFileData(sha256, config=temp_key_config)
+                self.temporary_data[sha256] = temporary_data
+                if self.submission.params.initial_data:
+                    try:
+                        for key, value in dict(json.loads(self.submission.params.initial_data)).items():
+                            if len(str(value)) > config.submission.max_temp_data_length:
+                                continue
+                            temporary_data.set_value(key, value)
+
+                    except (ValueError, TypeError):
+                        pass
 
         if results is not None:
             rescan = scheduler.expand_categories(self.submission.params.services.rescan)
+            result_keys = list(results.keys())
 
             # Replay the process of routing files for dispatcher internal state.
             for k, result in results.items():
@@ -282,24 +316,35 @@ class SubmissionTask:
                     self.forbid_for_children(sha256, service_name)
 
             # Replay the process of receiving results for dispatcher internal state
-            for k, result in results.items():
-                sha256, service, _ = k.split('.', 2)
-                if service not in rescan:
-                    extracted = result['response']['extracted']
-                    children: list[str] = [r['sha256'] for r in extracted]
-                    self.register_children(sha256, children)
-                    children_detail: list[tuple[str, str]] = [(r['sha256'], r['parent_relation']) for r in extracted]
-                    self.service_results[(sha256, service)] = ResultSummary(
-                        key=k, drop=result['drop_file'], score=result['result']['score'],
-                        children=children_detail, partial=result.get('partial', False))
-
-                tags = Result(result).scored_tag_dict()
-                for key, tag in tags.items():
-                    if key in self.file_tags[sha256].keys():
-                        # Sum score of already known tags
-                        self.file_tags[sha256][key]['score'] += tag['score']
-                    else:
-                        self.file_tags[sha256][key] = tag
+            # iterate through result based on file depth
+            for sha256, depth in sorted_file_depth:
+                results_to_process = list(filter(lambda k: sha256 in k, result_keys))
+                for result_key in results_to_process:
+                    result = results[result_key]
+                    sha256, service, _ = result_key.split(".", 2)
+
+                    if service not in rescan:
+                        extracted = result["response"]["extracted"]
+                        children: list[str] = [r["sha256"] for r in extracted]
+                        self.register_children(sha256, children)
+                        children_detail: list[tuple[str, str]] = [
+                            (r["sha256"], r["parent_relation"]) for r in extracted
+                        ]
+                        self.service_results[(sha256, service)] = ResultSummary(
+                            key=result_key,
+                            drop=result["drop_file"],
+                            score=result["result"]["score"],
+                            children=children_detail,
+                            partial=result.get("partial", False),
+                        )
+
+                    tags = Result(result).scored_tag_dict()
+                    for key, tag in tags.items():
+                        if key in self.file_tags[sha256].keys():
+                            # Sum score of already known tags
+                            self.file_tags[sha256][key]["score"] += tag["score"]
+                        else:
+                            self.file_tags[sha256][key] = tag
 
         if errors is not None:
             for e in errors:
@@ -334,6 +379,7 @@ class SubmissionTask:
         _parent_map is for dynamic recursion prevention
         temporary_data is for cascading the temp data to children
         """
+
         parent_temp = self.temporary_data[parent]
         for child in children:
             if child not in self.temporary_data:
@@ -706,7 +752,13 @@ class Dispatcher(ThreadedCoreBase):
                 # Start of process dispatcher transaction
                 with apm_span(self.apm_client, 'submission_message'):
                     # This is probably a complete task
-                    task = SubmissionTask(scheduler=self.scheduler, datastore=self.datastore, **message)
+
+                    task = SubmissionTask(
+                        scheduler=self.scheduler,
+                        datastore=self.datastore,
+                        config=self.config,
+                        **message,
+                    )
 
                     # Check the sid table
                     if task.sid in self.bad_sids:
@@ -739,6 +791,7 @@ class Dispatcher(ThreadedCoreBase):
 
         if not self.active_submissions.exists(sid):
             self.log.info("[%s] New submission received", sid)
+
             task.trace('submission_start')
             self.active_submissions.add(sid, {
                 'completed_queue': task.completed_queue,
@@ -760,21 +813,6 @@ class Dispatcher(ThreadedCoreBase):
         if submission.params.quota_item and submission.params.submitter:
             self.log.info(f"[{sid}] Submission counts towards {submission.params.submitter.upper()} quota")
 
-        # Apply initial data parameter
-        temp_key_config = dict(self.config.submission.default_temporary_keys)
-        temp_key_config.update(self.config.submission.temporary_keys)
-        temporary_data = TemporaryFileData(sha256, config=temp_key_config)
-        task.temporary_data[sha256] = temporary_data
-        if submission.params.initial_data:
-            try:
-                for key, value in dict(json.loads(submission.params.initial_data)).items():
-                    if len(str(value)) > self.config.submission.max_temp_data_length:
-                        continue
-                    temporary_data.set_value(key, value)
-
-            except (ValueError, TypeError) as err:
-                self.log.warning(f"[{sid}] could not process initialization data: {err}")
-
         self.tasks[sid] = task
         self._submission_timeouts.set(task.sid, SUBMISSION_TOTAL_TIMEOUT, None)
 
@@ -784,7 +822,10 @@ class Dispatcher(ThreadedCoreBase):
         # Initialize ancestry chain by identifying the root file
         file_info = self.get_fileinfo(task, sha256)
         file_type = file_info.type if file_info else 'NOT_FOUND'
-        temporary_data.local_values['ancestry'] = [[dict(type=file_type, parent_relation="ROOT", sha256=sha256)]]
+
+        task.temporary_data[sha256].local_values["ancestry"] = [
+            [dict(type=file_type, parent_relation="ROOT", sha256=sha256)]
+        ]
 
         # Start the file dispatching
         task.active_files.add(sha256)
@@ -875,7 +916,6 @@ class Dispatcher(ThreadedCoreBase):
             schedule_summary = [list(stage.keys()) for stage in task.file_schedules[sha256]]
             task.trace('schedule_built', sha256=sha256, message=str(schedule_summary))
 
-
         file_info = task.file_info[sha256]
         schedule: list = list(task.file_schedules[sha256])
         deep_scan, ignore_filtering = submission.params.deep_scan, submission.params.ignore_filtering

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/assemblyline_core/ingester/ingester.py

@@ -112,6 +112,7 @@ class IngestTask(odm.Model):
     ingest_id = odm.UUID()  # Ingestion Identifier
     ingest_time = odm.Date(default="NOW")  # Time at which the file was ingested
     notify_time = odm.Optional(odm.Date())  # Time at which the user is notify the submission is finished
+    to_ingest = odm.Boolean(default=False)
 
 
 class Ingester(ThreadedCoreBase):
@@ -250,7 +251,13 @@ class Ingester(ThreadedCoreBase):
                         submission=sub,
                         ingest_id=sub.sid,
                     ))
-                    task.submission.sid = None  # Reset to new random uuid
+
+                    # if this is a new task from imported bundle we want to use the same sid
+                    # because all the submission information are stored in the datastore
+                    # else create a new sid for this submission
+                    if "bundle.source" not in task.submission.metadata:
+                        task.submission.sid = None  # Reset to new random uuid
+
                     # Write all input to the traffic queue
                     self.traffic_queue.publish(SubmissionMessage({
                         'msg': sub,
@@ -920,10 +927,15 @@ class Ingester(ThreadedCoreBase):
         return reason
 
     def submit(self, task: IngestTask):
-        self.submit_client.submit(
-            submission_obj=task.submission,
-            completed_queue=COMPLETE_QUEUE_NAME,
-        )
+
+        if "bundle.source" in task.submission.metadata:
+            self.submit_client.send_bundle_to_dispatch(task.submission, completed_queue=COMPLETE_QUEUE_NAME)
+        else:
+
+            self.submit_client.submit(
+                submission_obj=task.submission,
+                completed_queue=COMPLETE_QUEUE_NAME,
+            )
 
         self.timeout_queue.push(int(now(_max_time)), task.submission.scan_key)
         self.log.info(f"[{task.ingest_id} :: {task.sha256}] Submitted to dispatcher for analysis")

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/assemblyline_core/replay/client.py

@@ -297,11 +297,14 @@ class APIClient(ClientBase):
         self.al_client.bundle.create(id, output=bundle_path, use_alert=use_alert)
 
     def load_bundle(self, bundle_path, min_classification, rescan_services, exist_ok=True, reclassification=None):
-        self.al_client.bundle.import_bundle(bundle_path,
-                                            min_classification=min_classification,
-                                            rescan_services=rescan_services,
-                                            exist_ok=exist_ok,
-                                            reclassification=reclassification)
+        self.al_client.bundle.import_bundle(
+            bundle_path,
+            min_classification=min_classification,
+            rescan_services=rescan_services,
+            exist_ok=exist_ok,
+            reclassification=reclassification,
+            to_ingest=True,  # send submissions to ingester
+        )
 
     def load_json(self, file_path, reclassification=None):
         from assemblyline_client import ClientError
@@ -412,11 +415,14 @@ class DirectClient(ClientBase):
         os.rename(temp_bundle_file, bundle_path)
 
     def load_bundle(self, bundle_path, min_classification, rescan_services, exist_ok=True, reclassification=None):
-        import_bundle(bundle_path,
-                      min_classification=min_classification,
-                      rescan_services=rescan_services,
-                      exist_ok=exist_ok,
-                      reclassification=reclassification)
+        import_bundle(
+            bundle_path,
+            min_classification=min_classification,
+            rescan_services=rescan_services,
+            exist_ok=exist_ok,
+            reclassification=reclassification,
+            to_ingest=True,  # send submissions to ingester
+        )
 
     def load_json(self, file_path, reclassification=None):
         # We're assuming all JSON that loaded has an "enabled" field
@@ -442,7 +448,6 @@ class DirectClient(ClientBase):
                         else:
                             raise
 
-
                 if collection == "workflow":
                     # If there has been any edits by another user, then preserve the enabled state
                     # Otherwise, the workflow will be synchronized with the origin system

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/assemblyline_core/submission_client.py

@@ -37,6 +37,8 @@ from assemblyline.odm.models.result import Result
 from assemblyline.odm.models.submission import File, Submission
 from assemblyline.odm.models.config import Config
 from assemblyline_core.dispatching.client import DispatchClient
+from assemblyline_core.ingester.constants import INGEST_QUEUE_NAME
+from assemblyline.remote.datatypes.queues.named import NamedQueue
 
 Classification = forge.get_classification()
 SECONDS_PER_DAY = 24 * 60 * 60
@@ -72,6 +74,7 @@ class SubmissionClient:
 
         # A client for interacting with the dispatcher
         self.dispatcher = DispatchClient(datastore, redis)
+        self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, redis)
 
     def __enter__(self):
         return self
@@ -84,8 +87,16 @@ class SubmissionClient:
             self.identify.stop()
 
     @elasticapm.capture_span(span_type='submission_client')
-    def rescan(self, submission: Submission, results: Dict[str, Result], file_infos: Dict[str, FileInfo],
-               file_tree, errors: List[str], rescan_services: List[str]):
+    def rescan(
+        self,
+        submission,
+        results: Dict[str, Result],
+        file_infos: Dict[str, FileInfo],
+        file_tree: dict,
+        errors: List[str],
+        rescan_services: List[str],
+        to_ingest: bool = False,
+    ):
         """
         Rescan a submission started on another system.
         """
@@ -114,8 +125,29 @@ class SubmissionClient:
         self.datastore.submission.save(submission_obj.sid, submission_obj)
 
         # Dispatch the submission
-        self.log.debug("Submission complete. Dispatching: %s", submission_obj.sid)
-        self.dispatcher.dispatch_bundle(submission_obj, results, file_infos, file_tree, errors)
+        if to_ingest:
+            self.log.debug("Submission complete. Submission sent to ingester: %s", submission_obj.sid)
+
+            submission_obj = SubmissionObject(
+                {
+                    "sid": submission["sid"],
+                    "files": submission.get("files", []),
+                    "metadata": submission.get("metadata", {}),
+                    "params": submission.get("params", {}),
+                    "notification": submission.get("notification", {}),
+                    "scan_key": submission.get("scan_key", None),
+                    "errors": errors,
+                    "file_infos": file_infos,
+                    "file_tree": file_tree,
+                    "results": results,
+                }
+            ).as_primitives()
+
+            self.ingest_queue.push(submission_obj)
+
+        else:
+            self.log.debug("Submission complete. Dispatching: %s", submission_obj.sid)
+            self.dispatcher.dispatch_bundle(submission_obj, results, file_infos, file_tree, errors)
 
         return submission
 
@@ -252,3 +284,22 @@ class SubmissionClient:
             if extracted_path:
                 if os.path.exists(extracted_path):
                     os.unlink(extracted_path)
+
+    @elasticapm.capture_span(span_type="submission_client")
+    def send_bundle_to_dispatch(
+        self,
+        submission_obj: SubmissionObject,
+        completed_queue: str = None,
+    ):
+
+        sid = submission_obj.sid
+        submission = self.datastore.submission.get(sid)
+
+        self.dispatcher.dispatch_bundle(
+            submission=submission,
+            results=submission_obj.results,
+            file_infos=submission_obj.file_infos,
+            file_tree=submission_obj.file_tree,
+            errors=submission_obj.errors,
+            completed_queue=completed_queue,
+        )

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/assemblyline_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline-core
-Version: 4.6.1.dev229
+Version: 4.6.1.dev231
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/test/test_dispatcher.py

@@ -5,7 +5,7 @@ from unittest import mock
 
 import pytest
 from assemblyline_core.dispatching.client import DISPATCH_RESULT_QUEUE, DispatchClient
-from assemblyline_core.dispatching.dispatcher import Dispatcher, ServiceTask, Submission
+from assemblyline_core.dispatching.dispatcher import Dispatcher, ServiceTask, Submission, SubmissionTask
 from assemblyline_core.dispatching.schedules import Scheduler as RealScheduler
 
 # noinspection PyUnresolvedReferences
@@ -18,8 +18,9 @@ from assemblyline.common.metrics import MetricsFactory
 from assemblyline.odm import models
 from assemblyline.odm.models.error import Error
 from assemblyline.odm.models.file import File
-from assemblyline.odm.models.result import Result
+from assemblyline.odm.models.result import Result, File as ResponseFile
 from assemblyline.odm.models.user import User
+from assemblyline.odm.models.submission import Submission as SubmissionModel
 from assemblyline.odm.randomizer import (
     get_random_hash,
     random_minimal_obj,
@@ -454,3 +455,148 @@ def test_prevent_result_overwrite(clean_redis, clean_datastore):
     msg_result_key = message['result_summary']['key']
 
     assert msg_result_key != result_key
+
+
+def test_create_submission_task(datastore_connection, config, filestore, clean_redis):
+
+    # create a test submission
+    scheduler = Scheduler(datastore_connection, config, clean_redis)
+    submission = random_model_obj(SubmissionModel)
+    sid = get_random_hash(64)
+    submission.sid = sid
+
+    # create files and results with file tree:
+    #     root_1
+    #       |
+    #     middle_1
+    #     |     |
+    #   leaf_1 leaf_2
+
+    leaf_1_sha = get_random_hash(64)
+    leaf_2_sha = get_random_hash(64)
+    middle_1_sha = get_random_hash(64)
+    root_1_sha = get_random_hash(64)
+
+    file_shas = [root_1_sha, middle_1_sha, leaf_2_sha, leaf_1_sha]
+
+    files = [
+        {"name": ["root_1"], "size": 1, "sha256": root_1_sha},
+        {"name": ["middle_1"], "size": 1, "sha256": middle_1_sha},
+        {"name": ["leaf-2"], "size": 1, "sha256": leaf_2_sha},
+        {
+            "name": ["leaf-1"],
+            "size": 1,
+            "sha256": leaf_1_sha,
+        },
+    ]
+
+    submission.files = files
+
+    leaf_1 = {
+        "name": ["leaf-1"],
+        "type": "test_type",
+        "sha256": leaf_1_sha,
+        "children": {},
+        "truncated": False,
+        "score": 0,
+    }
+
+    leaf_2 = {
+        "name": ["leaf-2"],
+        "type": "test_type",
+        "sha256": leaf_2_sha,
+        "children": {},
+        "truncated": False,
+        "score": 0,
+    }
+
+    middle_1 = {
+        "name": ["middle-1"],
+        "type": "test_type",
+        "sha256": middle_1_sha,
+        "children": {leaf_1_sha: leaf_1, leaf_2_sha: leaf_2},
+        "truncated": False,
+        "score": 0,
+    }
+
+    root_1 = {
+        root_1_sha: {
+            "name": ["root-1"],
+            "type": "test_type",
+            "sha256": root_1_sha,
+            "children": {middle_1_sha: middle_1},
+            "truncated": False,
+            "score": 0,
+        }
+    }
+
+    file_infos = {}
+    errors = []
+    for file_sha in file_shas:
+        file_infos[file_sha] = random_model_obj(File).as_primitives()
+        errors.append(f"{file_sha}.serviceName.v0_0_0.c0.e0")
+
+    submission.errors = errors
+
+    results = {}
+
+    result_root_1_key = f"{root_1_sha}.serviceName.v0_0_0.c0"
+    results[result_root_1_key] = random_model_obj(Result).as_primitives()
+    results[result_root_1_key]["sha256"] = root_1_sha
+    result_file_middle = random_model_obj(ResponseFile)
+    result_file_middle.sha256 = middle_1_sha
+    results[result_root_1_key]["response"]["extracted"] = [result_file_middle]
+
+    result_middle_1_key = f"{middle_1_sha}.serviceName.v0_0_0.c0"
+    results[result_middle_1_key] = random_model_obj(Result).as_primitives()
+    results[result_middle_1_key]["sha256"] = middle_1_sha
+    result_file_leaf1 = random_model_obj(ResponseFile)
+    result_file_leaf1.sha256 = leaf_1_sha
+    result_file_leaf2 = random_model_obj(ResponseFile)
+    result_file_leaf2.sha256 = leaf_2_sha
+
+    results[result_middle_1_key]["response"]["extracted"] = [result_file_leaf1, result_file_leaf2]
+
+    result_leaf_1_key = f"{leaf_1_sha}.serviceName.v0_0_0.c0"
+    results[result_leaf_1_key] = random_model_obj(Result).as_primitives()
+    results[result_leaf_1_key]["sha256"] = leaf_1_sha
+    results[result_leaf_1_key]["response"]["extracted"] = []
+
+    result_leaf_2_key = f"{leaf_2_sha}.serviceName.v0_0_0.c0"
+    results[result_leaf_2_key] = random_model_obj(Result).as_primitives()
+    results[result_leaf_2_key]["sha256"] = leaf_2_sha
+    results[result_leaf_2_key]["response"]["extracted"] = []
+
+    submission.results = results.keys()
+
+    # Create a message from submission queue
+    submission_message = {
+        "submission": submission.as_primitives(),
+        "results": results,
+        "file_infos": file_infos,
+        "file_tree": root_1,
+        "errors": errors,
+        "completed_queue": "test",
+    }
+
+    task = SubmissionTask(
+        scheduler=scheduler,
+        datastore=datastore_connection,
+        config=config,
+        **submission_message,
+    )
+
+    assert task.sid == sid
+
+    parent_map = {middle_1_sha: {root_1_sha}, leaf_1_sha: {middle_1_sha}, leaf_2_sha: {middle_1_sha}}
+
+    assert task._parent_map.keys() == parent_map.keys()
+    for key, val in parent_map.items():
+        assert val == task._parent_map[key]
+
+    file_depth = {root_1_sha: 0, middle_1_sha: 1, leaf_1_sha: 2, leaf_2_sha: 2}
+    assert task.file_depth.keys() == file_depth.keys()
+    for key, val in file_depth.items():
+        assert val == task.file_depth[key]
+
+    assert task.temporary_data.keys() == file_depth.keys()

{assemblyline_core-4.6.1.dev229 → assemblyline_core-4.6.1.dev231}/test/test_simulation.py

@@ -12,6 +12,7 @@ import threading
 import logging
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Any
+from assemblyline_core.submission_client import SubmissionClient
 
 import pytest
 
@@ -29,6 +30,7 @@ from assemblyline.odm.models.user import User
 from assemblyline.odm.randomizer import random_model_obj
 from assemblyline.odm.messages.submission import Submission as SubmissionInput
 from assemblyline.remote.datatypes.queues.named import NamedQueue
+from assemblyline.odm.random_data import create_submission
 
 import assemblyline_core
 from assemblyline_core.plumber.run_plumber import Plumber
@@ -84,7 +86,6 @@ class MockService(ServerBase):
             self.log.info(f"{self.service_name} has received a job {task.sid}")
 
             file = self.filestore.get(task.fileinfo.sha256)
-
             instructions = json.loads(file)
             instructions = instructions.get(self.service_name, {})
             self.log.info(f"{self.service_name} following instruction: {instructions}")
@@ -1326,7 +1327,7 @@ def test_final_partial(core: CoreSession, metrics):
             name='abc123'
         )]
     )).as_primitives())
- 
+
     # Wait until both of the services have started (so service b doesn't get the temp data a produces on its first run)
     start = time.time()
     while sum(s.hits.get(sha, 0) for s in core.services) != 2:
@@ -1342,7 +1343,7 @@ def test_final_partial(core: CoreSession, metrics):
     while sum(s.finish.get(sha, 0) for s in core_a) < 1:
         if time.time() - start > RESPONSE_TIMEOUT:
             pytest.fail()
-        time.sleep(0.01)    
+        time.sleep(0.01)
 
     # Let b finish, it should produce a partial result then rerun right away
     core_b.local_lock.set()
@@ -1453,3 +1454,48 @@ def test_complex_extracted(core: CoreSession, metrics):
         if result.partial:
             partial_results += 1
     assert partial_results == 0, 'partial_results'
+
+
+@pytest.mark.parametrize("to_ingest,submission_metadata", [(False, {}), (True, {}),(True, {"bundle.source": "test_source"})])
+def test_rescan_submission(core: CoreSession, metrics: MetricsCounter, to_ingest, submission_metadata):
+    # when a submission is rescan, the submission information should be stored in the database
+    # file_tree, results, errors are passed to ingestion queue first and then sent to dispatcher
+
+    # Create a new submission
+    submission = create_submission(core.ds, core.filestore, metadata=submission_metadata)
+
+    file_hashes = [x[:64] for x in submission["results"]]
+    file_hashes.extend([x[:64] for x in submission["errors"]])
+    file_hashes.extend([f["sha256"] for f in submission["files"]])
+    file_tree = core.ds.get_or_create_file_tree(submission, core.config.submission.max_extraction_depth)["tree"]
+    file_infos = core.ds.file.multiget(list(set(file_hashes)), as_dictionary=True, as_obj=False)
+    rescan_services = ["core-aaa"]
+    result_ids = list(filter(lambda x: not x.endswith(".e"), submission.results))
+    results = core.ds.result.multiget(result_ids, as_dictionary=True, as_obj=False)
+
+    with SubmissionClient(datastore=core.ds, filestore=core.filestore, config=core.config, identify=None) as sc:
+        sc.ingest_queue = core.ingest_queue
+        sc.rescan(
+            submission.as_primitives(),
+            results,
+            file_infos,
+            file_tree,
+            submission["errors"],
+            rescan_services,
+            to_ingest=to_ingest,
+        )
+
+    if to_ingest:
+
+        metrics.expect("ingester", "submissions_ingested", 1)
+        metrics.expect("dispatcher", "submissions_completed", 1)
+        metrics.expect("ingester", "submissions_completed", 1)
+        if submission_metadata:
+            metrics.expect("dispatcher", "files_completed", len(submission["files"]))
+
+    else:
+        # when to_ingest is false, the submission should be route to dispatcher directly
+        metrics.expect("dispatcher", "submissions_completed", 1)
+        metrics.expect("dispatcher", "files_completed", 1)
+        metrics.expect("ingester", "submissions_ingested", 0)
+        metrics.expect("ingester", "submissions_completed", 0)

assemblyline_core-4.6.1.dev229/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.6.1.dev229