assemblyline 4.7.5.dev13__tar.gz → 4.7.5.dev16__tar.gz

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

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

assemblyline-4.7.5.dev16/assemblyline/VERSION

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

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev16}/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.dev16}/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
@@ -143,7 +143,7 @@ def create_transport(url, connection_attempts=None):
         valid_bool_keys = ['allow_directory_access', 'use_default_credentials', 'initalize_container', 'read_only']
         extras = _get_extras(parse_qs(parsed.query), valid_str_keys=valid_str_keys, valid_bool_keys=valid_bool_keys)
 
-        t = TransportAzure(base=base, host=host, connection_attempts=connection_attempts, **extras)
+        t = TransportAzure(base=base, host=host, connection_attempts=connection_attempts, port=port, **extras)
 
     else:
         raise FileStoreException("Unknown transport: %s" % scheme)
@@ -200,6 +200,22 @@ class FileStore(object):
                     trace = get_stacktrace_info(ex)
                     self.log.info('Transport problem: %s', trace)
 
+    def delete_batch_chunk_size(self) -> int:
+        return min(transport.delete_batch_chunk_size() for transport in self.transports)
+
+    @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
+                t.delete_batch(file_list)
+
     @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.dev16}/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
 
@@ -35,7 +37,7 @@ class TransportAzure(Transport):
 
     def __init__(self, base=None, access_key=None, tenant_id=None, client_id=None, client_secret=None,
                  host=None, connection_attempts=None, allow_directory_access=False, use_default_credentials=False,
-                 initalize_container=True, read_only=False):
+                 initalize_container=True, read_only=False, port=None):
         self.log = logging.getLogger('assemblyline.transport.azure')
         self.read_only = False
         self.connection_attempts: Optional[int] = connection_attempts
@@ -44,6 +46,8 @@ class TransportAzure(Transport):
         # Get URL
         self.host = host
         self.endpoint_url = f"https://{self.host}"
+        if port:
+            self.endpoint_url += ":" + str(port)
 
         # Get container and base_path
         parts = base.strip("/").split("/", 1)
@@ -57,7 +61,7 @@ class TransportAzure(Transport):
         if use_default_credentials:
             if (tenant_id and client_id) and (not client_secret):
                 self.credential = WorkloadIdentityCredential(tenant_id=tenant_id,
-                                                            client_id=client_id)
+                                                             client_id=client_id)
             else:
                 # Service accounts will by default create the enviromental variables, and use them as params
                 self.credential = DefaultAzureCredential()
@@ -144,6 +148,21 @@ class TransportAzure(Transport):
             if not isinstance(error.cause, ResourceNotFoundError):
                 raise
 
+    def delete_batch_chunk_size(self) -> int:
+        return 256
+
+    def delete_batch(self, file_list: Iterable[str]):
+        """Deletes a batch of files."""
+        keys = [self.normalize(path) for path in file_list]
+        for key_chunk in chunk.chunk(keys, 256):
+            responses: Iterable[HttpResponse] = self.with_retries(
+                self.container_client.delete_blobs, *key_chunk, raise_on_any_failure=False
+            )
+            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.dev16}/assemblyline/filestore/transport/base.py

@@ -46,6 +46,16 @@ class Transport(object):
         """
         raise TransportException("Not Implemented")
 
+    def delete_batch_chunk_size(self) -> int:
+        return 1
+
+    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.dev16}/assemblyline/filestore/transport/ftp.py

@@ -154,7 +154,12 @@ class TransportFTP(Transport):
     @reconnect_retry_on_fail
     def delete(self, path):
         path = self.normalize(path)
-        self.ftp.delete(path)
+        try:
+            self.ftp.delete(path)
+        except ftplib.error_perm as e:
+            # If the file doesnt exist we get a 550.
+            if not str(e).startswith('550'):
+                raise
 
     @reconnect_retry_on_fail
     def exists(self, path) -> bool:

{assemblyline-4.7.5.dev13 → assemblyline-4.7.5.dev16}/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
@@ -131,7 +134,7 @@ class TransportS3(Transport):
                 return ret_val
 
             except (EndpointConnectionError, ConnectionClosedError):
-                self.log.warning(f"No connection to S3 transport {self.endpoint_url}, retrying...")
+                self.log.warning("No connection to S3 transport %s, retrying...", self.endpoint_url)
                 retries += 1
         raise ConnectionError(f"Couldn't connect to the requested S3 endpoint {self.endpoint_url} inside retry limit")
 
@@ -139,6 +142,31 @@ class TransportS3(Transport):
         key = self.normalize(path)
         self.with_retries(self.client.delete_object, Bucket=self.bucket, Key=key)
 
+    def delete_batch_chunk_size(self) -> int:
+        if self.use_batch_delete:
+            return 1000
+        else:
+            return 1
+
+    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.dev16}/assemblyline/filestore/transport/sftp.py

@@ -117,8 +117,11 @@ class TransportSFTP(Transport):
 
     @reconnect_retry_on_fail
     def delete(self, path):
-        path = self.normalize(path)
-        self.sftp.remove(path)
+        try:
+            path = self.normalize(path)
+            self.sftp.remove(path)
+        except FileNotFoundError:
+            pass
 
     @reconnect_retry_on_fail
     def exists(self, path):

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline
-Version: 4.7.5.dev13
+Version: 4.7.5.dev16
 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