assemblyline 4.7.5.dev13__tar.gz → 4.7.5.dev15__tar.gz

{assemblyline-4.7.5.dev13/assemblyline.egg-info → assemblyline-4.7.5.dev15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline
-Version: 4.7.5.dev13
+Version: 4.7.5.dev15
 Summary: Assemblyline 4 - Automated malware analysis framework
 Home-page: https://github.com/CybercentreCanada/assemblyline-base
 Author: CCCS Assemblyline development team

assemblyline-4.7.5.dev15/assemblyline/VERSION

@@ -0,0 +1 @@
+4.7.5.dev15

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev15}/assemblyline/datastore/collection.py

@@ -962,6 +962,8 @@ class ESCollection(Generic[ModelType]):
                     # If this index has an archive, check is the document was found in it.
                     if self.archive_name:
                         doc['_source']['from_archive'] = self.is_archive_index(index)
+                    elif 'from_archive' in doc['_source']:
+                        doc['_source']['from_archive'] = False
 
                     if version:
                         return self._normalize_output(doc['_source']), f"{doc['_seq_no']}---{doc['_primary_term']}"
@@ -1012,14 +1014,25 @@ class ESCollection(Generic[ModelType]):
 
     @typing.overload
     def get_if_exists(self, key: str, as_obj: typing.Literal[True] = True, index_type: typing.Optional[Index] = None,
-                      version=False) -> typing.Optional[ModelType]:
+                      version: typing.Literal[False] = False) -> typing.Optional[ModelType]:
         ...
 
     @typing.overload
     def get_if_exists(self, key: str, as_obj: typing.Literal[False], index_type: typing.Optional[Index] = None,
-                      version=False) -> typing.Optional[dict]:
+                      version: typing.Literal[False] = False) -> typing.Optional[dict]:
         ...
 
+    @typing.overload
+    def get_if_exists(self, key: str, as_obj: typing.Literal[True] = True, index_type: typing.Optional[Index] = None,
+                      version=True) -> tuple[typing.Optional[ModelType], str]:
+        ...
+
+    @typing.overload
+    def get_if_exists(self, key: str, as_obj: typing.Literal[False], index_type: typing.Optional[Index] = None,
+                      version=True) -> tuple[typing.Optional[dict], str]:
+        ...
+
+
     def get_if_exists(self, key, as_obj=True, index_type=None, version=False):
         """
         Get a document from the datastore but do not retry if not found.
@@ -1378,6 +1391,8 @@ class ESCollection(Generic[ModelType]):
         # If this index has an archive, check is the document was found in it.
         if self.archive_name:
             source_data['from_archive'] = self.is_archive_index(result['_index'])
+        elif 'from_archive' in source_data:
+            source_data['from_archive'] = False
 
         if self.model_class:
             if not fields:

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev15}/assemblyline/filestore/__init__.py

@@ -128,7 +128,7 @@ def create_transport(url, connection_attempts=None):
 
     elif scheme == 's3':
         valid_str_keys = ['aws_region', 's3_bucket']
-        valid_bool_keys = ['use_ssl', 'verify', 'boto_defaults', 'read_only']
+        valid_bool_keys = ['use_ssl', 'verify', 'boto_defaults', 'read_only', 'use_batch_delete']
         extras = _get_extras(parse_qs(parsed.query), valid_str_keys=valid_str_keys, valid_bool_keys=valid_bool_keys)
 
         # If user/password not specified, access might be dictated by IAM roles
@@ -200,6 +200,24 @@ class FileStore(object):
                     trace = get_stacktrace_info(ex)
                     self.log.info('Transport problem: %s', trace)
 
+    @elasticapm.capture_span(span_type='filestore')
+    def delete_batch(self, file_list: list[str], location='all'):
+        # Don't even call into the transports if there are no files given
+        if not file_list:
+            return
+
+        with elasticapm.capture_span(name='delete_batch', span_type='filestore', labels={'batch_size': len(file_list)}):
+            for t in self.slice(location):
+                if t.read_only:
+                    # Don't attempt to delete from read only transports
+                    continue
+
+                try:
+                    t.delete_batch(file_list)
+                except Exception as ex:
+                    trace = get_stacktrace_info(ex)
+                    self.log.info('Transport problem: %s', trace)
+
     @elasticapm.capture_span(span_type='filestore')
     def download(self, src_path: str, dest_path: str, location='any'):
         successful = False

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev15}/assemblyline/filestore/transport/azure.py

@@ -15,6 +15,7 @@ from azure.core.exceptions import (
     ServiceRequestError,
     TooManyRedirectsError,
 )
+from azure.core.pipeline.transport import HttpResponse
 from azure.identity import (
     ClientSecretCredential,
     DefaultAzureCredential,
@@ -22,6 +23,7 @@ from azure.identity import (
 )
 from azure.storage.blob import BlobServiceClient
 
+from assemblyline.common import chunk
 from assemblyline.common.exceptions import ChainAll
 from assemblyline.filestore.transport.base import Transport, TransportException
 
@@ -144,6 +146,17 @@ class TransportAzure(Transport):
             if not isinstance(error.cause, ResourceNotFoundError):
                 raise
 
+    def delete_batch(self, file_list: Iterable[str]):
+        """Deletes a batch of files."""
+        container_client = self.service_client.get_container_client(self.blob_container)
+        keys = [self.normalize(path) for path in file_list]
+        for key_chunk in chunk.chunk(keys, 256):
+            responses: Iterable[HttpResponse] = self.with_retries(container_client.delete_blobs(*key_chunk))
+            for response in responses:
+                if response.status_code == 404:
+                    continue
+                response.raise_for_status()
+
     def exists(self, path):
         key = self.normalize(path)
         blob_client = self.service_client.get_blob_client(self.blob_container, key)

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev15}/assemblyline/filestore/transport/base.py

@@ -46,6 +46,13 @@ class Transport(object):
         """
         raise TransportException("Not Implemented")
 
+    def delete_batch(self, file_list: Iterable[str]):
+        """
+        Deletes a batch of files
+        """
+        for path in file_list:
+            self.delete(path)
+
     def exists(self, path: str) -> bool:
         """
         Returns True if the path exists, False otherwise.

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev15}/assemblyline/filestore/transport/s3.py

@@ -12,6 +12,7 @@ from botocore.exceptions import (
     EndpointConnectionError,
 )
 
+from assemblyline.common import chunk
 from assemblyline.common.exceptions import ChainAll
 from assemblyline.filestore.transport.base import Transport, TransportException
 
@@ -35,12 +36,14 @@ class TransportS3(Transport):
     DEFAULT_HOST = "s3.amazonaws.com"
 
     def __init__(self, base=None, accesskey=None, secretkey=None, aws_region=None, s3_bucket="al-storage",
-                 host=None, port=None, use_ssl=None, verify=True, connection_attempts=None, boto_defaults=False, read_only=False):
+                 host=None, port=None, use_ssl=None, verify=True, connection_attempts=None,
+                 boto_defaults=False, read_only=False, use_batch_delete: bool = False):
         self.log = logging.getLogger('assemblyline.transport.s3')
         self.base = base
         self.bucket = s3_bucket
         self.accesskey = accesskey
         self.retry_limit: int = connection_attempts
+        self.use_batch_delete: bool = use_batch_delete
 
         if use_ssl is None:
             self.use_ssl = True
@@ -139,6 +142,25 @@ class TransportS3(Transport):
         key = self.normalize(path)
         self.with_retries(self.client.delete_object, Bucket=self.bucket, Key=key)
 
+    def delete_batch(self, file_list: Iterable[str]):
+        """Deletes a batch of files."""
+        if not self.use_batch_delete:
+            return super().delete_batch(file_list)
+
+        keys = [self.normalize(path) for path in file_list]
+        for key_chunk in chunk.chunk(keys, 1000):
+            response = self.with_retries(
+                self.client.delete_objects,
+                Bucket=self.bucket,
+                Delete={
+                    'Objects': [{'Key': key} for key in key_chunk],
+                    'Quiet': True,
+                },
+            )
+
+            if response['Errors']:
+                raise TransportException('Could not delete some member of batch', response['Errors'][0]['Message'])
+
     def exists(self, path):
         # checks to see if KEY exists
         key = self.normalize(path)

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev15/assemblyline.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline
-Version: 4.7.5.dev13
+Version: 4.7.5.dev15
 Summary: Assemblyline 4 - Automated malware analysis framework
 Home-page: https://github.com/CybercentreCanada/assemblyline-base
 Author: CCCS Assemblyline development team

assemblyline-4.7.5.dev13/assemblyline/VERSION

@@ -1 +0,0 @@
-4.7.5.dev13