assemblyline-core 4.6.1.dev163__tar.gz → 4.7.0.dev45__tar.gz

{assemblyline_core-4.6.1.dev163 → assemblyline_core-4.7.0.dev45}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: assemblyline-core
-Version: 4.6.1.dev163
+Version: 4.7.0.dev45
 Summary: Assemblyline 4 - Core components
 Home-page: https://github.com/CybercentreCanada/assemblyline-core/
 Author: CCCS Assemblyline development team

assemblyline_core-4.7.0.dev45/assemblyline_core/VERSION

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

{assemblyline_core-4.6.1.dev163 → 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.6.1.dev163 → 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.6.1.dev163 → assemblyline_core-4.7.0.dev45}/assemblyline_core/replay/client.py

@@ -297,11 +297,14 @@ class APIClient(ClientBase):
         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, reclassification=None):
-        self.al_client.bundle.import_bundle(bundle_path,
-                                            min_classification=min_classification,
-                                            rescan_services=rescan_services,
-                                            exist_ok=exist_ok,
-                                            reclassification=reclassification)
+        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, reclassification=None):
         from assemblyline_client import ClientError
@@ -412,11 +415,14 @@ class DirectClient(ClientBase):
         os.rename(temp_bundle_file, bundle_path)
 
     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)
+        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, reclassification=None):
         # We're assuming all JSON that loaded has an "enabled" field
@@ -442,7 +448,6 @@ class DirectClient(ClientBase):
                         else:
                             raise
 
-
                 if collection == "workflow":
                     # If there has been any edits by another user, then preserve the enabled state
                     # Otherwise, the workflow will be synchronized with the origin system

{assemblyline_core-4.6.1.dev163 → assemblyline_core-4.7.0.dev45}/assemblyline_core/scaler/controllers/kubernetes_ctl.py

@@ -14,6 +14,12 @@ from time import sleep
 from typing import List, Optional, Tuple
 
 import urllib3
+from assemblyline.odm.models.config import Selector
+from assemblyline.odm.models.service import (
+    DependencyConfig,
+    DockerConfig,
+    PersistentVolume,
+)
 from cryptography import x509
 from cryptography.hazmat.primitives import hashes, serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
@@ -22,6 +28,7 @@ from kubernetes import client, config, watch
 from kubernetes.client import (
     V1Affinity,
     V1Capabilities,
+    V1ConfigMap,
     V1ConfigMapVolumeSource,
     V1Container,
     V1Deployment,
@@ -62,12 +69,6 @@ from kubernetes.client import (
 )
 from kubernetes.client.rest import ApiException
 
-from assemblyline.odm.models.config import Selector
-from assemblyline.odm.models.service import (
-    DependencyConfig,
-    DockerConfig,
-    PersistentVolume,
-)
 from assemblyline_core.scaler.controllers.interface import ControllerInterface
 
 # RESERVE_MEMORY_PER_NODE = os.environ.get('RESERVE_MEMORY_PER_NODE')
@@ -390,7 +391,7 @@ class KubernetesController(ControllerInterface):
     def _dependency_name(self, service_name: str, container_name: str):
         return f"{self._deployment_name(service_name)}-{container_name}".lower()
 
-    def add_config_mount(self, name: str, config_map: str, key: str, target_path: str, read_only=True, core=False):
+    def add_config_mount(self, name: str, config_map: str, key: Optional[str], target_path: str, read_only=True, core=False):
         volumes, mounts = self.volumes, self.mounts
         if core:
             volumes, mounts = self.core_volumes, self.core_mounts
@@ -1379,3 +1380,19 @@ class KubernetesController(ControllerInterface):
         for np in (existing_netpol - {np.metadata.name for np in network_policies}):
             self.net_api.delete_namespaced_network_policy(namespace=self.namespace, name=np,
                                                           _request_timeout=API_TIMEOUT)
+
+    def update_config_map(self, data: dict, name: str):
+        """Update or create a ConfigMap in Kubernetes."""
+        config_map = V1ConfigMap(
+            metadata=V1ObjectMeta(name=name, namespace=self.namespace),
+            data=data
+        )
+        try:
+            self.api.patch_namespaced_config_map(name=name, namespace=self.namespace, body=config_map,
+                                                 _request_timeout=API_TIMEOUT)
+        except ApiException as error:
+            if error.status == 404:
+                self.api.create_namespaced_config_map(namespace=self.namespace, body=config_map,
+                                                      _request_timeout=API_TIMEOUT)
+            else:
+                raise

{assemblyline_core-4.6.1.dev163 → assemblyline_core-4.7.0.dev45}/assemblyline_core/scaler/scaler_server.py

@@ -19,7 +19,6 @@ from typing import Any, Dict, Optional
 
 import elasticapm
 import yaml
-
 from assemblyline.common.constants import (
     SCALER_TIMEOUT_QUEUE,
     SERVICE_STATE_HASH,
@@ -44,14 +43,13 @@ from assemblyline.remote.datatypes.hash import ExpiringHash, Hash
 from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline.remote.datatypes.queues.priority import PriorityQueue
 from assemblyline.remote.datatypes.queues.priority import length as pq_length
-from assemblyline_core.scaler.controllers import KubernetesController
+
+from assemblyline_core.scaler import collection
+from assemblyline_core.scaler.controllers import DockerController, KubernetesController
 from assemblyline_core.scaler.controllers.interface import ServiceControlError
 from assemblyline_core.server_base import ServiceStage, ThreadedCoreBase
 from assemblyline_core.updater.helper import get_registry_config
 
-from . import collection
-from .controllers import DockerController
-
 APM_SPAN_TYPE = 'scaler'
 
 # How often (in seconds) to download new service data, try to scale managed services,
@@ -298,6 +296,16 @@ class ScalerServer(ThreadedCoreBase):
                         # be shared with privileged services.
                         pass
 
+        # Create a configuration file specifically meant for privileged services to consume
+        # This should only contain the relevant information to connect to the databases
+        privileged_config = yaml.dump({
+            'datastore': self.config.datastore.as_primitives(),
+            'filestore': self.config.filestore.as_primitives(),
+            'core': {
+                'redis': self.config.core.redis.as_primitives()
+            }
+        })
+
         labels = {
             'app': 'assemblyline',
             'section': 'service',
@@ -340,7 +348,9 @@ class ScalerServer(ThreadedCoreBase):
                                                    )
 
             # Add global configuration for privileged services
-            self.controller.add_config_mount(KUBERNETES_AL_CONFIG, config_map=KUBERNETES_AL_CONFIG, key="config",
+            # Check if the ConfigMap already exists, if it does, update it
+            self.controller.update_config_map(data={'config': privileged_config}, name='privileged-service-config')
+            self.controller.add_config_mount(KUBERNETES_AL_CONFIG, config_map='privileged-service-config', key="config",
                                              target_path="/etc/assemblyline/config.yml", read_only=True, core=True)
 
             # If we're passed an override for server-server and it's defining an HTTPS connection, then add a global
@@ -382,7 +392,7 @@ class ScalerServer(ThreadedCoreBase):
 
                 with open(os.path.join(DOCKER_CONFIGURATION_PATH, 'config.yml'), 'w') as handle:
                     # Convert to JSON before converting to YAML to account for direct ODM representation errors
-                    yaml.dump(json.loads(self.config.json()), handle)
+                    handle.write(privileged_config)
 
                 with open(os.path.join(DOCKER_CONFIGURATION_PATH, 'classification.yml'), 'w') as handle:
                     yaml.dump(get_classification().original_definition, handle)