assemblyline-core 4.6.1.dev214__tar.gz → 4.6.1.dev218__tar.gz

{assemblyline_core-4.6.1.dev214 → assemblyline_core-4.6.1.dev218}/PKG-INFO

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

assemblyline_core-4.6.1.dev218/assemblyline_core/VERSION

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

{assemblyline_core-4.6.1.dev214 → assemblyline_core-4.6.1.dev218}/assemblyline_core/tasking_client.py

@@ -162,28 +162,27 @@ class TaskingClient:
                 self.datastore.service_delta.commit()
                 self.log.info(f"{log_prefix}{service.name} version ({service.version}) registered")
             else:
-                # Check for any updates to the service that should be applied to the service delta
-                service_delta = flatten(self.datastore.service_delta.get(service.name, as_obj=False))
-                if service_delta.get('version') is not None and service_delta['version'] != service.version:
-                    # Update the service delta based on changes in the new service version
-                    old_service = flatten(self.datastore.service.get(f'{service.name}_{service_delta["version"]}',
-                                                                    as_obj=False))
-                    new_service = flatten(service.as_primitives())
-
-                    # Check for cases elements from internal lists have been removed
-                    for key, id_field in self.datastore.service_list_keys.items():
-                        if key in service_delta:
-                            # Check to see what was removed from the list in the previous version
-                            old_ids = set(item[id_field] for item in old_service.get(key, []))
-                            new_ids = set(item[id_field] for item in service_delta[key])
-                            removed_ids = old_ids - new_ids
-
-                            # Add any new items to the delta so it can be included in the merged service configuration
-                            added_ids = set(item[id_field] for item in new_service[key]) - removed_ids
-                            if added_ids:
-                                service_delta[key].extend([{id_field: added_id} for added_id in added_ids])
-
-                    self.datastore.service_delta.save(service.name, unflatten(service_delta))
+                # Check for any changes that should be merged into the service delta
+                service_delta = self.datastore.service_delta.get(service.name, as_obj=False)
+
+                # Check for any new configuration keys that should be added to the service delta
+                if service_delta.get('config'):
+                    new_config = {k: v for k, v in service.config.items() if k not in service_delta['config']}
+                    if new_config:
+                        if 'config' not in service_delta:
+                            service_delta['config'] = {}
+                        service_delta['config'].update(new_config)
+
+                # Check for any new submission parameters that should be added to the service delta
+                if service_delta.get('submission_params'):
+                    old_submission_params = {param['name'] for param in service_delta['submission_params']}
+                    for param in service.submission_params:
+                        if param['name'] not in old_submission_params:
+                            # New parameter, add it to the old submission params
+                            service_delta['submission_params'].append(param.as_primitives())
+
+                # Save any changes to the service delta
+                self.datastore.service_delta.save(service.name, service_delta)
 
             new_heuristics = []
 

{assemblyline_core-4.6.1.dev214 → assemblyline_core-4.6.1.dev218}/assemblyline_core/updater/helper.py

@@ -1,22 +1,23 @@
 import os
+import requests
 import re
 import socket
 import string
 import time
+
+from assemblyline.common.version import FRAMEWORK_VERSION, SYSTEM_VERSION
+from assemblyline.odm.models.config import Config as SystemConfig, ServiceRegistry
+from assemblyline.odm.models.service import Service as ServiceConfig, DockerConfig
+
 from base64 import b64encode
 from collections import defaultdict
 from logging import Logger
 from typing import Dict, List
+from packaging.version import parse, Version
 from urllib.parse import urlencode
 
-import requests
-from assemblyline.common.version import FRAMEWORK_VERSION, SYSTEM_VERSION
-from assemblyline.odm.models.config import Config as SystemConfig
-from assemblyline.odm.models.config import ServiceRegistry
-from assemblyline.odm.models.service import DockerConfig
-from assemblyline.odm.models.service import Service as ServiceConfig
 from azure.identity import DefaultAzureCredential
-from packaging.version import Version, parse
+
 
 DEFAULT_DOCKER_REGISTRY = "hub.docker.com"
 

{assemblyline_core-4.6.1.dev214 → assemblyline_core-4.6.1.dev218}/assemblyline_core/updater/run_updater.py

@@ -11,12 +11,6 @@ from concurrent.futures import ThreadPoolExecutor
 from typing import Any, List, Optional
 
 import docker
-from assemblyline.common import isotime
-from assemblyline.odm.messages.changes import Operation, ServiceChange
-from assemblyline.odm.models.config import Mount, Selector
-from assemblyline.odm.models.service import DockerConfig, Service
-from assemblyline.remote.datatypes.events import EventSender, EventWatcher
-from assemblyline.remote.datatypes.hash import Hash
 from kubernetes import client, config
 from kubernetes.client import (
     V1Capabilities,
@@ -40,6 +34,12 @@ from kubernetes.client import (
 )
 from kubernetes.client.rest import ApiException
 
+from assemblyline.common import isotime
+from assemblyline.odm.messages.changes import Operation, ServiceChange
+from assemblyline.odm.models.config import Mount, Selector
+from assemblyline.odm.models.service import DockerConfig, Service
+from assemblyline.remote.datatypes.events import EventSender, EventWatcher
+from assemblyline.remote.datatypes.hash import Hash
 from assemblyline_core.scaler.controllers.kubernetes_ctl import (
     PRIVILEGED_SERVICE_ACCOUNT_NAME,
     create_docker_auth_config,
@@ -692,32 +692,14 @@ class ServiceUpdater(ThreadedCoreBase):
                         service_name, latest_tag = service_key.split("_")
 
                         if self.datastore.service.get_if_exists(service_key):
-                            service_image = self.datastore.service.get(service_key, as_obj=False)['docker_config']['image'].rsplit(':', 1)[0].split('/', 1)[1]
-
-                            service_delta = self.datastore.service_delta.get(service_name, as_obj=False)
-                            # Update the service version in the delta
-                            service_delta['version'] = latest_tag
-
-                            # If the container images related to the service have been changed, update the version
-                            if service_delta.get('docker_config', {}).get('image'):
-                                image_path, _ = service_delta['docker_config']['image'].rsplit(':', 1)
-                                service_delta['docker_config']['image'] = f"{image_path}:{latest_tag}"
-
-                            if service_delta.get('dependencies', {}).items():
-                                for _, dep_data in service_delta['dependencies'].items():
-                                    if dep_data.get('image'):
-                                        dep_image_path, _ = dep_data['image'].rsplit(':', 1)
-                                        if service_image in dep_image_path:
-                                            # Only update dependencies that share the same image as the main service
-                                            dep_data['image'] = f"{dep_image_path}:{latest_tag}"
+                            operations = [(self.datastore.service_delta.UPDATE_SET, 'version', latest_tag)]
 
                             # Check if a service waas previously disabled and re-enable it
                             if service_name in self.incompatible_services:
                                 self.incompatible_services.remove(service_name)
-                                service_delta['enabled'] = True
-
+                                operations.append((self.datastore.service_delta.UPDATE_SET, 'enabled', True))
 
-                            if self.datastore.service_delta.save(service_name, service_delta):
+                            if self.datastore.service_delta.update(service_name, operations):
                                 # Update completed, cleanup
                                 self.service_events.send(service_name, {
                                     'operation': Operation.Modified,

{assemblyline_core-4.6.1.dev214 → assemblyline_core-4.6.1.dev218}/assemblyline_core.egg-info/PKG-INFO

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

assemblyline_core-4.6.1.dev218/test/test_tasking_client.py

@@ -0,0 +1,77 @@
+from assemblyline_core.tasking_client import TaskingClient
+
+from assemblyline.datastore.helper import AssemblylineDatastore as Datastore
+from assemblyline.odm.models.heuristic import Heuristic
+from assemblyline.odm.models.result import Heuristic as SectionHeuristic
+from assemblyline.odm.models.result import Result, Section
+from assemblyline.odm.models.service import Service
+from assemblyline.odm.randomizer import random_minimal_obj
+
+
+def test_register_service(datastore_connection: Datastore):
+    client = TaskingClient(datastore_connection, register_only=True)
+
+    # Test service registration
+    service = random_minimal_obj(Service).as_primitives()
+    heuristics = [random_minimal_obj(Heuristic).as_primitives() for _ in range(2)]
+    service['heuristics'] = heuristics
+    assert client.register_service(service)
+    assert all([datastore_connection.heuristic.exists(h['heur_id']) for h in heuristics])
+
+    # Test registration with heuristics that were removed but still have related results
+    heuristic = heuristics.pop(0)
+    result = random_minimal_obj(Result)
+    section = random_minimal_obj(Section)
+    section.heuristic = SectionHeuristic(heuristic)
+    result.result.sections = [section]
+    datastore_connection.result.save('test_result', result)
+    datastore_connection.result.commit()
+
+    # Heuristics that were removed should still reside in the system if there are still associated data to it
+    service['heuristics'] = heuristics
+    assert client.register_service(service)
+    assert datastore_connection.heuristic.exists(heuristic['heur_id'])
+
+    # Test registration with removed heuristics that have no related results
+    datastore_connection.result.delete('test_result')
+    datastore_connection.result.commit()
+    assert client.register_service(service)
+    assert not datastore_connection.heuristic.exists(heuristic['heur_id'])
+
+def test_service_update(datastore_connection: Datastore):
+    client = TaskingClient(datastore_connection, register_only=True)
+
+    # Test service registration
+    service = random_minimal_obj(Service).as_primitives()
+    assert client.register_service(service)
+
+    # Test registering a service update where there's a new submission parameter and configuration
+    service['submission_params'].append({"name": 'new_param', 'type': 'str', 'default': 'default_value', 'value': 'default_value'})
+    service['version'] = "new_version"
+    service['config']  = {'new_config': 'value'}
+    assert client.register_service(service)
+
+    # Pretend I'm the updater acknowledging the new version has been registered
+    datastore_connection.service_delta.update(service['name'], [(datastore_connection.service_delta.UPDATE_SET, 'version', service['version'])])
+
+    # We should see the new submission parameter in the service while applying delta changes
+    delta = datastore_connection.get_service_with_delta(service['name'])
+    assert delta['submission_params'][-1]['name'] == 'new_param'
+    assert delta['config']['new_config'] == 'value'
+
+    # Test registering a service update where the user has changed a submission parameter prior to the update and new parameter was added
+    assert datastore_connection.service_delta.update(service['name'], [(datastore_connection.service_delta.UPDATE_APPEND, 'submission_params', {'name': 'new_param', 'type': 'str', 'default': 'custom_value', 'value': 'custom_value'})])
+    datastore_connection.service_delta.commit()
+
+    service['submission_params'].append({"name": 'new_new_param', 'type': 'str', 'default': 'default_value', 'value': 'default_value'})
+    service['version'] = "new_new_version"
+
+    assert client.register_service(service)
+
+    # Pretend I'm the updater acknowledging the new version has been registered
+    datastore_connection.service_delta.update(service['name'], [(datastore_connection.service_delta.UPDATE_SET, 'version', service['version'])])
+
+    # We expect to see both the updated submission parameter and the newly added one (while still keeping the custom value changes)
+    delta = datastore_connection.get_service_with_delta(service['name'])
+    assert delta['submission_params'][-2]['name'] == 'new_param' and delta['submission_params'][-2]['value'] == 'custom_value'
+    assert delta['submission_params'][-1]['name'] == 'new_new_param'

assemblyline_core-4.6.1.dev214/assemblyline_core/VERSION

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

assemblyline_core-4.6.1.dev214/test/test_tasking_client.py

@@ -1,89 +0,0 @@
-from assemblyline.odm.models.heuristic import Heuristic
-from assemblyline.odm.models.result import Heuristic as SectionHeuristic
-from assemblyline.odm.models.result import Result, Section
-from assemblyline.odm.models.service import Service
-from assemblyline.odm.randomizer import random_minimal_obj, random_model_obj
-from assemblyline_core.tasking_client import TaskingClient
-
-
-def test_register_service(datastore_connection):
-    client = TaskingClient(datastore_connection, register_only=True)
-
-    # Test service registration
-    service = random_model_obj(Service).as_primitives()
-    heuristics = [random_minimal_obj(Heuristic).as_primitives() for _ in range(2)]
-    service['heuristics'] = heuristics
-    assert client.register_service(service)
-    assert all([datastore_connection.heuristic.exists(h['heur_id']) for h in heuristics])
-
-    # Test registration with heuristics that were removed but still have related results
-    heuristic = heuristics.pop(0)
-    result = random_minimal_obj(Result)
-    section = random_minimal_obj(Section)
-    section.heuristic = SectionHeuristic(heuristic)
-    result.result.sections = [section]
-    datastore_connection.result.save('test_result', result)
-    datastore_connection.result.commit()
-
-    # Heuristics that were removed should still reside in the system if there are still associated data to it
-    service['heuristics'] = heuristics
-    assert client.register_service(service)
-    assert datastore_connection.heuristic.exists(heuristic['heur_id'])
-
-    # Test registration with removed heuristics that have no related results
-    datastore_connection.result.delete('test_result')
-    datastore_connection.result.commit()
-    assert client.register_service(service)
-    assert not datastore_connection.heuristic.exists(heuristic['heur_id'])
-
-    # Test registration with new and removed sources, submission parameters and configurations
-    datastore_connection.service_delta.save(service['name'], {
-        "version": service['version'],
-        # Let's imagine a user added a new configuration parameter between updates
-        "config": {"user_config": "user_value"},
-        "submission_params": [
-            {
-                # Let's imagine a user added a new submission parameter between updates
-                "name": "user_param",
-                "type": "str",
-                "default": "user_value",
-                "value": "user_value",
-            }
-        ],
-        "update_config": {
-            # Let's imagine a user removed a sources between updates
-            "sources": service['update_config']['sources'][1:],
-        }
-    })
-
-    # Now let's update the service with a new version that has some changes
-    service['config']['new_config'] = 'new_value'
-    service['submission_params'].append({
-        "name": "new_param",
-        "type": "str",
-        "default": "new_value",
-        "value": "new_value",
-    })
-    service['update_config']['sources'].append({"name": "new_source", "uri": "http://new_source"})
-    service['version'] = "new_version"
-    assert client.register_service(service)
-
-    datastore_connection.service_delta.update(service['name'],
-                                              [(datastore_connection.service_delta.UPDATE_SET, 'version', service['version'])])
-
-    merged_service = datastore_connection.get_service_with_delta(service['name'], as_obj=False)
-    # Update sources that have been removed should stay removed
-    sources = [s['name'] for s in merged_service['update_config']['sources']]
-    assert service['update_config']['sources'][0]['name'] not in sources
-
-    # Update sources that are new should be added
-    assert 'new_source' in sources
-
-    # Submission parameters that are new should be added and old user parameters should stay
-    submission_param_names = [p['name'] for p in merged_service['submission_params']]
-    assert 'new_param' in submission_param_names
-    assert 'user_param' in submission_param_names
-
-    # Configurations that are new should be added and old user configurations should stay
-    assert merged_service['config']['new_config'] == 'new_value'
-    assert merged_service['config']['user_config'] == 'user_value'