assemblyline-core 4.6.1.dev222__tar.gz → 4.6.1.dev225__tar.gz

{assemblyline_core-4.6.1.dev222 → assemblyline_core-4.6.1.dev225}/PKG-INFO

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

assemblyline_core-4.6.1.dev225/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.6.1.dev225

{assemblyline_core-4.6.1.dev222 → assemblyline_core-4.6.1.dev225}/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.dev222 → assemblyline_core-4.6.1.dev225}/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)

{assemblyline_core-4.6.1.dev222 → assemblyline_core-4.6.1.dev225}/assemblyline_core/updater/helper.py

@@ -1,23 +1,21 @@
-import os
-import requests
 import re
 import socket
 import string
 import time
-
-from assemblyline.common.version import FRAMEWORK_VERSION, SYSTEM_VERSION
-from assemblyline.odm.models.config import Config as SystemConfig, ServiceRegistry
-from assemblyline.odm.models.service import Service as ServiceConfig, DockerConfig
-
 from base64 import b64encode
 from collections import defaultdict
 from logging import Logger
 from typing import Dict, List
-from packaging.version import parse, Version
 from urllib.parse import urlencode
 
+import requests
+from assemblyline.common.version import FRAMEWORK_VERSION, SYSTEM_VERSION
+from assemblyline.odm.models.config import Config as SystemConfig
+from assemblyline.odm.models.config import ServiceRegistry
+from assemblyline.odm.models.service import DockerConfig
+from assemblyline.odm.models.service import Service as ServiceConfig
 from azure.identity import DefaultAzureCredential
-
+from packaging.version import Version, parse
 
 DEFAULT_DOCKER_REGISTRY = "hub.docker.com"
 

{assemblyline_core-4.6.1.dev222 → assemblyline_core-4.6.1.dev225}/assemblyline_core.egg-info/PKG-INFO

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

assemblyline_core-4.6.1.dev222/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.6.1.dev222