secator 0.10.1a2__py3-none-any.whl → 0.10.1a4__py3-none-any.whl

secator/celery.py

@@ -2,7 +2,6 @@ import gc
 import json
 import logging
 import os
-import sys
 import uuid
 
 from time import time
@@ -13,14 +12,13 @@ from celery.app import trace
 from rich.logging import RichHandler
 from retry import retry
 
-from secator.celery_signals import setup_handlers
+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)
 
-IN_CELERY_WORKER_PROCESS = sys.argv and ('secator.celery.app' in sys.argv or 'worker' in sys.argv)
 
 #---------#
 # Logging #
@@ -138,64 +136,11 @@ def chunker(seq, size):
 	return (seq[pos:pos + size] for pos in range(0, len(seq), size))
 
 
-@app.task(bind=True)
-def handle_runner_error(self, results, runner):
-	"""Handle errors in Celery workflows (chunked tasks or runners)."""
-	results = forward_results(results)
-	runner.results = results
-	runner.log_results()
-	runner.run_hooks('on_end')
-	return runner.results
-
-
-def break_task(task, task_opts, targets, results=[], chunk_size=1):
-	"""Break a task into multiple of the same type."""
-	chunks = targets
-	if chunk_size > 1:
-		chunks = list(chunker(targets, chunk_size))
-	debug(
-		'',
-		obj={task.unique_name: 'CHUNKED', 'chunk_size': chunk_size, 'chunks': len(chunks), 'target_count': len(targets)},
-		obj_after=False,
-		sub='celery.state',
-		verbose=True
-	)
-
-	# Clone opts
-	opts = task_opts.copy()
-
-	# Build signatures
-	sigs = []
-	task.ids_map = {}
-	for ix, chunk in enumerate(chunks):
-		if not isinstance(chunk, list):
-			chunk = [chunk]
-		if len(chunks) > 0:  # add chunk to task opts for tracking chunks exec
-			opts['chunk'] = ix + 1
-			opts['chunk_count'] = len(chunks)
-		task_id = str(uuid.uuid4())
-		opts['has_parent'] = True
-		opts['enable_duplicate_check'] = False
-		opts['results'] = results
-		sig = type(task).si(chunk, **opts).set(queue=type(task).profile, task_id=task_id)
-		full_name = f'{task.name}_{ix + 1}'
-		task.add_subtask(task_id, task.name, f'{task.name}_{ix + 1}')
-		info = Info(message=f'Celery chunked task created: {task_id}', _source=full_name, _uuid=str(uuid.uuid4()))
-		task.add_result(info)
-		sigs.append(sig)
-
-	# Build Celery workflow
-	workflow = chord(
-		tuple(sigs),
-		handle_runner_error.s(runner=task).set(queue='results')
-	)
-	return workflow
-
-
 @app.task(bind=True)
 def run_task(self, args=[], kwargs={}):
-	print('run task')
 	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()
@@ -204,6 +149,8 @@ def run_task(self, args=[], kwargs={}):
 @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()
@@ -241,12 +188,17 @@ def run_command(self, results, name, targets, opts={}):
 	sync = not IN_CELERY_WORKER_PROCESS
 	task_cls = Task.get_task_class(name)
 	task = task_cls(targets, **opts)
+	task.started = True
+	task.run_hooks('on_start')
 	update_state(self, task, force=True)
 
 	# Chunk task if needed
-	if task_cls.needs_chunking(targets, sync):
-		console.print(Info(message=f'Task {name} requires chunking, breaking into {len(targets)} tasks'))
-		return self.replace(break_task(task, opts, targets, results=results))
+	if task.needs_chunking(sync):
+		if IN_CELERY_WORKER_PROCESS:
+			console.print(Info(message=f'Task {name} requires chunking, breaking into {len(targets)} tasks'))
+		tasks = break_task(task, opts, results=results)
+		update_state(self, task, force=True)
+		return self.replace(tasks)
 
 	# Update state live
 	[update_state(self, task) for _ in task]
@@ -268,9 +220,56 @@ def forward_results(results):
 		results = results['results']
 	results = flatten(results)
 	results = deduplicate(results, attr='_uuid')
-	console.print(Info(message=f'Forwarding {len(results)} results ...'))
+	if IN_CELERY_WORKER_PROCESS:
+		console.print(Info(message=f'Forwarding {len(results)} results ...'))
 	return results
 
+
+@app.task
+def mark_runner_started(runner):
+	"""Mark a runner as started and run on_start hooks.
+
+	Args:
+		runner (Runner): Secator runner instance
+
+	Returns:
+		list: Runner results
+	"""
+	runner.started = True
+	# runner.start_time = time()
+	runner.run_hooks('on_start')
+	return runner.results
+
+
+@app.task
+def mark_runner_complete(results, runner):
+	"""Mark a runner as completed and run on_end hooks.
+
+	Args:
+		results (list): Task results
+		runner (Runner): Secator runner instance
+
+	Returns:
+		list: Final results
+	"""
+	results = forward_results(results)
+
+	# If sync mode, don't update the runner as it's already done
+	if runner.sync:
+		return results
+
+	# Run final processing
+	runner.results = results
+	runner.done = True
+	runner.progress = 100
+	if not runner.no_process:
+		runner.mark_duplicates()
+		runner.results = runner.filter_results()
+	runner.log_results()
+	runner.run_hooks('on_end')
+	return runner.results
+
+
 #--------------#
 # Celery utils #
 #--------------#
@@ -285,3 +284,50 @@ def is_celery_worker_alive():
 	else:
 		console.print(Info(message='No Celery worker available, running locally'))
 	return result
+
+
+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
+	opts = task_opts.copy()
+
+	# Build signatures
+	sigs = []
+	task.ids_map = {}
+	for ix, chunk in enumerate(chunks):
+		if not isinstance(chunk, list):
+			chunk = [chunk]
+		if len(chunks) > 0:  # add chunk to task opts for tracking chunks exec
+			opts['chunk'] = ix + 1
+			opts['chunk_count'] = len(chunks)
+		task_id = str(uuid.uuid4())
+		opts['has_parent'] = True
+		opts['enable_duplicate_check'] = False
+		opts['results'] = results
+		sig = type(task).si(chunk, **opts).set(queue=type(task).profile, task_id=task_id)
+		full_name = f'{task.name}_{ix + 1}'
+		task.add_subtask(task_id, task.name, f'{task.name}_{ix + 1}')
+		info = Info(message=f'Celery chunked task created: {task_id}', _source=full_name, _uuid=str(uuid.uuid4()))
+		task.add_result(info)
+		sigs.append(sig)
+
+	# Mark main task as async since it's being chunked
+	task.sync = False
+
+	# Build Celery workflow
+	workflow = chord(
+		tuple(sigs),
+		mark_runner_complete.s(runner=task).set(queue='results')
+	)
+	return workflow

secator/celery_signals.py

@@ -1,5 +1,6 @@
 import os
 import signal
+import sys
 import threading
 from pathlib import Path
 
@@ -10,6 +11,7 @@ from secator.output_types import Info
 from secator.rich import console
 
 IDLE_TIMEOUT = CONFIG.celery.worker_kill_after_idle_seconds
+IN_CELERY_WORKER_PROCESS = sys.argv and ('secator.celery.app' in sys.argv or 'worker' in sys.argv)
 
 # File-based state management system
 STATE_DIR = Path("/tmp/celery_state")

secator/celery_utils.py

@@ -12,7 +12,7 @@ from rich.padding import Padding
 from rich.progress import Progress as RichProgress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from secator.config import CONFIG
 from secator.definitions import STATE_COLORS
-from secator.output_types import Error
+from secator.output_types import Error, Info, State
 from secator.rich import console
 from secator.utils import debug, traceback_as_string
 
@@ -76,10 +76,31 @@ class CeleryData(object):
 
 			# Get live results and print progress
 			for data in CeleryData.poll(result, ids_map, refresh_interval):
-				yield from data['results']
+				for result in data['results']:
+
+					# Add dynamic subtask to ids_map
+					if isinstance(result, Info):
+						message = result.message
+						if message.startswith('Celery chunked task created: '):
+							task_id = message.split(' ')[-1]
+							ids_map[task_id] = {
+								'id': task_id,
+								'name': result._source,
+								'full_name': result._source,
+								'descr': '',
+								'state': 'PENDING',
+								'count': 0,
+								'progress': 0
+							}
+					yield result
 
 				if print_remote_info:
 					task_id = data['id']
+					if task_id not in progress_cache:
+						if CONFIG.runners.show_subtasks:
+							progress_cache[task_id] = progress.add_task('', advance=0, **data)
+						else:
+							continue
 					progress_id = progress_cache[task_id]
 					CeleryData.update_progress(progress, progress_id, data)
 
@@ -117,9 +138,24 @@ class CeleryData(object):
 		"""
 		while True:
 			try:
+				main_task = State(
+					task_id=result.id,
+					state=result.state,
+					_source='celery'
+				)
+				debug(f"Main task state: {result.id} - {result.state}", sub='celery.poll', verbose=True)
+				yield {'id': result.id, 'results': [main_task]}
 				yield from CeleryData.get_all_data(result, ids_map)
+
 				if result.ready():
 					debug('result is ready', sub='celery.poll', id=result.id)
+					main_task = State(
+						task_id=result.id,
+						state=result.state,
+						_source='celery'
+					)
+					debug(f"Final main task state: {result.id} - {result.state}", sub='celery.poll', verbose=True)
+					yield {'id': result.id, 'results': [main_task]}
 					yield from CeleryData.get_all_data(result, ids_map)
 					break
 			except (KeyboardInterrupt, GreenletExit):

secator/config.py

@@ -93,6 +93,7 @@ class Runners(StrictModel):
 	skip_exploit_search: bool = False
 	skip_cve_low_confidence: bool = False
 	remove_duplicates: bool = False
+	show_chunk_progress: bool = False
 
 
 class Security(StrictModel):

secator/output_types/__init__.py

@@ -7,6 +7,7 @@ __all__ = [
 	'Progress',
 	'Record',
 	'Stat',
+	'State',
 	'Subdomain',
 	'Url',
 	'UserAccount',
@@ -29,9 +30,10 @@ from secator.output_types.info import Info
 from secator.output_types.warning import Warning
 from secator.output_types.error import Error
 from secator.output_types.stat import Stat
+from secator.output_types.state import State
 
 EXECUTION_TYPES = [
-	Target, Progress, Info, Warning, Error
+	Target, Progress, Info, Warning, Error, State
 ]
 STAT_TYPES = [
 	Stat

secator/output_types/state.py

@@ -0,0 +1,29 @@
+import time
+from dataclasses import dataclass, field
+
+from secator.output_types._base import OutputType
+from secator.utils import rich_to_ansi
+
+
+@dataclass
+class State(OutputType):
+    """Represents the state of a Celery task."""
+
+    task_id: str
+    state: str
+    _type: str = field(default='state', repr=True)
+    _source: str = field(default='', repr=True)
+    _timestamp: int = field(default_factory=lambda: time.time(), compare=False)
+    _uuid: str = field(default='', repr=True, compare=False)
+    _context: dict = field(default_factory=dict, repr=True, compare=False)
+    _tagged: bool = field(default=False, repr=True, compare=False)
+    _duplicate: bool = field(default=False, repr=True, compare=False)
+    _related: list = field(default_factory=list, compare=False)
+    _icon = '📊'
+    _color = 'bright_blue'
+
+    def __str__(self) -> str:
+        return f"Task {self.task_id} is {self.state}"
+
+    def __repr__(self) -> str:
+        return rich_to_ansi(f"{self._icon} [bold {self._color}]{self.state}[/] {self.task_id}")

secator/runners/_base.py

@@ -12,7 +12,7 @@ import humanize
 from secator.definitions import ADDONS_ENABLED
 from secator.celery_utils import CeleryData
 from secator.config import CONFIG
-from secator.output_types import FINDING_TYPES, OutputType, Progress, Info, Warning, Error, Target
+from secator.output_types import FINDING_TYPES, OutputType, Progress, Info, Warning, Error, Target, State
 from secator.report import Report
 from secator.rich import console, console_stdout
 from secator.runners._helpers import (get_task_folder_id, process_extractor, run_extractors)
@@ -97,6 +97,7 @@ class Runner:
 		self.threads = []
 		self.no_poll = self.run_opts.get('no_poll', False)
 		self.quiet = self.run_opts.get('quiet', False)
+		self.started = False
 
 		# Runner process options
 		self.no_process = self.run_opts.get('no_process', False)
@@ -117,10 +118,21 @@ class Runner:
 		self.raise_on_error = self.run_opts.get('raise_on_error', False)
 		self.print_opts = {k: v for k, v in self.__dict__.items() if k.startswith('print_') if v}
 
+		# Chunks
+		self.has_parent = self.run_opts.get('has_parent', False)
+		self.has_children = self.run_opts.get('has_children', False)
+		self.chunk = self.run_opts.get('chunk', None)
+		self.chunk_count = self.run_opts.get('chunk_count', None)
+		self.unique_name = self.name.replace('/', '_')
+		self.unique_name = f'{self.unique_name}_{self.chunk}' if self.chunk else self.unique_name
+
+		# Add prior results to runner results
+		[self.add_result(result, print=False, output=False) for result in results]
+
 		# Determine inputs
 		inputs = [inputs] if not isinstance(inputs, list) else inputs
-		if results:
-			inputs, run_opts, errors = run_extractors(results, run_opts, inputs)
+		if not self.chunk and self.results:
+			inputs, run_opts, errors = run_extractors(self.results, run_opts, inputs)
 			for error in errors:
 				self.add_result(error, print=True)
 		self.inputs = inputs
@@ -163,18 +175,6 @@ class Runner:
 		self.validators = {name: [] for name in VALIDATORS + getattr(self, 'validators', [])}
 		self.register_validators(validators)
 
-		# Chunks
-		self.has_parent = self.run_opts.get('has_parent', False)
-		self.has_children = self.run_opts.get('has_children', False)
-		self.chunk = self.run_opts.get('chunk', None)
-		self.chunk_count = self.run_opts.get('chunk_count', None)
-		self.unique_name = self.name.replace('/', '_')
-		self.unique_name = f'{self.unique_name}_{self.chunk}' if self.chunk else self.unique_name
-
-		# Process prior results
-		for result in results:
-			list(self._process_item(result, print=False, output=False))
-
 		# Input post-process
 		self.run_hooks('before_init')
 
@@ -238,6 +238,8 @@ class Runner:
 
 	@property
 	def status(self):
+		if not self.started:
+			return 'PENDING'
 		if not self.done:
 			return 'RUNNING'
 		return 'FAILURE' if len(self.self_errors) > 0 else 'SUCCESS'
@@ -283,11 +285,8 @@ class Runner:
 				self.run_hooks('on_end')
 				return
 
-			# Choose yielder
-			yielder = self.yielder_celery if self.celery_result else self.yielder
-
 			# Loop and process items
-			for item in yielder():
+			for item in self.yielder():
 				yield from self._process_item(item)
 				self.run_hooks('on_interval')
 
@@ -326,16 +325,18 @@ class Runner:
 				self.add_result(error, print=True)
 				yield error
 
-	def add_result(self, item, print=False):
+	def add_result(self, item, print=False, output=True):
 		"""Add item to runner results.
 
 		Args:
 			item (OutputType): Item.
 			print (bool): Whether to print it or not.
+			output (bool): Whether to add it to the output or not.
 		"""
 		self.uuids.append(item._uuid)
 		self.results.append(item)
-		self.output += repr(item) + '\n'
+		if output:
+			self.output += repr(item) + '\n'
 		if print:
 			self._print_item(item)
 
@@ -481,16 +482,52 @@ class Runner:
 					dupe = self.run_hooks('on_duplicate', dupe)
 
 	def yielder(self):
-		"""Yield results. Should be implemented by derived classes."""
-		raise NotImplementedError()
+		"""Base yielder implementation.
+
+		This should be overridden by derived classes if they need custom behavior.
+		Otherwise, they can implement build_celery_workflow() and get standard behavior.
 
-	def yielder_celery(self):
-		"""Yield results from Celery result."""
-		yield from CeleryData.iter_results(
-			self.celery_result,
-			ids_map=self.celery_ids_map,
-			print_remote_info=False
-		)
+		Yields:
+			secator.output_types.OutputType: Secator output type.
+		"""
+		# Build Celery workflow
+		workflow = self.build_celery_workflow()
+
+		# Run workflow and get results
+		if self.sync:
+			self.print_item = False
+			self.started = True
+			results = workflow.apply().get()
+			yield from results
+		else:
+			self.celery_result = workflow()
+			self.celery_ids.append(str(self.celery_result.id))
+			yield Info(
+				message=f'Celery task created: {self.celery_result.id}',
+				task_id=self.celery_result.id
+			)
+			if self.no_poll:
+				return
+			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'
+			)
+
+		# Yield results
+		yield from results
+
+	def build_celery_workflow(self):
+		"""Build Celery workflow.
+
+		This should be implemented by derived classes.
+
+		Returns:
+			celery.Signature: Celery task signature.
+		"""
+		raise NotImplementedError("Derived classes must implement build_celery_workflow()")
 
 	def toDict(self):
 		"""Dict representation of the runner."""
@@ -651,6 +688,7 @@ class Runner:
 		"""Log runner results."""
 		if self.no_poll:
 			return
+		self.started = True
 		self.done = True
 		self.progress = 100
 		self.end_time = datetime.fromtimestamp(time())
@@ -831,20 +869,32 @@ class Runner:
 		# Update item context
 		item._context.update(self.context)
 
-		# Return if already seen
-		if item._uuid in self.uuids:
-			return
-
 		# Add uuid to item
 		if not item._uuid:
 			item._uuid = str(uuid.uuid4())
 
+		# Return if already seen
+		if item._uuid in self.uuids:
+			return
+
 		# Add source to item
 		if not item._source:
 			item._source = self.unique_name
 
+		# Check for state updates
+		if isinstance(item, State) and self.celery_result and item.task_id == self.celery_result.id:
+			self.debug(f'Updating runner state from Celery: {item.state}', sub='state')
+			if item.state in ['FAILURE', 'SUCCESS', 'REVOKED']:
+				self.started = True
+				self.done = True
+			elif item.state in ['RUNNING']:
+				self.started = True
+			self.debug(f'Runner {self.unique_name} is {self.status} (started: {self.started}, done: {self.done})', sub='state')
+			self.last_updated_celery = item._timestamp
+			return
+
 		# If progress item, update runner progress
-		if isinstance(item, Progress) and item._source == self.unique_name:
+		elif isinstance(item, Progress) and item._source == self.unique_name:
 			self.progress = item.percent
 			if not should_update(CONFIG.runners.progress_update_frequency, self.last_updated_progress, item._timestamp):
 				return

secator/runners/command.py

@@ -193,11 +193,10 @@ class Command(Runner):
 		})
 		return res
 
-	@classmethod
-	def needs_chunking(cls, targets, sync):
-		many_targets = len(targets) > 1
-		targets_over_chunk_size = cls.input_chunk_size and len(targets) > cls.input_chunk_size
-		has_file_flag = cls.file_flag is not None
+	def needs_chunking(self, sync):
+		many_targets = len(self.inputs) > 1
+		targets_over_chunk_size = self.input_chunk_size and len(self.inputs) > self.input_chunk_size
+		has_file_flag = self.file_flag is not None
 		chunk_it = (sync and many_targets and not has_file_flag) or (not sync and many_targets and targets_over_chunk_size)
 		return chunk_it
 

secator/runners/scan.py

@@ -1,9 +1,7 @@
 import logging
 
-from secator.template import TemplateLoader
 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.utils import merge_opts
 
@@ -19,25 +17,39 @@ class Scan(Runner):
 		from secator.celery import run_scan
 		return run_scan.delay(args=args, kwargs=kwargs)
 
-	def yielder(self):
-		"""Run scan.
+	def build_celery_workflow(self):
+		"""Build Celery workflow for scan execution.
 
-		Yields:
-			secator.output_types.OutputType: Secator output type.
+		Returns:
+			celery.Signature: Celery task signature.
 		"""
+		from celery import chain
+		from secator.celery import mark_runner_started, mark_runner_complete
+		from secator.template import TemplateLoader
+
 		scan_opts = self.config.options
-		self.print_item = False
+
+		# Build chain of workflows
+		sigs = []
 		for name, workflow_opts in self.config.workflows.items():
-			# Run workflow
 			run_opts = self.run_opts.copy()
 			opts = merge_opts(scan_opts, workflow_opts, run_opts)
+			config = TemplateLoader(name=f'workflows/{name}')
 			workflow = Workflow(
-				TemplateLoader(name=f'workflows/{name}'),
+				config,
 				self.inputs,
-				results=[],
+				results=self.results,
 				run_opts=opts,
 				hooks=self._hooks,
-				context=self.context.copy())
-
-			# Get results
-			yield from workflow
+				context=self.context.copy()
+			)
+			celery_workflow = workflow.build_celery_workflow()
+			for task_id, task_info in workflow.celery_ids_map.items():
+				self.add_subtask(task_id, task_info['name'], task_info['descr'])
+			sigs.append(celery_workflow)
+
+		return chain(
+			mark_runner_started.si(self).set(queue='results'),
+			*sigs,
+			mark_runner_complete.s(self).set(queue='results'),
+		)

secator/runners/task.py

@@ -1,8 +1,8 @@
+import uuid
 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
+from celery import chain
 
 
 class Task(Runner):
@@ -14,52 +14,46 @@ class Task(Runner):
 		from secator.celery import run_task
 		return run_task.apply_async(kwargs={'args': args, 'kwargs': kwargs}, queue='celery')
 
-	def yielder(self):
-		"""Run task.
+	def build_celery_workflow(self):
+		"""Build Celery workflow for task execution.
 
-		Yields:
-			secator.output_types.OutputType: Secator output type.
+		Args:
+			run_opts (dict): Run options.
+			results (list): Prior results.
+
+		Returns:
+			celery.Signature: Celery task signature.
 		"""
+		from secator.celery import run_command, mark_runner_started, mark_runner_complete
+
 		# Get task class
 		task_cls = Task.get_task_class(self.config.name)
 
 		# Run opts
-		run_opts = self.run_opts.copy()
-		run_opts.pop('output', None)
-		run_opts.pop('no_poll', False)
+		opts = self.run_opts.copy()
+		opts.pop('output', None)
+		opts.pop('no_poll', False)
 
 		# Set task output types
 		self.output_types = task_cls.output_types
 		self.enable_duplicate_check = False
 
 		# Get hooks
-		hooks = {task_cls: self.hooks}
-		run_opts['hooks'] = hooks
-		run_opts['context'] = self.context
+		hooks = self._hooks.get(Task, {})
+		opts['hooks'] = hooks
+		opts['context'] = self.context
 
-		# Run task
-		if self.sync:
-			self.print_item = False
-			result = task_cls.si(self.inputs, **run_opts)
-			results = result.apply().get()
-		else:
-			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
-			)
-			if self.no_poll:
-				return
-			results = CeleryData.iter_results(
-				self.celery_result,
-				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')
+		# Create task signature
+		task_id = str(uuid.uuid4())
+		sig = run_command.s(self.config.name, self.inputs, opts).set(queue=task_cls.profile, task_id=task_id)
+		self.add_subtask(task_id, self.config.name, self.config.description or '')
 
-		# Yield task results
-		yield from results
+		# Build signature chain with lifecycle management
+		return chain(
+			mark_runner_started.si(self).set(queue='results'),
+			sig,
+			mark_runner_complete.s(self).set(queue='results'),
+		)
 
 	@staticmethod
 	def get_task_class(name):

secator/runners/workflow.py

@@ -4,8 +4,6 @@ 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):
@@ -17,65 +15,38 @@ class Workflow(Runner):
 		from secator.celery import run_workflow
 		return run_workflow.delay(args=args, kwargs=kwargs)
 
-	def yielder(self):
-		"""Run workflow.
+	@classmethod
+	def s(cls, *args, **kwargs):
+		from secator.celery import run_workflow
+		return run_workflow.s(args=args, kwargs=kwargs)
 
-		Yields:
-			secator.output_types.OutputType: Secator output type.
-		"""
-		# Task opts
-		run_opts = self.run_opts.copy()
-		run_opts['hooks'] = self._hooks.get(Task, {})
-		run_opts.pop('no_poll', False)
-
-		# Build Celery workflow
-		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
-			yield Info(
-				message=f'Celery task created: {self.celery_result.id}',
-				task_id=self.celery_result.id
-			)
-			if self.no_poll:
-				return
-			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
-
-	def build_celery_workflow(self, run_opts={}, results=[]):
-		""""Build Celery workflow.
+	def build_celery_workflow(self):
+		"""Build Celery workflow for workflow execution.
 
 		Returns:
-			tuple(celery.chain, List[str]): Celery task chain, Celery task ids.
+			celery.Signature: Celery task signature.
 		"""
 		from celery import chain
-		from secator.celery import forward_results
+		from secator.celery import mark_runner_started, mark_runner_complete
+
+		# Prepare run options
+		opts = self.run_opts.copy()
+		opts['hooks'] = self._hooks.get(Task, {})
+		opts.pop('no_poll', False)
+
+		# Build task signatures
 		sigs = self.get_tasks(
 			self.config.tasks.toDict(),
 			self.inputs,
 			self.config.options,
-			run_opts)
-		sigs = [forward_results.si(results).set(queue='results')] + sigs + [forward_results.s().set(queue='results')]
-		workflow = chain(*sigs)
-		return workflow
+			opts)
+
+		# Build workflow chain with lifecycle management
+		return chain(
+			mark_runner_started.si(self).set(queue='results'),
+			*sigs,
+			mark_runner_complete.s(self).set(queue='results'),
+		)
 
 	def get_tasks(self, config, inputs, workflow_opts, run_opts):
 		"""Get tasks recursively as Celery chains / chords.

secator/scans/__init__.py

@@ -2,9 +2,9 @@ from secator.cli import ALL_SCANS
 
 
 def generate_class(config):
-	from secator.runners import Workflow
+	from secator.runners import Scan
 
-	class workflow(Workflow):
+	class scan(Scan):
 		def __init__(self, inputs=[], **run_opts):
 			hooks = run_opts.pop('hooks', {})
 			results = run_opts.pop('results', [])
@@ -16,12 +16,12 @@ def generate_class(config):
 				run_opts=run_opts,
 				hooks=hooks,
 				context=context)
-	return workflow, config.name
+	return scan, config.name
 
 
 DYNAMIC_SCANS = {}
-for workflow in ALL_SCANS:
-	cls, name = generate_class(workflow)
+for scan in ALL_SCANS:
+	cls, name = generate_class(scan)
 	DYNAMIC_SCANS[name] = cls
 
 globals().update(DYNAMIC_SCANS)

{secator-0.10.1a2.dist-info → secator-0.10.1a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: secator
-Version: 0.10.1a2
+Version: 0.10.1a4
 Summary: The pentester's swiss knife.
 Project-URL: Homepage, https://github.com/freelabz/secator
 Project-URL: Issues, https://github.com/freelabz/secator/issues

{secator-0.10.1a2.dist-info → secator-0.10.1a4.dist-info}/RECORD

@@ -1,10 +1,10 @@
 secator/.gitignore,sha256=da8MUc3hdb6Mo0WjZu2upn5uZMbXcBGvhdhTQ1L89HI,3093
 secator/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-secator/celery.py,sha256=9KXKv4EamJYJrHt_Ppn7aIp1AiFaTn2V0J_tZBwtWK0,8802
-secator/celery_signals.py,sha256=HobT7hCbVKPEHvCNwxCvQxFVUyocU1kkrTXi67b1DDw,4346
-secator/celery_utils.py,sha256=UWqLZpUaOXcztC_GD6uEDLiP8bGmD3WiTQN-u3lialg,7712
+secator/celery.py,sha256=8BV-G_r0gYZ48oYejMpvY4wvPJwYcjZ3NHnpjzOkB94,9706
+secator/celery_signals.py,sha256=iumfx7tTeoavAbHijBtij0JzeIqElxQldNZtuZmFY_U,4456
+secator/celery_utils.py,sha256=bW1yzMCjfIiesU4SOVNVuy0I8HukJyh8KmNB4w0woJM,8857
 secator/cli.py,sha256=3_tTTusW12MCejFgtOeYjiedjrJpyQj_gsCK8FkTMJA,43922
-secator/config.py,sha256=xItKM29yvMqzNZZygSNZXZ2V9vJbTdRuLTfIoRfP3XE,19653
+secator/config.py,sha256=CdVBh6d4k13SpkQKyHQfMFHgkLypUH07kAKLmCJJO1w,19688
 secator/decorators.py,sha256=3kYadCz6haIZtnjkFHSRfenTdc6Yu7bHd-0IVjhD72w,13902
 secator/definitions.py,sha256=gFtLT9fjNtX_1qkiCjNfQyCvYq07IhScsQzX4o20_SE,3084
 secator/installer.py,sha256=Q5qmGbxGmuhysEA9YovTpy-YY2TxxFskhrzSX44c42E,17971
@@ -50,7 +50,7 @@ secator/exporters/txt.py,sha256=oMtr22di6cqyE_5yJoiWP-KElrI5QgvK1cOUrj7H7js,730
 secator/hooks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 secator/hooks/gcs.py,sha256=MIhntyWYz9BZdTXhWl5JznaczSq1_7fl3TVqPufuTSo,1490
 secator/hooks/mongodb.py,sha256=XKbm_SrcSbQ2koILWvhzSg4tqdvHXgX5aU5x46Edu1s,7716
-secator/output_types/__init__.py,sha256=LxCW0K1f2vdgUapc4pIEsUpBfC0TQVvqo7T57rGuZGk,1159
+secator/output_types/__init__.py,sha256=L3q9NXPaW0TGeidx5YH-6dWhOXD1GizztAcL2lqIA8Q,1221
 secator/output_types/_base.py,sha256=OgS6ICt66TzPsqo1JZwRIIwbng2HRX1i_u5qbUECgNk,2820
 secator/output_types/error.py,sha256=39gpEJfKM2EuyOhD9lSkjjna2QicMvnLdFav6kHmhlg,1529
 secator/output_types/exploit.py,sha256=-BKTqPBg94rVgjw8YSmcYuBCI2x-73WwMd9ITP9qr3Y,1750
@@ -60,6 +60,7 @@ secator/output_types/port.py,sha256=JdqXnEF8XuwaWFMT8Vghj7fKLwtsImuUdRfMmITgmWM,
 secator/output_types/progress.py,sha256=MIbmnrLHNodLL42UgiaqLHL0OG5-w6mtUrhn0ZhksjA,1343
 secator/output_types/record.py,sha256=HnsKxlIhkgswA_Yjz7BZ1vDjP53l6OJ0BCOtCSDwCSY,1250
 secator/output_types/stat.py,sha256=90oN2Ghc4k0B0FOdp6MOWiNgmXMmLHYknjunDeEKKRE,1129
+secator/output_types/state.py,sha256=-kQs_P-v_d_J8mgMRJA9Pa0SaOVHN__Fq_ateDc0tiA,1038
 secator/output_types/subdomain.py,sha256=ivJ_2kmrJ8hdB8wmvRJYlKV1BcE3Cds_vAI_5wL7ES4,1344
 secator/output_types/tag.py,sha256=_XEqWAvAvmi7nd2ldfEE71zQx97jTSph2iDHkeqGTyk,1470
 secator/output_types/target.py,sha256=lmPw2aFOGIOFG4XXo6vNVZBBAZlnApJjyDVepDY54TU,871
@@ -68,14 +69,14 @@ secator/output_types/user_account.py,sha256=rm10somxyu30JHjj629IkR15Nhahylud_fVO
 secator/output_types/vulnerability.py,sha256=nF7OT9zGez8sZvLrkhjBOORjVi8hCqfCYUFq3eZ_ywo,2870
 secator/output_types/warning.py,sha256=47GtmG083GqGPb_R5JDFmARJ9Mqrme58UxwJhgdGPuI,853
 secator/runners/__init__.py,sha256=EBbOk37vkBy9p8Hhrbi-2VtM_rTwQ3b-0ggTyiD22cE,290
-secator/runners/_base.py,sha256=T9gjOqe-UPDHe5ZdVRBtUtxTefRgDcq9JV08F6UV5ZU,29596
+secator/runners/_base.py,sha256=3LNWH_jEfaZbtgHYzfl6mFmUJ9IoC62vKmLlv-MmTAw,31293
 secator/runners/_helpers.py,sha256=QhJmdmFdu5XSx3LBFf4Q4Hy2EXS6bLGnJUq8G7C6f68,2410
 secator/runners/celery.py,sha256=bqvDTTdoHiGRCt0FRvlgFHQ_nsjKMP5P0PzGbwfCj_0,425
-secator/runners/command.py,sha256=PqCOHDKJXvG4weB8mXDTElGxc8i8pK2RoyTKUBpHASU,25480
-secator/runners/scan.py,sha256=Pab_o_liI5fhlv2OOwYNmonz5JFYYVqtQFf9eyAQpiE,1071
-secator/runners/task.py,sha256=f2AduWpIy8JHK-Qitl_2Kh0fia573_YHAyAlV6MsJ50,2068
-secator/runners/workflow.py,sha256=XEhBfL-f3vGH0HgEPnj62d8ITxjH_tPXiNSVkaonuwQ,3862
-secator/scans/__init__.py,sha256=nlNLiRl7Vu--c_iXClFFcagMd_b_OWKitq8tX1-1krQ,641
+secator/runners/command.py,sha256=9AvjZgSXctP8D-ffPCtlnXEiGqTeaD2wVGhiGNuROb0,25469
+secator/runners/scan.py,sha256=9FjDsFmQrAWfA6crWkCJaVqG3-t2HBVjcsv4UQp_9b8,1500
+secator/runners/task.py,sha256=59jPXKSxFtSNXsm6VTAz8li2jxpM0Bkcgcn77HIDCrY,1869
+secator/runners/workflow.py,sha256=qldnRm7r_SCvRHJFkZ7eaml62RZkOeCdT18PU357grY,2982
+secator/scans/__init__.py,sha256=1EEbngbDbvWxmeDYC6uux00WWy1v5qHtSpk6NVz27rM,617
 secator/serializers/__init__.py,sha256=OP5cmFl77ovgSCW_IDcZ21St2mUt5UK4QHfrsK2KvH8,248
 secator/serializers/dataclass.py,sha256=RqICpfsYWGjHAACAA2h2jZ_69CFHim4VZwcBqowGMcQ,1010
 secator/serializers/json.py,sha256=UJwAymRzjF-yBKOgz1MTOyBhQcdQg7fOKRXgmHIu8fo,411
@@ -109,8 +110,8 @@ secator/tasks/searchsploit.py,sha256=gvtLZbL2hzAZ07Cf0cSj2Qs0GvWK94XyHvoPFsetXu8
 secator/tasks/subfinder.py,sha256=C6W5NnXT92OUB1aSS9IYseqdI3wDMAz70TOEl8X-o3U,1213
 secator/tasks/wpscan.py,sha256=036ywiEqZfX_Bt071U7qIm7bi6pNk7vodflmuslJurA,5550
 secator/workflows/__init__.py,sha256=ivpZHiYYlj4JqlXLRmB9cmAPUGdk8QcUrCRL34hIqEA,665
-secator-0.10.1a2.dist-info/METADATA,sha256=YgZWsub4cdLtNiz-2UKkPKhLAi8mx0wFIPwGjnYU794,14726
-secator-0.10.1a2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-secator-0.10.1a2.dist-info/entry_points.txt,sha256=lPgsqqUXWgiuGSfKy-se5gHdQlAXIwS_A46NYq7Acic,44
-secator-0.10.1a2.dist-info/licenses/LICENSE,sha256=19W5Jsy4WTctNkqmZIqLRV1gTDOp01S3LDj9iSgWaJ0,2867
-secator-0.10.1a2.dist-info/RECORD,,
+secator-0.10.1a4.dist-info/METADATA,sha256=EWh5NvMc8ZTE9xYmtbbE7eikaHPHgQfTasXPWlsyQdA,14726
+secator-0.10.1a4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+secator-0.10.1a4.dist-info/entry_points.txt,sha256=lPgsqqUXWgiuGSfKy-se5gHdQlAXIwS_A46NYq7Acic,44
+secator-0.10.1a4.dist-info/licenses/LICENSE,sha256=19W5Jsy4WTctNkqmZIqLRV1gTDOp01S3LDj9iSgWaJ0,2867
+secator-0.10.1a4.dist-info/RECORD,,