secator 0.22.0__py3-none-any.whl

secator/.gitignore

@@ -0,0 +1,162 @@
+# ---> Python
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+

secator/celery.py

@@ -0,0 +1,453 @@
+import gc
+import json
+import logging
+import os
+
+from time import time
+
+from celery import Celery, chord
+from celery.canvas import signature
+from celery.app import trace
+
+from rich.logging import RichHandler
+from retry import retry
+
+from secator.celery_signals import IN_CELERY_WORKER_PROCESS, setup_handlers
+from secator.config import CONFIG
+from secator.output_types import Info
+from secator.rich import console
+from secator.runners import Scan, Task, Workflow
+from secator.utils import (debug, deduplicate, flatten, should_update)
+
+
+#---------#
+# Logging #
+#---------#
+
+rich_handler = RichHandler(rich_tracebacks=True)
+rich_handler.setLevel(logging.INFO)
+logging.basicConfig(
+	level='NOTSET',
+	format="%(threadName)s:%(message)s",
+	datefmt="[%X]",
+	handlers=[rich_handler],
+	force=True)
+logging.getLogger('kombu').setLevel(logging.ERROR)
+logging.getLogger('celery').setLevel(logging.DEBUG if 'celery.debug' in CONFIG.debug or 'celery.*' in CONFIG.debug else logging.WARNING)  # noqa: E501
+logger = logging.getLogger(__name__)
+trace.LOG_SUCCESS = "Task %(name)s[%(id)s] succeeded in %(runtime)ss"
+
+
+#------------#
+# Celery app #
+#------------#
+
+app = Celery(__name__)
+app.conf.update({
+	# Content types
+	'accept_content': ['application/x-python-serialize', 'application/json'],
+
+	# Broker config
+	'broker_url': CONFIG.celery.broker_url,
+	'broker_transport_options': json.loads(CONFIG.celery.broker_transport_options) if CONFIG.celery.broker_transport_options else {  # noqa: E501
+		'data_folder_in': CONFIG.dirs.celery_data,
+		'data_folder_out': CONFIG.dirs.celery_data,
+		'control_folder': CONFIG.dirs.celery_data,
+		'visibility_timeout': CONFIG.celery.broker_visibility_timeout,
+	},
+	'broker_connection_retry_on_startup': True,
+	'broker_pool_limit': CONFIG.celery.broker_pool_limit,
+	'broker_connection_timeout': CONFIG.celery.broker_connection_timeout,
+
+	# Result backend config
+	'result_backend': CONFIG.celery.result_backend,
+	'result_expires': CONFIG.celery.result_expires,
+	'result_backend_transport_options': json.loads(CONFIG.celery.result_backend_transport_options) if CONFIG.celery.result_backend_transport_options else {},  # noqa: E501
+	'result_extended': not CONFIG.addons.mongodb.enabled,
+	'result_backend_thread_safe': True,
+	'result_serializer': 'pickle',
+	'result_accept_content': ['application/x-python-serialize'],
+
+	# Task config
+	'task_acks_late': CONFIG.celery.task_acks_late,
+	'task_compression': 'gzip',
+	'task_create_missing_queues': True,
+	'task_eager_propagates': False,
+	'task_reject_on_worker_lost': CONFIG.celery.task_reject_on_worker_lost,
+	'task_routes': {
+		'secator.celery.run_workflow': {'queue': 'celery'},
+		'secator.celery.run_scan': {'queue': 'celery'},
+		'secator.celery.run_task': {'queue': 'celery'},
+		'secator.celery.forward_results': {'queue': 'results'},
+		'secator.hooks.mongodb.*': {'queue': 'mongodb'}
+	},
+	'task_store_eager_result': True,
+	'task_send_sent_event': CONFIG.celery.task_send_sent_event,
+	'task_serializer': 'pickle',
+	'task_accept_content': ['application/x-python-serialize'],
+
+	# Event config
+	'event_serializer': 'pickle',
+	'event_accept_content': ['application/x-python-serialize'],
+
+	# Worker config
+	# 'worker_direct': True,  # TODO: consider enabling this to allow routing to specific workers
+	'worker_max_tasks_per_child': CONFIG.celery.worker_max_tasks_per_child,
+	# 'worker_max_memory_per_child': 100000  # TODO: consider enabling this
+	'worker_pool_restarts': True,
+	'worker_prefetch_multiplier': CONFIG.celery.worker_prefetch_multiplier,
+	'worker_send_task_events': CONFIG.celery.worker_send_task_events
+})
+app.autodiscover_tasks(['secator.hooks.mongodb'], related_name=None)
+if IN_CELERY_WORKER_PROCESS:
+	setup_handlers()
+
+
+@retry(Exception, tries=3, delay=2)
+def update_state(celery_task, task, force=False):
+	"""Update task state to add metadata information."""
+	if not IN_CELERY_WORKER_PROCESS:
+		return
+	if task.no_live_updates:
+		return
+	if not force and not should_update(CONFIG.runners.backend_update_frequency, task.last_updated_celery):
+		return
+	task.last_updated_celery = time()
+	debug(
+		'',
+		sub='celery.state',
+		id=celery_task.request.id,
+		obj={task.unique_name: task.status, 'count': task.self_findings_count},
+		obj_after=False,
+		verbose=True
+	)
+	return celery_task.update_state(
+		state='RUNNING',
+		meta=task.celery_state
+	)
+
+
+def revoke_task(task_id, task_name=None):
+	message = f'Revoked task {task_id}'
+	if task_name:
+		message += f' ({task_name})'
+	app.control.revoke(task_id, terminate=True)
+	console.print(Info(message=message))
+
+
+#--------------#
+# Celery tasks #
+#--------------#
+
+
+def chunker(seq, size):
+	return (seq[pos:pos + size] for pos in range(0, len(seq), size))
+
+
+@app.task(bind=True)
+def run_task(self, args=[], kwargs={}):
+	console.print(Info(message=f'Running task {self.request.id}'))
+	if 'context' not in kwargs:
+		kwargs['context'] = {}
+	kwargs['context']['celery_id'] = self.request.id
+	task = Task(*args, **kwargs)
+	task.run()
+
+
+@app.task(bind=True)
+def run_workflow(self, args=[], kwargs={}):
+	console.print(Info(message=f'Running workflow {self.request.id}'))
+	if 'context' not in kwargs:
+		kwargs['context'] = {}
+	kwargs['context']['celery_id'] = self.request.id
+	workflow = Workflow(*args, **kwargs)
+	workflow.run()
+
+
+@app.task(bind=True)
+def run_scan(self, args=[], kwargs={}):
+	console.print(Info(message=f'Running scan {self.request.id}'))
+	if 'context' not in kwargs:
+		kwargs['context'] = {}
+	kwargs['context']['celery_id'] = self.request.id
+	scan = Scan(*args, **kwargs)
+	scan.run()
+
+
+@app.task(bind=True)
+def run_command(self, results, name, targets, opts={}):
+	# Set Celery request id in context
+	context = opts.get('context', {})
+	context['celery_id'] = self.request.id
+	context['worker_name'] = os.environ.get('WORKER_NAME', 'unknown')
+
+	# Set routing key in context
+	if IN_CELERY_WORKER_PROCESS:
+		quiet = not CONFIG.cli.worker_command_verbose
+		opts.update({
+			'print_item': True,
+			'print_line': True,
+			'print_cmd': True,
+			'print_target': True,
+			'print_profiles': True,
+			'quiet': quiet
+		})
+		routing_key = self.request.delivery_info['routing_key']
+		context['routing_key'] = routing_key
+		debug(f'Task "{name}" running with routing key "{routing_key}"', sub='celery.state')
+
+	# Flatten + dedupe + filter results
+	results = forward_results(results)
+
+	# Set task opts
+	opts['context'] = context
+	opts['results'] = results
+	opts['sync'] = True
+
+	# Initialize task
+	sync = not IN_CELERY_WORKER_PROCESS
+	task_cls = Task.get_task_class(name)
+	task = task_cls(targets, **opts)
+	chunk_it = task.needs_chunking(sync)
+	task.has_children = chunk_it
+	task.mark_started()
+	update_state(self, task, force=True)
+
+	# Chunk task if needed
+	if chunk_it:
+		if IN_CELERY_WORKER_PROCESS:
+			console.print(Info(message=f'Task {name} requires chunking'))
+		workflow = break_task(task, opts, results=results)
+		if IN_CELERY_WORKER_PROCESS:
+			console.print(Info(message=f'Task {name} successfully broken into {len(workflow)} chunks'))
+		update_state(self, task, force=True)
+		console.print(Info(message=f'Task {name} updated state, replacing task with Celery chord workflow'))
+		return replace(self, workflow)
+
+	# Update state live
+	for _ in task:
+		update_state(self, task)
+	update_state(self, task, force=True)
+
+	if CONFIG.addons.mongodb.enabled:
+		return [r._uuid for r in task.results]
+	return task.results
+
+
+@app.task
+def forward_results(results):
+	"""Forward results to the next task (bridge task).
+
+	Args:
+		results (list): Results to forward.
+
+	Returns:
+		list: List of uuids.
+	"""
+	if isinstance(results, list):
+		for ix, item in enumerate(results):
+			if isinstance(item, dict) and 'results' in item:
+				results[ix] = item['results']
+	elif 'results' in results:
+		results = results['results']
+
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Deduplicating {len(results)} results'))
+
+	results = flatten(results)
+	if IN_CELERY_WORKER_PROCESS and CONFIG.addons.mongodb.enabled:
+		console.print(Info(message=f'Extracting uuids from {len(results)} results'))
+		uuids = [r._uuid for r in results if hasattr(r, '_uuid')]
+		uuids.extend([r for r in results if isinstance(r, str)])
+		results = list(set(uuids))
+	else:
+		results = deduplicate(results, attr='_uuid')
+
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Forwarded {len(results)} flattened and deduplicated results'))
+
+	return results
+
+
+@app.task
+def mark_runner_started(results, runner, enable_hooks=True):
+	"""Mark a runner as started and run on_start hooks.
+
+	Args:
+		results (List): Previous results.
+		runner (Runner): Secator runner instance.
+		enable_hooks (bool): Enable hooks.
+
+	Returns:
+		list: Runner results
+	"""
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Runner {runner.unique_name} has started, running mark_started'))
+	debug(f'Runner {runner.unique_name} has started, running mark_started', sub='celery')
+	if results:
+		results = forward_results(results)
+	runner.enable_hooks = enable_hooks
+	if IN_CELERY_WORKER_PROCESS and CONFIG.addons.mongodb.enabled:
+		from secator.hooks.mongodb import get_results
+		results = get_results(results)
+	for item in results:
+		runner.add_result(item, print=False)
+	runner.mark_started()
+	if IN_CELERY_WORKER_PROCESS and CONFIG.addons.mongodb.enabled:
+		return [r._uuid for r in runner.results]
+	return runner.results
+
+
+@app.task
+def mark_runner_completed(results, runner, enable_hooks=True):
+	"""Mark a runner as completed and run on_end hooks.
+
+	Args:
+		results (list): Task results
+		runner (Runner): Secator runner instance
+		enable_hooks (bool): Enable hooks.
+
+	Returns:
+		list: Final results
+	"""
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Runner {runner.unique_name} has finished, running mark_completed'))
+	debug(f'Runner {runner.unique_name} has finished, running mark_completed', sub='celery')
+	results = forward_results(results)
+	runner.enable_hooks = enable_hooks
+	if IN_CELERY_WORKER_PROCESS and CONFIG.addons.mongodb.enabled:
+		from secator.hooks.mongodb import get_results
+		results = get_results(results)
+	for item in results:
+		runner.add_result(item, print=False)
+	runner.mark_completed()
+	if IN_CELERY_WORKER_PROCESS and CONFIG.addons.mongodb.enabled:
+		return [r._uuid for r in runner.results]
+	return runner.results
+
+
+#--------------#
+# Celery utils #
+#--------------#
+
+
+def is_celery_worker_alive():
+	"""Check if a Celery worker is available."""
+	result = app.control.broadcast('ping', reply=True, limit=1, timeout=1)
+	result = bool(result)
+	if result:
+		console.print(Info(message='Celery worker is available, running remotely'))
+	else:
+		console.print(Info(message='No Celery worker available, running locally'))
+	return result
+
+
+def replace(task_instance, sig):
+	"""Replace this task, with a new task inheriting the task id.
+
+	Execution of the host task ends immediately and no subsequent statements
+	will be run.
+
+	.. versionadded:: 4.0
+
+	Arguments:
+		sig (Signature): signature to replace with.
+		visitor (StampingVisitor): Visitor API object.
+
+	Raises:
+		~@Ignore: This is always raised when called in asynchronous context.
+		It is best to always use ``return self.replace(...)`` to convey
+		to the reader that the task won't continue after being replaced.
+	"""
+	chord = task_instance.request.chord
+	sig.freeze(task_instance.request.id)
+	replaced_task_nesting = task_instance.request.get('replaced_task_nesting', 0) + 1
+	sig.set(
+		chord=chord,
+		group_id=task_instance.request.group,
+		group_index=task_instance.request.group_index,
+		root_id=task_instance.request.root_id,
+		replaced_task_nesting=replaced_task_nesting
+	)
+	import psutil
+	import os
+	process = psutil.Process(os.getpid())
+	length = len(task_instance.request.chain) if task_instance.request.chain else 0
+	# console.print(f'Adding {length} chain tasks from request chain')
+	for ix, t in enumerate(reversed(task_instance.request.chain or [])):
+		console.print(Info(message=f'Adding chain task {t.name} from request chain ({ix + 1}/{length})'))
+		chain_task = signature(t, app=task_instance.app)
+		chain_task.set(replaced_task_nesting=replaced_task_nesting)
+		sig |= chain_task
+		del chain_task
+		del t
+		memory_bytes = process.memory_info().rss
+		console.print(Info(message=f'Memory usage: {memory_bytes / 1024 / 1024:.2f} MB (chain task {ix + 1}/{length})'))
+	gc.collect()
+	return task_instance.on_replace(sig)
+
+
+def break_task(task, task_opts, results=[]):
+	"""Break a task into multiple of the same type."""
+	chunks = task.inputs
+	if task.input_chunk_size > 1:
+		chunks = list(chunker(task.inputs, task.input_chunk_size))
+	debug(
+		'',
+		obj={task.unique_name: 'CHUNKED', 'chunk_size': task.input_chunk_size, 'chunks': len(chunks), 'target_count': len(task.inputs)},  # noqa: E501
+		obj_after=False,
+		sub='celery.state',
+		verbose=True
+	)
+
+	# Clone opts
+	base_opts = task_opts.copy()
+
+	# Build signatures
+	sigs = []
+	task.ids_map = {}
+	for ix, chunk in enumerate(chunks):
+		if not isinstance(chunk, list):
+			chunk = [chunk]
+
+		# Add chunk info to opts
+		opts = base_opts.copy()
+		opts.update({'chunk': ix + 1, 'chunk_count': len(chunks)})
+		debug('', obj={
+			task.unique_name: 'CHUNK',
+			'chunk': f'{ix + 1} / {len(chunks)}',
+			'target_count': len(chunk),
+			'targets': chunk
+		}, sub='celery.state')  # noqa: E501
+
+		# Construct chunked signature
+		opts['has_parent'] = True
+		opts['enable_duplicate_check'] = False
+		opts['results'] = results
+		if 'targets_' in opts:
+			del opts['targets_']
+		sig = type(task).si(chunk, **opts)
+		task_id = sig.freeze().task_id
+		full_name = f'{task.name}_{ix + 1}'
+		task.add_subtask(task_id, task.name, full_name)
+		if IN_CELERY_WORKER_PROCESS:
+			info = Info(message=f'Celery chunked task created ({ix + 1} / {len(chunks)}): {task_id}')
+			task.add_result(info)
+		sigs.append(sig)
+
+	# Mark main task as async since it's being chunked
+	task.sync = False
+	task.results = []
+	task.uuids = set()
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Task {task.unique_name} is now async, building chord with {len(sigs)} chunks'))
+	# console.print(Info(message=f'Results: {results}'))
+
+	# Build Celery workflow
+	workflow = chord(
+		tuple(sigs),
+		mark_runner_completed.s(runner=task).set(queue='results')
+	)
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Task {task.unique_name} chord built with {len(sigs)} chunks, returning workflow'))
+	return workflow