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

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

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

assemblyline_core-4.6.1.dev214/assemblyline_core/VERSION

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

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

@@ -3,7 +3,6 @@ import time
 from typing import Any, Dict, Optional
 
 import elasticapm
-
 from assemblyline.common import forge
 from assemblyline.common.constants import SERVICE_STATE_HASH, ServiceStatus
 from assemblyline.common.dict_utils import flatten, unflatten
@@ -22,6 +21,7 @@ from assemblyline.odm.models.service import Service
 from assemblyline.odm.models.tagging import Tagging
 from assemblyline.remote.datatypes.events import EventSender, EventWatcher
 from assemblyline.remote.datatypes.hash import ExpiringHash
+
 from assemblyline_core.dispatching.client import DispatchClient
 
 
@@ -161,9 +161,32 @@ class TaskingClient:
                 self.datastore.service_delta.save(service.name, {'version': service.version})
                 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))
 
             new_heuristics = []
-            
+
             plan = self.datastore.heuristic.get_bulk_plan()
             for index, heuristic in enumerate(heuristics):
                 heuristic_id = f'#{index}'  # Set heuristic id to it's position in the list for logging purposes

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

@@ -1,23 +1,22 @@
 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.dev202 → assemblyline_core-4.6.1.dev214}/assemblyline_core/updater/run_updater.py

@@ -11,6 +11,12 @@ 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,
@@ -34,12 +40,6 @@ 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,14 +692,32 @@ class ServiceUpdater(ThreadedCoreBase):
                         service_name, latest_tag = service_key.split("_")
 
                         if self.datastore.service.get_if_exists(service_key):
-                            operations = [(self.datastore.service_delta.UPDATE_SET, 'version', latest_tag)]
+                            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}"
 
                             # Check if a service waas previously disabled and re-enable it
                             if service_name in self.incompatible_services:
                                 self.incompatible_services.remove(service_name)
-                                operations.append((self.datastore.service_delta.UPDATE_SET, 'enabled', True))
+                                service_delta['enabled'] = True
 
-                            if self.datastore.service_delta.update(service_name, operations):
+
+                            if self.datastore.service_delta.save(service_name, service_delta):
                                 # Update completed, cleanup
                                 self.service_events.send(service_name, {
                                     'operation': Operation.Modified,

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

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

assemblyline_core-4.6.1.dev214/test/test_tasking_client.py

@@ -0,0 +1,89 @@
+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'

assemblyline_core-4.6.1.dev202/assemblyline_core/VERSION

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

assemblyline_core-4.6.1.dev202/test/test_tasking_client.py

@@ -1,37 +0,0 @@
-from assemblyline_core.tasking_client import TaskingClient
-
-from assemblyline.odm.models.service import Service
-from assemblyline.odm.models.heuristic import Heuristic
-from assemblyline.odm.models.result import Result, Section, Heuristic as SectionHeuristic
-
-from assemblyline.odm.randomizer import random_minimal_obj
-
-def test_register_service(datastore_connection):
-    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'])