secator 0.10.0__tar.gz → 0.10.1a1__tar.gz

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

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

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

{secator-0.10.0 → secator-0.10.1a1}/secator/celery.py

@@ -7,7 +7,7 @@ import uuid
 
 from time import time
 
-from celery import Celery, chain, chord
+from celery import Celery, chord
 from celery.app import trace
 
 from rich.logging import RichHandler
@@ -15,10 +15,9 @@ from retry import retry
 
 from secator.celery_signals import setup_handlers
 from secator.config import CONFIG
-from secator.output_types import Info, Error
+from secator.output_types import Info
 from secator.rich import console
 from secator.runners import Scan, Task, Workflow
-from secator.runners._helpers import run_extractors
 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)
@@ -103,7 +102,7 @@ if IN_CELERY_WORKER_PROCESS:
 @retry(Exception, tries=3, delay=2)
 def update_state(celery_task, task, force=False):
 	"""Update task state to add metadata information."""
-	if task.sync:
+	if not IN_CELERY_WORKER_PROCESS:
 		return
 	if not force and not should_update(CONFIG.runners.backend_update_frequency, task.last_updated_celery):
 		return
@@ -139,6 +138,16 @@ 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
@@ -167,7 +176,8 @@ def break_task(task, task_opts, targets, results=[], chunk_size=1):
 		task_id = str(uuid.uuid4())
 		opts['has_parent'] = True
 		opts['enable_duplicate_check'] = False
-		sig = type(task).s(chunk, **opts).set(queue=type(task).profile, task_id=task_id)
+		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()))
@@ -175,23 +185,16 @@ def break_task(task, task_opts, targets, results=[], chunk_size=1):
 		sigs.append(sig)
 
 	# Build Celery workflow
-	workflow = chain(
-		forward_results.s(results).set(queue='results'),
-		chord(
-			tuple(sigs),
-			forward_results.s().set(queue='results'),
-		)
+	workflow = chord(
+		tuple(sigs),
+		handle_runner_error.s(runner=task).set(queue='results')
 	)
-	if task.sync:
-		task.print_item = False
-		task.results = workflow.apply().get()
-	else:
-		result = workflow.apply_async()
-		task.celery_result = result
+	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}'))
 	kwargs['context']['celery_id'] = self.request.id
 	task = Task(*args, **kwargs)
@@ -218,111 +221,41 @@ def run_scan(self, args=[], kwargs={}):
 
 @app.task(bind=True)
 def run_command(self, results, name, targets, opts={}):
-	chunk = opts.get('chunk')
-	sync = opts.get('sync', True)
+	if IN_CELERY_WORKER_PROCESS:
+		opts.update({'print_item': True, 'print_line': True, 'print_cmd': True})
+		routing_key = self.request.delivery_info['routing_key']
+		console.print(Info(message=f'Task "{name}" running with routing key "{routing_key}"'))
+
+	# Flatten + dedupe + filter results
+	results = forward_results(results)
 
 	# 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')
 	opts['context'] = context
-	opts['print_remote_info'] = False
 	opts['results'] = results
+	opts['sync'] = True
 
-	# If we are in a Celery worker, print everything, always
-	if IN_CELERY_WORKER_PROCESS:
-		opts.update({
-			'print_item': True,
-			'print_line': True,
-			'print_cmd': True
-		})
-		routing_key = self.request.delivery_info['routing_key']
-		console.print(Info(message=f'Task "{name}" running with routing key "{routing_key}"'))
+	# Initialize task
+	sync = not IN_CELERY_WORKER_PROCESS
+	task_cls = Task.get_task_class(name)
+	task = task_cls(targets, **opts)
+	update_state(self, task, force=True)
 
-	# Flatten + dedupe results
-	results = flatten(results)
-	results = deduplicate(results, attr='_uuid')
+	# 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))
+
+	# Update state live
+	[update_state(self, task) for _ in task]
+	update_state(self, task, force=True)
+
+	# Garbage collection to save RAM
+	gc.collect()
 
-	# Get expanded targets
-	if not chunk and results:
-		targets, opts = run_extractors(results, opts, targets)
-		debug('after extractors', obj={'targets': targets, 'opts': opts}, sub='celery.state')
-
-	task = None
-
-	try:
-
-		# Get task class
-		task_cls = Task.get_task_class(name)
-
-		# Check if chunkable
-		many_targets = len(targets) > 1
-		targets_over_chunk_size = task_cls.input_chunk_size and len(targets) > task_cls.input_chunk_size
-		has_file_flag = task_cls.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)
-		task_opts = opts.copy()
-		task_opts.update({
-			'print_remote_info': False,
-			'has_children': chunk_it,
-		})
-
-		if IN_CELERY_WORKER_PROCESS and chunk_it and routing_key != 'poll':
-			console.print(Info(message=f'Task {name} is chunkable but not running on "poll" queue, re-routing to "poll" queue'))
-			raise self.replace(run_command.si(results, name, targets, opts=opts).set(queue='poll', task_id=self.request.id))
-
-		if chunk_it:
-			task_opts['print_cmd'] = False
-
-		task = task_cls(targets, **task_opts)
-		debug(
-			'',
-			obj={
-				f'{task.unique_name}': 'CHUNK STATUS',
-				'chunk_it': chunk_it,
-				'sync': task.sync,
-				'many_targets': many_targets,
-				'targets_over_chunk_size': targets_over_chunk_size,
-			},
-			obj_after=False,
-			id=self.request.id,
-			sub='celery.state',
-			verbose=True
-		)
-
-		# Chunk task if needed
-		if chunk_it:
-			chunk_size = task_cls.input_chunk_size if has_file_flag else 1
-			break_task(
-				task,
-				opts,
-				targets,
-				results=results,
-				chunk_size=chunk_size)
-			console.print(Info(message=f'Task "{name}" starts polling for chunked results'))
-
-		# Update state before starting
-		update_state(self, task)
-
-		# Update state for each item found
-		for _ in task:
-			update_state(self, task)
-
-	except BaseException as e:
-		if not task:
-			raise e
-		error = Error.from_exception(e)
-		error._source = task.unique_name
-		error._uuid = str(uuid.uuid4())
-		task.add_result(error, print=True)
-		task.stop_celery_tasks()
-
-	finally:
-		if not task:
-			raise
-		update_state(self, task, force=True)
-		gc.collect()
-		debug('', obj={task.unique_name: task.status, 'results': task.results}, sub='celery.results', verbose=True)
-		return task.results
+	return task.results
 
 
 @app.task
@@ -335,6 +268,7 @@ def forward_results(results):
 		results = results['results']
 	results = flatten(results)
 	results = deduplicate(results, attr='_uuid')
+	console.print(Info(message=f'Forwarding {len(results)} results ...'))
 	return results
 
 #--------------#

secator-0.10.1a1/secator/celery_signals.py

@@ -0,0 +1,141 @@
+import os
+import signal
+import threading
+from pathlib import Path
+
+from celery import signals
+
+from secator.config import CONFIG
+from secator.output_types import Info
+from secator.rich import console
+
+IDLE_TIMEOUT = CONFIG.celery.worker_kill_after_idle_seconds
+
+# File-based state management system
+STATE_DIR = Path("/tmp/celery_state")
+STATE_DIR.mkdir(exist_ok=True, parents=True)
+
+
+def get_lock_file_path():
+    worker_name = os.environ.get("WORKER_NAME", f"unknown_{os.getpid()}")
+    return Path(f"/tmp/celery_worker_{worker_name}.lock")
+
+
+def set_task_running(task_id):
+    """Mark that a task is running in current worker"""
+    with open(get_lock_file_path(), "w") as f:
+        f.write(task_id)
+
+
+def clear_task_running():
+    """Clear the task running state"""
+    lock_file = get_lock_file_path()
+    if lock_file.exists():
+        lock_file.unlink()
+
+
+def is_task_running():
+    """Check if a task is currently running"""
+    return get_lock_file_path().exists()
+
+
+def kill_worker(parent=False):
+    """Kill current worker using its pid by sending a SIGTERM to Celery master process."""
+    worker_name = os.environ.get('WORKER_NAME', 'unknown')
+
+    # Check if a task is running via the lock file
+    if not is_task_running():
+        pid = os.getppid() if parent else os.getpid()
+        console.print(Info(message=f'Sending SIGTERM to worker {worker_name} with pid {pid}'))
+        os.kill(pid, signal.SIGTERM)
+    else:
+        console.print(Info(message=f'Cancelling worker shutdown of {worker_name} since a task is running'))
+
+
+def setup_idle_timer(timeout):
+    """Setup a timer to kill the worker after being idle"""
+    if timeout == -1:
+        return
+
+    console.print(Info(message=f'Starting inactivity timer for {timeout} seconds ...'))
+    timer = threading.Timer(timeout, kill_worker)
+    timer.daemon = True  # Make sure timer is killed when worker exits
+    timer.start()
+
+
+def maybe_override_logging():
+    def decorator(func):
+        if CONFIG.celery.override_default_logging:
+            return signals.setup_logging.connect(func)
+        else:
+            return func
+    return decorator
+
+
+@maybe_override_logging()
+def setup_logging(*args, **kwargs):
+    """Override celery's logging setup to prevent it from altering our settings.
+    github.com/celery/celery/issues/1867
+    """
+    pass
+
+
+def capture_worker_name(sender, instance, **kwargs):
+    os.environ["WORKER_NAME"] = '{0}'.format(sender)
+
+
+def worker_init_handler(**kwargs):
+    if IDLE_TIMEOUT != -1:
+        setup_idle_timer(IDLE_TIMEOUT)
+
+
+def task_prerun_handler(task_id, **kwargs):
+    # Mark that a task is running
+    set_task_running(task_id)
+
+
+def task_postrun_handler(**kwargs):
+    # Mark that no task is running
+    clear_task_running()
+
+    # Get sender name from kwargs
+    sender_name = kwargs['sender'].name
+
+    if CONFIG.celery.worker_kill_after_task and sender_name.startswith('secator.'):
+        worker_name = os.environ.get('WORKER_NAME', 'unknown')
+        console.print(Info(message=f'Shutdown worker {worker_name} since config celery.worker_kill_after_task is set.'))
+        kill_worker(parent=True)
+        return
+
+    # Set up a new idle timer
+    if IDLE_TIMEOUT != -1:
+        console.print(Info(message=f'Reset inactivity timer to {IDLE_TIMEOUT} seconds'))
+        setup_idle_timer(IDLE_TIMEOUT)
+
+
+def task_revoked_handler(request=None, **kwargs):
+    """Handle revoked tasks by clearing the task running state"""
+    console.print(Info(message='Task was revoked, clearing running state'))
+    clear_task_running()
+
+    # Set up a new idle timer
+    if IDLE_TIMEOUT != -1:
+        console.print(Info(message=f'Reset inactivity timer to {IDLE_TIMEOUT} seconds after task revocation'))
+        setup_idle_timer(IDLE_TIMEOUT)
+
+
+def worker_shutdown_handler(**kwargs):
+    """Cleanup lock files when worker shuts down"""
+    lock_file = get_lock_file_path()
+    if lock_file.exists():
+        lock_file.unlink()
+
+
+def setup_handlers():
+    signals.celeryd_after_setup.connect(capture_worker_name)
+    signals.setup_logging.connect(setup_logging)
+    signals.task_prerun.connect(task_prerun_handler)
+    signals.task_postrun.connect(task_postrun_handler)
+    signals.task_revoked.connect(task_revoked_handler)
+    signals.worker_ready.connect(worker_init_handler)
+    signals.worker_shutdown.connect(worker_shutdown_handler)

{secator-0.10.0 → secator-0.10.1a1}/secator/celery_utils.py

@@ -264,5 +264,5 @@ class CeleryData(object):
 				CeleryData.get_task_ids(result.parent, ids=ids)
 
 		except kombu.exceptions.DecodeError:
-			debug('kombu decode error', sub='celery.data.get_task_ids')
+			debug('kombu decode error', sub='celery.data')
 			return

{secator-0.10.0 → secator-0.10.1a1}/secator/cli.py

@@ -171,7 +171,8 @@ def worker(hostname, concurrency, reload, queue, pool, check, dev, stop, show):
 		patterns = "celery.py;tasks/*.py;runners/*.py;serializers/*.py;output_types/*.py;hooks/*.py;exporters/*.py"
 		cmd = f'watchmedo auto-restart --directory=./ --patterns="{patterns}" --recursive -- {cmd}'
 
-	Command.execute(cmd, name='secator_worker')
+	ret = Command.execute(cmd, name='secator_worker')
+	sys.exit(ret.return_code)
 
 
 #-------#

{secator-0.10.0 → secator-0.10.1a1}/secator/hooks/mongodb.py

@@ -20,11 +20,19 @@ MONGODB_MAX_POOL_SIZE = CONFIG.addons.mongodb.max_pool_size
 
 logger = logging.getLogger(__name__)
 
-client = pymongo.MongoClient(
-	escape_mongodb_url(MONGODB_URL),
-	maxPoolSize=MONGODB_MAX_POOL_SIZE,
-	serverSelectionTimeoutMS=MONGODB_CONNECT_TIMEOUT
-)
+_mongodb_client = None
+
+
+def get_mongodb_client():
+	"""Get or create MongoDB client"""
+	global _mongodb_client
+	if _mongodb_client is None:
+		_mongodb_client = pymongo.MongoClient(
+			escape_mongodb_url(MONGODB_URL),
+			maxPoolSize=MONGODB_MAX_POOL_SIZE,
+			serverSelectionTimeoutMS=MONGODB_CONNECT_TIMEOUT
+		)
+	return _mongodb_client
 
 
 def get_runner_dbg(runner):
@@ -39,6 +47,7 @@ def get_runner_dbg(runner):
 
 
 def update_runner(self):
+	client = get_mongodb_client()
 	db = client.main
 	type = self.config.type
 	collection = f'{type}s'
@@ -72,6 +81,7 @@ def update_finding(self, item):
 	if type(item) not in FINDING_TYPES:
 		return item
 	start_time = time.time()
+	client = get_mongodb_client()
 	db = client.main
 	update = item.toDict()
 	_type = item._type
@@ -97,15 +107,14 @@ def update_finding(self, item):
 
 
 def find_duplicates(self):
+	from secator.celery import IN_CELERY_WORKER_PROCESS
 	ws_id = self.toDict().get('context', {}).get('workspace_id')
 	if not ws_id:
 		return
-	if self.sync:
-		debug(f'running duplicate check on workspace {ws_id}', sub='hooks.mongodb')
+	if not IN_CELERY_WORKER_PROCESS:
 		tag_duplicates(ws_id)
 	else:
-		celery_id = tag_duplicates.delay(ws_id)
-		debug(f'running duplicate check on workspace {ws_id}', id=celery_id, sub='hooks.mongodb')
+		tag_duplicates.delay(ws_id)
 
 
 def load_finding(obj):
@@ -132,6 +141,8 @@ def tag_duplicates(ws_id: str = None):
 	Args:
 		ws_id (str): Workspace id.
 	"""
+	debug(f'running duplicate check on workspace {ws_id}', sub='hooks.mongodb')
+	client = get_mongodb_client()
 	db = client.main
 	workspace_query = list(
 		db.findings.find({'_context.workspace_id': str(ws_id), '_tagged': True}).sort('_timestamp', -1))
@@ -172,19 +183,19 @@ def tag_duplicates(ws_id: str = None):
 				'seen dupes': len(seen_dupes)
 			},
 			id=ws_id,
-			sub='hooks.mongodb.duplicates',
+			sub='hooks.mongodb',
 			verbose=True)
 		tmp_duplicates_ids = list(dict.fromkeys([i._uuid for i in tmp_duplicates]))
-		debug(f'duplicate ids: {tmp_duplicates_ids}', id=ws_id, sub='hooks.mongodb.duplicates', verbose=True)
+		debug(f'duplicate ids: {tmp_duplicates_ids}', id=ws_id, sub='hooks.mongodb', verbose=True)
 
 		# Update latest object as non-duplicate
 		if tmp_duplicates:
 			duplicates.extend([f for f in tmp_duplicates])
 			db.findings.update_one({'_id': ObjectId(item._uuid)}, {'$set': {'_related': tmp_duplicates_ids}})
-			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb.duplicates', verbose=True)
+			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb', verbose=True)
 			non_duplicates.append(item)
 		else:
-			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb.duplicates', verbose=True)
+			debug(f'adding {item._uuid} as non-duplicate', id=ws_id, sub='hooks.mongodb', verbose=True)
 			non_duplicates.append(item)
 
 	# debug(f'found {len(duplicates)} total duplicates')
@@ -208,7 +219,7 @@ def tag_duplicates(ws_id: str = None):
 			'duplicates': len(duplicates_ids),
 			'non-duplicates': len(non_duplicates_ids)
 		},
-		sub='hooks.mongodb.duplicates')
+		sub='hooks.mongodb')
 
 
 HOOKS = {
@@ -232,6 +243,6 @@ HOOKS = {
 		'on_item': [update_finding],
 		'on_duplicate': [update_finding],
 		'on_interval': [update_runner],
-		'on_end': [update_runner, find_duplicates]
+		'on_end': [update_runner]
 	}
 }

{secator-0.10.0 → secator-0.10.1a1}/secator/output_types/error.py

@@ -21,10 +21,13 @@ class Error(OutputType):
 	_sort_by = ('_timestamp',)
 
 	def from_exception(e, **kwargs):
-		message = type(e).__name__
+		errtype = type(e).__name__
+		message = errtype
 		if str(e):
 			message += f': {str(e)}'
-		return Error(message=message, traceback=traceback_as_string(e), **kwargs)
+		traceback = traceback_as_string(e) if errtype not in ['KeyboardInterrupt', 'GreenletExit'] else ''
+		error = Error(message=message, traceback=traceback, **kwargs)
+		return error
 
 	def __str__(self):
 		return self.message