secator 0.10.1a1__tar.gz → 0.10.1a3__tar.gz

{secator-0.10.1a1 → secator-0.10.1a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: secator
-Version: 0.10.1a1
+Version: 0.10.1a3
 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.1a1 → secator-0.10.1a3}/pyproject.toml

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

{secator-0.10.1a1 → secator-0.10.1a3}/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 #
@@ -148,14 +146,14 @@ def handle_runner_error(self, results, runner):
 	return runner.results
 
 
-def break_task(task, task_opts, targets, results=[], chunk_size=1):
+def break_task(task, task_opts, results=[]):
 	"""Break a task into multiple of the same type."""
-	chunks = targets
-	if chunk_size > 1:
-		chunks = list(chunker(targets, chunk_size))
+	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': chunk_size, 'chunks': len(chunks), 'target_count': len(targets)},
+		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
@@ -196,6 +194,8 @@ def break_task(task, task_opts, targets, results=[], chunk_size=1):
 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 +204,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()
@@ -244,9 +246,11 @@ def run_command(self, results, name, targets, opts={}):
 	update_state(self, task, force=True)
 
 	# Chunk task if needed
-	if task_cls.needs_chunking(targets, sync):
+	if task.needs_chunking(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))
+		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,7 +272,8 @@ 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
 
 #--------------#

{secator-0.10.1a1 → secator-0.10.1a3}/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.1a1 → secator-0.10.1a3}/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
 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)
 

{secator-0.10.1a1 → secator-0.10.1a3}/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.1a1 → secator-0.10.1a3}/secator/runners/_base.py

@@ -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'
@@ -326,16 +328,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)
 
@@ -640,6 +644,7 @@ class Runner:
 
 	def log_start(self):
 		"""Log runner start."""
+		self.started = True
 		if not self.print_remote_info:
 			return
 		remote_str = 'starting' if self.sync else 'sent to Celery worker'
@@ -831,14 +836,14 @@ 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

{secator-0.10.1a1 → secator-0.10.1a3}/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.1a1 → secator-0.10.1a3}/secator/runners/scan.py

@@ -3,8 +3,8 @@ 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.output_types import Error
 from secator.utils import merge_opts
 
 logger = logging.getLogger(__name__)
@@ -28,16 +28,16 @@ class Scan(Runner):
 		scan_opts = self.config.options
 		self.print_item = False
 		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)
 			workflow = Workflow(
 				TemplateLoader(name=f'workflows/{name}'),
-				targets,
-				results=[],
+				self.inputs,
+				results=self.results,
 				run_opts=opts,
 				hooks=self._hooks,
 				context=self.context.copy())
-
-			# Get results
 			yield from workflow
+			if len(self.errors) > 0:
+				self.add_result(Error(message=f'Stopping scan since workflow {name} has errors'))
+				return

{secator-0.10.1a1 → secator-0.10.1a3}/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)