assemblyline-core 4.5.1.dev426__tar.gz → 4.7.0.dev45__tar.gz

{assemblyline_core-4.5.1.dev426 → assemblyline_core-4.7.0.dev45}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: assemblyline-core
-Version: 4.5.1.dev426
+Version: 4.7.0.dev45
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team
@@ -33,6 +33,7 @@ Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
 Dynamic: license
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: summary

assemblyline_core-4.7.0.dev45/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.7.0.dev45

{assemblyline_core-4.5.1.dev426 → assemblyline_core-4.7.0.dev45}/assemblyline_core/dispatching/client.py

@@ -16,6 +16,7 @@ from assemblyline.common.constants import DISPATCH_RUNNING_TASK_HASH, SUBMISSION
     make_watcher_list_name, DISPATCH_TASK_HASH
 from assemblyline.common.forge import CachedObject, get_service_queue
 from assemblyline.common.isotime import now_as_iso
+from assemblyline.common.dispatcher import Dispatcher
 from assemblyline.datastore.exceptions import VersionConflictException
 from assemblyline.odm.base import DATEFORMAT
 from assemblyline.odm.messages.dispatching import DispatcherCommandMessage, CREATE_WATCH, \
@@ -30,7 +31,7 @@ from assemblyline.remote.datatypes.hash import ExpiringHash, Hash
 from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline.remote.datatypes.set import ExpiringSet, Set
 from assemblyline_core.dispatching.dispatcher import DISPATCH_START_EVENTS, DISPATCH_RESULT_QUEUE, \
-    DISPATCH_COMMAND_QUEUE, QUEUE_EXPIRY, BAD_SID_HASH, ServiceTask, Dispatcher
+    DISPATCH_COMMAND_QUEUE, QUEUE_EXPIRY, BAD_SID_HASH, ServiceTask
 
 
 MAX_CANCEL_RESPONSE_WAIT = 10

assemblyline_core-4.7.0.dev45/assemblyline_core/dispatching/dispatcher.py

@@ -0,0 +1,327 @@
+from __future__ import annotations
+
+import dataclasses
+import enum
+import os
+import threading
+import time
+import uuid
+from contextlib import contextmanager
+from copy import deepcopy
+from queue import Empty, PriorityQueue, Queue
+from typing import TYPE_CHECKING, Any, Iterable, Optional
+
+import elasticapm
+
+from assemblyline.common.constants import (
+    DISPATCH_RUNNING_TASK_HASH,
+    DISPATCH_TASK_HASH,
+    SCALER_TIMEOUT_QUEUE,
+    SUBMISSION_QUEUE,
+    make_watcher_list_name,
+)
+from assemblyline.common.forge import (
+    get_apm_client,
+    get_classification,
+    get_service_queue,
+)
+from assemblyline.common.isotime import now_as_iso
+from assemblyline.common.metrics import MetricsFactory
+from assemblyline.common.postprocess import ActionWorker
+from assemblyline.datastore.helper import AssemblylineDatastore
+from assemblyline.odm.messages.changes import Operation, ServiceChange
+from assemblyline.odm.messages.dispatcher_heartbeat import Metrics
+from assemblyline.odm.messages.dispatching import (
+    CREATE_WATCH,
+    LIST_OUTSTANDING,
+    UPDATE_BAD_SID,
+    CreateWatch,
+    DispatcherCommandMessage,
+    ListOutstanding,
+    WatchQueueMessage,
+)
+from assemblyline.odm.messages.service_heartbeat import Metrics as ServiceMetrics
+from assemblyline.odm.messages.submission import (
+    SubmissionMessage,
+    from_datastore_submission,
+)
+from assemblyline.odm.messages.task import FileInfo
+from assemblyline.odm.messages.task import Task as ServiceTask
+from assemblyline.odm.models.error import Error
+from assemblyline.odm.models.result import Result
+from assemblyline.odm.models.service import Service
+from assemblyline.odm.models.submission import Submission, TraceEvent
+from assemblyline.odm.models.user import User
+from assemblyline.remote.datatypes.events import EventWatcher
+from assemblyline.remote.datatypes.exporting_counter import export_metrics_once
+from assemblyline.remote.datatypes.hash import Hash
+from assemblyline.remote.datatypes.queues.comms import CommsQueue
+from assemblyline.remote.datatypes.queues.named import NamedQueue
+from assemblyline.remote.datatypes.set import ExpiringSet, Set
+from assemblyline.remote.datatypes.user_quota_tracker import UserQuotaTracker
+from assemblyline_core.server_base import ThreadedCoreBase
+
+from ..ingester.constants import COMPLETE_QUEUE_NAME
+from .schedules import Scheduler
+
+if TYPE_CHECKING:
+    from redis import Redis
+
+    from assemblyline.odm.models.file import File
+    from assemblyline.odm.models.config import Config
+
+
+APM_SPAN_TYPE = 'handle_message'
+
+AL_SHUTDOWN_GRACE = int(os.environ.get('AL_SHUTDOWN_GRACE', '60'))
+AL_SHUTDOWN_QUIT = 60
+FINALIZING_WINDOW = max(AL_SHUTDOWN_GRACE - AL_SHUTDOWN_QUIT, 0)
+RESULT_BATCH_SIZE = int(os.environ.get('DISPATCHER_RESULT_BATCH_SIZE', '50'))
+ERROR_BATCH_SIZE = int(os.environ.get('DISPATCHER_ERROR_BATCH_SIZE', '50'))
+DAY_IN_SECONDS = 24 * 60 * 60
+
+
+class KeyType(enum.Enum):
+    OVERWRITE = 'overwrite'
+    UNION = 'union'
+
+
+class Action(enum.IntEnum):
+    start = 0
+    result = 1
+    dispatch_file = 2
+    service_timeout = 3
+    check_submission = 4
+    bad_sid = 5
+
+
+@dataclasses.dataclass(order=True)
+class DispatchAction:
+    kind: Action
+    sid: str = dataclasses.field(compare=False)
+    sha: Optional[str] = dataclasses.field(compare=False, default=None)
+    service_name: Optional[str] = dataclasses.field(compare=False, default=None)
+    worker_id: Optional[str] = dataclasses.field(compare=False, default=None)
+    data: Any = dataclasses.field(compare=False, default=None)
+    event: Optional[threading.Event] = dataclasses.field(compare=False, default=None)
+
+
+
+@contextmanager
+def apm_span(client, span_name: str):
+    try:
+        if client:
+            client.begin_transaction(APM_SPAN_TYPE)
+        yield None
+        if client:
+            client.end_transaction(span_name, 'success')
+    except Exception:
+        if client:
+            client.end_transaction(span_name, 'exception')
+        raise
+
+
+DISPATCH_TASK_ASSIGNMENT = 'dispatcher-tasks-assigned-to-'
+TASK_ASSIGNMENT_PATTERN = DISPATCH_TASK_ASSIGNMENT + '*'
+DISPATCH_START_EVENTS = 'dispatcher-start-events-'
+DISPATCH_RESULT_QUEUE = 'dispatcher-results-'
+DISPATCH_COMMAND_QUEUE = 'dispatcher-commands-'
+DISPATCH_DIRECTORY = 'dispatchers-directory'
+DISPATCH_DIRECTORY_FINALIZE = 'dispatchers-directory-finalizing'
+BAD_SID_HASH = 'bad-sid-hash'
+QUEUE_EXPIRY = 60*60
+SERVICE_VERSION_EXPIRY_TIME = 30 * 60  # How old service version info can be before we ignore it
+GUARD_TIMEOUT = 60*2
+GLOBAL_TASK_CHECK_INTERVAL = 60*10
+TIMEOUT_EXTRA_TIME = 5
+TIMEOUT_TEST_INTERVAL = 5
+MAX_RESULT_BUFFER = 64
+RESULT_THREADS = max(1, int(os.getenv('DISPATCHER_RESULT_THREADS', '2')))
+FINALIZE_THREADS = max(1, int(os.getenv('DISPATCHER_FINALIZE_THREADS', '2')))
+
+# After 20 minutes, check if a submission is still making progress.
+# In the case of a crash somewhere else in the system, we may not have
+# gotten a message we are expecting. This should prompt a retry in most
+# cases.
+SUBMISSION_TOTAL_TIMEOUT = 60 * 20
+
+
+class Dispatcher(ThreadedCoreBase):
+    # @staticmethod
+    # def all_instances(persistent_redis: Redis):
+    #     return Hash(DISPATCH_DIRECTORY, host=persistent_redis).keys()
+
+    # @staticmethod
+    # def instance_assignment_size(persistent_redis, instance_id):
+    #     return Hash(DISPATCH_TASK_ASSIGNMENT + instance_id, host=persistent_redis).length()
+
+    # @staticmethod
+    # def instance_assignment(persistent_redis, instance_id) -> list[str]:
+    #     return Hash(DISPATCH_TASK_ASSIGNMENT + instance_id, host=persistent_redis).keys()
+
+    # @staticmethod
+    # def all_queue_lengths(redis, instance_id):
+    #     return {
+    #         'start': NamedQueue(DISPATCH_START_EVENTS + instance_id, host=redis).length(),
+    #         'result': NamedQueue(DISPATCH_RESULT_QUEUE + instance_id, host=redis).length(),
+    #         'command': NamedQueue(DISPATCH_COMMAND_QUEUE + instance_id, host=redis).length()
+    #     }
+
+    def __init__(self, datastore=None, redis=None, redis_persist=None, logger=None,
+                 config=None, counter_name: str = 'dispatcher'):
+        super().__init__('assemblyline.dispatcher', config=config, datastore=datastore,
+                         redis=redis, redis_persist=redis_persist, logger=logger)
+
+        # Load the datastore collections that we are going to be using
+        self.instance_id = uuid.uuid4().hex
+        self.tasks: dict[str, SubmissionTask] = {}
+        self.finalizing = threading.Event()
+        self.finalizing_start = 0.0
+
+        # Build some utility classes
+        self.scheduler = Scheduler(self.datastore, self.config, self.redis)
+        self.running_tasks: Hash[dict] = Hash(DISPATCH_RUNNING_TASK_HASH, host=self.redis)
+        self.scaler_timeout_queue = NamedQueue(SCALER_TIMEOUT_QUEUE, host=self.redis_persist)
+
+        self.classification_engine = get_classification()
+
+        # Output. Duplicate our input traffic into this queue so it may be cloned by other systems
+        self.traffic_queue = CommsQueue('submissions', self.redis)
+        self.quota_tracker = UserQuotaTracker('submissions', timeout=60 * 60, host=self.redis_persist)
+        self.submission_queue = NamedQueue(SUBMISSION_QUEUE, self.redis)
+
+        # Table to track the running dispatchers
+        self.dispatchers_directory: Hash[int] = Hash(DISPATCH_DIRECTORY, host=self.redis_persist)
+        self.dispatchers_directory_finalize: Hash[int] = Hash(DISPATCH_DIRECTORY_FINALIZE, host=self.redis_persist)
+        self.running_dispatchers_estimate = 1
+
+        # Tables to track what submissions are running where
+        self.active_submissions = Hash(DISPATCH_TASK_ASSIGNMENT+self.instance_id, host=self.redis_persist)
+        self.submissions_assignments = Hash(DISPATCH_TASK_HASH, host=self.redis_persist)
+        self.ingester_scanning = Hash('m-scanning-table', self.redis_persist)
+
+        # Communications queues
+        self.start_queue: NamedQueue[tuple[str, str, str, str]] =\
+            NamedQueue(DISPATCH_START_EVENTS+self.instance_id, host=self.redis, ttl=QUEUE_EXPIRY)
+        self.result_queue: NamedQueue[dict] =\
+            NamedQueue(DISPATCH_RESULT_QUEUE+self.instance_id, host=self.redis, ttl=QUEUE_EXPIRY)
+        self.command_queue: NamedQueue[dict] =\
+            NamedQueue(DISPATCH_COMMAND_QUEUE+self.instance_id, host=self.redis, ttl=QUEUE_EXPIRY)
+
+        # Publish counters to the metrics sink.
+        self.counter = MetricsFactory(metrics_type='dispatcher', schema=Metrics, name=counter_name,
+                                      redis=self.redis, config=self.config)
+
+        self.apm_client = None
+        if self.config.core.metrics.apm_server.server_url:
+            elasticapm.instrument()
+            self.apm_client = get_apm_client("dispatcher")
+
+        self._service_timeouts: TimeoutTable[tuple[str, str, str], str] = TimeoutTable()
+        self._submission_timeouts: TimeoutTable[str, None] = TimeoutTable()
+
+        # Setup queues for work to be divided into
+        self.process_queues: list[PriorityQueue[DispatchAction]] = [PriorityQueue() for _ in range(RESULT_THREADS)]
+        self.queue_ready_signals: list[threading.Semaphore] = [threading.Semaphore(MAX_RESULT_BUFFER)
+                                                               for _ in range(RESULT_THREADS)]
+
+        # Queue of finished submissions/errors waiting to be saved into elastic
+        self.finalize_queue = Queue()
+        self.error_queue: Queue[tuple[str, Error]] = Queue()
+
+        # Queue to hold of service timeouts that need to be processed
+        # They will be held in this queue until results in redis are
+        # already processed
+        self.timeout_queue: Queue[DispatchAction] = Queue()
+
+        # Utility object to handle post-processing actions
+        self.postprocess_worker = ActionWorker(cache=False, config=self.config, datastore=self.datastore,
+                                               redis_persist=self.redis_persist)
+
+        # Update bad sid list
+        self.redis_bad_sids = Set(BAD_SID_HASH, host=self.redis_persist)
+        self.bad_sids: set[str] = set(self.redis_bad_sids.members())
+
+        # Event Watchers
+        self.service_change_watcher = EventWatcher(self.redis, deserializer=ServiceChange.deserialize)
+        self.service_change_watcher.register('changes.services.*', self._handle_service_change_event)
+
+    def stop(self):
+        super().stop()
+        self.service_change_watcher.stop()
+        self.postprocess_worker.stop()
+
+    def try_run(self):
+        self.log.info(f'Using dispatcher id {self.instance_id}')
+        self.service_change_watcher.start()
+        threads = {
+            # Process to protect against old dead tasks timing out
+            'Global Timeout Backstop': self.timeout_backstop,
+        }
+
+        for ii in range(RESULT_THREADS):
+            # Process results
+            threads[f'Service Update Worker #{ii}'] = self.service_worker_factory(ii)
+
+        self.maintain_threads(threads)
+
+        # If the dispatcher is exiting cleanly remove as many tasks from the service queues as we can
+        service_queues = {}
+        for task in self.tasks.values():
+            for (_sha256, service_name), dispatch_key in task.queue_keys.items():
+                try:
+                    s_queue = service_queues[service_name]
+                except KeyError:
+                    s_queue = get_service_queue(service_name, self.redis)
+                    service_queues[service_name] = s_queue
+                s_queue.remove(dispatch_key)
+
+
+    def timeout_backstop(self):
+        while self.running:
+            cpu_mark = time.process_time()
+            time_mark = time.time()
+
+            # Start of process dispatcher transaction
+            with apm_span(self.apm_client, 'timeout_backstop'):
+                dispatcher_instances = set(Dispatcher.all_instances(persistent_redis=self.redis_persist))
+                error_tasks = []
+
+                # iterate running tasks
+                for _task_key, task_body in self.running_tasks:
+                    task = ServiceTask(task_body)
+                    # Its a bad task if it's dispatcher isn't running
+                    if task.metadata['dispatcher__'] not in dispatcher_instances:
+                        error_tasks.append(task)
+                    # Its a bad task if its OUR task, but we aren't tracking that submission anymore
+                    if task.metadata['dispatcher__'] == self.instance_id and task.sid not in self.tasks:
+                        error_tasks.append(task)
+
+                # Refresh our dispatcher list.
+                dispatcher_instances = set(Dispatcher.all_instances(persistent_redis=self.redis_persist))
+                other_dispatcher_instances = dispatcher_instances - {self.instance_id}
+
+                # The remaining running tasks (probably) belong to dead dispatchers and can be killed
+                for task in error_tasks:
+                    # Check against our refreshed dispatcher list in case it changed during the previous scan
+                    if task.metadata['dispatcher__'] in other_dispatcher_instances:
+                        continue
+
+                    # If its already been handled, we don't need to
+                    if not self.running_tasks.pop(task.key()):
+                        continue
+
+                    # Kill the task that would report to a dead dispatcher
+                    self.log.warning(f"[{task.sid}]Task killed by backstop {task.service_name} {task.fileinfo.sha256}")
+                    self.scaler_timeout_queue.push({
+                        'service': task.service_name,
+                        'container': task.metadata['worker__']
+                    })
+
+                    # Report to the metrics system that a recoverable error has occurred for that service
+                    export_metrics_once(task.service_name, ServiceMetrics, dict(fail_recoverable=1),
+                                        host=task.metadata['worker__'], counter_type='service', redis=self.redis)
+
+            self.counter.increment_execution_time('cpu_seconds', time.process_time() - cpu_mark)
+            self.counter.increment_execution_time('busy_seconds', time.time() - time_mark)
+            self.sleep(GLOBAL_TASK_CHECK_INTERVAL)

assemblyline_core-4.7.0.dev45/assemblyline_core/ingester/ingester.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python
+"""
+Ingester
+
+Ingester is responsible for monitoring for incoming submission requests,
+sending submissions, waiting for submissions to complete, sending a message
+to a notification queue as specified by the submission and, based on the
+score received, possibly sending a message to indicate that an alert should
+be created.
+"""
+
+import logging
+import threading
+import time
+from os import environ
+from random import random
+from typing import Any, Iterable, List, Optional, Tuple
+
+import elasticapm
+
+from assemblyline import odm
+from assemblyline.common import exceptions, forge, isotime
+from assemblyline.common.constants import DROP_PRIORITY
+from assemblyline.common.exceptions import get_stacktrace_info
+from assemblyline.common.importing import load_module_by_path
+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.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.queues.named import NamedQueue
+from assemblyline.remote.datatypes.queues.priority import PriorityQueue
+from assemblyline.remote.datatypes.user_quota_tracker import UserQuotaTracker
+from assemblyline_core.dispatching.dispatcher import Dispatcher
+from assemblyline_core.server_base import ThreadedCoreBase
+from assemblyline_core.submission_client import SubmissionClient
+
+from .constants import COMPLETE_QUEUE_NAME, INGEST_QUEUE_NAME, drop_chance
+
+_dup_prefix = 'w-m-'
+_notification_queue_prefix = 'nq-'
+_max_retries = 10
+_retry_delay = 60 * 4  # Wait 4 minutes to retry
+_max_time = 2 * 24 * 60 * 60  # Wait 2 days for responses.
+HOUR_IN_SECONDS = 60 * 60
+COMPLETE_THREADS = int(environ.get('INGESTER_COMPLETE_THREADS', 4))
+INGEST_THREADS = int(environ.get('INGESTER_INGEST_THREADS', 1))
+SUBMIT_THREADS = int(environ.get('INGESTER_SUBMIT_THREADS', 4))
+
+
+def must_drop(length: int, maximum: int) -> bool:
+    """
+    To calculate the probability of dropping an incoming submission we compare
+    the number returned by random() which will be in the range [0,1) and the
+    number returned by tanh() which will be in the range (-1,1).
+
+    If length is less than maximum the number returned by tanh will be negative
+    and so drop will always return False since the value returned by random()
+    cannot be less than 0.
+
+    If length is greater than maximum, drop will return False with a probability
+    that increases as the distance between maximum and length increases:
+
+        Length           Chance of Dropping
+
+        <= maximum       0
+        1.5 * maximum    0.76
+        2 * maximum      0.96
+        3 * maximum      0.999
+    """
+    return random() < drop_chance(length, maximum)
+
+
+@odm.model()
+class IngestTask(odm.Model):
+    # Submission Parameters
+    submission: MessageSubmission = odm.compound(MessageSubmission)
+
+    # Shortcut for properties of the submission
+    @property
+    def file_size(self) -> int:
+        return sum(file.size for file in self.submission.files)
+
+    @property
+    def params(self) -> SubmissionParams:
+        return self.submission.params
+
+    @property
+    def sha256(self) -> str:
+        return self.submission.files[0].sha256
+
+    # Information about the ingestion itself, parameters irrelevant
+    retries = odm.Integer(default=0)
+
+    # Fields added after a submission is complete for notification/bookkeeping processes
+    failure = odm.Text(default='')  # If the ingestion has failed for some reason, what is it?
+    score = odm.Optional(odm.Integer())  # Score from previous processing of this file
+    extended_scan = odm.Enum(EXTENDED_SCAN_VALUES, default="skipped")  # Status of the extended scan
+    ingest_id = odm.UUID()  # Ingestion Identifier
+    ingest_time = odm.Date(default="NOW")  # Time at which the file was ingested
+    notify_time = odm.Optional(odm.Date())  # Time at which the user is notify the submission is finished
+    to_ingest = odm.Boolean(default=False)
+

{assemblyline_core-4.5.1.dev426 → assemblyline_core-4.7.0.dev45}/assemblyline_core/metrics/heartbeat_formatter.py

@@ -6,6 +6,7 @@ from assemblyline.common import forge, metrics
 from assemblyline.common.archiving import ARCHIVE_QUEUE_NAME
 from assemblyline.common.constants import DISPATCH_TASK_HASH, SUBMISSION_QUEUE, \
     SERVICE_STATE_HASH, ServiceStatus
+from assemblyline.common.dispatcher import Dispatcher
 from assemblyline.datastore.exceptions import SearchException
 from assemblyline.odm.messages.retrohunt_heartbeat import RetrohuntMessage
 from assemblyline.odm.messages.scaler_heartbeat import ScalerMessage
@@ -25,7 +26,6 @@ from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline.remote.datatypes.queues.priority import PriorityQueue
 
 from assemblyline_core.alerter.run_alerter import ALERT_QUEUE_NAME, ALERT_RETRY_QUEUE_NAME
-from assemblyline_core.dispatching.dispatcher import Dispatcher
 from assemblyline_core.ingester import INGEST_QUEUE_NAME, drop_chance
 from assemblyline_core.ingester.constants import COMPLETE_QUEUE_NAME
 

{assemblyline_core-4.5.1.dev426 → assemblyline_core-4.7.0.dev45}/assemblyline_core/replay/client.py

@@ -4,13 +4,13 @@ import time
 
 from assemblyline.common import forge
 from assemblyline.common.bundling import create_bundle, import_bundle
+from assemblyline.common.classification import InvalidClassification
 from assemblyline.odm import Model
-from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline.remote.datatypes.hash import Hash
-from assemblyline_core.replay.replay import INPUT_TYPES
+from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline_core.badlist_client import BadlistClient
+from assemblyline_core.replay.replay import INPUT_TYPES
 from assemblyline_core.safelist_client import SafelistClient
-from assemblyline_core.signature_client import SignatureClient
 
 EMPTY_WAIT_TIME = int(os.environ.get('EMPTY_WAIT_TIME', '30'))
 REPLAY_REQUESTED = 'requested'
@@ -296,13 +296,17 @@ class APIClient(ClientBase):
     def create_al_bundle(self, id, bundle_path, use_alert=False):
         self.al_client.bundle.create(id, output=bundle_path, use_alert=use_alert)
 
-    def load_bundle(self, bundle_path, min_classification, rescan_services, exist_ok=True):
-        self.al_client.bundle.import_bundle(bundle_path,
-                                            min_classification=min_classification,
-                                            rescan_services=rescan_services,
-                                            exist_ok=exist_ok)
+    def load_bundle(self, bundle_path, min_classification, rescan_services, exist_ok=True, reclassification=None):
+        self.al_client.bundle.import_bundle(
+            bundle_path,
+            min_classification=min_classification,
+            rescan_services=rescan_services,
+            exist_ok=exist_ok,
+            reclassification=reclassification,
+            to_ingest=True,  # send submissions to ingester
+        )
 
-    def load_json(self, file_path):
+    def load_json(self, file_path, reclassification=None):
         from assemblyline_client import ClientError
 
         # We're assuming all JSON that loaded has an "enabled" field
@@ -374,6 +378,7 @@ class DirectClient(ClientBase):
         # Initialize connection to redis-persistent for checkpointing
         redis_persist = get_client(config.core.redis.persistent.host,
                                    config.core.redis.persistent.port, False)
+        self.classification = forge.get_classification()
         self.datastore = forge.get_datastore(config=config)
         self.queues = {
             queue_type: NamedQueue(f"replay_{queue_type}", host=redis)
@@ -409,13 +414,17 @@ class DirectClient(ClientBase):
         temp_bundle_file = create_bundle(id, working_dir=os.path.dirname(bundle_path), use_alert=use_alert)
         os.rename(temp_bundle_file, bundle_path)
 
-    def load_bundle(self, bundle_path, min_classification, rescan_services, exist_ok=True):
-        import_bundle(bundle_path,
-                      min_classification=min_classification,
-                      rescan_services=rescan_services,
-                      exist_ok=exist_ok)
+    def load_bundle(self, bundle_path, min_classification, rescan_services, exist_ok=True, reclassification=None):
+        import_bundle(
+            bundle_path,
+            min_classification=min_classification,
+            rescan_services=rescan_services,
+            exist_ok=exist_ok,
+            reclassification=reclassification,
+            to_ingest=True,  # send submissions to ingester
+        )
 
-    def load_json(self, file_path):
+    def load_json(self, file_path, reclassification=None):
         # We're assuming all JSON that loaded has an "enabled" field
         collection = os.path.basename(file_path).split('_', 1)[0]
         with open(file_path) as fp:
@@ -428,6 +437,16 @@ class DirectClient(ClientBase):
 
                 # Let's see if there's an existing document with the same ID in the collection
                 obj = es_collection.get_if_exists(id, as_obj=False)
+                if obj:
+                    # Check if the classification of the object is compatible with the system's classification
+                    try:
+                        self.classification.normalize_classification(obj['classification'])
+                    except InvalidClassification:
+                        if reclassification:
+                            # If reclassification is requested, then we can change the classification
+                            obj['classification'] = reclassification
+                        else:
+                            raise
 
                 if collection == "workflow":
                     # If there has been any edits by another user, then preserve the enabled state

{assemblyline_core-4.5.1.dev426 → assemblyline_core-4.7.0.dev45}/assemblyline_core/replay/loader/run.py

@@ -86,8 +86,8 @@ class ReplayLoader(ReplayBase):
         self.maintain_threads(threads)
 
     def stop(self):
-        super().stop()
         self.cache.close()
+        return super().stop()
 
 
 if __name__ == '__main__':

{assemblyline_core-4.5.1.dev426 → assemblyline_core-4.7.0.dev45}/assemblyline_core/replay/loader/run_worker.py

@@ -1,5 +1,5 @@
-import shutil
 import os
+import shutil
 
 from cart import unpack_file
 
@@ -32,15 +32,16 @@ class ReplayLoaderWorker(ReplayBase):
                     if file_path.endswith(".al_bundle"):
                         self.client.load_bundle(file_path,
                                                 min_classification=self.replay_config.loader.min_classification,
-                                                rescan_services=self.replay_config.loader.rescan)
+                                                rescan_services=self.replay_config.loader.rescan,
+                                                reclassification=self.replay_config.loader.reclassification)
                     elif file_path.endswith(".al_json"):
-                        self.client.load_json(file_path)
+                        self.client.load_json(file_path, reclassification=self.replay_config.loader.reclassification)
 
                     elif file_path.endswith(".al_json.cart"):
                         cart_path = file_path
                         file_path = file_path[:-5]
                         unpack_file(cart_path, file_path)
-                        self.client.load_json(file_path)
+                        self.client.load_json(file_path, reclassification=self.replay_config.loader.reclassification)
                         os.unlink(cart_path)
 
                     if os.path.exists(file_path):
@@ -55,11 +56,11 @@ class ReplayLoaderWorker(ReplayBase):
                         # Terminate on NFS-related error
                         self.log.warning("'Invalid cross-device link' exception detected. Terminating..")
                         self.stop()
-                except Exception:
+                except Exception as e:
                     # Make sure failed directory exists
                     os.makedirs(self.replay_config.loader.failed_directory, exist_ok=True)
 
-                    self.log.error(f"Failed to load the bundle file {file_path}, moving it to the failed directory.")
+                    self.log.error(f"Failed to load the bundle file {file_path}, moving it to the failed directory. Reason: {e}")
                     failed_path = os.path.join(self.replay_config.loader.failed_directory, os.path.basename(file_path))
                     shutil.move(file_path, failed_path)