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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline-core
-Version: 4.6.1.dev215
+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.dev215 → assemblyline_core-4.6.1.dev218}/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,31 @@ 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 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 = []
-            
+
             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.dev215 → 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.dev215
+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.dev215/assemblyline_core/VERSION

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

assemblyline_core-4.6.1.dev215/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'])