secator 0.10.1a2__tar.gz → 0.10.1a4__tar.gz

{secator-0.10.1a2 → secator-0.10.1a4}/PKG-INFO

@@ -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 → secator-0.10.1a4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = 'hatchling.build'
 
 [project]
 name = 'secator'
-version = "0.10.1a2"
+version = "0.10.1a4"
 authors = [{ name = 'FreeLabz', email = 'sales@freelabz.com' }]
 readme = 'README.md'
 description = "The pentester's swiss knife."

{secator-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a4/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-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a2 → secator-0.10.1a4}/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-0.10.1a4/secator/runners/scan.py

@@ -0,0 +1,55 @@
+import logging
+
+from secator.config import CONFIG
+from secator.runners._base import Runner
+from secator.runners.workflow import Workflow
+from secator.utils import merge_opts
+
+logger = logging.getLogger(__name__)
+
+
+class Scan(Runner):
+
+	default_exporters = CONFIG.scans.exporters
+
+	@classmethod
+	def delay(cls, *args, **kwargs):
+		from secator.celery import run_scan
+		return run_scan.delay(args=args, kwargs=kwargs)
+
+	def build_celery_workflow(self):
+		"""Build Celery workflow for scan execution.
+
+		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
+
+		# Build chain of workflows
+		sigs = []
+		for name, workflow_opts in self.config.workflows.items():
+			run_opts = self.run_opts.copy()
+			opts = merge_opts(scan_opts, workflow_opts, run_opts)
+			config = TemplateLoader(name=f'workflows/{name}')
+			workflow = Workflow(
+				config,
+				self.inputs,
+				results=self.results,
+				run_opts=opts,
+				hooks=self._hooks,
+				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'),
+		)