assemblyline-core 4.5.0.27__tar.gz → 4.5.0.29__tar.gz

{assemblyline-core-4.5.0.27 → assemblyline-core-4.5.0.29}/PKG-INFO

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

assemblyline-core-4.5.0.29/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.0.29

{assemblyline-core-4.5.0.27 → assemblyline-core-4.5.0.29}/assemblyline_core/badlist_client.py

@@ -52,9 +52,9 @@ class BadlistClient:
             data.setdefault('file', {})
 
         # Ensure expiry_ts is set on tag-related items
-        dtl = data.pop('dtl', None) or self.config.core.expiry.badlisted_tag_dtl * 24 * 3600
+        dtl = data.pop('dtl', None) or self.config.core.expiry.badlisted_tag_dtl
         if dtl:
-            data['expiry_ts'] = now_as_iso(dtl)
+            data['expiry_ts'] = now_as_iso(dtl * 24 * 3600)
 
         # Set last updated
         data['added'] = data['updated'] = now_as_iso()

{assemblyline-core-4.5.0.27 → assemblyline-core-4.5.0.29}/assemblyline_core/ingester/ingester.py

@@ -595,7 +595,6 @@ class Ingester(ThreadedCoreBase):
                 task.params.submitter) if g in str(task.params.classification)]
 
         # Check if this file is already being processed
-        self.stamp_filescore_key(task)
         pprevious, previous, score = None, None, None
         if not param.ignore_cache:
             pprevious, previous, score, _ = self.check(task, count_miss=False)
@@ -627,6 +626,33 @@ class Ingester(ThreadedCoreBase):
         # (So we don't end up dropping the resubmission).
         if previous:
             self.counter.increment('duplicates')
+
+            if self.config.core.ingester.always_create_submission:
+                # Create a submission record based on the cache hit
+                submission = self.datastore.submission.get(previous, as_obj=False)
+
+                # Assign the current submission as the PSID for the new submission
+                pprevious = previous
+                previous = task.ingest_id
+                task.params.psid = pprevious
+
+                submission.update({
+                    'archived': False,
+                    'archive_ts': None,
+                    'classification': task.params.classification,
+                    'expiry_ts':now_as_iso(task.params.ttl * 24 * 60 * 60),
+                    'from_archive': False,
+                    'metadata': task.submission.metadata,
+                    'params': task.params.as_primitives(),
+                    'sid': previous,
+                    'to_be_deleted': False,
+                    'times': {
+                        'submitted': task.ingest_time,
+                        'completed': "NOW"
+                    },
+                })
+                self.datastore.submission.save(previous, submission)
+
             self.finalize(pprevious, previous, score, task, cache=True)
 
             # On cache hits of any kind we want to send out a completed message

{assemblyline-core-4.5.0.27 → assemblyline-core-4.5.0.29}/assemblyline_core/safelist_client.py

@@ -69,9 +69,9 @@ class SafelistClient:
             data.setdefault('file', {})
 
         # Ensure expiry_ts is set on tag-related items
-        dtl = data.pop('dtl', None) or self.config.core.expiry.safelisted_tag_dtl * 24 * 3600
+        dtl = data.pop('dtl', None) or self.config.core.expiry.safelisted_tag_dtl
         if dtl:
-            data['expiry_ts'] = now_as_iso(dtl)
+            data['expiry_ts'] = now_as_iso(dtl * 24 * 3600)
 
         # Set last updated
         data['added'] = data['updated'] = now_as_iso()

{assemblyline-core-4.5.0.27 → assemblyline-core-4.5.0.29}/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.27 → assemblyline-core-4.5.0.29}/assemblyline_core/scaler/scaler_server.py

@@ -17,6 +17,7 @@ import copy
 from contextlib import contextmanager
 
 import elasticapm
+import json
 import yaml
 
 from assemblyline.remote.datatypes.queues.named import NamedQueue
@@ -285,11 +286,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 +307,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 +318,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 +327,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:
@@ -356,7 +361,8 @@ class ScalerServer(ThreadedCoreBase):
                 self.controller.core_mounts.append((DOCKER_CONFIGURATION_VOLUME, '/etc/assemblyline/'))
 
                 with open(os.path.join(DOCKER_CONFIGURATION_PATH, 'config.yml'), 'w') as handle:
-                    yaml.dump(self.config.as_primitives(), handle)
+                    # Convert to JSON before converting to YAML to account for direct ODM representation errors
+                    yaml.dump(json.loads(self.config.json()), handle)
 
                 with open(os.path.join(DOCKER_CONFIGURATION_PATH, 'classification.yml'), 'w') as handle:
                     yaml.dump(get_classification().original_definition, handle)
@@ -365,7 +371,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.27 → assemblyline-core-4.5.0.29}/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.27 → assemblyline-core-4.5.0.29}/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.27 → assemblyline-core-4.5.0.29}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline-core-4.5.0.27 → assemblyline-core-4.5.0.29}/test/test_worker_ingest.py

@@ -193,3 +193,56 @@ def test_ingest_size_error(ingest_harness):
     queue = ingester.notification_queues[queue_name]
     message = queue.pop()
     assert message is not None
+
+def test_ingest_always_create_submission(ingest_harness):
+    datastore, ingester, in_queue = ingest_harness
+
+    # Simulate configuration where we'll always create a submission
+    ingester.config.core.ingester.always_create_submission = True
+    get_if_exists = datastore.filescore.get_if_exists
+    try:
+        # Add a valid file score for all files
+        from assemblyline.odm.models.filescore import FileScore
+        from assemblyline.odm.models.submission import Submission
+        datastore.filescore.get_if_exists = mock.MagicMock(
+            return_value=FileScore(dict(psid='000', expiry_ts=0, errors=0, score=10, sid='001', time=time.time()))
+        )
+        # Create a submission for cache hit
+        old_sub = random_minimal_obj(Submission)
+        old_sub.sid = '001'
+        old_sub.params.psid = '000'
+        old_sub = old_sub.as_primitives()
+        datastore.submission.save('001', old_sub)
+
+        # Ingest a file
+        submission_msg = make_message(message={'sid': '002', 'metadata': {'blah': 'blah'}})
+        submission_msg['sid'] = '002'
+        in_queue.push(submission_msg)
+        ingester.handle_ingest()
+
+        # No file has made it into the internal buffer => cache hit and drop
+        datastore.filescore.get_if_exists.assert_called_once()
+        ingester.counter.increment.assert_any_call('cache_hit')
+        ingester.counter.increment.assert_any_call('duplicates')
+        assert ingester.unique_queue.length() == 0
+        assert ingester.ingest_queue.length() == 0
+
+        # Check to see if new submission was created
+        new_sub = datastore.submission.get_if_exists('002', as_obj=False)
+        assert new_sub and new_sub['params']['psid'] == old_sub['sid']
+
+        # Check to see if certain properties are same (anything relating to analysis)
+        assert all([old_sub.get(attr) == new_sub.get(attr) \
+                    for attr in ['error_count', 'errors', 'file_count', 'files', 'max_score', 'results', 'state', 'verdict']])
+
+        # Check to see if certain properties are different
+        # (anything that isn't related to analysis but can be set at submission time)
+        assert all([old_sub.get(attr) != new_sub.get(attr) \
+                    for attr in ['expiry_ts', 'metadata', 'params', 'times']])
+
+        # Check to see if certain properties have been nullified
+        # (properties that are set outside of submission)
+        assert not all([new_sub.get(attr) \
+                        for attr in ['archived', 'archive_ts', 'to_be_deleted', 'from_archive']])
+    finally:
+        datastore.filescore.get_if_exists = get_if_exists

assemblyline-core-4.5.0.27/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.0.27