assemblyline-core 4.5.1.dev501__tar.gz → 4.6.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/PKG-INFO +1 -1
  2. assemblyline_core-4.6.0.0/assemblyline_core/VERSION +1 -0
  3. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/dispatching/dispatcher.py +48 -40
  4. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/ingester/ingester.py +6 -5
  5. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/controllers/kubernetes_ctl.py +57 -22
  6. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/scaler_server.py +34 -34
  7. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/updater/run_updater.py +28 -19
  8. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/vacuum/worker.py +29 -26
  9. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core.egg-info/PKG-INFO +1 -1
  10. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_dispatcher.py +14 -14
  11. assemblyline_core-4.5.1.dev501/assemblyline_core/VERSION +0 -1
  12. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/LICENCE.md +0 -0
  13. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/README.md +0 -0
  14. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/__init__.py +0 -0
  15. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/alerter/__init__.py +0 -0
  16. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/alerter/processing.py +0 -0
  17. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/alerter/run_alerter.py +0 -0
  18. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/archiver/__init__.py +0 -0
  19. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/archiver/run_archiver.py +0 -0
  20. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/badlist_client.py +0 -0
  21. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/dispatching/__init__.py +0 -0
  22. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/dispatching/__main__.py +0 -0
  23. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/dispatching/client.py +0 -0
  24. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/dispatching/schedules.py +0 -0
  25. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/dispatching/timeout.py +0 -0
  26. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/expiry/__init__.py +0 -0
  27. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/expiry/run_expiry.py +0 -0
  28. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/ingester/__init__.py +0 -0
  29. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/ingester/__main__.py +0 -0
  30. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/ingester/constants.py +0 -0
  31. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/__init__.py +0 -0
  32. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/es_metrics.py +0 -0
  33. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/heartbeat_formatter.py +0 -0
  34. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/helper.py +0 -0
  35. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/metrics_server.py +0 -0
  36. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/run_heartbeat_manager.py +0 -0
  37. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/run_metrics_aggregator.py +0 -0
  38. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/metrics/run_statistics_aggregator.py +0 -0
  39. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/plumber/__init__.py +0 -0
  40. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/plumber/run_plumber.py +0 -0
  41. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/__init__.py +0 -0
  42. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/client.py +0 -0
  43. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/creator/__init__.py +0 -0
  44. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/creator/run.py +0 -0
  45. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/creator/run_worker.py +0 -0
  46. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/loader/__init__.py +0 -0
  47. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/loader/run.py +0 -0
  48. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/loader/run_worker.py +0 -0
  49. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/replay/replay.py +0 -0
  50. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/safelist_client.py +0 -0
  51. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/__init__.py +0 -0
  52. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/collection.py +0 -0
  53. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/controllers/__init__.py +0 -0
  54. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/controllers/docker_ctl.py +0 -0
  55. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/controllers/interface.py +0 -0
  56. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/scaler/run_scaler.py +0 -0
  57. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/server_base.py +0 -0
  58. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/signature_client.py +0 -0
  59. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/submission_client.py +0 -0
  60. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/tasking_client.py +0 -0
  61. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/updater/__init__.py +0 -0
  62. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/updater/helper.py +0 -0
  63. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/vacuum/__init__.py +0 -0
  64. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/vacuum/crawler.py +0 -0
  65. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/vacuum/department_map.py +0 -0
  66. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/vacuum/safelist.py +0 -0
  67. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/vacuum/stream_map.py +0 -0
  68. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/workflow/__init__.py +0 -0
  69. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core/workflow/run_workflow.py +0 -0
  70. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core.egg-info/SOURCES.txt +0 -0
  71. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core.egg-info/dependency_links.txt +0 -0
  72. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core.egg-info/requires.txt +0 -0
  73. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/assemblyline_core.egg-info/top_level.txt +0 -0
  74. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/setup.cfg +0 -0
  75. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/setup.py +0 -0
  76. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_alerter.py +0 -0
  77. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_badlist_client.py +0 -0
  78. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_expiry.py +0 -0
  79. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_plumber.py +0 -0
  80. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_replay.py +0 -0
  81. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_safelist_client.py +0 -0
  82. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_scaler.py +0 -0
  83. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_scheduler.py +0 -0
  84. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_signature_client.py +0 -0
  85. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_simulation.py +0 -0
  86. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_tasking_client.py +0 -0
  87. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_vacuum.py +0 -0
  88. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_worker_ingest.py +0 -0
  89. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_worker_submit.py +0 -0
  90. {assemblyline_core-4.5.1.dev501 → assemblyline_core-4.6.0.0}/test/test_workflow.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: assemblyline-core
3
- Version: 4.5.1.dev501
3
+ Version: 4.6.0.0
4
4
  Summary: Assemblyline 4 - Core components
5
5
  Home-page: https://github.com/CybercentreCanada/assemblyline-core/
6
6
  Author: CCCS Assemblyline development team
@@ -0,0 +1 @@
1
+ 4.6.0.0
@@ -1,42 +1,63 @@
1
1
  from __future__ import annotations
2
- import uuid
2
+
3
+ import dataclasses
4
+ import enum
5
+ import json
3
6
  import os
4
7
  import threading
5
8
  import time
9
+ import typing
10
+ import uuid
6
11
  from collections import defaultdict
7
12
  from contextlib import contextmanager
8
- import typing
9
- from typing import Optional, Any, TYPE_CHECKING, Iterable
10
- import json
11
- import enum
12
- from queue import PriorityQueue, Empty, Queue
13
- import dataclasses
14
13
  from copy import deepcopy
14
+ from queue import Empty, PriorityQueue, Queue
15
+ from typing import TYPE_CHECKING, Any, Iterable, Optional
15
16
 
16
17
  import elasticapm
17
18
 
18
19
  from assemblyline.common import isotime
19
- from assemblyline.common.constants import make_watcher_list_name, SUBMISSION_QUEUE, \
20
- DISPATCH_RUNNING_TASK_HASH, SCALER_TIMEOUT_QUEUE, DISPATCH_TASK_HASH
21
- from assemblyline.common.forge import get_service_queue, get_apm_client, get_classification
20
+ from assemblyline.common.constants import (
21
+ DISPATCH_RUNNING_TASK_HASH,
22
+ DISPATCH_TASK_HASH,
23
+ SCALER_TIMEOUT_QUEUE,
24
+ SUBMISSION_QUEUE,
25
+ make_watcher_list_name,
26
+ )
27
+ from assemblyline.common.forge import (
28
+ get_apm_client,
29
+ get_classification,
30
+ get_service_queue,
31
+ )
22
32
  from assemblyline.common.isotime import now_as_iso
23
33
  from assemblyline.common.metrics import MetricsFactory
24
34
  from assemblyline.common.postprocess import ActionWorker
25
35
  from assemblyline.datastore.helper import AssemblylineDatastore
26
- from assemblyline.odm.messages.changes import ServiceChange, Operation
36
+ from assemblyline.odm.messages.changes import Operation, ServiceChange
27
37
  from assemblyline.odm.messages.dispatcher_heartbeat import Metrics
38
+ from assemblyline.odm.messages.dispatching import (
39
+ CREATE_WATCH,
40
+ LIST_OUTSTANDING,
41
+ UPDATE_BAD_SID,
42
+ CreateWatch,
43
+ DispatcherCommandMessage,
44
+ ListOutstanding,
45
+ WatchQueueMessage,
46
+ )
28
47
  from assemblyline.odm.messages.service_heartbeat import Metrics as ServiceMetrics
29
- from assemblyline.odm.messages.dispatching import WatchQueueMessage, CreateWatch, DispatcherCommandMessage, \
30
- CREATE_WATCH, LIST_OUTSTANDING, UPDATE_BAD_SID, ListOutstanding
31
- from assemblyline.odm.messages.submission import SubmissionMessage, from_datastore_submission
32
- from assemblyline.odm.messages.task import FileInfo, Task as ServiceTask
48
+ from assemblyline.odm.messages.submission import (
49
+ SubmissionMessage,
50
+ from_datastore_submission,
51
+ )
52
+ from assemblyline.odm.messages.task import FileInfo
53
+ from assemblyline.odm.messages.task import Task as ServiceTask
33
54
  from assemblyline.odm.models.error import Error
34
55
  from assemblyline.odm.models.result import Result
35
56
  from assemblyline.odm.models.service import Service
36
57
  from assemblyline.odm.models.submission import Submission
37
58
  from assemblyline.odm.models.user import User
38
- from assemblyline.remote.datatypes.exporting_counter import export_metrics_once
39
59
  from assemblyline.remote.datatypes.events import EventWatcher
60
+ from assemblyline.remote.datatypes.exporting_counter import export_metrics_once
40
61
  from assemblyline.remote.datatypes.hash import Hash
41
62
  from assemblyline.remote.datatypes.queues.comms import CommsQueue
42
63
  from assemblyline.remote.datatypes.queues.named import NamedQueue
@@ -44,14 +65,15 @@ from assemblyline.remote.datatypes.set import ExpiringSet, Set
44
65
  from assemblyline.remote.datatypes.user_quota_tracker import UserQuotaTracker
45
66
  from assemblyline_core.server_base import ThreadedCoreBase
46
67
 
68
+ from ..ingester.constants import COMPLETE_QUEUE_NAME
47
69
  from .schedules import Scheduler
48
70
  from .timeout import TimeoutTable
49
- from ..ingester.constants import COMPLETE_QUEUE_NAME
50
71
 
51
72
  if TYPE_CHECKING:
52
- from assemblyline.odm.models.file import File
53
73
  from redis import Redis
54
74
 
75
+ from assemblyline.odm.models.file import File
76
+
55
77
 
56
78
  APM_SPAN_TYPE = 'handle_message'
57
79
 
@@ -62,9 +84,6 @@ RESULT_BATCH_SIZE = int(os.environ.get('DISPATCHER_RESULT_BATCH_SIZE', '50'))
62
84
  ERROR_BATCH_SIZE = int(os.environ.get('DISPATCHER_ERROR_BATCH_SIZE', '50'))
63
85
  DAY_IN_SECONDS = 24 * 60 * 60
64
86
 
65
- # TODO: DYNAMIC_ANALYSIS_CATEGORY can be removed after assemblyline version 4.6+
66
- DYNAMIC_ANALYSIS_CATEGORY = 'Dynamic Analysis'
67
-
68
87
 
69
88
  class KeyType(enum.Enum):
70
89
  OVERWRITE = 'overwrite'
@@ -161,17 +180,17 @@ class TemporaryFileData:
161
180
 
162
181
  def set_value(self, key: str, value: Any) -> bool:
163
182
  """Set the value of a temporary data key using the appropriate method for the key.
164
-
183
+
165
184
  Return true if this change could mean partial results should be reevaluated.
166
185
  """
167
186
  if self.config.get(key) == KeyType.UNION.value:
168
187
  return self._union_shared_value(key, value)
169
-
188
+
170
189
  if self.config.get(key) == KeyType.OVERWRITE.value:
171
190
  change = self.shared_values.get(key) != value
172
191
  self.shared_values[key] = value
173
192
  return change
174
-
193
+
175
194
  self.local_values[key] = value
176
195
  return False
177
196
 
@@ -257,10 +276,6 @@ class SubmissionTask:
257
276
  if not service:
258
277
  continue
259
278
 
260
- # TODO: the following 2 lines can be removed when assemblyline changed to version 4.6+
261
- if service.category == DYNAMIC_ANALYSIS_CATEGORY:
262
- self.forbid_for_children(sha256, service.name)
263
-
264
279
  prevented_services = scheduler.expand_categories(service.recursion_prevention)
265
280
 
266
281
  for service_name in prevented_services:
@@ -393,7 +408,7 @@ class SubmissionTask:
393
408
  entry.dispatch_needed = True
394
409
  else:
395
410
  # If there are results and there is a monitoring entry, the result was partial
396
- # so redispatch it immediately. If there are not partial results the monitoring
411
+ # so redispatch it immediately. If there are not partial results the monitoring
397
412
  # entry will have been cleared.
398
413
  self.redispatch_service(sha256, service)
399
414
  changed.append(sha256)
@@ -836,11 +851,9 @@ class Dispatcher(ThreadedCoreBase):
836
851
 
837
852
  forbidden_services = None
838
853
 
839
- # If Dynamic Recursion Prevention is in effect and the file is not part of the bypass list,
854
+ # If Recursion Prevention is in effect and the file is not part of the bypass list,
840
855
  # Find the list of services this file is forbidden from being sent to.
841
- # TODO: remove "or submission.params.ignore_dynamic_recursion_prevention" after assemblyline upgrade to version 4.6+
842
- ignore_drp = submission.params.ignore_recursion_prevention or submission.params.ignore_dynamic_recursion_prevention
843
- if not ignore_drp and sha256 not in task.dynamic_recursion_bypass:
856
+ if not submission.params.ignore_recursion_prevention and sha256 not in task.dynamic_recursion_bypass:
844
857
  forbidden_services = task.find_recursion_excluded_services(sha256)
845
858
 
846
859
  task.file_schedules[sha256] = self.scheduler.build_schedule(submission, file_info.type,
@@ -938,10 +951,6 @@ class Dispatcher(ThreadedCoreBase):
938
951
  tag_fields.append('score')
939
952
 
940
953
  # Mark this routing for the purposes of recursion prevention
941
- # TODO: The following 2 lines can be removed after assemblyline upgrade to version 4.6+
942
- if service.category == DYNAMIC_ANALYSIS_CATEGORY:
943
- task.forbid_for_children(sha256, service_name)
944
-
945
954
  prevented_services = self.scheduler.expand_categories(service.recursion_prevention)
946
955
 
947
956
  for service_name in prevented_services:
@@ -961,8 +970,7 @@ class Dispatcher(ThreadedCoreBase):
961
970
  max_files=task.submission.params.max_extracted,
962
971
  ttl=submission.params.ttl,
963
972
  ignore_cache=submission.params.ignore_cache,
964
- # TODO: remove "or submission.params.ignore_dynamic_recursion_prevention" after assemblyline upgrade to version 4.6+
965
- ignore_recursion_prevention=submission.params.ignore_recursion_prevention or submission.params.ignore_dynamic_recursion_prevention ,
973
+ ignore_recursion_prevention=submission.params.ignore_recursion_prevention,
966
974
  ignore_filtering=ignore_filtering,
967
975
  tags=[{field: x[field] for field in tag_fields} for x in tags],
968
976
  temporary_submission_data=[
@@ -1574,7 +1582,7 @@ class Dispatcher(ThreadedCoreBase):
1574
1582
  # Not worth running if we know we have services in queue
1575
1583
  if not any(_s == sha256 for _s, _ in task.queue_keys.keys()):
1576
1584
  force_redispatch.add(sha256)
1577
-
1585
+
1578
1586
  # Try to run the next stage
1579
1587
  for sha256 in force_redispatch:
1580
1588
  self.dispatch_file(task, sha256)
@@ -408,12 +408,13 @@ class Ingester(ThreadedCoreBase):
408
408
  self.apm_client.end_transaction('ingest_submit', 'exception')
409
409
 
410
410
  def handle_submission_delete(self, sid: Optional[str]):
411
- if not sid:
412
- return
413
-
414
- # Upon submission deletion, ensure to cleanup the local cache of filescores relative to the SID
415
411
  with self.cache_lock:
416
- self.cache = {k: v for k, v in self.cache.items() if v.sid != sid}
412
+ if not sid:
413
+ # Clear the entire local cache
414
+ self.cache = {}
415
+ else:
416
+ # Ensure to cleanup the local cache of filescores relative to the SID
417
+ self.cache = {k: v for k, v in self.cache.items() if v.sid != sid}
417
418
 
418
419
  def handle_complete(self):
419
420
  while self.running:
@@ -1,36 +1,73 @@
1
1
  from __future__ import annotations
2
+
2
3
  import base64
3
4
  import functools
4
5
  import json
5
- import uuid
6
6
  import os
7
7
  import threading
8
+ import uuid
8
9
  import weakref
9
- import urllib3
10
-
11
10
  from base64 import b64encode
12
- from cryptography import x509
13
- from cryptography.hazmat.primitives.asymmetric import rsa
14
- from cryptography.hazmat.primitives import serialization, hashes
15
11
  from collections import OrderedDict, defaultdict
16
12
  from datetime import datetime, timedelta
17
- from dateutil.tz import tzlocal
18
- from typing import List, Optional, Tuple
19
13
  from time import sleep
20
- from assemblyline.odm.models.config import Selector
14
+ from typing import List, Optional, Tuple
21
15
 
16
+ import urllib3
17
+ from cryptography import x509
18
+ from cryptography.hazmat.primitives import hashes, serialization
19
+ from cryptography.hazmat.primitives.asymmetric import rsa
20
+ from dateutil.tz import tzlocal
22
21
  from kubernetes import client, config, watch
23
- from kubernetes.client import V1Deployment, V1DeploymentSpec, V1PodTemplateSpec, V1DeploymentStrategy, \
24
- V1PodSpec, V1ObjectMeta, V1Volume, V1Container, V1VolumeMount, V1EnvVar, V1ConfigMapVolumeSource, \
25
- V1PersistentVolumeClaimVolumeSource, V1LabelSelector, V1ResourceRequirements, V1PersistentVolumeClaim, \
26
- V1PersistentVolumeClaimSpec, V1NetworkPolicy, V1NetworkPolicySpec, V1NetworkPolicyEgressRule, V1NetworkPolicyPeer, \
27
- V1NetworkPolicyIngressRule, V1Secret, V1SecretVolumeSource, V1LocalObjectReference, V1Service, \
28
- V1ServiceSpec, V1ServicePort, V1PodSecurityContext, V1Probe, V1ExecAction, V1SecurityContext, \
29
- V1Affinity, V1NodeAffinity, V1NodeSelector, V1NodeSelectorTerm, V1NodeSelectorRequirement, V1Toleration, \
30
- V1Capabilities, V1SeccompProfile
22
+ from kubernetes.client import (
23
+ V1Affinity,
24
+ V1Capabilities,
25
+ V1ConfigMapVolumeSource,
26
+ V1Container,
27
+ V1Deployment,
28
+ V1DeploymentSpec,
29
+ V1DeploymentStrategy,
30
+ V1EnvVar,
31
+ V1ExecAction,
32
+ V1LabelSelector,
33
+ V1LocalObjectReference,
34
+ V1NetworkPolicy,
35
+ V1NetworkPolicyEgressRule,
36
+ V1NetworkPolicyIngressRule,
37
+ V1NetworkPolicyPeer,
38
+ V1NetworkPolicySpec,
39
+ V1NodeAffinity,
40
+ V1NodeSelector,
41
+ V1NodeSelectorRequirement,
42
+ V1NodeSelectorTerm,
43
+ V1ObjectMeta,
44
+ V1PersistentVolumeClaim,
45
+ V1PersistentVolumeClaimSpec,
46
+ V1PersistentVolumeClaimVolumeSource,
47
+ V1PodSecurityContext,
48
+ V1PodSpec,
49
+ V1PodTemplateSpec,
50
+ V1Probe,
51
+ V1ResourceRequirements,
52
+ V1SeccompProfile,
53
+ V1Secret,
54
+ V1SecretVolumeSource,
55
+ V1SecurityContext,
56
+ V1Service,
57
+ V1ServicePort,
58
+ V1ServiceSpec,
59
+ V1Toleration,
60
+ V1Volume,
61
+ V1VolumeMount,
62
+ )
31
63
  from kubernetes.client.rest import ApiException
32
- from assemblyline.odm.models.service import DependencyConfig, DockerConfig, PersistentVolume
33
64
 
65
+ from assemblyline.odm.models.config import Selector
66
+ from assemblyline.odm.models.service import (
67
+ DependencyConfig,
68
+ DockerConfig,
69
+ PersistentVolume,
70
+ )
34
71
  from assemblyline_core.scaler.controllers.interface import ControllerInterface
35
72
 
36
73
  # RESERVE_MEMORY_PER_NODE = os.environ.get('RESERVE_MEMORY_PER_NODE')
@@ -250,7 +287,7 @@ def parse_cpu(string: str) -> float:
250
287
  class KubernetesController(ControllerInterface):
251
288
  def __init__(self, logger, namespace: str, prefix: str, priority: str, dependency_priority: str,
252
289
  cpu_reservation: float, linux_node_selector: Selector, labels=None, log_level="INFO", core_env={},
253
- default_service_account=None, cluster_pod_list=True, enable_pod_security=False,
290
+ cluster_pod_list=True, enable_pod_security=False,
254
291
  default_service_tolerations=[],
255
292
  priv_labels=None):
256
293
  # Try loading a kubernetes connection from either the fact that we are running
@@ -295,7 +332,6 @@ class KubernetesController(ControllerInterface):
295
332
  self.core_mounts: dict[str, V1VolumeMount] = {}
296
333
  self._external_profiles = weakref.WeakValueDictionary()
297
334
  self._service_limited_env: dict[str, dict[str, str]] = defaultdict(dict)
298
- self.default_service_account: Optional[str] = default_service_account
299
335
  self.cluster_pod_list = cluster_pod_list
300
336
  self.security_policy = RESTRICTED_POD_SECURITY_CONTEXT if enable_pod_security else None
301
337
  self.default_service_tolerations = [V1Toleration(**toleration.as_primitives()) for toleration in default_service_tolerations]
@@ -836,8 +872,7 @@ class KubernetesController(ControllerInterface):
836
872
  metadata = V1ObjectMeta(name=deployment_name, labels=all_labels, annotations={CHANGE_KEY_NAME: change_key})
837
873
 
838
874
  # Figure out which (if any) service account to use
839
- service_account = self.default_service_account or \
840
- (PRIVILEGED_SERVICE_ACCOUNT_NAME if core_mounts else UNPRIVILEGED_SERVICE_ACCOUNT_NAME)
875
+ service_account = PRIVILEGED_SERVICE_ACCOUNT_NAME if core_mounts else UNPRIVILEGED_SERVICE_ACCOUNT_NAME
841
876
  if docker_config.service_account:
842
877
  service_account = docker_config.service_account
843
878
 
@@ -2,46 +2,55 @@
2
2
  An auto-scaling service specific to Assemblyline services.
3
3
  """
4
4
  from __future__ import annotations
5
+
6
+ import concurrent.futures
7
+ import copy
5
8
  import functools
6
- import threading
7
- from collections import defaultdict
8
- from string import Template
9
- from typing import Dict, Optional, Any
10
- import os
11
- import re
9
+ import json
12
10
  import math
13
- import time
11
+ import os
14
12
  import platform
15
- import concurrent.futures
16
- import copy
13
+ import threading
14
+ import time
15
+ from collections import defaultdict
17
16
  from contextlib import contextmanager
17
+ from string import Template
18
+ from typing import Any, Dict, Optional
18
19
 
19
20
  import elasticapm
20
- import json
21
21
  import yaml
22
22
 
23
- from assemblyline.remote.datatypes.queues.named import NamedQueue
24
- from assemblyline.remote.datatypes.queues.priority import PriorityQueue, length as pq_length
25
- from assemblyline.remote.datatypes.exporting_counter import export_metrics_once
26
- from assemblyline.remote.datatypes.hash import ExpiringHash, Hash
27
- from assemblyline.remote.datatypes.events import EventWatcher, EventSender
28
- from assemblyline.odm.models.service import Service, DockerConfig, EnvironmentVariable
29
- from assemblyline.odm.models.config import Mount
30
- from assemblyline.odm.messages.scaler_heartbeat import Metrics
31
- from assemblyline.odm.messages.scaler_status_heartbeat import Status
32
- from assemblyline.odm.messages.changes import ServiceChange, Operation
33
- from assemblyline.common.dict_utils import get_recursive_sorted_tuples, flatten
23
+ from assemblyline.common.constants import (
24
+ SCALER_TIMEOUT_QUEUE,
25
+ SERVICE_STATE_HASH,
26
+ ServiceStatus,
27
+ )
28
+ from assemblyline.common.dict_utils import flatten, get_recursive_sorted_tuples
29
+ from assemblyline.common.forge import (
30
+ get_apm_client,
31
+ get_classification,
32
+ get_service_queue,
33
+ )
34
34
  from assemblyline.common.uid import get_id_from_data
35
- from assemblyline.common.forge import get_classification, get_service_queue, get_apm_client
36
- from assemblyline.common.constants import SCALER_TIMEOUT_QUEUE, SERVICE_STATE_HASH, ServiceStatus
37
35
  from assemblyline.common.version import FRAMEWORK_VERSION, SYSTEM_VERSION
38
- from assemblyline_core.updater.helper import get_registry_config
36
+ from assemblyline.odm.messages.changes import Operation, ServiceChange
37
+ from assemblyline.odm.messages.scaler_heartbeat import Metrics
38
+ from assemblyline.odm.messages.scaler_status_heartbeat import Status
39
+ from assemblyline.odm.models.config import Mount
40
+ from assemblyline.odm.models.service import DockerConfig, EnvironmentVariable, Service
41
+ from assemblyline.remote.datatypes.events import EventSender, EventWatcher
42
+ from assemblyline.remote.datatypes.exporting_counter import export_metrics_once
43
+ from assemblyline.remote.datatypes.hash import ExpiringHash, Hash
44
+ from assemblyline.remote.datatypes.queues.named import NamedQueue
45
+ from assemblyline.remote.datatypes.queues.priority import PriorityQueue
46
+ from assemblyline.remote.datatypes.queues.priority import length as pq_length
39
47
  from assemblyline_core.scaler.controllers import KubernetesController
40
48
  from assemblyline_core.scaler.controllers.interface import ServiceControlError
41
49
  from assemblyline_core.server_base import ServiceStage, ThreadedCoreBase
50
+ from assemblyline_core.updater.helper import get_registry_config
42
51
 
43
- from .controllers import DockerController
44
52
  from . import collection
53
+ from .controllers import DockerController
45
54
 
46
55
  APM_SPAN_TYPE = 'scaler'
47
56
 
@@ -326,7 +335,6 @@ class ScalerServer(ThreadedCoreBase):
326
335
  core_env=core_env,
327
336
  cluster_pod_list=self.config.core.scaler.cluster_pod_list,
328
337
  enable_pod_security=self.config.core.scaler.enable_pod_security,
329
- default_service_account=self.config.services.service_account,
330
338
  default_service_tolerations=service_defaults_config.tolerations,
331
339
  priv_labels=priv_labels
332
340
  )
@@ -348,14 +356,6 @@ class ScalerServer(ThreadedCoreBase):
348
356
 
349
357
  # Add default mounts for (non-)privileged services
350
358
  for mount in service_defaults_config.mounts:
351
- # Deprecated configuration for mounting ConfigMap
352
- # TODO: Deprecate code on next major change
353
- if mount.config_map:
354
- self.controller.add_config_mount(mount.name, config_map=mount.config_map, key=mount.key,
355
- target_path=mount.path, read_only=mount.read_only,
356
- core=mount.privileged_only)
357
- continue
358
-
359
359
  if mount.resource_type == 'configmap':
360
360
  # ConfigMap-based mount
361
361
  self.controller.add_config_mount(mount.name, config_map=mount.resource_name, key=mount.resource_key,
@@ -7,17 +7,31 @@ import os
7
7
  import re
8
8
  import time
9
9
  import uuid
10
-
11
10
  from concurrent.futures import ThreadPoolExecutor
12
11
  from typing import Any, List, Optional
13
12
 
14
13
  import docker
15
-
16
- from kubernetes.client import V1Job, V1ObjectMeta, V1JobSpec, V1PodTemplateSpec, V1PodSpec, V1Volume, \
17
- V1VolumeMount, V1EnvVar, V1Container, V1ResourceRequirements, \
18
- V1ConfigMapVolumeSource, V1Secret, V1SecretVolumeSource, V1LocalObjectReference, V1Toleration, V1SecurityContext, \
19
- V1Capabilities, V1SeccompProfile
20
14
  from kubernetes import client, config
15
+ from kubernetes.client import (
16
+ V1Capabilities,
17
+ V1ConfigMapVolumeSource,
18
+ V1Container,
19
+ V1EnvVar,
20
+ V1Job,
21
+ V1JobSpec,
22
+ V1LocalObjectReference,
23
+ V1ObjectMeta,
24
+ V1PodSpec,
25
+ V1PodTemplateSpec,
26
+ V1ResourceRequirements,
27
+ V1SeccompProfile,
28
+ V1Secret,
29
+ V1SecretVolumeSource,
30
+ V1SecurityContext,
31
+ V1Toleration,
32
+ V1Volume,
33
+ V1VolumeMount,
34
+ )
21
35
  from kubernetes.client.rest import ApiException
22
36
 
23
37
  from assemblyline.common import isotime
@@ -26,7 +40,11 @@ from assemblyline.odm.models.config import Mount, Selector
26
40
  from assemblyline.odm.models.service import DockerConfig, Service
27
41
  from assemblyline.remote.datatypes.events import EventSender, EventWatcher
28
42
  from assemblyline.remote.datatypes.hash import Hash
29
- from assemblyline_core.scaler.controllers.kubernetes_ctl import create_docker_auth_config, selector_to_node_affinity, PRIVILEGED_SERVICE_ACCOUNT_NAME
43
+ from assemblyline_core.scaler.controllers.kubernetes_ctl import (
44
+ PRIVILEGED_SERVICE_ACCOUNT_NAME,
45
+ create_docker_auth_config,
46
+ selector_to_node_affinity,
47
+ )
30
48
  from assemblyline_core.server_base import ThreadedCoreBase
31
49
  from assemblyline_core.updater.helper import get_latest_tag_for_service
32
50
 
@@ -157,7 +175,7 @@ class DockerUpdateInterface:
157
175
 
158
176
  class KubernetesUpdateInterface:
159
177
  def __init__(self, logger, prefix, namespace, priority_class, extra_labels, linux_node_selector: Selector,
160
- log_level="INFO", default_service_account=None, default_service_tolerations=[], enable_pod_security=False):
178
+ log_level="INFO", default_service_tolerations=[], enable_pod_security=False):
161
179
  # Try loading a kubernetes connection from either the fact that we are running
162
180
  # inside of a cluster, or we have a configuration in the normal location
163
181
  try:
@@ -187,7 +205,6 @@ class KubernetesUpdateInterface:
187
205
  self.priority_class = priority_class
188
206
  self.extra_labels = extra_labels
189
207
  self.log_level = log_level
190
- self.default_service_account = default_service_account
191
208
  self.secret_env = []
192
209
  self.linux_node_selector = linux_node_selector
193
210
  self.default_service_tolerations = [V1Toleration(**toleration.as_primitives()) for toleration in default_service_tolerations]
@@ -268,13 +285,7 @@ class KubernetesUpdateInterface:
268
285
  read_only=mount.read_only,
269
286
  )
270
287
 
271
- if mount.config_map:
272
- # Deprecated configuration for mounting ConfigMap
273
- # TODO: Deprecate code on next major change
274
- vol_kwargs.update(dict(config_map=V1ConfigMapVolumeSource(name=mount.config_map, optional=False)))
275
- vol_mount_kwargs.update(dict(sub_path=mount.key))
276
-
277
- elif mount.resource_type == 'secret':
288
+ if mount.resource_type == 'secret':
278
289
  # Secret-based source
279
290
  vol_kwargs.update(dict(secret=V1SecretVolumeSource(secret_name=mount.resource_name)))
280
291
  vol_mount_kwargs.update(dict(sub_path=mount.resource_key))
@@ -346,7 +357,7 @@ class KubernetesUpdateInterface:
346
357
  restart_policy='Never',
347
358
  containers=[container],
348
359
  priority_class_name=self.priority_class,
349
- service_account_name=docker_config.service_account or self.default_service_account or PRIVILEGED_SERVICE_ACCOUNT_NAME,
360
+ service_account_name=docker_config.service_account or PRIVILEGED_SERVICE_ACCOUNT_NAME,
350
361
  affinity=selector_to_node_affinity(self.linux_node_selector),
351
362
  tolerations=self.default_service_tolerations
352
363
  )
@@ -487,7 +498,6 @@ class ServiceUpdater(ThreadedCoreBase):
487
498
  priority_class='al-core-priority',
488
499
  extra_labels=extra_labels,
489
500
  log_level=self.config.logging.log_level,
490
- default_service_account=self.config.services.service_account,
491
501
  linux_node_selector=self.config.core.scaler.linux_node_selector,
492
502
  default_service_tolerations=self.config.core.scaler.service_defaults.tolerations,
493
503
  enable_pod_security=self.config.core.scaler.enable_pod_security)
@@ -512,7 +522,6 @@ class ServiceUpdater(ThreadedCoreBase):
512
522
  tag = 'stable'
513
523
  else:
514
524
  tag = 'latest'
515
- service_key = None
516
525
  try:
517
526
  service = Service(
518
527
  {'name': service_name,
@@ -1,52 +1,55 @@
1
1
  import contextlib
2
2
  import datetime
3
3
  import json
4
- import tempfile
5
4
  import logging
6
5
  import os
7
- import time
8
- import signal
9
6
  import shutil
10
- from copy import deepcopy
11
- from typing import Optional, Any
12
- from multiprocessing import Lock, Event
7
+ import signal
8
+ import tempfile
13
9
  import threading
10
+ import time
11
+ from copy import deepcopy
12
+ from multiprocessing import Event, Lock
13
+ from typing import Any, Optional
14
14
 
15
- import elasticapm
16
15
  import arrow
16
+ import elasticapm
17
17
 
18
- from assemblyline.common.forge import CachedObject, get_classification, get_config, get_datastore, get_filestore, \
19
- get_apm_client
18
+ from assemblyline.common import identify
20
19
  from assemblyline.common.codec import decode_file
21
20
  from assemblyline.common.dict_utils import flatten
21
+ from assemblyline.common.forge import (
22
+ CachedObject,
23
+ get_apm_client,
24
+ get_classification,
25
+ get_config,
26
+ get_datastore,
27
+ get_filestore,
28
+ )
29
+ from assemblyline.common.isotime import now_as_iso
22
30
  from assemblyline.common.log import init_logging
23
31
  from assemblyline.common.metrics import MetricsFactory
24
- from assemblyline.datastore.helper import AssemblylineDatastore, MetadataValidator
25
- from assemblyline.common import identify
26
- from assemblyline.common.isotime import now_as_iso
32
+ from assemblyline.common.str_utils import safe_str
27
33
  from assemblyline.common.uid import get_random_id
34
+ from assemblyline.datastore.helper import AssemblylineDatastore, MetadataValidator
35
+ from assemblyline.filestore import FileStore
36
+ from assemblyline.odm.messages.submission import Submission
37
+ from assemblyline.odm.messages.vacuum_heartbeat import Metrics
28
38
  from assemblyline.odm.models import user
29
39
  from assemblyline.odm.models.config import Config
30
40
  from assemblyline.odm.models.submission import DEFAULT_SRV_SEL
31
41
  from assemblyline.odm.models.user_settings import UserSettings
32
- from assemblyline.remote.datatypes.queues.comms import CommsQueue
33
- from assemblyline.odm.messages.vacuum_heartbeat import Metrics
34
-
35
- from assemblyline.filestore import FileStore
36
- from assemblyline.common.str_utils import safe_str
37
42
  from assemblyline.remote.datatypes import get_client as get_redis_client
38
- from assemblyline.odm.messages.submission import Submission
39
- from assemblyline.remote.datatypes.queues.named import NamedQueue
40
43
  from assemblyline.remote.datatypes.hash import Hash
41
-
42
- from assemblyline_core.vacuum.crawler import VACUUM_BUFFER_NAME
44
+ from assemblyline.remote.datatypes.queues.comms import CommsQueue
45
+ from assemblyline.remote.datatypes.queues.named import NamedQueue
43
46
  from assemblyline_core.ingester.constants import INGEST_QUEUE_NAME
47
+ from assemblyline_core.vacuum.crawler import VACUUM_BUFFER_NAME
44
48
 
45
- from .safelist import VacuumSafelist
46
- from .department_map import DepartmentMap
47
- from .stream_map import StreamMap, Stream
48
49
  from .crawler import heartbeat
49
-
50
+ from .department_map import DepartmentMap
51
+ from .safelist import VacuumSafelist
52
+ from .stream_map import Stream, StreamMap
50
53
 
51
54
  # init_logging('assemblyline.vacuum.worker')
52
55
  logger = logging.getLogger('assemblyline.vacuum.worker')
@@ -410,7 +413,7 @@ class FileProcessor(threading.Thread):
410
413
  'deep_scan': False,
411
414
  "priority": 150,
412
415
  "ignore_cache": False,
413
- "ignore_dynamic_recursion_prevention": False,
416
+ "ignore_recursion_prevention": False,
414
417
  "ignore_filtering": False,
415
418
  "type": "INGEST"
416
419
  })
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: assemblyline-core
3
- Version: 4.5.1.dev501
3
+ Version: 4.6.0.0
4
4
  Summary: Assemblyline 4 - Core components
5
5
  Home-page: https://github.com/CybercentreCanada/assemblyline-core/
6
6
  Author: CCCS Assemblyline development team
@@ -1,20 +1,10 @@
1
+ import json
1
2
  import logging
2
3
  import time
3
4
  from unittest import mock
4
5
 
5
- import json
6
6
  import pytest
7
-
8
- from assemblyline.common.forge import get_service_queue, get_classification
9
- from assemblyline.odm.models.error import Error
10
- from assemblyline.odm.models.file import File
11
- from assemblyline.odm.models.result import Result
12
- from assemblyline.odm.models.user import User
13
- from assemblyline.odm.randomizer import random_model_obj, random_minimal_obj, get_random_hash
14
- from assemblyline.odm import models
15
- from assemblyline.common.metrics import MetricsFactory
16
-
17
- from assemblyline_core.dispatching.client import DispatchClient, DISPATCH_RESULT_QUEUE
7
+ from assemblyline_core.dispatching.client import DISPATCH_RESULT_QUEUE, DispatchClient
18
8
  from assemblyline_core.dispatching.dispatcher import Dispatcher, ServiceTask, Submission
19
9
  from assemblyline_core.dispatching.schedules import Scheduler as RealScheduler
20
10
 
@@ -23,6 +13,18 @@ from assemblyline_core.dispatching.timeout import TimeoutTable
23
13
  from mocking import ToggleTrue
24
14
  from test_scheduler import dummy_service
25
15
 
16
+ from assemblyline.common.forge import get_classification, get_service_queue
17
+ from assemblyline.common.metrics import MetricsFactory
18
+ from assemblyline.odm import models
19
+ from assemblyline.odm.models.error import Error
20
+ from assemblyline.odm.models.file import File
21
+ from assemblyline.odm.models.result import Result
22
+ from assemblyline.odm.models.user import User
23
+ from assemblyline.odm.randomizer import (
24
+ get_random_hash,
25
+ random_minimal_obj,
26
+ random_model_obj,
27
+ )
26
28
 
27
29
  logger = logging.getLogger('assemblyline.test')
28
30
 
@@ -311,8 +313,6 @@ def test_dispatch_extracted_bypass_drp(clean_redis, clean_datastore):
311
313
  submission = random_model_obj(Submission)
312
314
  submission.to_be_deleted = False
313
315
 
314
- # the following 1 line can be removed after assemblyline upgrade to version 4.6+
315
- submission.params.ignore_dynamic_recursion_prevention = False
316
316
  submission.params.ignore_recursion_prevention = False
317
317
  submission.params.services.selected = ['extract', 'sandbox']
318
318
  submission.files = [dict(name='./file', sha256=file_hash)]
@@ -1 +0,0 @@
1
- 4.5.1.dev501