def __init__(self, component_name: str, logger: logging.Logger = None, shutdown_timeout: float = None, config=None, datastore=None, redis=None, redis_persist=None): super().__init__(component_name=component_name, logger=logger, shutdown_timeout=shutdown_timeout, config=config) self.datastore: AssemblylineDatastore = datastore or forge.get_datastore( self.config) # Connect to all of our persistent redis structures self.redis: Redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.redis_persist: Redis = redis_persist or get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) # Create a cached service data object, and access to the service status self.service_info = typing.cast(typing.Dict[str, Service], forge.CachedObject(self._get_services)) self._service_stage_hash = get_service_stage_hash(self.redis)
def __init__(self, datastore=None, redis=None, redis_persist=None, logger=None): self.config = forge.get_config() self.redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) redis_persist = redis_persist or get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) self.timeout_watcher = WatcherClient(redis_persist) self.submission_queue = NamedQueue(SUBMISSION_QUEUE, self.redis) self.file_queue = NamedQueue(FILE_QUEUE, self.redis) self.ds = datastore or forge.get_datastore(self.config) self.log = logger or logging.getLogger( "assemblyline.dispatching.client") self.results = datastore.result self.errors = datastore.error self.files = datastore.file self.active_submissions = ExpiringHash(DISPATCH_TASK_HASH, host=redis_persist) self.running_tasks = ExpiringHash(DISPATCH_RUNNING_TASK_HASH, host=self.redis) self.service_data = cast(Dict[str, Service], CachedObject(self._get_services))
def __init__(self, shutdown_timeout: int = SHUTDOWN_SECONDS_LIMIT): super(RunPrivilegedService, self).__init__(f'assemblyline.service.{SERVICE_NAME}', shutdown_timeout=shutdown_timeout) self.client_id = os.environ.get('HOSTNAME', 'dev-service') self.redis = get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.redis_persist = get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) self.tasking_client = TaskingClient(redis=self.redis, redis_persist=self.redis_persist) self.tasking_dir = os.environ.get('TASKING_DIR', tempfile.gettempdir()) self.filestore = forge.get_filestore() self.service = None self.service_config = {} self.service_name = None self.service_tool_version = None self.status = STATUSES.INITIALIZING self.metric_factory = None self.log.setLevel(LOG_LEVEL)
def __init__(self, sender, log, config=None, redis=None): self.sender = sender self.log = log self.config = config or forge.get_config() self.datastore = forge.get_datastore(self.config) self.redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.redis_persist = get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) self.status_queue = CommsQueue(STATUS_QUEUE, self.redis) self.dispatch_active_hash = Hash(DISPATCH_TASK_HASH, self.redis_persist) self.dispatcher_submission_queue = NamedQueue(SUBMISSION_QUEUE, self.redis) self.ingest_scanning = Hash('m-scanning-table', self.redis_persist) self.ingest_unique_queue = PriorityQueue('m-unique', self.redis_persist) self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.redis_persist) self.ingest_complete_queue = NamedQueue(COMPLETE_QUEUE_NAME, self.redis) self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.redis_persist) constants = forge.get_constants(self.config) self.c_rng = constants.PRIORITY_RANGES['critical'] self.h_rng = constants.PRIORITY_RANGES['high'] self.m_rng = constants.PRIORITY_RANGES['medium'] self.l_rng = constants.PRIORITY_RANGES['low'] self.c_s_at = self.config.core.ingester.sampling_at['critical'] self.h_s_at = self.config.core.ingester.sampling_at['high'] self.m_s_at = self.config.core.ingester.sampling_at['medium'] self.l_s_at = self.config.core.ingester.sampling_at['low'] self.to_expire = {k: 0 for k in metrics.EXPIRY_METRICS} if self.config.core.expiry.batch_delete: self.delete_query = f"expiry_ts:[* TO {self.datastore.ds.now}-{self.config.core.expiry.delay}" \ f"{self.datastore.ds.hour}/DAY]" else: self.delete_query = f"expiry_ts:[* TO {self.datastore.ds.now}-{self.config.core.expiry.delay}" \ f"{self.datastore.ds.hour}]" self.scheduler = BackgroundScheduler(daemon=True) self.scheduler.add_job( self._reload_expiry_queues, 'interval', seconds=self.config.core.metrics.export_interval * 4) self.scheduler.start()
def __init__(self, names, host=None, port=None, private=False): self.c = get_client(host, port, private) self.p = retry_call(self.c.pubsub) if not isinstance(names, list): names = [names] self.names = names self._connected = False
def __init__(self, prefix="counter", host=None, port=None, track_counters=False): self.c = get_client(host, port, False) self.prefix = prefix if track_counters: self.tracker = Hash("c-tracker-%s" % prefix, host=host, port=port) else: self.tracker = None
def redis(): config = forge.get_config() client = get_client(config.core.metrics.redis.host, config.core.metrics.redis.port, False) client.flushdb() yield client client.flushdb()
def __init__(self, name, host=None, port=None, private=False): self.c = get_client(host, port, private) self.r = self.c.register_script(pq_pop_script) self.s = self.c.register_script(pq_push_script) self.t = self.c.register_script(pq_unpush_script) self._deque_range = self.c.register_script(pq_dequeue_range_script) self.name = name
def test_create_single_alert(config, datastore): persistent_redis = get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, ) alerter = Alerter() # Swap our alerter onto a private queue so our test doesn't get intercepted alerter.alert_queue = alert_queue = NamedQueue(uuid.uuid4().hex, persistent_redis) # Get a random submission submission = random.choice(all_submissions) all_submissions.remove(submission) # Generate a task for the submission ingest_msg = random_model_obj(IngestTask) ingest_msg.submission.sid = submission.sid ingest_msg.submission.metadata = submission.metadata ingest_msg.submission.params = submission.params ingest_msg.submission.files = submission.files alert_queue.push(ingest_msg.as_primitives()) alert_type = alerter.run_once() assert alert_type == 'create' datastore.alert.commit() res = datastore.alert.search("id:*", as_obj=False) assert res['total'] == 1 alert = datastore.alert.get(res['items'][0]['alert_id']) assert alert.sid == submission.sid
def __init__(self, name, timeout, host=None, port=None): self.uuid = get_random_id() self.c = get_client(host, port, False) self.lock_release = '-'.join(('lock', str(timeout), name, 'released')) self.lock_holder = '-'.join(('lock', str(timeout), name, 'holder')) self.timeout = timeout self._acquire = self.c.register_script(lock_acquire_script) self._release = self.c.register_script(lock_release_script)
def __init__(self, log, host, user, apikey, verify, alert_fqs=None, submission_fqs=None, lookback_time='*'): # Import assemblyline client from assemblyline_client import get_client # Setup AL client self.al_client = get_client(host, apikey=(user, apikey), verify=verify) super().__init__(log, alert_fqs=alert_fqs, submission_fqs=submission_fqs, lookback_time=lookback_time)
def __init__(self, host=None, port=None, private=None, deserializer: Callable[[str], MessageType] = json.loads): client: Redis[Any] = get_client(host, port, private) self.pubsub = retry_call(client.pubsub) self.worker: Optional[threading.Thread] = None self.deserializer = deserializer
def __init__(self, name: str, host=None, port=None, private: bool = False, ttl: int = 0): self.c = get_client(host, port, private) self.name: str = name self.ttl: int = ttl self.last_expire_time = 0
def __init__(self, log, alert_fqs=None, submission_fqs=None, lookback_time='*'): # Setup datastore config = forge.get_config() redis = get_client(config.core.redis.nonpersistent.host, config.core.redis.nonpersistent.port, False) self.datastore = forge.get_datastore(config=config) self.alert_queue = NamedQueue("replay_alert", host=redis) self.file_queue = NamedQueue("replay_file", host=redis) self.submission_queue = NamedQueue("replay_submission", host=redis) super().__init__(log, alert_fqs=alert_fqs, submission_fqs=submission_fqs, lookback_time=lookback_time)
def __init__(self, prefix: str, host=None, port=None, private=None, serializer: Callable[[MessageType], str] = json.dumps): self.client: Redis[Any] = get_client(host, port, private) self.prefix = prefix.lower() if not self.prefix.endswith('.'): self.prefix += '.' self.serializer = serializer
def get_statistics_cache(config=None, redis=None): from assemblyline.remote.datatypes import get_client from assemblyline.remote.datatypes.hash import Hash if not redis: if not config: config = get_config() redis = get_client(config.core.redis.persistent.host, config.core.redis.persistent.port, False) return Hash("cached_statistics", redis)
def __init__(self, redis_persist=None): config = forge.get_config() self.redis = redis_persist or get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, ) self.hash = ExpiringHash(name=WATCHER_HASH, ttl=MAX_TIMEOUT, host=redis_persist) self.queue = UniquePriorityQueue(WATCHER_QUEUE, redis_persist)
def redis_connection(): from assemblyline.remote.datatypes import get_client c = get_client(None, None, False) try: ret_val = c.ping() if ret_val: return c except ConnectionError: pass return pytest.skip( "Connection to the Redis server failed. This test cannot be performed..." )
def export_metrics_once(name, schema, metrics, host=None, counter_type=None, config=None, redis=None): """Manually publish metric counts to the metrics system. This was built for when the service server is reporting metrics for execution and caching on behalf of many services. At the moment the metrics system uses the hosts to count the number of instances of each service. This could be done with a single auto exporting counter for the service server, but that may require significant downstream changes in the metrics system. """ config = config or forge.get_config() redis = redis or get_client(config.core.metrics.redis.host, config.core.metrics.redis.port, False) # Separate out the timers and normal counters timer_schema = set() counter_schema = set() for _k, field_type in schema.fields().items(): if isinstance(field_type, PerformanceTimer): timer_schema.add(_k) else: counter_schema.add(_k) for _k in timer_schema: counter_schema.discard(_k + '_count') channel = forge.get_metrics_sink(redis) counts = Counters({key: 0 for key in counter_schema}) counts.update({key + '.t': 0 for key in timer_schema}) counts.update({key + '.c': 0 for key in timer_schema}) for metric, value in metrics.items(): if metric in counter_schema: counts[metric] += value elif metric in timer_schema: counts[metric + ".c"] += 1 counts[metric + ".t"] += value else: raise ValueError(f"{metric} is not an accepted counter") counts['type'] = counter_type or name counts['name'] = name counts['host'] = host channel.publish(dict(counts.items()))
def __init__(self, datastore, redis, redis_persist, logger, counter_name='dispatcher'): # Load the datastore collections that we are going to be using self.datastore: AssemblylineDatastore = datastore self.log: logging.Logger = logger self.submissions: Collection = datastore.submission self.results: Collection = datastore.result self.errors: Collection = datastore.error self.files: Collection = datastore.file # Create a config cache that will refresh config values periodically self.config: Config = forge.get_config() # Connect to all of our persistent redis structures self.redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.redis_persist = redis_persist or get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) # Build some utility classes self.scheduler = Scheduler(datastore, self.config, self.redis) self.classification_engine = forge.get_classification() self.timeout_watcher = WatcherClient(self.redis_persist) self.submission_queue = NamedQueue(SUBMISSION_QUEUE, self.redis) self.file_queue = NamedQueue(FILE_QUEUE, self.redis) self._nonper_other_queues = {} self.active_submissions = ExpiringHash(DISPATCH_TASK_HASH, host=self.redis_persist) self.running_tasks = ExpiringHash(DISPATCH_RUNNING_TASK_HASH, host=self.redis) # Publish counters to the metrics sink. self.counter = MetricsFactory(metrics_type='dispatcher', schema=Metrics, name=counter_name, redis=self.redis, config=self.config)
def redis_connection(config): try: from assemblyline.remote.datatypes import get_client c = get_client(config.core.redis.nonpersistent.host, config.core.redis.nonpersistent.port, False) ret_val = c.ping() if ret_val: return c except redis.ConnectionError: pass pytest.skip( "Connection to the Redis server failed. This test cannot be performed..." )
def _reset_service_updates(signature_type): service_updates = Hash( 'service-updates', get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, )) for svc in service_updates.items(): if svc.lower() == signature_type.lower(): update_data = service_updates.get(svc) update_data['next_update'] = now_as_iso(120) update_data['previous_update'] = now_as_iso(-10**10) service_updates.set(svc, update_data) break
def remove_service(servicename, **_): """ Remove a service configuration Variables: servicename => Name of the service to remove Arguments: None Data Block: None Result example: {"success": true} # Has the deletion succeeded """ svc = STORAGE.service_delta.get(servicename) if svc: success = True if not STORAGE.service_delta.delete(servicename): success = False if not STORAGE.service.delete_by_query(f"id:{servicename}*"): success = False STORAGE.heuristic.delete_by_query(f"id:{servicename.upper()}*") STORAGE.signature.delete_by_query(f"type:{servicename.lower()}*") # Notify components watching for service config changes event_sender.send(servicename, { 'operation': Operation.Removed, 'name': servicename }) # Clear potentially unused keys from Redis related to service Hash( f'service-updates-{servicename}', get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, )).delete() return make_api_response({"success": success}) else: return make_api_response({"success": False}, err=f"Service {servicename} does not exist", status_code=404)
def __init__(self, metrics_type, schema, name=None, redis=None, config=None, export_zero=True): self.config = config or forge.get_config() self.redis = redis or get_client(self.config.core.metrics.redis.host, self.config.core.metrics.redis.port, False) # Separate out the timers and normal counters timer_schema = set() counter_schema = set() for _k, field_type in schema.fields().items(): if isinstance(field_type, PerformanceTimer): timer_schema.add(_k) else: counter_schema.add(_k) for _k in timer_schema: counter_schema.discard(_k + '_count') self.type = metrics_type self.name = name or metrics_type # Initialize legacy metrics self.metrics_handler = AutoExportingCounters( self.name, redis=self.redis, config=self.config, counter_type=metrics_type, timer_names=timer_schema, counter_names=counter_schema, export_zero=export_zero) self.metrics_handler.start()
def __init__(self): super().__init__('assemblyline.alerter') # Publish counters to the metrics sink. self.counter = MetricsFactory('alerter', Metrics) self.datastore = forge.get_datastore(self.config) self.persistent_redis = get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) self.process_alert_message = forge.get_process_alert_message() self.running = False self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.persistent_redis) if self.config.core.metrics.apm_server.server_url is not None: self.log.info( f"Exporting application metrics to: {self.config.core.metrics.apm_server.server_url}" ) elasticapm.instrument() self.apm_client = elasticapm.Client( server_url=self.config.core.metrics.apm_server.server_url, service_name="alerter") else: self.apm_client = None
def __init__(self, host=None, port=None, private=False): self.c = get_client(host, port, private)
import os import threading from assemblyline.common import forge from assemblyline.common import log as al_log from assemblyline.common.version import BUILD_MINOR, FRAMEWORK_VERSION, SYSTEM_VERSION from assemblyline.remote.datatypes.counters import Counters from assemblyline.remote.datatypes import get_client from assemblyline_core.safelist_client import SafelistClient from assemblyline_core.tasking_client import TaskingClient config = forge.get_config() redis = get_client( host=config.core.redis.nonpersistent.host, port=config.core.redis.nonpersistent.port, private=False, ) redis_persist = get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, ) ################################################################# # Configuration CLASSIFICATION = forge.get_classification() DEBUG = config.ui.debug VERSION = os.environ.get(
def __init__(self, datastore, logger, classification=None, redis=None, persistent_redis=None, metrics_name='ingester'): self.datastore = datastore self.log = logger # Cache the user groups self.cache_lock = threading.RLock( ) # TODO are middle man instances single threaded now? self._user_groups = {} self._user_groups_reset = time.time() // HOUR_IN_SECONDS self.cache = {} self.notification_queues = {} self.whitelisted = {} self.whitelisted_lock = threading.RLock() # Create a config cache that will refresh config values periodically self.config = forge.CachedObject(forge.get_config) # Module path parameters are fixed at start time. Changing these involves a restart self.is_low_priority = load_module_by_path( self.config.core.ingester.is_low_priority) self.get_whitelist_verdict = load_module_by_path( self.config.core.ingester.get_whitelist_verdict) self.whitelist = load_module_by_path( self.config.core.ingester.whitelist) # Constants are loaded based on a non-constant path, so has to be done at init rather than load constants = forge.get_constants(self.config) self.priority_value = constants.PRIORITIES self.priority_range = constants.PRIORITY_RANGES self.threshold_value = constants.PRIORITY_THRESHOLDS # Connect to the redis servers self.redis = redis or get_client( host=self.config.core.redis.nonpersistent.host, port=self.config.core.redis.nonpersistent.port, private=False, ) self.persistent_redis = persistent_redis or get_client( host=self.config.core.redis.persistent.host, port=self.config.core.redis.persistent.port, private=False, ) # Classification engine self.ce = classification or forge.get_classification() # Metrics gathering factory self.counter = MetricsFactory(metrics_type='ingester', schema=Metrics, redis=self.redis, config=self.config, name=metrics_name) # State. The submissions in progress are stored in Redis in order to # persist this state and recover in case we crash. self.scanning = Hash('m-scanning-table', self.persistent_redis) # Input. The dispatcher creates a record when any submission completes. self.complete_queue = NamedQueue(_completeq_name, self.redis) # Internal. Dropped entries are placed on this queue. # self.drop_queue = NamedQueue('m-drop', self.persistent_redis) # Input. An external process places submission requests on this queue. self.ingest_queue = NamedQueue(INGEST_QUEUE_NAME, self.persistent_redis) # Output. Duplicate our input traffic into this queue so it may be cloned by other systems self.traffic_queue = CommsQueue('submissions', self.redis) # Internal. Unique requests are placed in and processed from this queue. self.unique_queue = PriorityQueue('m-unique', self.persistent_redis) # Internal, delay queue for retrying self.retry_queue = PriorityQueue('m-retry', self.persistent_redis) # Internal, timeout watch queue self.timeout_queue = PriorityQueue('m-timeout', self.redis) # Internal, queue for processing duplicates # When a duplicate file is detected (same cache key => same file, and same # submission parameters) the file won't be ingested normally, but instead a reference # will be written to a duplicate queue. Whenever a file is finished, in the complete # method, not only is the original ingestion finalized, but all entries in the duplicate queue # are finalized as well. This has the effect that all concurrent ingestion of the same file # are 'merged' into a single submission to the system. self.duplicate_queue = MultiQueue(self.persistent_redis) # Output. submissions that should have alerts generated self.alert_queue = NamedQueue(ALERT_QUEUE_NAME, self.persistent_redis) # Utility object to help submit tasks to dispatching self.submit_client = SubmissionClient(datastore=self.datastore, redis=self.redis)
from assemblyline.remote.datatypes.events import EventSender from assemblyline.remote.datatypes.hash import Hash from assemblyline_core.updater.helper import get_latest_tag_for_service from assemblyline_ui.api.base import api_login, make_api_response, make_file_response, make_subapi_blueprint from assemblyline_ui.api.v4.signature import _reset_service_updates from assemblyline_ui.config import LOGGER, STORAGE, config, CLASSIFICATION as Classification SUB_API = 'service' service_api = make_subapi_blueprint(SUB_API, api_version=4) service_api._doc = "Manage the different services" latest_service_tags = Hash( 'service-tags', get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, )) service_update = Hash( 'container-update', get_client( host=config.core.redis.persistent.host, port=config.core.redis.persistent.port, private=False, )) event_sender = EventSender('changes.services', host=config.core.redis.nonpersistent.host, port=config.core.redis.nonpersistent.port)
def __init__(self, name, host=None, port=None): self.c = get_client(host, port, False) self.name = name self._drop_card = self.c.register_script(_drop_card_script) self._limited_add = self.c.register_script(_limited_add)