assemblyline-core 4.5.1.dev493__tar.gz → 4.5.1.dev497__tar.gz

{assemblyline_core-4.5.1.dev493 → assemblyline_core-4.5.1.dev497}/PKG-INFO

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

assemblyline_core-4.5.1.dev497/assemblyline_core/VERSION

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

{assemblyline_core-4.5.1.dev493 → assemblyline_core-4.5.1.dev497}/assemblyline_core/dispatching/dispatcher.py

@@ -351,15 +351,17 @@ class SubmissionTask:
             values=values,
         )
 
-    def partial_result(self, sha256, service_name):
+    def partial_result(self, sha256, service_name) -> bool:
         """Note that a partial result has been recieved. If a dispatch was requested process that now."""
         try:
             entry = self.monitoring[(sha256, service_name)]
         except KeyError:
-            return
+            return False
 
         if entry.dispatch_needed:
             self.redispatch_service(sha256, service_name)
+            return True
+        return False
 
     def clear_monitoring_entry(self, sha256, service_name):
         """A service has completed normally. If the service is monitoring clear out the record."""
@@ -1426,9 +1428,11 @@ class Dispatcher(ThreadedCoreBase):
         self.clear_timeout(task, sha256, service_name)
         task.service_logs.pop((sha256, service_name), None)
 
+        force_redispatch = set()
         if summary.partial:
             self.log.info("[%s/%s] %s returned partial results", sid, sha256, service_name)
-            task.partial_result(sha256, service_name)
+            if task.partial_result(sha256, service_name):
+                force_redispatch.add(sha256)
         else:
             task.clear_monitoring_entry(sha256, service_name)
 
@@ -1475,11 +1479,11 @@ class Dispatcher(ThreadedCoreBase):
             summary.children = [(c, 'EXTRACTED') for c in old_children]
 
         # Record the result as a summary
-        task.service_results[(sha256, service_name)] = summary
+        if not force_redispatch:
+            task.service_results[(sha256, service_name)] = summary
         task.register_children(sha256, [c for c, _ in summary.children])
 
         # Update the temporary data table for this file
-        force_redispatch = set()
         for key, value in (temporary_data or {}).items():
             if len(str(value)) <= self.config.submission.max_temp_data_length:
                 if task.temporary_data[sha256].set_value(key, value):

{assemblyline_core-4.5.1.dev493 → assemblyline_core-4.5.1.dev497}/assemblyline_core/ingester/ingester.py

@@ -18,34 +18,37 @@ from typing import Any, Iterable, List, Optional, Tuple
 
 import elasticapm
 
-from assemblyline.common.postprocess import ActionWorker
-from assemblyline_core.server_base import ThreadedCoreBase
+from assemblyline import odm
+from assemblyline.common import exceptions, forge, isotime
 from assemblyline.common.constants import DROP_PRIORITY
-from assemblyline.common.metrics import MetricsFactory
-from assemblyline.common.str_utils import dotdump, safe_str
 from assemblyline.common.exceptions import get_stacktrace_info
-from assemblyline.common.isotime import now, now_as_iso
 from assemblyline.common.importing import load_module_by_path
-from assemblyline.common import forge, exceptions, isotime
+from assemblyline.common.isotime import now, now_as_iso
+from assemblyline.common.metrics import MetricsFactory
+from assemblyline.common.postprocess import ActionWorker
+from assemblyline.common.str_utils import dotdump, safe_str
 from assemblyline.datastore.exceptions import DataStoreException
 from assemblyline.filestore import CorruptedFileStoreException, FileStoreException
+from assemblyline.odm.messages.ingest_heartbeat import Metrics
+from assemblyline.odm.messages.submission import Submission as MessageSubmission
+from assemblyline.odm.messages.submission import SubmissionMessage
+from assemblyline.odm.models.alert import EXTENDED_SCAN_VALUES
 from assemblyline.odm.models.filescore import FileScore
+from assemblyline.odm.models.submission import Submission as DatabaseSubmission
+from assemblyline.odm.models.submission import SubmissionParams
 from assemblyline.odm.models.user import User
-from assemblyline.odm.messages.ingest_heartbeat import Metrics
-from assemblyline.remote.datatypes.queues.named import NamedQueue
-from assemblyline.remote.datatypes.queues.priority import PriorityQueue
+from assemblyline.remote.datatypes.events import EventWatcher
+from assemblyline.remote.datatypes.hash import Hash
 from assemblyline.remote.datatypes.queues.comms import CommsQueue
 from assemblyline.remote.datatypes.queues.multi import MultiQueue
-from assemblyline.remote.datatypes.hash import Hash
+from assemblyline.remote.datatypes.queues.named import NamedQueue
+from assemblyline.remote.datatypes.queues.priority import PriorityQueue
 from assemblyline.remote.datatypes.user_quota_tracker import UserQuotaTracker
-from assemblyline import odm
-from assemblyline.odm.models.submission import SubmissionParams, Submission as DatabaseSubmission
-from assemblyline.odm.models.alert import EXTENDED_SCAN_VALUES
-from assemblyline.odm.messages.submission import Submission as MessageSubmission, SubmissionMessage
-
 from assemblyline_core.dispatching.dispatcher import Dispatcher
+from assemblyline_core.server_base import ThreadedCoreBase
 from assemblyline_core.submission_client import SubmissionClient
-from .constants import INGEST_QUEUE_NAME, drop_chance, COMPLETE_QUEUE_NAME
+
+from .constants import COMPLETE_QUEUE_NAME, INGEST_QUEUE_NAME, drop_chance
 
 _dup_prefix = 'w-m-'
 _notification_queue_prefix = 'nq-'
@@ -185,6 +188,11 @@ class Ingester(ThreadedCoreBase):
         self.async_submission_tracker = UserQuotaTracker('async_submissions', timeout=24 * 60 * 60,  # 1 day timeout
                                                          redis=self.redis_persist)
 
+        # Watchers
+        self.submission_delete_watcher = EventWatcher(self.redis)
+        self.submission_delete_watcher.register("delete.submission", self.handle_submission_delete)
+        self.submission_delete_watcher.start()
+
         if self.config.core.metrics.apm_server.server_url is not None:
             self.log.info(f"Exporting application metrics to: {self.config.core.metrics.apm_server.server_url}")
             elasticapm.instrument()
@@ -399,6 +407,14 @@ class Ingester(ThreadedCoreBase):
                 if self.apm_client:
                     self.apm_client.end_transaction('ingest_submit', 'exception')
 
+    def handle_submission_delete(self, sid: Optional[str]):
+        if not sid:
+            return
+
+        # Upon submission deletion, ensure to cleanup the local cache of filescores relative to the SID
+        with self.cache_lock:
+            self.cache = {k: v for k, v in self.cache.items() if v.sid != sid}
+
     def handle_complete(self):
         while self.running:
             result = self.complete_queue.pop(timeout=3)

{assemblyline_core-4.5.1.dev493 → assemblyline_core-4.5.1.dev497}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline_core-4.5.1.dev493 → assemblyline_core-4.5.1.dev497}/test/test_simulation.py

@@ -72,6 +72,9 @@ class MockService(ServerBase):
         self.dispatch_client = DispatchClient(self.datastore, redis)
         self.hits = dict()
         self.drops = dict()
+        self.finish = dict()
+        self.local_lock = threading.Event()
+        self.local_lock.set()
 
     def try_run(self):
         while self.running:
@@ -97,6 +100,9 @@ class MockService(ServerBase):
             if instructions.get('lock', False):
                 _global_semaphore.acquire(blocking=True, timeout=instructions['lock'])
 
+            if instructions.get('local_lock', False):
+                self.local_lock.wait(instructions['local_lock'])
+
             if 'drop' in instructions:
                 if instructions['drop'] >= hits:
                     self.drops[task.fileinfo.sha256] = self.drops.get(task.fileinfo.sha256, 0) + 1
@@ -124,8 +130,10 @@ class MockService(ServerBase):
                         partial = True
                         break
 
-            if partial:
-                print(self.service_name, "will produce partial results")
+                if partial:
+                    print(self.service_name, "will produce partial results")
+                else:
+                    print(self.service_name, "will produce complete results")
 
             result_data = {
                 'archive_ts': None,
@@ -151,15 +159,17 @@ class MockService(ServerBase):
                 result_key = result.build_key(get_random_id())
             self.dispatch_client.service_finished(task.sid, result_key, result,
                                                   temporary_data=instructions.get('temporary_data'))
+            self.finish[task.fileinfo.sha256] = self.finish.get(task.fileinfo.sha256, 0) + 1
 
 
 class CoreSession:
-    def __init__(self, config, ingest):
+    def __init__(self, config, ingest, services):
         self.ds: typing.Optional[AssemblylineDatastore] = None
         self.filestore = None
         self.redis = None
         self.config: Config = config
         self.ingest: Ingester = ingest
+        self.services: list[MockService] = services
         self.dispatcher: Dispatcher
 
     @property
@@ -253,7 +263,7 @@ def core(request, redis, filestore, config, clean_datastore: AssemblylineDatasto
     service_config: list[tuple[str, int, str, dict]] = [
         ('pre', 1, 'EXTRACT', {'extra_data': True, 'monitored_keys': ['passwords']}),
         ('core-a', 2, 'CORE', {}),
-        ('core-b', 1, 'CORE', {}),
+        ('core-b', 1, 'CORE', {'extra_data': True, 'monitored_keys': ['passwords']}),
         ('finish', 1, 'POST', {'extra_data': True})
     ]
 
@@ -282,7 +292,7 @@ def core(request, redis, filestore, config, clean_datastore: AssemblylineDatasto
 
     ingester = Ingester(datastore=ds, redis=redis, persistent_redis=redis, config=config)
 
-    fields = CoreSession(config, ingester)
+    fields = CoreSession(config, ingester, services)
     fields.redis = redis
     fields.ds = ds
 
@@ -1275,6 +1285,97 @@ def test_temp_data_monitoring(core: CoreSession, metrics):
     assert partial_results == 0, 'partial_results'
 
 
+def test_final_partial(core: CoreSession, metrics):
+    # This test was written to cover an error where a partial result produced as the final
+    # result of a submission would not trigger dispatching when it should due to data
+    # that was produced while it was running.
+
+    # Both services run at the same time, but one requires info from the other.
+    # We lock down the timing of the service completion so that:
+    # a) both run at the same time so that in its first run core-b does not have the
+    #    temp data it wants and produces a partial result.
+    # b) core-a finishes before core-b adding the temporary data to the dispatcher
+    # c) core-b finishes and should trigger a rerun with the data to produce a full result
+    sha, size = ready_body(core, {
+        'core-a': {'local_lock': 10, 'temporary_data': {'passwords': ['test_temp_data_monitoring']}},
+        'core-b': {'local_lock': 10, 'partial': {'passwords': 'test_temp_data_monitoring'}},
+    })
+
+    core_a = [s for s in core.services if s.service_name == 'core-a']
+    core_b = [s for s in core.services if s.service_name == 'core-b'][0]
+
+    for service in core_a + [core_b]:
+        service.local_lock.clear()
+
+    core.ingest_queue.push(SubmissionInput(dict(
+        metadata={},
+        params=dict(
+            description="file abc123",
+            services=dict(selected=['core-a', 'core-b']),
+            submitter='user',
+            groups=['user'],
+            max_extracted=10000
+        ),
+        notification=dict(
+            queue='temp-final-partial',
+            threshold=0
+        ),
+        files=[dict(
+            sha256=sha,
+            size=size,
+            name='abc123'
+        )]
+    )).as_primitives())
+ 
+    # Wait until both of the services have started (so service b doesn't get the temp data a produces on its first run)
+    start = time.time()
+    while sum(s.hits.get(sha, 0) for s in core.services) != 2:
+        if time.time() - start > RESPONSE_TIMEOUT:
+            pytest.fail()
+        time.sleep(0.01)
+
+    # Release a
+    for service in core_a:
+        service.local_lock.set()
+
+    # Let a finish so that the temporary data is added in the dispatcher
+    while sum(s.finish.get(sha, 0) for s in core_a) < 1:
+        if time.time() - start > RESPONSE_TIMEOUT:
+            pytest.fail()
+        time.sleep(0.01)    
+
+    # Let b finish, it should produce a partial result then rerun right away
+    core_b.local_lock.set()
+
+    notification_queue = NamedQueue('nq-temp-final-partial', core.redis)
+    dropped_task = notification_queue.pop(timeout=RESPONSE_TIMEOUT)
+    assert dropped_task
+    dropped_task = IngestTask(dropped_task)
+    sub: Submission = core.ds.submission.get(dropped_task.submission.sid)
+
+    # The submission should produce no errors and two results
+    assert len(sub.errors) == 0
+    assert len(sub.results) == 2, sub.results
+
+    # b service should have run twice to produce the results
+    assert core_b.hits[sha] >= 2, f'core_b.hits {core_b.hits[sha]}'
+
+    # Wait until we get feedback from the metrics channel
+    metrics.expect('ingester', 'submissions_ingested', 1)
+    metrics.expect('ingester', 'submissions_completed', 1)
+    metrics.expect('dispatcher', 'submissions_completed', 1)
+    metrics.expect('dispatcher', 'files_completed', 1)
+
+    # Verify thath there are no partial results in the final submission
+    partial_results = 0
+    for res in sub.results:
+        result = core.ds.get_single_result(res, as_obj=True)
+        assert result is not None, res
+        if result.partial:
+            partial_results += 1
+    assert partial_results == 0, 'partial_results'
+
+
 def test_complex_extracted(core: CoreSession, metrics):
     # stages to this processing when everything goes well
     # 1. extract a file that will process to produce a partial result

assemblyline_core-4.5.1.dev493/assemblyline_core/VERSION

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