assemblyline-core 4.6.0.dev11__tar.gz → 4.6.0.2__tar.gz

{assemblyline_core-4.6.0.dev11 → assemblyline_core-4.6.0.2}/PKG-INFO

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

assemblyline_core-4.6.0.2/assemblyline_core/VERSION

@@ -0,0 +1 @@
+4.6.0.2

{assemblyline_core-4.6.0.dev11 → assemblyline_core-4.6.0.2}/assemblyline_core/dispatching/dispatcher.py

@@ -54,7 +54,7 @@ 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
+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
@@ -311,6 +311,16 @@ class SubmissionTask:
         """Shortcut to read submission SID"""
         return self.submission.sid
 
+    def trace(self, event_type: str, sha256: Optional[str] = None,
+              service: Optional[str] = None, message: Optional[str] = None) -> None:
+        if self.submission.params.trace:
+            self.submission.tracing_events.append(TraceEvent({
+                'event_type': event_type,
+                'service': service,
+                'file': sha256,
+                'message': message,
+            }))
+
     def forbid_for_children(self, sha256: str, service_name: str):
         """Mark that children of a given file should not be routed to a service."""
         try:
@@ -729,6 +739,7 @@ class Dispatcher(ThreadedCoreBase):
 
         if not self.active_submissions.exists(sid):
             self.log.info("[%s] New submission received", sid)
+            task.trace('submission_start')
             self.active_submissions.add(sid, {
                 'completed_queue': task.completed_queue,
                 'submission': submission.as_primitives()
@@ -742,6 +753,7 @@ class Dispatcher(ThreadedCoreBase):
             }).as_primitives())
 
         else:
+            task.trace('submission_start', message='Received a pre-existing submission')
             self.log.warning(f"[{sid}] Received a pre-existing submission, check if it is complete")
 
         # Refresh the quota hold
@@ -840,6 +852,7 @@ class Dispatcher(ThreadedCoreBase):
         sid = submission.sid
         if self.apm_client:
             elasticapm.label(sid=sid, sha256=sha256)
+        task.trace('dispatch_file', sha256=sha256)
 
         file_depth: int = task.file_depth[sha256]
         # If its the first time we've seen this file, we won't have a schedule for it
@@ -859,6 +872,9 @@ class Dispatcher(ThreadedCoreBase):
             task.file_schedules[sha256] = self.scheduler.build_schedule(submission, file_info.type,
                                                                         file_depth, forbidden_services,
                                                                         task.service_access_control)
+            schedule_summary = [list(stage.keys()) for stage in task.file_schedules[sha256]]
+            task.trace('schedule_built', sha256=sha256, message=str(schedule_summary))
+
 
         file_info = task.file_info[sha256]
         schedule: list = list(task.file_schedules[sha256])
@@ -990,9 +1006,9 @@ class Dispatcher(ThreadedCoreBase):
 
             if sent or enqueued or running:
                 # If we have confirmed that we are waiting, or have taken an action, log that.
-                self.log.info(f"[{sid}] File {sha256} sent to: {sent} "
-                              f"already in queue for: {enqueued} "
-                              f"running on: {running}")
+                logs = f"sent to: {sent} already in queue for: {enqueued} running on: {running}"
+                self.log.info(f"[{sid}] File {sha256} {logs}")
+                task.trace('dispatch_file_result', sha256=sha256, message=logs)
                 return False
             elif skipped:
                 # Not waiting for anything, and have started skipping what is left over
@@ -1007,10 +1023,12 @@ class Dispatcher(ThreadedCoreBase):
 
         self.counter.increment('files_completed')
         if len(task.queue_keys) > 0 or len(task.running_services) > 0:
-            self.log.info(f"[{sid}] Finished processing file '{sha256}', submission incomplete "
-                          f"(queued: {len(task.queue_keys)} running: {len(task.running_services)})")
+            logs = f"queued: {len(task.queue_keys)} running: {len(task.running_services)}"
+            self.log.info(f"[{sid}] Finished processing file '{sha256}', submission incomplete ({logs})")
+            task.trace('file_finished', sha256=sha256, message=logs)
         else:
             self.log.info(f"[{sid}] Finished processing file '{sha256}', checking if submission complete")
+            task.trace('file_finished', sha256=sha256)
             return self.check_submission(task)
         return False
 
@@ -1104,14 +1122,17 @@ class Dispatcher(ThreadedCoreBase):
         # file isn't done yet, and hasn't been filtered by any of the previous few steps
         # poke those files.
         if pending_files:
+            task.trace('submission_check', message=f"Dispatching {list(pending_files)}")
             self.log.debug(f"[{task.submission.sid}] Dispatching {len(pending_files)} files: {list(pending_files)}")
             for file_hash in pending_files:
                 if self.dispatch_file(task, file_hash):
                     return True
         elif processing_files:
+            task.trace('submission_check', message=f"Waiting for {list(processing_files)}")
             self.log.debug("[%s] Not finished waiting on %d files: %s",
                            task.submission.sid, len(processing_files), list(processing_files))
         else:
+            task.trace('submission_check', message="Finished")
             self.log.debug("[%s] Finalizing submission.", task.submission.sid)
             max_score = max(file_scores.values()) if file_scores else 0  # Submissions with no results have no score
             if self.tasks.pop(task.sid, None):
@@ -1341,6 +1362,8 @@ class Dispatcher(ThreadedCoreBase):
                         self.log.warning(f'[{message.sid}] Service started missing data.')
                         continue
 
+                    task.trace('service_start', sha256=message.sha, service=message.service_name,
+                               message=message.worker_id)
                     key = (message.sha, message.service_name)
                     if task.queue_keys.pop(key, None) is not None:
                         # If this task is already finished (result message processed before start
@@ -1384,8 +1407,9 @@ class Dispatcher(ThreadedCoreBase):
                     continue
 
                 if task:
-                    task.service_logs[(message.sha, message.service_name)].append(
-                        f'Service timeout at {now_as_iso()} on worker {message.worker_id}')
+                    log = f'Service timeout at {now_as_iso()} on worker {message.worker_id}'
+                    task.service_logs[(message.sha, message.service_name)].append(log)
+                    task.trace('service_timeout', sha256=message.sha, service=message.service_name, message=log)
                     self.timeout_service(task, message.sha, message.service_name, message.worker_id)
 
             elif kind == Action.dispatch_file:
@@ -1429,6 +1453,8 @@ class Dispatcher(ThreadedCoreBase):
             self.log.exception(f"Malformed result message, missing key: {missing}")
             return
 
+        task.trace('process_result', sha256=sha256, service=service_name, message="Processing result " + summary.key)
+
         # Add SHA256s of files that allowed to run regardless of Dynamic Recursion Prevention
         task.dynamic_recursion_bypass = task.dynamic_recursion_bypass.union(set(dynamic_recursion_bypass))
 
@@ -1599,6 +1625,8 @@ class Dispatcher(ThreadedCoreBase):
     @elasticapm.capture_span(span_type='dispatcher')
     def process_service_error(self, task: SubmissionTask, error_key, error: Error):
         self.log.info(f'[{task.submission.sid}] Error from service {error.response.service_name} on {error.sha256}')
+        task.trace('process_error', sha256=error.sha256, service=error.response.service_name,
+                   message="Service error " + error_key)
         self.clear_timeout(task, error.sha256, error.response.service_name)
         key = (error.sha256, error.response.service_name)
         if error.response.status == "FAIL_NONRECOVERABLE":

{assemblyline_core-4.6.0.dev11 → assemblyline_core-4.6.0.2}/assemblyline_core/plumber/run_plumber.py

@@ -7,22 +7,22 @@ disabled or deleted for which a service queue exists, the dispatcher will be inf
 had an error.
 """
 import threading
-import warnings
+from copy import deepcopy
 from typing import Optional
 
 from assemblyline.common.constants import service_queue_name
-from assemblyline.common.forge import get_service_queue, get_config
+from assemblyline.common.forge import get_config, get_service_queue
 from assemblyline.common.isotime import DAY_IN_SECONDS, now_as_iso
 from assemblyline.odm.models.apikey import get_apikey_id
 from assemblyline.odm.models.error import Error
 from assemblyline.odm.models.service import Service
 from assemblyline.odm.models.user import load_roles, load_roles_form_acls
+from assemblyline.odm.models.user_settings import SubmissionProfileParams
 from assemblyline.remote.datatypes import retry_call
 from assemblyline.remote.datatypes.queues.named import NamedQueue
 from assemblyline_core.dispatching.client import DispatchClient
 from assemblyline_core.server_base import CoreBase, ServiceStage
 
-
 DAY = 60 * 60 * 24
 TASK_DELETE_CHUNK = 10000
 
@@ -56,8 +56,14 @@ class Plumber(CoreBase):
                                      name="redis_notification_queue_cleanup")
         nq_thread.start()
 
+        # Start a user apikey cleanup thread
         ua_thread = threading.Thread(target=self.user_apikey_cleanup, daemon=True, name="user_apikey_cleanup")
         ua_thread.start()
+
+        # Start a user setting migration thread
+        usm_thread = threading.Thread(target=self.migrate_user_settings, daemon=True, name="migrate_user_settings")
+        usm_thread.start()
+
         self.service_queue_plumbing()
 
     def service_queue_plumbing(self):
@@ -193,60 +199,127 @@ class Plumber(CoreBase):
             self.sleep(2)
         self.log.info(f"Done watching {service_name} service queue")
 
-
-
     def user_apikey_cleanup(self):
-        query = "id:*"
-        offset = 0
-        rows = 100
-        total = 1
-        cur_total = 0
-
-        config = get_config()
-        apikey_max_dtl = config.auth.apikey_max_dtl
-
-        expiry_ts = now_as_iso(apikey_max_dtl * DAY_IN_SECONDS) if apikey_max_dtl is not None else None
-
-        while cur_total < total:
-            result = self.datastore.user.search(query, offset=offset, rows=rows)
-            total = result.get('total', 0)
-            cur_total = cur_total + (result.get("count", total))
-
-            # check for API keys in total
-            users = result.get('items', [])
+        expiry_ts = None
+        if self.config.auth.apikey_max_dtl is not None:
+            expiry_ts = now_as_iso(self.config.auth.apikey_max_dtl * DAY_IN_SECONDS)
 
-            for u in users:
-                uname = u['uname']
-                user = self.datastore.user.get(uname)
-                apikeys = user.apikeys
+        for user in self.datastore.user.stream_search(query="*", fl="*", as_obj=False):
+            uname = user['uname']
+            apikeys = user['apikeys']
 
-                for key in apikeys:
-                    old_apikey = apikeys[key]
-                    key_id = get_apikey_id(key, uname)
+            for key in apikeys:
+                old_apikey = apikeys[key]
+                key_id = get_apikey_id(key, uname)
 
-                    roles = None
-                    if old_apikey['acl'] == ["C"]:
+                roles = None
+                if old_apikey['acl'] == ["C"]:
 
-                        roles = [r for r in old_apikey['roles']
-                                    if r in load_roles(user['type'], user['roles'])]
-
-                    else:
-                        roles = [r for r in load_roles_form_acls(old_apikey['acl'], roles)
+                    roles = [r for r in old_apikey['roles']
                                 if r in load_roles(user['type'], user['roles'])]
-                    new_apikey = {
-                        "password": old_apikey['password'],
-                        "acl": old_apikey['acl'],
-                        "uname": uname,
-                        "key_name": key,
-                        "roles": roles,
-                        "expiry_ts": expiry_ts
-                    }
-                    self.datastore.apikey.save(key_id, new_apikey)
-
-                user['apikeys'] = {}
-                self.datastore.user.save(uname, user)
-
 
+                else:
+                    roles = [r for r in load_roles_form_acls(old_apikey['acl'], roles)
+                            if r in load_roles(user['type'], user['roles'])]
+                new_apikey = {
+                    "password": old_apikey['password'],
+                    "acl": old_apikey['acl'],
+                    "uname": uname,
+                    "key_name": key,
+                    "roles": roles,
+                    "expiry_ts": expiry_ts
+                }
+                self.datastore.apikey.save(key_id, new_apikey)
+
+            user['apikeys'] = {}
+            self.datastore.user.save(uname, user)
+
+        # Commit changes made to indices
+        self.datastore.user.commit()
+        self.datastore.apikey.commit()
+
+    def migrate_user_settings(self):
+        service_list = self.datastore.list_all_services(as_obj=False)
+
+        # Migrate user settings to the new format
+        for doc in self.datastore.user_settings.scan_with_search_after(query={
+            "bool": {
+                "must": {
+                    "query_string": {
+                        "query": "*",
+                    }
+                }
+            }
+        }):
+            user_settings = doc["_source"]
+            if user_settings.get('submission_profiles'):
+                # User settings already migrated, skip
+                continue
+
+            # Create the list of submission profiles
+            submission_profiles = {
+                # Grant everyone a default of which all settings from before 4.6 should be transferred
+                "default" : SubmissionProfileParams(user_settings).as_primitives(strip_null=True)
+            }
+
+            profile_error_checks = {}
+            # Try to apply the original settings to the system-defined profiles
+            for profile in self.config.submission.profiles:
+                updates = deepcopy(user_settings)
+                profile_error_checks.setdefault(profile.name, 0)
+                validated_profile = profile.params.as_primitives(strip_null=True)
+
+                updates.setdefault("services", {})
+                updates["services"].setdefault("selected", [])
+                updates["services"].setdefault("excluded", [])
+
+                # Append the exclusion list set by the profile
+                updates['services']['excluded'] = updates['services']['excluded'] + \
+                    list(validated_profile.get("services", {}).get("excluded", []))
+
+                # Ensure the selected services are not in the excluded list
+                for service in service_list:
+                    if service['name'] in updates['services']['selected'] and service['name'] in updates['services']['excluded']:
+                        updates['services']['selected'].remove(service['name'])
+                        profile_error_checks[profile.name] += 1
+                    elif service['category'] in updates['services']['selected'] and service['category'] in updates['services']['excluded']:
+                        updates['services']['selected'].remove(service['category'])
+                        profile_error_checks[profile.name] += 1
+
+
+                # Check the services parameters
+                for param_type, list_of_params in profile.restricted_params.items():
+                    # Check if there are restricted submission parameters
+                    if param_type == "submission":
+                        requested_params = (set(list_of_params) & set(updates.keys())) - set({'services', 'service_spec'})
+                        if requested_params:
+                            # Track the number of errors for each profile
+                            profile_error_checks[profile.name] += len(requested_params)
+                            for param in requested_params:
+                                # Remove the parameter from the updates
+                                updates.pop(param, None)
+
+                    # Check if there are restricted service parameters
+                    else:
+                        service_spec = updates.get('service_spec', {}).get(param_type, {})
+                        requested_params = set(list_of_params) & set(service_spec)
+                        if requested_params:
+                            # Track the number of errors for each profile
+                            profile_error_checks[profile.name] += len(requested_params)
+                            for param in requested_params:
+                                # Remove the parameter from the updates
+                                service_spec.pop(param, None)
+
+                            if not service_spec:
+                                # Remove the service spec if empty
+                                updates['service_spec'].pop(param_type, None)
+                submission_profiles[profile.name] = SubmissionProfileParams(updates).as_primitives(strip_null=True)
+
+            # Assign the profile with the least number of errors
+            user_settings['submission_profiles'] = submission_profiles
+            user_settings['preferred_submission_profile'] = sorted(profile_error_checks.items(), key=lambda x: x[1])[0][0]
+
+            self.datastore.user_settings.save(doc["_id"], user_settings)
 
 
 

{assemblyline_core-4.6.0.dev11 → assemblyline_core-4.6.0.2}/assemblyline_core/scaler/controllers/docker_ctl.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
-import docker
+
 import os
 import threading
 import time
-from collections import defaultdict
-from typing import List, Optional, Tuple, Dict
 import uuid
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import docker
 
 from assemblyline.odm.models.service import DependencyConfig, DockerConfig
+
 from .interface import ControllerInterface, ServiceControlError
 
 # Where to find the update directory inside this container.
@@ -382,7 +385,7 @@ class DockerController(ControllerInterface):
                            container.labels]
                 running = running[0:-delta]
                 for container in running:
-                    container.kill()
+                    container.kill(signal='SIGTERM')
 
             if delta > 0:
                 # Start delta instances of the service
@@ -531,6 +534,7 @@ class DockerController(ControllerInterface):
         dynamically rather than in prepare_network.
         """
         from docker.errors import NotFound
+
         # Create network for service
         network_name = f'{COMPOSE_PROJECT}_service-net-{service_name}'
         try:
@@ -553,6 +557,7 @@ class DockerController(ControllerInterface):
         This lets us override the auth_config on a per image basis.
         """
         from docker.errors import ImageNotFound
+
         # Split the image string into "[registry/]image_name" and "tag"
         repository, _, tag = service.container_config.image.rpartition(':')
         if '/' in tag:

{assemblyline_core-4.6.0.dev11 → assemblyline_core-4.6.0.2}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline_core-4.6.0.dev11 → assemblyline_core-4.6.0.2}/test/test_plumber.py

@@ -1,14 +1,16 @@
 from time import sleep
 from unittest import mock
 
-from assemblyline.odm.messages.task import Task
-from assemblyline.odm.models.service import Service
-from assemblyline.odm.random_data import random_model_obj
 from assemblyline_core.plumber.run_plumber import Plumber
 from assemblyline_core.server_base import ServiceStage
 from mocking import TrueCountTimes
 from redis import Redis
 
+from assemblyline.odm.messages.task import Task
+from assemblyline.odm.models.service import Service
+from assemblyline.odm.models.user import ApiKey, User
+from assemblyline.odm.random_data import random_model_obj
+
 
 def test_expire_missing_service():
     redis = mock.MagicMock(spec=Redis)
@@ -107,3 +109,50 @@ def test_cleanup_old_tasks(datastore_connection):
                                               q="task.start_time_in_millis:0",
                                               track_total_hits=True,
                                               size=0)['hits']['total']['value'] == 0
+
+def test_user_setting_migrations(datastore_connection):
+    from assemblyline.odm.models.config import SubmissionProfileParams
+
+    SubmissionProfileParams.fields().keys()
+    # Create a bunch of random "old" tasks and clean them up
+    redis = mock.MagicMock(spec=Redis)
+    redis_persist = mock.MagicMock(spec=Redis)
+    plumber = Plumber(redis=redis, redis_persist=redis_persist, datastore=datastore_connection, delay=1)
+
+    # Create a user with old settings (format prior to 4.6)
+    settings = {'classification': 'TLP:CLEAR', 'deep_scan': False, 'description': '', 'download_encoding': 'cart', 'default_external_sources': ['Malware Bazaar', 'VirusTotal'], 'default_zip_password': 'zippy', 'executive_summary': False, 'expand_min_score': 500, 'generate_alert': False, 'ignore_cache': False, 'ignore_dynamic_recursion_prevention': False, 'ignore_recursion_prevention': False, 'ignore_filtering': False, 'malicious': False, 'priority': 369, 'profile': False, 'service_spec': {'AVClass': {'include_malpedia_dataset': False}}, 'services': {'selected': ['Extraction', 'ConfigExtractor', 'YARA'], 'excluded': [], 'rescan': [], 'resubmit': [], 'runtime_excluded': []}, 'submission_view': 'report', 'ttl': 0}
+
+    user_account = random_model_obj(User, as_json=True)
+    user_account['uname'] = "admin"
+    user_account['apikeys'] = {'test': random_model_obj(ApiKey, as_json=True)}
+    datastore_connection.ds.client.index(index="user_settings", id="admin", document=settings)
+    datastore_connection.ds.client.index(index="user", id="admin", document=user_account)
+
+    datastore_connection.user_settings.commit()
+    datastore_connection.user.commit()
+
+    # Initiate the migration
+    plumber.user_apikey_cleanup()
+    plumber.migrate_user_settings()
+
+    # Check that the settings have been migrated
+    migrated_settings = datastore_connection.user_settings.get("admin", as_obj=False)
+
+    # Check to see if API keys for the user were transferred to the new index
+    assert datastore_connection.apikey.search('uname:admin', rows=0)['total'] > 0
+
+    # Deprecated settings should be removed
+    assert "ignore_dynamic_recursion_prevention" not in migrated_settings
+
+    # All former submission settings at the root-level should be moved to submission profiles
+    assert all([key not in migrated_settings for key in SubmissionProfileParams.fields().keys()] )
+
+    for settings in migrated_settings['submission_profiles'].values():
+        assert settings['classification'] == 'TLP:C'
+        assert settings['deep_scan'] is False
+        assert settings['generate_alert'] is False
+        assert settings['ignore_cache'] is False
+        assert settings['priority'] == 369
+        # Full service spec should be preserved in default profile (along with others by default if there's no restricted parameters)
+        assert settings['service_spec'] == {'AVClass': {'include_malpedia_dataset': False}}
+        assert settings['ttl'] == 0

assemblyline_core-4.6.0.dev11/assemblyline_core/VERSION

@@ -1 +0,0 @@
-4.6.0.dev11