assemblyline-core 4.5.0.26__tar.gz → 4.5.0.28__tar.gz

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/PKG-INFO

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

assemblyline-core-4.5.0.28/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.0.28

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/archiver/run_archiver.py

@@ -64,7 +64,7 @@ class Archiver(ServerBase):
             try:
                 if len(message) == 3:
                     archive_type, type_id, delete_after = message
-                    metadata = None
+                    metadata = {}
                     use_alternate_dtl = False
                 elif len(message) == 4:
                     archive_type, type_id, delete_after, metadata = message
@@ -90,8 +90,9 @@ class Archiver(ServerBase):
                         submission, version = self.datastore.submission.get_if_exists(type_id, version=True)
 
                         # If we have metadata passed in the message, we need to apply it before archiving the submission
-                        if metadata and self.config.core.archiver.use_metadata:
-                            submission.metadata.update({f"archive.{k}": v for k, v in metadata.items()})
+                        if metadata and self.config.submission.metadata.archive:
+                            submission.metadata.update({k: v for k, v in metadata.items()
+                                                        if k not in submission.metadata})
                             self.datastore.submission.save(type_id, submission, version=version)
 
                         break

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/dispatching/client.py

@@ -279,6 +279,17 @@ class DispatchClient:
                     else:
                         result.expiry_ts = None
                 try:
+                    if self.ds.result.exists(result_key):
+                        # A result already exists for this key
+                        # Regenerate entire result key based on result and modified task (ignore caching)
+                        task.ignore_cache = True
+                        result_key = Result.help_build_key(sha256=task.fileinfo.sha256,
+                                                           service_name=result.response.service_name,
+                                                           service_version=result.response.service_version,
+                                                           service_tool_version=result.response.service_tool_version,
+                                                           is_empty=False,
+                                                           task=task)
+                        version = "create"
                     self.ds.result.save(result_key, result, version=version)
                     break
                 except VersionConflictException as vce:

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/expiry/run_expiry.py

@@ -4,38 +4,48 @@ from __future__ import annotations
 import concurrent.futures
 import threading
 import functools
-import elasticapm
 import time
-
 from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, Future, as_completed
 from concurrent.futures.process import BrokenProcessPool
-from datemath import dm
 from typing import Callable, Optional, TYPE_CHECKING
 
-from assemblyline.common.isotime import epoch_to_iso, now_as_iso
-from assemblyline.datastore.collection import Index
+import elasticapm
+from datemath import dm
+
 from assemblyline_core.server_base import ServerBase
 from assemblyline_core.dispatching.dispatcher import BAD_SID_HASH
 from assemblyline.common import forge
+from assemblyline.common.isotime import epoch_to_iso, now_as_iso
 from assemblyline.common.metrics import MetricsFactory
 from assemblyline.filestore import FileStore
 from assemblyline.odm.messages.expiry_heartbeat import Metrics
 from assemblyline.remote.datatypes import get_client
+from assemblyline.datastore.collection import Index
 from assemblyline.remote.datatypes.set import Set
 
 if TYPE_CHECKING:
     from assemblyline.datastore.collection import ESCollection
 
 
-def file_delete_worker(logger, filestore_urls, file_batch) -> list[str]:
+def file_delete_worker(logger, filestore_urls, file_batch, archive_filestore_urls=None) -> list[tuple[str, bool]]:
     try:
         filestore = FileStore(*filestore_urls)
-
-        def filestore_delete(sha256: str) -> Optional[str]:
-            filestore.delete(sha256)
-            if not filestore.exists(sha256):
-                return sha256
-            return None
+        if archive_filestore_urls and filestore_urls != archive_filestore_urls:
+            archivestore = FileStore(*archive_filestore_urls)
+        else:
+            archivestore = filestore
+
+        def filestore_delete(item: tuple[str, bool]) -> tuple[Optional[str], Optional[bool]]:
+            sha256, from_archive = item
+            if from_archive:
+                archivestore.delete(sha256)
+                if not archivestore.exists(sha256):
+                    return sha256, True
+            else:
+                filestore.delete(sha256)
+                if not filestore.exists(sha256):
+                    return sha256, False
+            return None, None
 
         return _file_delete_worker(logger, filestore_delete, file_batch)
 
@@ -44,8 +54,11 @@ def file_delete_worker(logger, filestore_urls, file_batch) -> list[str]:
     return []
 
 
-def _file_delete_worker(logger, delete_action: Callable[[str], Optional[str]], file_batch) -> list[str]:
-    finished_files: list[str] = []
+ActionSignature = Callable[[tuple[str, bool]], tuple[Optional[str], Optional[bool]]]
+
+
+def _file_delete_worker(logger, delete_action: ActionSignature, file_batch) -> list[tuple[str, bool]]:
+    finished_files: list[tuple[str, bool]] = []
     try:
         futures = []
 
@@ -55,9 +68,9 @@ def _file_delete_worker(logger, delete_action: Callable[[str], Optional[str]], f
 
             for future in as_completed(futures):
                 try:
-                    erased_name = future.result()
-                    if erased_name:
-                        finished_files.append(erased_name)
+                    erased_name, from_archive = future.result()
+                    if erased_name and from_archive is not None:
+                        finished_files.append((erased_name, from_archive))
                 except Exception as error:
                     logger.exception("Error in filestore worker: " + str(error))
 
@@ -67,17 +80,29 @@ def _file_delete_worker(logger, delete_action: Callable[[str], Optional[str]], f
 
 
 class ExpiryManager(ServerBase):
-    def __init__(self, redis_persist=None):
-        self.config = forge.get_config()
+    def __init__(self, redis_persist=None, datastore=None, filestore=None, config=None, classification=None):
+        self.config = config or forge.get_config()
 
         super().__init__('assemblyline.expiry', shutdown_timeout=self.config.core.expiry.sleep_time + 5)
-        self.datastore = forge.get_datastore(config=self.config)
-        self.filestore = forge.get_filestore(config=self.config)
-        self.classification = forge.get_classification()
+
+        # Set Archive related configs
+        if self.config.datastore.archive.enabled:
+            self.archive_access = True
+            self.index_type = Index.HOT_AND_ARCHIVE
+        else:
+            self.archive_access = False
+            self.index_type = Index.HOT
+
+        self.datastore = datastore or forge.get_datastore(config=self.config, archive_access=self.archive_access)
+        self.filestore = filestore or forge.get_filestore(config=self.config)
+        self.classification = classification or forge.get_classification()
         self.expirable_collections: list[ESCollection] = []
         self.counter = MetricsFactory('expiry', Metrics)
         self.file_delete_worker = ProcessPoolExecutor(self.config.core.expiry.delete_workers)
-        self.same_storage = self.config.filestore.storage == self.config.filestore.archive
+        if self.config.filestore.archive:
+            self.same_storage = self.config.filestore.storage == self.config.filestore.archive
+        else:
+            self.same_storage = True
         self.current_submission_cleanup = set()
 
         self.redis_persist = redis_persist or get_client(
@@ -127,14 +152,15 @@ class ExpiryManager(ServerBase):
     def filestore_delete(self, file_batch, _):
         return self.file_delete_worker.submit(file_delete_worker, logger=self.log,
                                               filestore_urls=list(self.config.filestore.storage),
-                                              file_batch=file_batch)
+                                              file_batch=file_batch,
+                                              archive_filestore_urls=list(self.config.filestore.archive))
 
     def cachestore_delete(self, file_batch, _):
         return self.file_delete_worker.submit(file_delete_worker, logger=self.log,
                                               filestore_urls=list(self.config.filestore.cache),
                                               file_batch=file_batch)
 
-    def _finish_delete(self, collection: ESCollection, task: Future, expire_only: list[str]):
+    def _finish_delete(self, collection: ESCollection, task: Future, expire_only: list[tuple[str, bool]]):
         # Wait until the worker process finishes deleting files
         file_list: list[str] = []
         while self.running:
@@ -145,25 +171,41 @@ class ExpiryManager(ServerBase):
             except concurrent.futures.TimeoutError:
                 pass
 
-        file_list.extend(expire_only)
-
-        # build a batch delete job for all the removed files
-        bulk = collection.get_bulk_plan()
-        for sha256 in file_list:
-            bulk.add_delete_operation(sha256)
-
-        if len(file_list) > 0:
+        if file_list:
             self.log.info(f'[{collection.name}] Deleted associated files from the '
                           f'{"cachestore" if "cache" in collection.name else "filestore"}...')
-            collection.bulk(bulk)
-            self.counter.increment(f'{collection.name}', increment_by=len(file_list))
-            self.log.info(f"[{collection.name}] Deleted {len(file_list)} items from the datastore...")
         else:
+            self.log.info(f'[{collection.name}] Nothing was deleted from the '
+                          f'{"cachestore" if "cache" in collection.name else "filestore"}...')
+
+        # From the files to be deleted, check which are from the hot index
+        hot_file_list = [x[0] for x in file_list if not x[1]]
+        hot_file_list.extend([x[0] for x in expire_only if not x[1]])
+
+        # From the files to be deleted, check which are from the archive index
+        archive_file_list = [x[0] for x in file_list if x[1]]
+        archive_file_list.extend([x[0] for x in expire_only if x[1]])
+
+        for cur_file_list, index_type in [(hot_file_list, Index.HOT), (archive_file_list, Index.ARCHIVE)]:
+            if not cur_file_list:
+                # Nothing to delete from this index type
+                continue
+
+            # build a batch delete job for all the removed files
+            bulk = collection.get_bulk_plan(index_type=index_type)
+            for sha256 in cur_file_list:
+                bulk.add_delete_operation(sha256)
+
+            collection.bulk(bulk)
+            self.counter.increment(f'{collection.name}', increment_by=len(cur_file_list))
+            self.log.info(f"[{collection.name}] Deleted {len(cur_file_list)} items from the datastore...")
+
+        if not hot_file_list and not archive_file_list:
             self.log.warning(f'[{collection.name}] Expiry unable to clean up any of the files in filestore.')
 
-    def _simple_delete(self, collection, delete_query, number_to_delete):
+    def _simple_delete(self, collection: ESCollection, delete_query, number_to_delete):
         self.heartbeat()
-        collection.delete_by_query(delete_query)
+        collection.delete_by_query(delete_query, index_type=self.index_type)
         self.counter.increment(f'{collection.name}', increment_by=number_to_delete)
         self.log.info(f"[{collection.name}] Deleted {number_to_delete} items from the datastore...")
 
@@ -195,22 +237,45 @@ class ExpiryManager(ServerBase):
         # check if we are dealing with an index that needs file cleanup
         if self.config.core.expiry.delete_storage and collection.name in self.fs_hashmap:
             # Delete associated files
-            delete_objects: list[str] = []
-            for item in collection.stream_search(delete_query, fl='id', as_obj=False):
+            delete_objects: list[tuple[str, bool]] = []
+            for item in collection.stream_search(
+                    delete_query, fl='id,from_archive', as_obj=False, index_type=self.index_type):
                 self.heartbeat()
-                delete_objects.append(item['id'])
+                delete_objects.append((item['id'], item.get('from_archive', False)))
 
             # Filter archived documents if archive filestore is the same as the filestore
-            expire_only = []
-            if self.same_storage and self.config.datastore.archive.enabled and collection.name == 'file':
-                archived_files = self.datastore.file.multiexists(delete_objects, index_type=Index.ARCHIVE)
-                delete_objects = [k for k, v in archived_files.items() if not v]
-                expire_only = [k for k, v in archived_files.items() if v]
+            expire_only: list[tuple[str, bool]] = []
+            if self.same_storage and self.archive_access and collection.name == 'file':
+                # Separate hot and archive files
+                delete_from_archive = [i[0] for i in delete_objects if i[1]]
+                delete_from_hot = [i[0] for i in delete_objects if not i[1]]
+
+                # Check for overlap
+                overlap = set(delete_from_archive).intersection(set(delete_from_hot))
+                delete_from_archive = list(set(delete_from_archive)-overlap)
+                delete_from_hot = list(set(delete_from_hot)-overlap)
+
+                # Create the original delete_object form the overlap
+                delete_objects = [(k, False) for k in overlap]
+                delete_objects.extend([(k, True) for k in overlap])
+
+                if delete_from_hot:
+                    # Check hot objects to delete if they are in archive
+                    archived_files = self.datastore.file.multiexists(delete_from_hot, index_type=Index.ARCHIVE)
+                    delete_objects.extend([(k, False) for k, v in archived_files.items() if not v])
+                    expire_only.extend([(k, False) for k, v in archived_files.items() if v])
+
+                if delete_from_archive:
+                    # Check hot objects to delete if they are in archive
+                    hot_files = self.datastore.file.multiexists(delete_from_archive, index_type=Index.HOT)
+                    delete_objects.extend([(k, True) for k, v in hot_files.items() if not v])
+                    expire_only.extend([(k, True) for k, v in hot_files.items() if v])
 
             delete_tasks = self.fs_hashmap[collection.name](delete_objects, final_date)
 
             # Proceed with deletion, but only after all the scheduled deletes for this
-            self.log.info(f"[{collection.name}] Scheduled {len(delete_objects)}/{number_to_delete} files to be removed")
+            self.log.info(f"[{collection.name}] Scheduled {len(delete_objects)}/{number_to_delete} files to be "
+                          f"removed from the {'cachestore' if 'cache' in collection.name else 'filestore'}")
             self._finish_delete(collection, delete_tasks, expire_only)
 
         else:
@@ -264,7 +329,7 @@ class ExpiryManager(ServerBase):
            will be affected in between start date and the date found"""
         rows = collection.search(f"expiry_ts: {{{start} TO {final_date}]", rows=1,
                                  offset=self.expiry_size - 1, sort='expiry_ts asc',
-                                 as_obj=False, fl='expiry_ts')
+                                 as_obj=False, fl='expiry_ts', index_type=self.index_type)
         if rows['items']:
             return rows['items'][0]['expiry_ts'], self.expiry_size
         return final_date, rows['total']

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/scaler/controllers/kubernetes_ctl.py

@@ -26,7 +26,7 @@ from kubernetes.client import V1Deployment, V1DeploymentSpec, V1PodTemplateSpec,
     V1PersistentVolumeClaimSpec, V1NetworkPolicy, V1NetworkPolicySpec, V1NetworkPolicyEgressRule, V1NetworkPolicyPeer, \
     V1NetworkPolicyIngressRule, V1Secret, V1SecretVolumeSource, V1LocalObjectReference, V1Service, \
     V1ServiceSpec, V1ServicePort, V1PodSecurityContext, V1Probe, V1ExecAction, V1SecurityContext, \
-    V1Affinity, V1NodeAffinity, V1NodeSelector, V1NodeSelectorTerm, V1NodeSelectorRequirement
+    V1Affinity, V1NodeAffinity, V1NodeSelector, V1NodeSelectorTerm, V1NodeSelectorRequirement, V1Toleration
 from kubernetes.client.rest import ApiException
 from assemblyline.odm.models.service import DependencyConfig, DockerConfig, PersistentVolume
 
@@ -241,7 +241,7 @@ def parse_cpu(string: str) -> float:
 class KubernetesController(ControllerInterface):
     def __init__(self, logger, namespace: str, prefix: str, priority: str, dependency_priority: str,
                  cpu_reservation: float, linux_node_selector: Selector, labels=None, log_level="INFO", core_env={},
-                 default_service_account=None, cluster_pod_list=True):
+                 default_service_account=None, cluster_pod_list=True, default_service_tolerations = []):
         # Try loading a kubernetes connection from either the fact that we are running
         # inside of a cluster, or have a config file that tells us how
         try:
@@ -285,6 +285,7 @@ class KubernetesController(ControllerInterface):
         self._service_limited_env: dict[str, dict[str, str]] = defaultdict(dict)
         self.default_service_account: Optional[str] = default_service_account
         self.cluster_pod_list = cluster_pod_list
+        self.default_service_tolerations = [V1Toleration(**toleration.as_primitives()) for toleration in default_service_tolerations]
 
         # A record of previously reported events so that we don't report the same message repeatedly, fill it with
         # existing messages so we don't have a huge dump of duplicates on restart
@@ -849,6 +850,7 @@ class KubernetesController(ControllerInterface):
             security_context=V1PodSecurityContext(fs_group=1000),
             service_account_name=service_account,
             affinity=selector_to_node_affinity(self.linux_node_selector),
+            tolerations=self.default_service_tolerations
         )
 
         if use_pull_secret:

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/scaler/scaler_server.py

@@ -285,11 +285,13 @@ class ScalerServer(ThreadedCoreBase):
             'privilege': 'service'
         }
 
+        service_defaults_config = self.config.core.scaler.service_defaults
+
         # If Scaler has envs that set service-server env, then that should override configured values
         if SERVICE_API_HOST:
-            self.config.core.scaler.service_defaults.environment = \
+            service_defaults_config.environment = \
                 [EnvironmentVariable(dict(name="SERVICE_API_HOST", value=SERVICE_API_HOST))] + \
-                [env for env in self.config.core.scaler.service_defaults.environment if env.name != "SERVICE_API_HOST"]
+                [env for env in service_defaults_config.environment if env.name != "SERVICE_API_HOST"]
 
         if self.config.core.scaler.additional_labels:
             labels.update({k: v for k, v in (_l.split("=") for _l in self.config.core.scaler.additional_labels)})
@@ -304,7 +306,9 @@ class ScalerServer(ThreadedCoreBase):
                                                    log_level=self.config.logging.log_level,
                                                    core_env=core_env,
                                                    cluster_pod_list=self.config.core.scaler.cluster_pod_list,
-                                                   default_service_account=self.config.services.service_account)
+                                                   default_service_account=self.config.services.service_account,
+                                                   default_service_tolerations=service_defaults_config.tolerations
+                                                   )
 
             # Add global configuration for privileged services
             self.controller.add_config_mount(KUBERNETES_AL_CONFIG, config_map=KUBERNETES_AL_CONFIG, key="config",
@@ -313,7 +317,7 @@ class ScalerServer(ThreadedCoreBase):
             # If we're passed an override for server-server and it's defining an HTTPS connection, then add a global
             # mount for the Root CA that needs to be mounted
             if INTERNAL_ENCRYPT:
-                self.config.core.scaler.service_defaults.mounts.append(Mount(dict(
+                service_defaults_config.mounts.append(Mount(dict(
                     name="root-ca",
                     path="/etc/assemblyline/ssl/al_root-ca.crt",
                     resource_type="secret",
@@ -322,7 +326,7 @@ class ScalerServer(ThreadedCoreBase):
                 )))
 
             # Add default mounts for (non-)privileged services
-            for mount in self.config.core.scaler.service_defaults.mounts:
+            for mount in service_defaults_config.mounts:
                 # Deprecated configuration for mounting ConfigMap
                 # TODO: Deprecate code on next major change
                 if mount.config_map:
@@ -365,7 +369,7 @@ class ScalerServer(ThreadedCoreBase):
             if CLASSIFICATION_HOST_PATH:
                 self.controller.global_mounts.append((CLASSIFICATION_HOST_PATH, '/etc/assemblyline/classification.yml'))
 
-            for mount in self.config.core.scaler.service_defaults.mounts:
+            for mount in service_defaults_config.mounts:
                 # Mounts are all storage-based since there's no equivalent to ConfigMaps in Docker
                 if mount.privileged_only:
                     self.controller.core_mounts.append((mount.name, mount.path))

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/tasking_client.py

@@ -91,7 +91,7 @@ class TaskingClient:
             self.event_listener.stop()
 
     @elasticapm.capture_span(span_type='tasking_client')
-    def upload_file(self, file_path, classification, ttl, is_section_image, expected_sha256=None):
+    def upload_file(self, file_path, classification, ttl, is_section_image, is_supplementary, expected_sha256=None):
         # Identify the file info of the uploaded file
         file_info = self.identify.fileinfo(file_path)
 
@@ -105,8 +105,12 @@ class TaskingClient:
                 file_info['expiry_ts'] = None
 
             # Update the datastore with the uploaded file
-            self.datastore.save_or_freshen_file(file_info['sha256'], file_info, file_info['expiry_ts'],
-                                                file_info['classification'], is_section_image=is_section_image)
+            self.datastore.save_or_freshen_file(
+                file_info['sha256'],
+                file_info, file_info['expiry_ts'],
+                file_info['classification'],
+                is_section_image=is_section_image,
+                is_supplementary=is_supplementary)
 
             # Upload file to the filestore (upload already checks if the file exists)
             self.filestore.upload(file_path, file_info['sha256'])
@@ -349,7 +353,8 @@ class TaskingClient:
                 file_info['classification'] = item['classification']
                 self.datastore.save_or_freshen_file(item['sha256'], file_info,
                                                     file_info['expiry_ts'], file_info['classification'],
-                                                    is_section_image=item.get('is_section_image', False))
+                                                    is_section_image=item.get('is_section_image', False),
+                                                    is_supplementary=item.get('is_supplementary', False))
             return False
 
         if task.ttl:

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core/updater/run_updater.py

@@ -15,7 +15,7 @@ import docker
 
 from kubernetes.client import V1Job, V1ObjectMeta, V1JobSpec, V1PodTemplateSpec, V1PodSpec, V1Volume, \
     V1VolumeMount, V1EnvVar, V1Container, V1ResourceRequirements, \
-    V1ConfigMapVolumeSource, V1Secret, V1SecretVolumeSource, V1LocalObjectReference
+    V1ConfigMapVolumeSource, V1Secret, V1SecretVolumeSource, V1LocalObjectReference, V1Toleration
 from kubernetes import client, config
 from kubernetes.client.rest import ApiException
 
@@ -148,7 +148,7 @@ class DockerUpdateInterface:
 
 class KubernetesUpdateInterface:
     def __init__(self, logger, prefix, namespace, priority_class, extra_labels, linux_node_selector: Selector,
-                 log_level="INFO", default_service_account=None):
+                 log_level="INFO", default_service_account=None, default_service_tolerations=[]):
         # Try loading a kubernetes connection from either the fact that we are running
         # inside of a cluster, or we have a configuration in the normal location
         try:
@@ -181,6 +181,8 @@ class KubernetesUpdateInterface:
         self.default_service_account = default_service_account
         self.secret_env = []
         self.linux_node_selector = linux_node_selector
+        self.default_service_tolerations = [V1Toleration(**toleration.as_primitives()) for toleration in default_service_tolerations]
+
 
         # Get the deployment of this process. Use that information to fill out the secret info
         deployment = self.apps_api.read_namespaced_deployment(name='updater', namespace=self.namespace)
@@ -465,7 +467,8 @@ class ServiceUpdater(ThreadedCoreBase):
                                                         extra_labels=extra_labels,
                                                         log_level=self.config.logging.log_level,
                                                         default_service_account=self.config.services.service_account,
-                                                        linux_node_selector=self.config.core.scaler.linux_node_selector)
+                                                        linux_node_selector=self.config.core.scaler.linux_node_selector,
+                                                        default_service_tolerations=self.config.core.scaler.service_defaults.tolerations)
             # Add all additional mounts to privileged services
             self.mounts = self.config.core.scaler.service_defaults.mounts
         else:

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/test/test_dispatcher.py

@@ -14,8 +14,8 @@ from assemblyline.odm.randomizer import random_model_obj, random_minimal_obj, ge
 from assemblyline.odm import models
 from assemblyline.common.metrics import MetricsFactory
 
-from assemblyline_core.dispatching.client import DispatchClient
-from assemblyline_core.dispatching.dispatcher import Dispatcher, Submission
+from assemblyline_core.dispatching.client import DispatchClient, DISPATCH_RESULT_QUEUE
+from assemblyline_core.dispatching.dispatcher import Dispatcher, ServiceTask, Submission
 from assemblyline_core.dispatching.schedules import Scheduler as RealScheduler
 
 # noinspection PyUnresolvedReferences
@@ -123,7 +123,7 @@ def test_simple(clean_redis, clean_datastore):
     user: User = random_model_obj(User)
     ds.user.save(user.uname, user)
 
-    sub: Submission = random_model_obj(models.submission.Submission)
+    sub: Submission = random_model_obj(Submission)
     sub.sid = sid = 'first-submission'
     sub.params.ignore_cache = False
     sub.params.max_extracted = 5
@@ -242,7 +242,7 @@ def test_dispatch_extracted(clean_redis, clean_datastore):
         ds.file.save(fh, obj)
 
     # Inject the fake submission
-    submission = random_model_obj(models.submission.Submission)
+    submission = random_model_obj(Submission)
     submission.to_be_deleted = False
     submission.files = [dict(name='./file', sha256=file_hash)]
     submission.params.submitter = user.uname
@@ -284,8 +284,8 @@ def test_dispatch_extracted(clean_redis, clean_datastore):
 @mock.patch('assemblyline_core.dispatching.dispatcher.MetricsFactory', mock.MagicMock())
 @mock.patch('assemblyline_core.dispatching.dispatcher.Scheduler', DRPScheduler)
 def test_dispatch_extracted_bypass_drp(clean_redis, clean_datastore):
-    # Dynamic Recursion Prevention is to prevent services belonging to the 'Dynamic Analysis' from analyzing the children
-    # of files they've analyzed.
+    # Dynamic Recursion Prevention is to prevent services belonging to the 'Dynamic Analysis'
+    # from analyzing the children of files they've analyzed.
 
     # The bypass should allow services to specify files to run through Dynamic Analysis regardless of the
     # Dynamic Recursion Prevention parameter.
@@ -308,7 +308,7 @@ def test_dispatch_extracted_bypass_drp(clean_redis, clean_datastore):
         ds.file.save(fh, obj)
 
     # Inject the fake submission
-    submission = random_model_obj(models.submission.Submission)
+    submission = random_model_obj(Submission)
     submission.to_be_deleted = False
     submission.params.ignore_dynamic_recursion_prevention = False
     submission.params.services.selected = ['extract', 'sandbox']
@@ -372,7 +372,7 @@ def test_dispatch_extracted_bypass_drp(clean_redis, clean_datastore):
     disp.service_worker(disp.process_queue_index(sid))
 
     # 'sandbox' should have a task for the extracted file
-    #disp.dispatch_file(disp.tasks.get(sid), second_file_hash)
+    # disp.dispatch_file(disp.tasks.get(sid), second_file_hash)
     job = client.request_work('0', 'sandbox', '0')
     assert job.fileinfo.sha256 == second_file_hash
     assert job.filename == 'second-*'
@@ -417,3 +417,37 @@ def test_timeout():
 
     # Expire nothing
     assert len(table.timeouts()) == 0
+
+
+def test_prevent_result_overwrite(clean_redis, clean_datastore):
+    client = DispatchClient(clean_datastore, clean_redis, clean_redis)
+    dispatcher_name = "test"
+    result_queue = client._get_queue_from_cache(DISPATCH_RESULT_QUEUE + dispatcher_name)
+
+    # Create a task and add it to set of running tasks
+    task = random_model_obj(ServiceTask)
+    task.metadata['dispatcher__'] = dispatcher_name
+
+    # Create a result that's not "empty"
+    result = random_model_obj(Result)
+    result.response.service_name = task.service_name
+    result.sha256 = task.fileinfo.sha256
+    result.result.score = 1
+    result_key = result.build_key()
+
+    # Submit result to be saved
+    client.running_tasks.add(task.key(), task.as_primitives())
+    client.service_finished(task.sid, result_key, result)
+
+    # Pop result from queue, we expect to get the same result key as earlier
+    message = result_queue.pop(blocking=False)
+    msg_result_key = message['result_summary']['key']
+    assert msg_result_key == result_key
+
+    # Save the same result again but we expect to be saved under another key
+    client.running_tasks.add(task.key(), task.as_primitives())
+    client.service_finished(task.sid, result_key, result)
+    message = result_queue.pop(blocking=False)
+    msg_result_key = message['result_summary']['key']
+
+    assert msg_result_key != result_key

{assemblyline-core-4.5.0.26 → assemblyline-core-4.5.0.28}/test/test_expiry.py

@@ -4,6 +4,7 @@ import random
 import concurrent.futures
 
 from assemblyline.common.isotime import now_as_iso
+from assemblyline.datastore.helper import AssemblylineDatastore
 from assemblyline.odm.randomizer import random_model_obj
 
 from assemblyline_core.expiry.run_expiry import ExpiryManager
@@ -14,34 +15,31 @@ expiry_collections_len = {}
 archive_collections_len = {}
 
 
-@pytest.fixture(scope='module')
-def datastore(archive_connection):
-    return archive_connection
-
-
-def purge_data(datastore):
-    for name, definition in datastore.ds.get_models().items():
+def purge_data(datastore_connection: AssemblylineDatastore):
+    for name, definition in datastore_connection.ds.get_models().items():
         if hasattr(definition, 'expiry_ts'):
-            getattr(datastore, name).wipe()
+            getattr(datastore_connection, name).wipe()
 
 
 @pytest.fixture(scope="function")
-def ds_expiry(request, datastore):
-    for name, definition in datastore.ds.get_models().items():
+def ds_expiry(request, datastore_connection):
+    for name, definition in datastore_connection.ds.get_models().items():
         if hasattr(definition, 'expiry_ts'):
-            collection = getattr(datastore, name)
+            collection = getattr(datastore_connection, name)
             collection.wipe()
             expiry_len = random.randint(MIN_OBJECTS, MAX_OBJECTS)
             for x in range(expiry_len):
                 obj = random_model_obj(collection.model_class)
+                if hasattr(definition, 'from_archive'):
+                    obj.from_archive = False
                 obj.expiry_ts = now_as_iso(-10000)
                 collection.save('longer_name'+str(x), obj)
 
             expiry_collections_len[name] = expiry_len
             collection.commit()
 
-    request.addfinalizer(lambda: purge_data(datastore))
-    return datastore
+    request.addfinalizer(lambda: purge_data(datastore_connection))
+    return datastore_connection
 
 
 class FakeCounter(object):
@@ -58,8 +56,8 @@ class FakeCounter(object):
         return self.counts.get(name, 0)
 
 
-def test_expire_all(ds_expiry):
-    expiry = ExpiryManager()
+def test_expire_all(config, ds_expiry, filestore):
+    expiry = ExpiryManager(config=config, datastore=ds_expiry, filestore=filestore)
     expiry.running = True
     expiry.counter = FakeCounter()
     with concurrent.futures.ThreadPoolExecutor(5) as pool:

assemblyline-core-4.5.0.26/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.0.26