secator 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

secator/runners/scan.py

@@ -5,8 +5,7 @@ from secator.config import CONFIG
 from secator.runners._base import Runner
 from secator.runners._helpers import run_extractors
 from secator.runners.workflow import Workflow
-from secator.rich import console
-from secator.output_types import Target
+from secator.utils import merge_opts
 
 logger = logging.getLogger(__name__)
 
@@ -24,38 +23,23 @@ class Scan(Runner):
 		"""Run scan.
 
 		Yields:
-			dict: Item yielded from individual workflow tasks.
+			secator.output_types.OutputType: Secator output type.
 		"""
-		# Yield targets
-		for target in self.targets:
-			yield Target(name=target, _source=self.config.name, _type='target', _context=self.context)
-
-		# Run workflows
+		scan_opts = self.config.options
+		self.print_item = False
 		for name, workflow_opts in self.config.workflows.items():
 
 			# Extract opts and and expand target from previous workflows results
-			targets, workflow_opts = run_extractors(self.results, workflow_opts or {}, self.targets)
-			if not targets:
-				console.log(f'No targets were specified for workflow {name}. Skipping.')
-				continue
-
-			# Workflow opts
-			run_opts = self.run_opts.copy()
-			fmt_opts = {
-				'json': run_opts.get('json', False),
-				'print_item': False,
-				'print_start': True,
-				'print_run_summary': True,
-				'print_progress': self.sync
-			}
-			run_opts.update(fmt_opts)
+			targets, workflow_opts = run_extractors(self.results, workflow_opts or {}, self.inputs)
 
 			# Run workflow
+			run_opts = self.run_opts.copy()
+			opts = merge_opts(scan_opts, workflow_opts, run_opts)
 			workflow = Workflow(
 				TemplateLoader(name=f'workflows/{name}'),
 				targets,
 				results=[],
-				run_opts=run_opts,
+				run_opts=opts,
 				hooks=self._hooks,
 				context=self.context.copy())
 

secator/runners/task.py

@@ -1,14 +1,15 @@
-from secator.definitions import DEBUG
-from secator.output_types import Target
 from secator.config import CONFIG
 from secator.runners import Runner
 from secator.utils import discover_tasks
+from secator.celery_utils import CeleryData
+from secator.output_types import Info
 
 
 class Task(Runner):
 	default_exporters = CONFIG.tasks.exporters
 	enable_hooks = False
 
+	@classmethod
 	def delay(cls, *args, **kwargs):
 		from secator.celery import run_task
 		return run_task.apply_async(kwargs={'args': args, 'kwargs': kwargs}, queue='celery')
@@ -16,11 +17,8 @@ class Task(Runner):
 	def yielder(self):
 		"""Run task.
 
-		Args:
-			sync (bool): Run in sync mode (main thread). If False, run in Celery worker in distributed mode.
-
-		Returns:
-			list: List of results.
+		Yields:
+			secator.output_types.OutputType: Secator output type.
 		"""
 		# Get task class
 		task_cls = Task.get_task_class(self.config.name)
@@ -28,26 +26,10 @@ class Task(Runner):
 		# Run opts
 		run_opts = self.run_opts.copy()
 		run_opts.pop('output', None)
-		dry_run = run_opts.get('show', False)
-		if dry_run:
-			self.print_item_count = False
-
-		# Fmt opts
-		fmt_opts = {
-			'json': run_opts.get('json', False),
-			'print_cmd': True,
-			'print_cmd_prefix': not self.sync,
-			'print_input_file': DEBUG > 0,
-			'print_item': True,
-			'print_item_count': not self.sync and not dry_run,
-			'print_line': True
-			# 'print_line': self.sync and not self.output_quiet,
-		}
-		# self.print_item = not self.sync  # enable print_item for base Task only if running remote
-		run_opts.update(fmt_opts)
 
 		# Set task output types
 		self.output_types = task_cls.output_types
+		self.enable_duplicate_check = False
 
 		# Get hooks
 		hooks = {task_cls: self.hooks}
@@ -56,23 +38,25 @@ class Task(Runner):
 
 		# Run task
 		if self.sync:
-			task = task_cls(self.targets, **run_opts)
-			if dry_run:  # don't run
-				return
+			self.print_item = False
+			result = task_cls.si(self.inputs, **run_opts)
+			results = result.apply().get()
 		else:
-			self.celery_result = task_cls.delay(self.targets, **run_opts)
-			task = self.process_live_tasks(
+			self.celery_result = task_cls.delay(self.inputs, **run_opts)
+			self.add_subtask(self.celery_result.id, self.config.name, self.config.description or '')
+			yield Info(
+				message=f'Celery task created: {self.celery_result.id}',
+				task_id=self.celery_result.id
+			)
+			results = CeleryData.iter_results(
 				self.celery_result,
-				description=False,
-				results_only=True,
-				print_remote_status=self.print_remote_status)
+				ids_map=self.celery_ids_map,
+				description=True,
+				print_remote_info=False,
+				print_remote_title=f'[bold gold3]{self.__class__.__name__.capitalize()}[/] [bold magenta]{self.name}[/] results')
 
 		# Yield task results
-		yield from task
-
-		# Yield targets
-		for target in self.targets:
-			yield Target(name=target, _source=self.config.name, _type='target', _context=self.context)
+		yield from results
 
 	@staticmethod
 	def get_task_class(name):
@@ -88,21 +72,3 @@ class Task(Runner):
 			if task_cls.__name__ == name:
 				return task_cls
 		raise ValueError(f'Task {name} not found. Aborting.')
-
-	@staticmethod
-	def get_tasks_from_conf(config):
-		"""Get task names from config. Ignore hierarchy and keywords.
-
-		TODO: Add hierarchy tree / add make flow diagrams.
-		"""
-		tasks = []
-		for name, opts in config.items():
-			if name == '_group':
-				tasks.extend(Task.get_tasks_from_conf(opts))
-			elif name == '_chain':
-				tasks.extend(Task.get_tasks_from_conf(opts))
-			else:
-				if '/' in name:
-					name = name.split('/')[0]
-				tasks.append(name)
-		return tasks

secator/runners/workflow.py

@@ -1,9 +1,11 @@
-from secator.definitions import DEBUG
-from secator.output_types import Target
+import uuid
+
 from secator.config import CONFIG
 from secator.runners._base import Runner
 from secator.runners.task import Task
 from secator.utils import merge_opts
+from secator.celery_utils import CeleryData
+from secator.output_types import Info
 
 
 class Workflow(Runner):
@@ -18,44 +20,39 @@ class Workflow(Runner):
 	def yielder(self):
 		"""Run workflow.
 
-		Args:
-			sync (bool): Run in sync mode (main thread). If False, run in Celery worker in distributed mode.
-
-		Returns:
-			list: List of results.
+		Yields:
+			secator.output_types.OutputType: Secator output type.
 		"""
-		# Yield targets
-		for target in self.targets:
-			yield Target(name=target, _source=self.config.name, _type='target', _context=self.context)
-
 		# Task opts
-		task_run_opts = self.run_opts.copy()
-		task_fmt_opts = {
-			'json': task_run_opts.get('json', False),
-			'print_cmd': True,
-			'print_cmd_prefix': not self.sync,
-			'print_description': self.sync,
-			'print_input_file': DEBUG > 0,
-			'print_item': True,
-			'print_item_count': True,
-			'print_line': not self.sync,
-			'print_progress': self.sync,
-		}
-
-		# Construct run opts
-		task_run_opts['hooks'] = self._hooks.get(Task, {})
-		task_run_opts.update(task_fmt_opts)
+		run_opts = self.run_opts.copy()
+		run_opts['hooks'] = self._hooks.get(Task, {})
 
 		# Build Celery workflow
-		workflow = self.build_celery_workflow(run_opts=task_run_opts, results=self.results)
+		workflow = self.build_celery_workflow(
+			run_opts=run_opts,
+			results=self.results
+		)
+		self.celery_ids = list(self.celery_ids_map.keys())
 
 		# Run Celery workflow and get results
 		if self.sync:
+			self.print_item = False
 			results = workflow.apply().get()
 		else:
 			result = workflow()
+			self.celery_ids.append(str(result.id))
 			self.celery_result = result
-			results = self.process_live_tasks(result, results_only=True, print_remote_status=self.print_remote_status)
+			yield Info(
+				message=f'Celery task created: {self.celery_result.id}',
+				task_id=self.celery_result.id
+			)
+			results = CeleryData.iter_results(
+				self.celery_result,
+				ids_map=self.celery_ids_map,
+				description=True,
+				print_remote_info=self.print_remote_info,
+				print_remote_title=f'[bold gold3]{self.__class__.__name__.capitalize()}[/] [bold magenta]{self.name}[/] results'
+			)
 
 		# Get workflow results
 		yield from results
@@ -64,44 +61,44 @@ class Workflow(Runner):
 		""""Build Celery workflow.
 
 		Returns:
-			celery.chain: Celery task chain.
+			tuple(celery.chain, List[str]): Celery task chain, Celery task ids.
 		"""
 		from celery import chain
 		from secator.celery import forward_results
 		sigs = self.get_tasks(
 			self.config.tasks.toDict(),
-			self.targets,
+			self.inputs,
 			self.config.options,
 			run_opts)
 		sigs = [forward_results.si(results).set(queue='io')] + sigs + [forward_results.s().set(queue='io')]
 		workflow = chain(*sigs)
 		return workflow
 
-	def get_tasks(self, obj, targets, workflow_opts, run_opts):
+	def get_tasks(self, config, inputs, workflow_opts, run_opts):
 		"""Get tasks recursively as Celery chains / chords.
 
 		Args:
-			obj (secator.config.TemplateLoader): Config.
-			targets (list): List of targets.
+			config (dict): Tasks config dict.
+			inputs (list): Inputs.
 			workflow_opts (dict): Workflow options.
 			run_opts (dict): Run options.
 			sync (bool): Synchronous mode (chain of tasks, no chords).
 
 		Returns:
-			list: List of signatures.
+			tuple (List[celery.Signature], List[str]): Celery signatures, Celery task ids.
 		"""
 		from celery import chain, chord
 		from secator.celery import forward_results
 		sigs = []
-		for task_name, task_opts in obj.items():
+		for task_name, task_opts in config.items():
 			# Task opts can be None
 			task_opts = task_opts or {}
 
 			# If it's a group, process the sublevel tasks as a Celery chord.
-			if task_name == '_group':
+			if task_name.startswith('_group'):
 				tasks = self.get_tasks(
 					task_opts,
-					targets,
+					inputs,
 					workflow_opts,
 					run_opts
 				)
@@ -109,7 +106,7 @@ class Workflow(Runner):
 			elif task_name == '_chain':
 				tasks = self.get_tasks(
 					task_opts,
-					targets,
+					inputs,
 					workflow_opts,
 					run_opts
 				)
@@ -125,9 +122,13 @@ class Workflow(Runner):
 				opts['hooks'] = {task: self._hooks.get(Task, {})}
 				opts['context'] = self.context.copy()
 				opts['name'] = task_name
+				opts['has_parent'] = True
+				opts['skip_if_no_inputs'] = True
 
 				# Create task signature
-				sig = task.s(targets, **opts).set(queue=task.profile)
+				task_id = str(uuid.uuid4())
+				sig = task.s(inputs, **opts).set(queue=task.profile, task_id=task_id)
+				self.add_subtask(task_id, task_name, task_opts.get('description', ''))
 				self.output_types.extend(task.output_types)
 			sigs.append(sig)
 		return sigs

secator/scans/__init__.py

@@ -0,0 +1,28 @@
+from secator.cli import ALL_SCANS
+
+
+def generate_class(config):
+	from secator.runners import Workflow
+
+	class workflow(Workflow):
+		def __init__(self, inputs=[], **run_opts):
+			hooks = run_opts.pop('hooks', {})
+			results = run_opts.pop('results', [])
+			context = run_opts.pop('context', {})
+			super().__init__(
+				config=config,
+				inputs=inputs,
+				results=results,
+				run_opts=run_opts,
+				hooks=hooks,
+				context=context)
+	return workflow, config.name
+
+
+DYNAMIC_SCANS = {}
+for workflow in ALL_SCANS:
+	cls, name = generate_class(workflow)
+	DYNAMIC_SCANS[name] = cls
+
+globals().update(DYNAMIC_SCANS)
+__all__ = list(DYNAMIC_SCANS)

secator/serializers/dataclass.py

@@ -1,4 +1,6 @@
+from datetime import date, datetime
 import json
+from pathlib import PosixPath
 from secator.output_types import OUTPUT_TYPES
 
 
@@ -6,6 +8,10 @@ class DataclassEncoder(json.JSONEncoder):
     def default(self, obj):
         if hasattr(obj, 'toDict'):
             return obj.toDict()
+        elif isinstance(obj, PosixPath):
+            return str(obj)
+        elif isinstance(obj, (datetime, date)):
+            return obj.isoformat()
         else:
             return json.JSONEncoder.default(self, obj)
 

secator/serializers/json.py

@@ -1,15 +1,20 @@
-import yaml
+import json
 
 
 class JSONSerializer:
 
+	def __init__(self, strict=False):
+		self.strict = strict
+
 	def run(self, line):
 		start_index = line.find('{')
 		end_index = line.rfind('}')
 		if start_index == -1 or end_index == -1:
-			return None
+			return
+		if start_index != 0 and self.strict:
+			return
 		try:
 			json_obj = line[start_index:end_index+1]
-			return yaml.safe_load(json_obj)
-		except yaml.YAMLError:
-			return None
+			yield json.loads(json_obj)
+		except json.decoder.JSONDecodeError:
+			return

secator/serializers/regex.py

@@ -3,15 +3,23 @@ import re
 
 class RegexSerializer:
 
-	def __init__(self, regex, fields=[]):
+	def __init__(self, regex, fields=[], findall=False):
 		self.regex = re.compile(regex)
 		self.fields = fields
+		self.findall = findall
 
 	def run(self, line):
-		match = self.regex.match(line)
+		if self.findall:
+			match = self.regex.findall(line)
+			yield from match
+			return
 		output = {}
+		match = self.regex.match(line)
 		if not match:
-			return None
+			return
+		if not self.fields:
+			yield match.group(0)
+			return
 		for field in self.fields:
 			output[field] = match.group(field)
-		return output
+		yield output