assemblyline-core 4.5.1.dev145__tar.gz → 4.5.1.dev151__tar.gz

{assemblyline-core-4.5.1.dev145 → assemblyline-core-4.5.1.dev151}/PKG-INFO

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

assemblyline-core-4.5.1.dev151/assemblyline_core/VERSION

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

{assemblyline-core-4.5.1.dev145 → assemblyline-core-4.5.1.dev151}/assemblyline_core/archiver/run_archiver.py

@@ -64,7 +64,7 @@ class Archiver(ServerBase):
             try:
                 if len(message) == 3:
                     archive_type, type_id, delete_after = message
-                    metadata = None
+                    metadata = {}
                     use_alternate_dtl = False
                 elif len(message) == 4:
                     archive_type, type_id, delete_after, metadata = message
@@ -90,8 +90,9 @@ class Archiver(ServerBase):
                         submission, version = self.datastore.submission.get_if_exists(type_id, version=True)
 
                         # If we have metadata passed in the message, we need to apply it before archiving the submission
-                        if metadata and self.config.core.archiver.use_metadata:
-                            submission.metadata.update({f"archive.{k}": v for k, v in metadata.items()})
+                        if metadata and self.config.submission.metadata.archive:
+                            submission.metadata.update({k: v for k, v in metadata.items()
+                                                        if k not in submission.metadata})
                             self.datastore.submission.save(type_id, submission, version=version)
 
                         break

{assemblyline-core-4.5.1.dev145 → assemblyline-core-4.5.1.dev151}/assemblyline_core/expiry/run_expiry.py

@@ -4,38 +4,48 @@ from __future__ import annotations
 import concurrent.futures
 import threading
 import functools
-import elasticapm
 import time
-
 from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, Future, as_completed
 from concurrent.futures.process import BrokenProcessPool
-from datemath import dm
 from typing import Callable, Optional, TYPE_CHECKING
 
-from assemblyline.common.isotime import epoch_to_iso, now_as_iso
-from assemblyline.datastore.collection import Index
+import elasticapm
+from datemath import dm
+
 from assemblyline_core.server_base import ServerBase
 from assemblyline_core.dispatching.dispatcher import BAD_SID_HASH
 from assemblyline.common import forge
+from assemblyline.common.isotime import epoch_to_iso, now_as_iso
 from assemblyline.common.metrics import MetricsFactory
 from assemblyline.filestore import FileStore
 from assemblyline.odm.messages.expiry_heartbeat import Metrics
 from assemblyline.remote.datatypes import get_client
+from assemblyline.datastore.collection import Index
 from assemblyline.remote.datatypes.set import Set
 
 if TYPE_CHECKING:
     from assemblyline.datastore.collection import ESCollection
 
 
-def file_delete_worker(logger, filestore_urls, file_batch) -> list[str]:
+def file_delete_worker(logger, filestore_urls, file_batch, archive_filestore_urls=None) -> list[tuple[str, bool]]:
     try:
         filestore = FileStore(*filestore_urls)
-
-        def filestore_delete(sha256: str) -> Optional[str]:
-            filestore.delete(sha256)
-            if not filestore.exists(sha256):
-                return sha256
-            return None
+        if archive_filestore_urls and filestore_urls != archive_filestore_urls:
+            archivestore = FileStore(*archive_filestore_urls)
+        else:
+            archivestore = filestore
+
+        def filestore_delete(item: tuple[str, bool]) -> tuple[Optional[str], Optional[bool]]:
+            sha256, from_archive = item
+            if from_archive:
+                archivestore.delete(sha256)
+                if not archivestore.exists(sha256):
+                    return sha256, True
+            else:
+                filestore.delete(sha256)
+                if not filestore.exists(sha256):
+                    return sha256, False
+            return None, None
 
         return _file_delete_worker(logger, filestore_delete, file_batch)
 
@@ -44,8 +54,11 @@ def file_delete_worker(logger, filestore_urls, file_batch) -> list[str]:
     return []
 
 
-def _file_delete_worker(logger, delete_action: Callable[[str], Optional[str]], file_batch) -> list[str]:
-    finished_files: list[str] = []
+ActionSignature = Callable[[tuple[str, bool]], tuple[Optional[str], Optional[bool]]]
+
+
+def _file_delete_worker(logger, delete_action: ActionSignature, file_batch) -> list[tuple[str, bool]]:
+    finished_files: list[tuple[str, bool]] = []
     try:
         futures = []
 
@@ -55,9 +68,9 @@ def _file_delete_worker(logger, delete_action: Callable[[str], Optional[str]], f
 
             for future in as_completed(futures):
                 try:
-                    erased_name = future.result()
-                    if erased_name:
-                        finished_files.append(erased_name)
+                    erased_name, from_archive = future.result()
+                    if erased_name and from_archive is not None:
+                        finished_files.append((erased_name, from_archive))
                 except Exception as error:
                     logger.exception("Error in filestore worker: " + str(error))
 
@@ -67,17 +80,29 @@ def _file_delete_worker(logger, delete_action: Callable[[str], Optional[str]], f
 
 
 class ExpiryManager(ServerBase):
-    def __init__(self, redis_persist=None):
-        self.config = forge.get_config()
+    def __init__(self, redis_persist=None, datastore=None, filestore=None, config=None, classification=None):
+        self.config = config or forge.get_config()
 
         super().__init__('assemblyline.expiry', shutdown_timeout=self.config.core.expiry.sleep_time + 5)
-        self.datastore = forge.get_datastore(config=self.config)
-        self.filestore = forge.get_filestore(config=self.config)
-        self.classification = forge.get_classification()
+
+        # Set Archive related configs
+        if self.config.datastore.archive.enabled:
+            self.archive_access = True
+            self.index_type = Index.HOT_AND_ARCHIVE
+        else:
+            self.archive_access = False
+            self.index_type = Index.HOT
+
+        self.datastore = datastore or forge.get_datastore(config=self.config, archive_access=self.archive_access)
+        self.filestore = filestore or forge.get_filestore(config=self.config)
+        self.classification = classification or forge.get_classification()
         self.expirable_collections: list[ESCollection] = []
         self.counter = MetricsFactory('expiry', Metrics)
         self.file_delete_worker = ProcessPoolExecutor(self.config.core.expiry.delete_workers)
-        self.same_storage = self.config.filestore.storage == self.config.filestore.archive
+        if self.config.filestore.archive:
+            self.same_storage = self.config.filestore.storage == self.config.filestore.archive
+        else:
+            self.same_storage = True
         self.current_submission_cleanup = set()
 
         self.redis_persist = redis_persist or get_client(
@@ -127,14 +152,15 @@ class ExpiryManager(ServerBase):
     def filestore_delete(self, file_batch, _):
         return self.file_delete_worker.submit(file_delete_worker, logger=self.log,
                                               filestore_urls=list(self.config.filestore.storage),
-                                              file_batch=file_batch)
+                                              file_batch=file_batch,
+                                              archive_filestore_urls=list(self.config.filestore.archive))
 
     def cachestore_delete(self, file_batch, _):
         return self.file_delete_worker.submit(file_delete_worker, logger=self.log,
                                               filestore_urls=list(self.config.filestore.cache),
                                               file_batch=file_batch)
 
-    def _finish_delete(self, collection: ESCollection, task: Future, expire_only: list[str]):
+    def _finish_delete(self, collection: ESCollection, task: Future, expire_only: list[tuple[str, bool]]):
         # Wait until the worker process finishes deleting files
         file_list: list[str] = []
         while self.running:
@@ -145,25 +171,41 @@ class ExpiryManager(ServerBase):
             except concurrent.futures.TimeoutError:
                 pass
 
-        file_list.extend(expire_only)
-
-        # build a batch delete job for all the removed files
-        bulk = collection.get_bulk_plan()
-        for sha256 in file_list:
-            bulk.add_delete_operation(sha256)
-
-        if len(file_list) > 0:
+        if file_list:
             self.log.info(f'[{collection.name}] Deleted associated files from the '
                           f'{"cachestore" if "cache" in collection.name else "filestore"}...')
-            collection.bulk(bulk)
-            self.counter.increment(f'{collection.name}', increment_by=len(file_list))
-            self.log.info(f"[{collection.name}] Deleted {len(file_list)} items from the datastore...")
         else:
+            self.log.info(f'[{collection.name}] Nothing was deleted from the '
+                          f'{"cachestore" if "cache" in collection.name else "filestore"}...')
+
+        # From the files to be deleted, check which are from the hot index
+        hot_file_list = [x[0] for x in file_list if not x[1]]
+        hot_file_list.extend([x[0] for x in expire_only if not x[1]])
+
+        # From the files to be deleted, check which are from the archive index
+        archive_file_list = [x[0] for x in file_list if x[1]]
+        archive_file_list.extend([x[0] for x in expire_only if x[1]])
+
+        for cur_file_list, index_type in [(hot_file_list, Index.HOT), (archive_file_list, Index.ARCHIVE)]:
+            if not cur_file_list:
+                # Nothing to delete from this index type
+                continue
+
+            # build a batch delete job for all the removed files
+            bulk = collection.get_bulk_plan(index_type=index_type)
+            for sha256 in cur_file_list:
+                bulk.add_delete_operation(sha256)
+
+            collection.bulk(bulk)
+            self.counter.increment(f'{collection.name}', increment_by=len(cur_file_list))
+            self.log.info(f"[{collection.name}] Deleted {len(cur_file_list)} items from the datastore...")
+
+        if not hot_file_list and not archive_file_list:
             self.log.warning(f'[{collection.name}] Expiry unable to clean up any of the files in filestore.')
 
-    def _simple_delete(self, collection, delete_query, number_to_delete):
+    def _simple_delete(self, collection: ESCollection, delete_query, number_to_delete):
         self.heartbeat()
-        collection.delete_by_query(delete_query)
+        collection.delete_by_query(delete_query, index_type=self.index_type)
         self.counter.increment(f'{collection.name}', increment_by=number_to_delete)
         self.log.info(f"[{collection.name}] Deleted {number_to_delete} items from the datastore...")
 
@@ -195,22 +237,45 @@ class ExpiryManager(ServerBase):
         # check if we are dealing with an index that needs file cleanup
         if self.config.core.expiry.delete_storage and collection.name in self.fs_hashmap:
             # Delete associated files
-            delete_objects: list[str] = []
-            for item in collection.stream_search(delete_query, fl='id', as_obj=False):
+            delete_objects: list[tuple[str, bool]] = []
+            for item in collection.stream_search(
+                    delete_query, fl='id,from_archive', as_obj=False, index_type=self.index_type):
                 self.heartbeat()
-                delete_objects.append(item['id'])
+                delete_objects.append((item['id'], item.get('from_archive', False)))
 
             # Filter archived documents if archive filestore is the same as the filestore
-            expire_only = []
-            if self.same_storage and self.config.datastore.archive.enabled and collection.name == 'file':
-                archived_files = self.datastore.file.multiexists(delete_objects, index_type=Index.ARCHIVE)
-                delete_objects = [k for k, v in archived_files.items() if not v]
-                expire_only = [k for k, v in archived_files.items() if v]
+            expire_only: list[tuple[str, bool]] = []
+            if self.same_storage and self.archive_access and collection.name == 'file':
+                # Separate hot and archive files
+                delete_from_archive = [i[0] for i in delete_objects if i[1]]
+                delete_from_hot = [i[0] for i in delete_objects if not i[1]]
+
+                # Check for overlap
+                overlap = set(delete_from_archive).intersection(set(delete_from_hot))
+                delete_from_archive = list(set(delete_from_archive)-overlap)
+                delete_from_hot = list(set(delete_from_hot)-overlap)
+
+                # Create the original delete_object form the overlap
+                delete_objects = [(k, False) for k in overlap]
+                delete_objects.extend([(k, True) for k in overlap])
+
+                if delete_from_hot:
+                    # Check hot objects to delete if they are in archive
+                    archived_files = self.datastore.file.multiexists(delete_from_hot, index_type=Index.ARCHIVE)
+                    delete_objects.extend([(k, False) for k, v in archived_files.items() if not v])
+                    expire_only.extend([(k, False) for k, v in archived_files.items() if v])
+
+                if delete_from_archive:
+                    # Check hot objects to delete if they are in archive
+                    hot_files = self.datastore.file.multiexists(delete_from_archive, index_type=Index.HOT)
+                    delete_objects.extend([(k, True) for k, v in hot_files.items() if not v])
+                    expire_only.extend([(k, True) for k, v in hot_files.items() if v])
 
             delete_tasks = self.fs_hashmap[collection.name](delete_objects, final_date)
 
             # Proceed with deletion, but only after all the scheduled deletes for this
-            self.log.info(f"[{collection.name}] Scheduled {len(delete_objects)}/{number_to_delete} files to be removed")
+            self.log.info(f"[{collection.name}] Scheduled {len(delete_objects)}/{number_to_delete} files to be "
+                          f"removed from the {'cachestore' if 'cache' in collection.name else 'filestore'}")
             self._finish_delete(collection, delete_tasks, expire_only)
 
         else:
@@ -264,7 +329,7 @@ class ExpiryManager(ServerBase):
            will be affected in between start date and the date found"""
         rows = collection.search(f"expiry_ts: {{{start} TO {final_date}]", rows=1,
                                  offset=self.expiry_size - 1, sort='expiry_ts asc',
-                                 as_obj=False, fl='expiry_ts')
+                                 as_obj=False, fl='expiry_ts', index_type=self.index_type)
         if rows['items']:
             return rows['items'][0]['expiry_ts'], self.expiry_size
         return final_date, rows['total']

{assemblyline-core-4.5.1.dev145 → assemblyline-core-4.5.1.dev151}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline-core-4.5.1.dev145 → assemblyline-core-4.5.1.dev151}/test/test_dispatcher.py

@@ -284,8 +284,8 @@ def test_dispatch_extracted(clean_redis, clean_datastore):
 @mock.patch('assemblyline_core.dispatching.dispatcher.MetricsFactory', mock.MagicMock())
 @mock.patch('assemblyline_core.dispatching.dispatcher.Scheduler', DRPScheduler)
 def test_dispatch_extracted_bypass_drp(clean_redis, clean_datastore):
-    # Dynamic Recursion Prevention is to prevent services belonging to the 'Dynamic Analysis' from analyzing the children
-    # of files they've analyzed.
+    # Dynamic Recursion Prevention is to prevent services belonging to the 'Dynamic Analysis'
+    # from analyzing the children of files they've analyzed.
 
     # The bypass should allow services to specify files to run through Dynamic Analysis regardless of the
     # Dynamic Recursion Prevention parameter.
@@ -372,7 +372,7 @@ def test_dispatch_extracted_bypass_drp(clean_redis, clean_datastore):
     disp.service_worker(disp.process_queue_index(sid))
 
     # 'sandbox' should have a task for the extracted file
-    #disp.dispatch_file(disp.tasks.get(sid), second_file_hash)
+    # disp.dispatch_file(disp.tasks.get(sid), second_file_hash)
     job = client.request_work('0', 'sandbox', '0')
     assert job.fileinfo.sha256 == second_file_hash
     assert job.filename == 'second-*'

{assemblyline-core-4.5.1.dev145 → assemblyline-core-4.5.1.dev151}/test/test_expiry.py

@@ -4,6 +4,7 @@ import random
 import concurrent.futures
 
 from assemblyline.common.isotime import now_as_iso
+from assemblyline.datastore.helper import AssemblylineDatastore
 from assemblyline.odm.randomizer import random_model_obj
 
 from assemblyline_core.expiry.run_expiry import ExpiryManager
@@ -14,34 +15,31 @@ expiry_collections_len = {}
 archive_collections_len = {}
 
 
-@pytest.fixture(scope='module')
-def datastore(archive_connection):
-    return archive_connection
-
-
-def purge_data(datastore):
-    for name, definition in datastore.ds.get_models().items():
+def purge_data(datastore_connection: AssemblylineDatastore):
+    for name, definition in datastore_connection.ds.get_models().items():
         if hasattr(definition, 'expiry_ts'):
-            getattr(datastore, name).wipe()
+            getattr(datastore_connection, name).wipe()
 
 
 @pytest.fixture(scope="function")
-def ds_expiry(request, datastore):
-    for name, definition in datastore.ds.get_models().items():
+def ds_expiry(request, datastore_connection):
+    for name, definition in datastore_connection.ds.get_models().items():
         if hasattr(definition, 'expiry_ts'):
-            collection = getattr(datastore, name)
+            collection = getattr(datastore_connection, name)
             collection.wipe()
             expiry_len = random.randint(MIN_OBJECTS, MAX_OBJECTS)
             for x in range(expiry_len):
                 obj = random_model_obj(collection.model_class)
+                if hasattr(definition, 'from_archive'):
+                    obj.from_archive = False
                 obj.expiry_ts = now_as_iso(-10000)
                 collection.save('longer_name'+str(x), obj)
 
             expiry_collections_len[name] = expiry_len
             collection.commit()
 
-    request.addfinalizer(lambda: purge_data(datastore))
-    return datastore
+    request.addfinalizer(lambda: purge_data(datastore_connection))
+    return datastore_connection
 
 
 class FakeCounter(object):
@@ -58,8 +56,8 @@ class FakeCounter(object):
         return self.counts.get(name, 0)
 
 
-def test_expire_all(ds_expiry):
-    expiry = ExpiryManager()
+def test_expire_all(config, ds_expiry, filestore):
+    expiry = ExpiryManager(config=config, datastore=ds_expiry, filestore=filestore)
     expiry.running = True
     expiry.counter = FakeCounter()
     with concurrent.futures.ThreadPoolExecutor(5) as pool:

assemblyline-core-4.5.1.dev145/assemblyline_core/VERSION

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