assemblyline-core 4.4.1.dev278__tar.gz → 4.4.1.dev281__tar.gz

{assemblyline-core-4.4.1.dev278 → assemblyline-core-4.4.1.dev281}/PKG-INFO

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

assemblyline-core-4.4.1.dev281/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.4.1.dev281

{assemblyline-core-4.4.1.dev278 → assemblyline-core-4.4.1.dev281}/assemblyline_core/archiver/run_archiver.py

@@ -86,9 +86,9 @@ class Archiver(ServerBase):
                                                      index_type=Index.HOT)
 
                 # Gather list of files and archives them
-                files = {f.sha256 for f in submission.files}
-                files.update(self.datastore.get_file_list_from_keys(submission.results, supplementary=True))
-                for sha256 in files:
+                files = {(f.sha256, False) for f in submission.files}
+                files.update(self.datastore.get_file_list_from_keys(submission.results))
+                for sha256, supplementary in files:
                     self.counter.increment('file')
 
                     # Get the tags for this file
@@ -111,7 +111,7 @@ class Archiver(ServerBase):
                     operations += [(self.datastore.file.UPDATE_APPEND_IF_MISSING, 'labels', x) for x in techniques]
                     operations += [(self.datastore.file.UPDATE_APPEND_IF_MISSING, 'labels', x) for x in infos]
 
-                    # create type specific labels
+                    # Create type specific labels
                     operations += [
                         (self.datastore.file.UPDATE_APPEND_IF_MISSING, 'label_categories.attribution', x)
                         for x in attributions]
@@ -122,6 +122,9 @@ class Archiver(ServerBase):
                         (self.datastore.file.UPDATE_APPEND_IF_MISSING, 'label_categories.info', x)
                         for x in infos]
 
+                    # Set the is_supplementary property
+                    operations += [(self.datastore.file.UPDATE_SET, 'is_supplementary', supplementary)]
+
                     # Apply auto-created labels
                     self.datastore.file.update(sha256, operations=operations, index_type=Index.ARCHIVE)
                     self.datastore.file.update(sha256, operations=operations, index_type=Index.HOT)

{assemblyline-core-4.4.1.dev278 → assemblyline-core-4.4.1.dev281}/assemblyline_core/plumber/run_plumber.py

@@ -7,18 +7,21 @@ disabled or deleted for which a service queue exists, the dispatcher will be inf
 had an error.
 """
 import threading
+import warnings
 from typing import Optional
-from assemblyline.common.forge import get_service_queue
 
-from assemblyline.odm.models.error import Error
-from assemblyline.common.isotime import now_as_iso
 from assemblyline.common.constants import service_queue_name
+from assemblyline.common.forge import get_service_queue
+from assemblyline.common.isotime import now, now_as_iso
+from assemblyline.common.security import generate_random_secret
+from assemblyline.datastore.store import TRANSPORT_TIMEOUT
+from assemblyline.odm.models.error import Error
 from assemblyline.odm.models.service import Service
 from assemblyline.remote.datatypes import retry_call
 from assemblyline.remote.datatypes.queues.named import NamedQueue
-
 from assemblyline_core.dispatching.client import DispatchClient
 from assemblyline_core.server_base import CoreBase, ServiceStage
+from elasticsearch import Elasticsearch
 
 DAY = 60 * 60 * 24
 TASK_DELETE_CHUNK = 10000
@@ -37,6 +40,31 @@ class Plumber(CoreBase):
         self.stop_signals: dict[str, threading.Event] = {}
         self.service_limit: dict[str, int] = {}
 
+        # Ensure roles for "plumber" user are created
+        self.datastore.ds.with_retries(
+            self.datastore.ds.client.security.put_role,
+            name="manage_tasks",
+            indices=[{"names": [".tasks"], "privileges": ["all"], "allow_restricted_indices": True}])
+
+        # Initialize/update 'plumber' user in Elasticsearch to perform cleanup
+        password = generate_random_secret()
+        self.datastore.ds.with_retries(
+            self.datastore.ds.client.security.put_user,
+            username="plumber",
+            password=password,
+            roles=["manage_tasks", "superuser"]
+        )
+
+        # Close existing connection and re-connect to the datastore as "plumber" user
+        self.datastore.ds.client.close()
+        self.datastore.ds.client = Elasticsearch(hosts=self.datastore.ds.get_hosts(),
+                                                 basic_auth=("plumber", password),
+                                                 max_retries=0,
+                                                 request_timeout=TRANSPORT_TIMEOUT,
+                                                 ca_certs=self.datastore.ds.ca_certs)
+        if not self.datastore.ds.ping():
+            raise Exception("Unable to connect to datastore as 'plumber'")
+
     def stop(self):
         for sig in self.stop_signals.values():
             sig.set()
@@ -151,7 +179,21 @@ class Plumber(CoreBase):
     def cleanup_old_tasks(self):
         self.log.info("Cleaning up task index for old completed tasks...")
         while self.running:
-            deleted = self.datastore.task_cleanup(deleteable_task_age=DAY, max_tasks=TASK_DELETE_CHUNK)
+            # Create a new task to delete expired tasks
+            # NOTE: This will delete up to 10000 completed tasks older then a day
+            q = f"completed:true AND task.start_time_in_millis:<{now(-1 * DAY) * 1000}"
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                task = self.datastore.ds.with_retries(self.datastore.ds.client.delete_by_query, index='.tasks',
+                                                      q=q, wait_for_completion=False, conflicts='proceed',
+                                                      max_docs=TASK_DELETE_CHUNK)
+
+            # Wait until the tasks deletion task is over
+            res = self.datastore.ds._get_task_results(task)
+
+            # Get the number of deleted items
+            deleted = res['deleted']
+
             if not deleted:
                 self.sleep(self.delay)
             else:

{assemblyline-core-4.4.1.dev278 → assemblyline-core-4.4.1.dev281}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline-core-4.4.1.dev278 → assemblyline-core-4.4.1.dev281}/test/test_plumber.py

@@ -1,15 +1,13 @@
+from time import sleep
 from unittest import mock
-from redis import Redis
 
 from assemblyline.odm.messages.task import Task
-
-from assemblyline.odm.random_data import random_model_obj
 from assemblyline.odm.models.service import Service
-
+from assemblyline.odm.random_data import random_model_obj
 from assemblyline_core.plumber.run_plumber import Plumber
 from assemblyline_core.server_base import ServiceStage
-
 from mocking import TrueCountTimes
+from redis import Redis
 
 
 def test_expire_missing_service():
@@ -24,6 +22,9 @@ def test_expire_missing_service():
     service_a.enabled = True
 
     datastore.list_all_services.return_value = [service_a]
+    datastore.ds.ca_certs = None
+    datastore.ds.get_hosts.return_value = ["http://localhost:9200"]
+
     plumber = Plumber(redis=redis, redis_persist=redis_persist, datastore=datastore, delay=1)
     plumber.get_service_stage = mock.MagicMock(return_value=ServiceStage.Running)
     plumber.dispatch_client = mock.MagicMock()
@@ -51,6 +52,8 @@ def test_flush_paused_queues():
     service_a.enabled = True
 
     datastore.list_all_services.return_value = [service_a]
+    datastore.ds.ca_certs = None
+    datastore.ds.get_hosts.return_value = ["http://localhost:9200"]
 
     plumber = Plumber(redis=redis, redis_persist=redis_persist, datastore=datastore, delay=1)
     plumber.get_service_stage = mock.MagicMock(return_value=ServiceStage.Running)
@@ -72,3 +75,35 @@ def test_flush_paused_queues():
     assert plumber.dispatch_client.service_failed.call_count == 1
     args = plumber.dispatch_client.service_failed.call_args
     assert args[0][0] == task.sid
+
+
+def test_cleanup_old_tasks(datastore_connection):
+    # Create a bunch of random "old" tasks and clean them up
+    redis = mock.MagicMock(spec=Redis)
+    redis_persist = mock.MagicMock(spec=Redis)
+    plumber = Plumber(redis=redis, redis_persist=redis_persist, datastore=datastore_connection, delay=1)
+
+    # Generate new documents in .tasks index
+    num_old_tasks = 10
+    [plumber.datastore.ds.client.index(index=".tasks", document={
+        "completed": True,
+        "task": {
+            "start_time_in_millis": 0
+        }
+    }) for _ in range(num_old_tasks)]
+    sleep(1)
+
+    # Assert that these have been indeed committed to the tasks index
+    assert plumber.datastore.ds.client.search(index='.tasks',
+                                              q="task.start_time_in_millis:0",
+                                              track_total_hits=True,
+                                              size=0)['hits']['total']['value'] == num_old_tasks
+
+    # Run task cleanup, we should return to no more "old" completed tasks
+    plumber.running = TrueCountTimes(count=1)
+    plumber.cleanup_old_tasks()
+    sleep(1)
+    assert plumber.datastore.ds.client.search(index='.tasks',
+                                              q="task.start_time_in_millis:0",
+                                              track_total_hits=True,
+                                              size=0)['hits']['total']['value'] == 0

assemblyline-core-4.4.1.dev278/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.4.1.dev278