assemblyline-core 4.6.1.dev28__tar.gz → 4.6.1.dev31__tar.gz

{assemblyline_core-4.6.1.dev28 → assemblyline_core-4.6.1.dev31}/PKG-INFO

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

assemblyline_core-4.6.1.dev31/assemblyline_core/VERSION

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

{assemblyline_core-4.6.1.dev28 → assemblyline_core-4.6.1.dev31}/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,8 +199,6 @@ 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
@@ -246,7 +250,88 @@ class Plumber(CoreBase):
                 user['apikeys'] = {}
                 self.datastore.user.save(uname, user)
 
+    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']['included'].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']['included'].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.1.dev28 → assemblyline_core-4.6.1.dev31}/assemblyline_core.egg-info/PKG-INFO

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

{assemblyline_core-4.6.1.dev28 → assemblyline_core-4.6.1.dev31}/test/test_plumber.py

@@ -1,14 +1,15 @@
 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.random_data import random_model_obj
+
 
 def test_expire_missing_service():
     redis = mock.MagicMock(spec=Redis)
@@ -107,3 +108,39 @@ 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}
+    datastore_connection.ds.client.index(index="user_settings", id="admin", document=settings)
+    datastore_connection.user_settings.commit()
+
+    # Initiate the migration
+    plumber.migrate_user_settings()
+
+    # Check that the settings have been migrated
+    migrated_settings = datastore_connection.user_settings.get("admin", as_obj=False)
+
+    # 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 name, settings in migrated_settings['submission_profiles'].items():
+        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, but not in pre-defined ones
+        assert settings['service_spec'] == {} if name != "default" else {'AVClass': {'include_malpedia_dataset': False}}
+        assert settings['ttl'] == 0

assemblyline_core-4.6.1.dev28/assemblyline_core/VERSION

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