assemblyline-core 4.5.0.76__tar.gz → 4.5.0.79__tar.gz

{assemblyline_core-4.5.0.76 → assemblyline_core-4.5.0.79}/PKG-INFO

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

assemblyline_core-4.5.0.79/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.5.0.79

{assemblyline_core-4.5.0.76 → assemblyline_core-4.5.0.79}/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(
@@ -126,7 +127,7 @@ class TaskingClient:
 
         try:
             # Get heuristics list
-            heuristics = service_data.pop('heuristics', None)
+            heuristics = service_data.pop('heuristics', [])
 
             # Patch update_channel, registry_type before Service registration object creation
             service_data['update_channel'] = service_data.get(
@@ -162,44 +163,62 @@ class TaskingClient:
                 self.log.info(f"{log_prefix}{service.name} version ({service.version}) registered")
 
             new_heuristics = []
-            if 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
-                    try:
-                        # Append service name to heuristic ID
-                        heuristic['heur_id'] = f"{service.name.upper()}.{str(heuristic['heur_id'])}"
-
-                        # Attack_id field is now a list, make it a list if we receive otherwise
-                        attack_id = heuristic.get('attack_id', None)
-                        if isinstance(attack_id, str):
-                            heuristic['attack_id'] = [attack_id]
-
-                        heuristic = Heuristic(heuristic)
-                        heuristic_id = heuristic.heur_id
-                        existing_heuristic_obj = self.datastore.heuristic.get_if_exists(heuristic_id)
-                        if existing_heuristic_obj:
-                            # Ensure statistics of heuristic are preserved
-                            heuristic.stats = existing_heuristic_obj.stats
-                        plan.add_upsert_operation(heuristic_id, heuristic)
-                    except Exception as e:
-                        msg = f"{service.name} has an invalid heuristic ({heuristic_id}): {str(e)}"
-                        self.log.exception(f"{log_prefix}{msg}")
-                        raise ValueError(msg)
-
+            
+            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
+                try:
+                    # Append service name to heuristic ID
+                    heuristic['heur_id'] = f"{service.name.upper()}.{str(heuristic['heur_id'])}"
+
+                    # Attack_id field is now a list, make it a list if we receive otherwise
+                    attack_id = heuristic.get('attack_id', None)
+                    if isinstance(attack_id, str):
+                        heuristic['attack_id'] = [attack_id]
+
+                    heuristic = Heuristic(heuristic)
+                    heuristic_id = heuristic.heur_id
+                    existing_heuristic_obj = self.datastore.heuristic.get_if_exists(heuristic_id)
+                    if existing_heuristic_obj:
+                        # Ensure statistics of heuristic are preserved
+                        heuristic.stats = existing_heuristic_obj.stats
+                    plan.add_upsert_operation(heuristic_id, heuristic)
+                except Exception as e:
+                    msg = f"{service.name} has an invalid heuristic ({heuristic_id}): {str(e)}"
+                    self.log.exception(f"{log_prefix}{msg}")
+                    raise ValueError(msg)
+
+            if plan.operations:
                 for item in self.datastore.heuristic.bulk(plan)['items']:
                     if item['update']['result'] != "noop":
                         new_heuristics.append(item['update']['_id'])
                         self.log.info(f"{log_prefix}{service.name} "
-                                      f"heuristic {item['update']['_id']}: {item['update']['result'].upper()}")
+                                    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)
 
-                self.datastore.heuristic.commit()
+            for heuristic in removed_heuristics:
+                # Only remove heuristics that aren't actively referenced in a result
+                if not self.datastore.result.search(f"result.sections.heuristic.heur_id:{heuristic}",
+                                                    rows=0, track_total_hits=True)['total']:
+                    self.datastore.heuristic.delete(heuristic)
 
-                # Notify components watching for heuristic config changes
-                self.event_sender.send('heuristics', {
-                    'operation': Operation.Modified,
-                    'service_name': service.name
-                })
+
+
+            self.datastore.heuristic.commit()
+
+            # Notify components watching for heuristic config changes
+            self.event_sender.send('heuristics', {
+                'operation': Operation.Modified,
+                'service_name': service.name
+            })
 
             service_config = self.datastore.get_service_with_delta(service.name, as_obj=False)
 

{assemblyline_core-4.5.0.76 → assemblyline_core-4.5.0.79}/assemblyline_core/workflow/run_workflow.py

@@ -59,7 +59,7 @@ class WorkflowManager(ServerBase):
 
         return ret_val
 
-    def try_run(self):
+    def try_run(self, run_once=False):
         self.datastore.alert.commit()
         while self.running:
             self.heartbeat()
@@ -109,6 +109,9 @@ class WorkflowManager(ServerBase):
                     if not workflow.enabled:
                         continue
 
+                    # Trigger a heartbeat to let the system know the workflow manager is still alive between tasks
+                    self.heartbeat()
+
                     # Start of transaction
                     if self.apm_client:
                         self.apm_client.begin_transaction("Execute workflows")
@@ -217,6 +220,8 @@ class WorkflowManager(ServerBase):
             else:
                 self.log.info("Skipping all workflows since there where no new alerts in the specified time period.")
 
+            if run_once:
+                break
             time.sleep(30)
             self.start_ts = end_ts
 

{assemblyline_core-4.5.0.76 → assemblyline_core-4.5.0.79}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline_core-4.5.0.76 → assemblyline_core-4.5.0.79}/assemblyline_core.egg-info/SOURCES.txt

@@ -81,6 +81,8 @@ test/test_scaler.py
 test/test_scheduler.py
 test/test_signature_client.py
 test/test_simulation.py
+test/test_tasking_client.py
 test/test_vacuum.py
 test/test_worker_ingest.py
-test/test_worker_submit.py
+test/test_worker_submit.py
+test/test_workflow.py

assemblyline_core-4.5.0.79/test/test_tasking_client.py

@@ -0,0 +1,37 @@
+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'])

assemblyline_core-4.5.0.79/test/test_workflow.py

@@ -0,0 +1,45 @@
+import pytest
+import random
+from assemblyline_core.workflow.run_workflow import WorkflowManager
+
+from assemblyline.common.isotime import now_as_iso
+from assemblyline.odm.models.workflow import Workflow
+from assemblyline.odm.random_data import create_alerts, wipe_alerts, wipe_workflows
+from assemblyline.odm.randomizer import random_minimal_obj
+
+
+@pytest.fixture(scope="module")
+def manager(datastore_connection):
+    try:
+        create_alerts(datastore_connection)
+        wipe_workflows(datastore_connection)
+        datastore_connection.alert.update_by_query("*", [(datastore_connection.alert.UPDATE_SET, 'reporting_ts', now_as_iso())])
+        datastore_connection.alert.commit()
+        yield WorkflowManager()
+    finally:
+        wipe_alerts(datastore_connection)
+
+def test_workflow(manager, datastore_connection):
+    # Create workflow that targets alerts based on YARA rule association
+    workflow = random_minimal_obj(Workflow)
+
+    yara_rule = random.choice(list(datastore_connection.alert.facet("al.yara").keys()))   
+    workflow.query = f'al.yara:"{yara_rule}"'
+    workflow.workflow_id = "AL_TEST"
+    workflow.labels = ["AL_TEST"]
+    workflow.priority = "LOW"
+    workflow.status = "MALICIOUS"    
+    datastore_connection.workflow.save(workflow.workflow_id, workflow)
+    datastore_connection.workflow.commit()
+
+    # Run Workflow manager to process new workflow against existing alerts
+    manager.running = True
+    manager.get_last_reporting_ts = lambda x: "now/d+1d"
+    manager.try_run(run_once=True)
+    datastore_connection.alert.commit()
+    
+    # Assert that custom labels were applied to alerts
+    assert datastore_connection.alert.search("label:AL_TEST", track_total_hits=True)['total']
+
+    # Assert that the change has been record in the alerts' event history
+    assert datastore_connection.alert.search(f"events.entity_id:{workflow.workflow_id}", track_total_hits=True)['total']

assemblyline_core-4.5.0.76/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.5.0.76