assemblyline-core 4.5.1.dev461__tar.gz → 4.5.1.dev462__tar.gz

{assemblyline_core-4.5.1.dev461 → assemblyline_core-4.5.1.dev462}/PKG-INFO

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

assemblyline_core-4.5.1.dev462/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.1.dev462

{assemblyline_core-4.5.1.dev461 → assemblyline_core-4.5.1.dev462}/assemblyline_core/tasking_client.py

@@ -49,6 +49,8 @@ class TaskingClient:
         self.event_sender = EventSender('changes', redis)
         self.cleanup = False
         self.event_listener = None
+        self.heuristics: dict[str, Heuristic] = {}
+        self.register_only = register_only
 
         # If we're performing service registration, we only need a connection to the datastore
         if not register_only:
@@ -56,7 +58,6 @@ class TaskingClient:
             self.event_listener = EventWatcher(redis)
             self.filestore = filestore or forge.get_filestore(self.config)
             self.heuristic_handler = HeuristicHandler(self.datastore)
-            self.heuristics: dict[str, Heuristic] = {}
             self.reload_heuristics({})
             self.status_table = ExpiringHash(SERVICE_STATE_HASH, ttl=60*30, host=redis)
             self.tag_safelister = forge.CachedObject(forge.get_tag_safelister, kwargs=dict(
@@ -195,6 +196,10 @@ class TaskingClient:
                                     f"heuristic {item['update']['_id']}: {item['update']['result'].upper()}")
 
             # Look for heuristics that are no longer managed by the service and clean them up
+            if self.register_only:
+                # Fetch current heuristics for the service during registration
+                self.reload_heuristics({'service_name': service.name})
+
             all_heuristics = set(h_id for h_id in self.heuristics.keys()
                                 if h_id.startswith(f"{service.name.upper()}."))
             removed_heuristics = all_heuristics - set(h['heur_id'] for h in heuristics)

{assemblyline_core-4.5.1.dev461 → assemblyline_core-4.5.1.dev462}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline_core-4.5.1.dev461 → assemblyline_core-4.5.1.dev462}/test/test_tasking_client.py

@@ -5,26 +5,12 @@ 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.random_data import (
-    create_badlists,
-    create_users,
-    wipe_badlist,
-    wipe_users,
-)
-from assemblyline.odm.randomizer import random_minimal_obj
 
+from assemblyline.odm.randomizer import random_minimal_obj
 
-@pytest.fixture(scope="module")
-def client(datastore_connection):
-    try:
-        create_users(datastore_connection)
-        create_badlists(datastore_connection)
-        yield TaskingClient(datastore_connection)
-    finally:
-        wipe_users(datastore_connection)
-        wipe_badlist(datastore_connection)
+def test_register_service(datastore_connection):
+    client = TaskingClient(datastore_connection, register_only=True)
 
-def test_register_service(client, datastore_connection):
     # Test service registration
     service = random_minimal_obj(Service).as_primitives()
     heuristics = [random_minimal_obj(Heuristic).as_primitives() for _ in range(2)]

assemblyline_core-4.5.1.dev461/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.1.dev461