secator 0.15.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

secator/tree.py

@@ -0,0 +1,196 @@
+from typing import List, Optional, Union
+from secator.template import TemplateLoader
+from dotmap import DotMap
+
+
+DEFAULT_RENDER_OPTS = {
+	'group': lambda x: f"[dim]group {x.name.split('/')[-1] if '/' in x.name else ''}[/]",
+	'task': lambda x: f"[bold gold3]:wrench: {x.name}[/]",
+	'workflow': lambda x: f"[bold dark_orange3]:gear: {x.name}[/]",
+	'scan': lambda x: f"[bold red]:magnifying_glass_tilted_left: {x.name}[/]",
+	'condition': lambda x: f"[dim cyan]# if {x}[/]" if x else ''
+}
+
+
+class TaskNode:
+	"""Represents a node in the workflow/scan task tree."""
+	def __init__(self, name: str, type_: str, id: str, opts: Optional[dict] = None, default_opts: Optional[dict] = None, condition: Optional[str] = None, description: Optional[str] = None, parent=None, ancestor=None):  # noqa: E501
+		self.name = name
+		self.type = type_
+		self.id = id
+		self.opts = opts or {}
+		self.default_opts = default_opts or {}
+		self.description = description
+		self.condition = condition
+		self.children: List[TaskNode] = []
+		self.parent = parent
+		self.ancestor = ancestor
+
+	def add_child(self, child: 'TaskNode') -> None:
+		"""Add a child node to this node."""
+		self.children.append(child)
+
+	def remove(self):
+		"""Remove this node from its parent."""
+		if self.parent:
+			self.parent.children.remove(self)
+
+	def __str__(self) -> str:
+		"""String representation with condition if present."""
+		if self.condition:
+			return f"{self.name} # if {self.condition}"
+		return self.name
+
+
+class RunnerTree:
+	"""Represents a tree of workflow/scan tasks."""
+	def __init__(self, name: str, type_: str, render_opts: Optional[dict] = DEFAULT_RENDER_OPTS):
+		self.name = name
+		self.type = type_
+		self.root_nodes: List[TaskNode] = []
+		self.render_opts = render_opts
+
+	def add_root_node(self, node: TaskNode) -> None:
+		"""Add a root-level node to the tree."""
+		self.root_nodes.append(node)
+
+	def render_tree(self) -> str:
+		"""Render the tree as a console-friendly string."""
+		lines = []
+		for node in self.root_nodes:
+			node_str = self.render_opts.get(node.type, lambda x: str(x))(node)
+			condition_str = self.render_opts.get('condition', lambda x: str(x) if x else '')(node.condition)
+			if condition_str:
+				node_str = f"{node_str} {condition_str}"
+			lines.append(node_str)
+			self._render_children(node, "", lines)
+		return "\n".join(lines)
+
+	def _render_children(self, node: TaskNode, prefix: str, lines: List[str]) -> None:
+		"""Helper method to recursively render child nodes."""
+		children_count = len(node.children)
+		for i, child in enumerate(node.children):
+			is_last = i == children_count - 1
+			branch = "└─ " if is_last else "├─ "
+			child_str = self.render_opts.get(child.type, lambda x: str(x))
+			condition_str = self.render_opts.get('condition', lambda x: str(x) if x else '')(child.condition)
+			render_str = f"{prefix}{branch}{child_str(child)}"
+			if child.description:
+				render_str += f" - [dim]{child.description}[/]"
+			if condition_str:
+				render_str += f" {condition_str}"
+			lines.append(render_str)
+			if child.children:
+				new_prefix = prefix + ("   " if is_last else "│  ")
+				self._render_children(child, new_prefix, lines)
+
+	def get_subtree(self, node: TaskNode) -> 'RunnerTree':
+		"""Get the subtree of this node."""
+		subtree = RunnerTree(node.name, node.type)
+		for child in node.children:
+			subtree.add_root_node(child)
+		return subtree
+
+
+def build_runner_tree(config: DotMap, condition: Optional[str] = None, parent: Optional[TaskNode] = None, ancestor: Optional[TaskNode] = None) -> Union[RunnerTree, str]:  # noqa: E501
+	"""
+	Build a tree representation from a runner config.
+
+	Args:
+		config (DotMap): The runner config.
+
+	Returns:
+		A RunnerTree object or an error message string
+	"""
+	tree = RunnerTree(config.name, config.type)
+
+	if config.type == 'workflow':
+		root_node = TaskNode(config.name, 'workflow', config.name, opts=config.options, default_opts=config.default_options, condition=condition, parent=parent, ancestor=ancestor)  # noqa: E501
+		tree.add_root_node(root_node)
+
+		# Add tasks to the tree
+		for task_name, task_details in config.tasks.items():
+			id = f'{config.name}.{task_name}'
+			if task_name.startswith('_group'):
+				group_node = TaskNode(task_name, 'group', id, parent=root_node, ancestor=root_node)
+				root_node.add_child(group_node)
+				for subtask_name, subtask_details in task_details.items():
+					subtask_details = subtask_details or {}
+					id = f'{config.name}.{subtask_name}'
+					condition = subtask_details.get('if')
+					description = subtask_details.get('description')
+					subtask_node = TaskNode(subtask_name, 'task', id, opts=subtask_details, condition=condition, description=description, parent=group_node, ancestor=root_node)  # noqa: E501
+					group_node.add_child(subtask_node)
+			else:
+				condition = task_details.get('if') if task_details else None
+				description = task_details.get('description') if task_details else None
+				task_node = TaskNode(task_name, 'task', id, opts=task_details, condition=condition, description=description, parent=root_node, ancestor=root_node)  # noqa: E501
+				root_node.add_child(task_node)
+
+	elif config.type == 'scan':
+		id = f'{config.name}'
+		root_node = TaskNode(config.name, 'scan', id, opts=config.options, parent=parent)
+		tree.add_root_node(root_node)
+
+		# Add workflows to the tree
+		for workflow_name, workflow_details in config.workflows.items():
+			id = f'{config.name}.{workflow_name}'
+			condition = workflow_details.get('if') if isinstance(workflow_details, dict) else None
+			split_name = workflow_name.split('/')
+			wf_name = split_name[0]
+			wf_config = TemplateLoader(name=f'workflow/{wf_name}')
+			wf_config.name = workflow_name
+			wf_tree = build_runner_tree(wf_config, condition, parent=root_node, ancestor=root_node)
+			if isinstance(wf_tree, RunnerTree):
+				for wf_root_node in wf_tree.root_nodes:
+					root_node.add_child(wf_root_node)
+
+	elif config.type == 'task':
+		root_node = TaskNode(config.name, 'task', config.name, opts={}, parent=parent, ancestor=ancestor)
+		tree.add_root_node(root_node)
+
+	return tree
+
+
+def walk_runner_tree(tree: RunnerTree, visit_func):
+	"""
+	Walk the RunnerTree and visit each node.
+
+	Args:
+		tree (RunnerTree): The RunnerTree to walk.
+		visit_func (function): A function to call on each node.
+	"""
+	for root_node in tree.root_nodes:
+		_walk_node(root_node, visit_func)
+
+
+def _walk_node(node: TaskNode, visit_func):
+	"""
+	Recursively walk the node and its children.
+
+	Args:
+		node (TaskNode): The node to walk.
+		visit_func (function): A function to call on each node.
+	"""
+	visit_func(node)
+	for child in node.children:
+		_walk_node(child, visit_func)
+
+
+def get_flat_node_list(tree: RunnerTree) -> List[TaskNode]:
+	"""
+	Get the flat list of all nodes in the RunnerTree.
+
+	Args:
+		tree (RunnerTree): The RunnerTree to traverse.
+
+	Returns:
+		List[TaskNode]: The list of all nodes in the tree.
+	"""
+	nodes = []
+
+	def collect_node(node: TaskNode):
+		nodes.append(node)
+
+	walk_runner_tree(tree, collect_node)
+	return nodes

secator/utils.py

@@ -1,6 +1,6 @@
 import fnmatch
-import inspect
 import importlib
+import ipaddress
 import itertools
 import json
 import logging
@@ -15,9 +15,7 @@ import warnings
 
 from datetime import datetime, timedelta
 from functools import reduce
-from inspect import isclass
 from pathlib import Path
-from pkgutil import iter_modules
 from time import time
 import traceback
 from urllib.parse import urlparse, quote
@@ -26,7 +24,8 @@ import humanize
 import ifaddr
 import yaml
 
-from secator.definitions import (DEBUG_COMPONENT, VERSION, DEV_PACKAGE)
+from secator.definitions import (DEBUG, VERSION, DEV_PACKAGE, IP, HOST, CIDR_RANGE,
+								 MAC_ADDRESS, SLUG, UUID, EMAIL, IBAN, URL, PATH, HOST_PORT)
 from secator.config import CONFIG, ROOT_FOLDER, LIB_FOLDER, download_file
 from secator.rich import console
 
@@ -73,18 +72,21 @@ def expand_input(input, ctx):
 	Returns:
 		str: Input.
 	"""
+	piped_input = ctx.obj['piped_input']
+	dry_run = ctx.obj['dry_run']
 	if input is None:  # read from stdin
-		if not ctx.obj['piped_input']:
-			console.print('Waiting for input on stdin ...', style='bold yellow')
-		rlist, _, _ = select.select([sys.stdin], [], [], CONFIG.cli.stdin_timeout)
-		if rlist:
-			data = sys.stdin.read().splitlines()
-		else:
-			console.print(
-				'No input passed on stdin. Showing help page.',
-				style='bold red')
-			return None
-		return data
+		if not piped_input and not dry_run:
+			console.print('No input passed on stdin. Showing help page.', style='bold red')
+			ctx.get_help()
+			sys.exit(1)
+		elif piped_input:
+			rlist, _, _ = select.select([sys.stdin], [], [], CONFIG.cli.stdin_timeout)
+			if rlist:
+				data = sys.stdin.read().splitlines()
+				return data
+			else:
+				console.print('No input passed on stdin.', style='bold red')
+				sys.exit(1)
 	elif os.path.exists(input):
 		if os.path.isfile(input):
 			with open(input, 'r') as f:
@@ -99,6 +101,9 @@ def expand_input(input, ctx):
 	if isinstance(input, list) and len(input) == 1:
 		return input[0]
 
+	if ctx.obj['dry_run'] and not input:
+		return ['TARGET']
+
 	return input
 
 
@@ -140,75 +145,6 @@ def deduplicate(array, attr=None):
 	return sorted(list(dict.fromkeys(array)))
 
 
-def discover_internal_tasks():
-	"""Find internal secator tasks."""
-	from secator.runners import Runner
-	package_dir = Path(__file__).resolve().parent / 'tasks'
-	task_classes = []
-	for (_, module_name, _) in iter_modules([str(package_dir)]):
-		if module_name.startswith('_'):
-			continue
-		try:
-			module = importlib.import_module(f'secator.tasks.{module_name}')
-		except ImportError as e:
-			console.print(f'[bold red]Could not import secator.tasks.{module_name}:[/]')
-			console.print(f'\t[bold red]{type(e).__name__}[/]: {str(e)}')
-			continue
-		for attribute_name in dir(module):
-			attribute = getattr(module, attribute_name)
-			if isclass(attribute):
-				bases = inspect.getmro(attribute)
-				if Runner in bases and hasattr(attribute, '__task__'):
-					task_classes.append(attribute)
-
-	# Sort task_classes by category
-	task_classes = sorted(
-		task_classes,
-		# key=lambda x: (get_command_category(x), x.__name__)
-		key=lambda x: x.__name__)
-
-	return task_classes
-
-
-def discover_external_tasks():
-	"""Find external secator tasks."""
-	output = []
-	sys.dont_write_bytecode = True
-	for path in CONFIG.dirs.templates.glob('**/*.py'):
-		try:
-			task_name = path.stem
-			module_name = f'secator.tasks.{task_name}'
-
-			# console.print(f'Importing module {module_name} from {path}')
-			spec = importlib.util.spec_from_file_location(module_name, path)
-			module = importlib.util.module_from_spec(spec)
-			# console.print(f'Adding module "{module_name}" to sys path')
-			sys.modules[module_name] = module
-
-			# console.print(f'Executing module "{module}"')
-			spec.loader.exec_module(module)
-
-			# console.print(f'Checking that {module} contains task {task_name}')
-			if not hasattr(module, task_name):
-				console.print(f'[bold orange1]Could not load external task "{task_name}" from module {path.name}[/] ({path})')
-				continue
-			cls = getattr(module, task_name)
-			console.print(f'[bold green]Successfully loaded external task "{task_name}"[/] ({path})')
-			output.append(cls)
-		except Exception as e:
-			console.print(f'[bold red]Could not load external module {path.name}. Reason: {str(e)}.[/] ({path})')
-	sys.dont_write_bytecode = False
-	return output
-
-
-def discover_tasks():
-	"""Find all secator tasks (internal + external)."""
-	global _tasks
-	if not _tasks:
-		_tasks = discover_internal_tasks() + discover_external_tasks()
-	return _tasks
-
-
 def import_dynamic(path, name=None):
 	"""Import class or module dynamically from path.
 
@@ -238,22 +174,6 @@ def import_dynamic(path, name=None):
 		return None
 
 
-def get_command_cls(cls_name):
-	"""Get secator command by class name.
-
-	Args:
-		cls_name (str): Class name to load.
-
-	Returns:
-		cls: Class.
-	"""
-	tasks_classes = discover_tasks()
-	for task_cls in tasks_classes:
-		if task_cls.__name__ == cls_name:
-			return task_cls
-	return None
-
-
 def get_command_category(command):
 	"""Get the category of a command.
 
@@ -383,7 +303,7 @@ def rich_to_ansi(text):
 			tmp_console.print(text, end='', soft_wrap=True)
 		return capture.get()
 	except Exception:
-		console.print(f'[bold red]Could not convert rich text to ansi: {text}[/]', highlight=False, markup=False)
+		print(f'Could not convert rich text to ansi: {text}[/]', file=sys.stderr)
 		return text
 
 
@@ -397,11 +317,11 @@ def rich_escape(obj):
 		any: Initial object, or escaped Rich string.
 	"""
 	if isinstance(obj, str):
-		return obj.replace('[', r'\[').replace(']', r'\]')
+		return obj.replace('[', r'\[').replace(']', r'\]').replace(r'\[/', r'\[\/')
 	return obj
 
 
-def format_object(obj, obj_breaklines=False):
+def format_debug_object(obj, obj_breaklines=False):
 	"""Format the debug object for printing.
 
 	Args:
@@ -413,7 +333,7 @@ def format_object(obj, obj_breaklines=False):
 	"""
 	sep = '\n ' if obj_breaklines else ', '
 	if isinstance(obj, dict):
-		return sep.join(f'[dim cyan]{k}[/] [dim yellow]->[/] [dim green]{v}[/]' for k, v in obj.items() if v is not None)  # noqa: E501
+		return sep.join(f'[bold blue]{k}[/] [yellow]->[/] [blue]{v}[/]' for k, v in obj.items() if v is not None)  # noqa: E501
 	elif isinstance(obj, list):
 		return f'[dim green]{sep.join(obj)}[/]'
 	return ''
@@ -421,21 +341,25 @@ def format_object(obj, obj_breaklines=False):
 
 def debug(msg, sub='', id='', obj=None, lazy=None, obj_after=True, obj_breaklines=False, verbose=False):
 	"""Print debug log if DEBUG >= level."""
-	if not DEBUG_COMPONENT == ['all']:
-		if not DEBUG_COMPONENT or DEBUG_COMPONENT == [""]:
+	if not DEBUG == ['all'] and not DEBUG == ['1']:
+		if not DEBUG or DEBUG == [""]:
 			return
-
 		if sub:
-			if verbose and sub not in DEBUG_COMPONENT:
-				sub = f'debug.{sub}'
-			if not any(sub.startswith(s) for s in DEBUG_COMPONENT):
+			for s in DEBUG:
+				if '*' in s and re.match(s + '$', sub):
+					break
+				elif not verbose and sub.startswith(s):
+					break
+				elif verbose and sub == s:
+					break
+			else:
 				return
 
 	if lazy:
 		msg = lazy(msg)
 
 	formatted_msg = f'[yellow4]{sub:13s}[/] ' if sub else ''
-	obj_str = format_object(obj, obj_breaklines) if obj else ''
+	obj_str = format_debug_object(obj, obj_breaklines) if obj else ''
 
 	# Constructing the message string based on object position
 	if obj_str and not obj_after:
@@ -446,7 +370,12 @@ def debug(msg, sub='', id='', obj=None, lazy=None, obj_after=True, obj_breakline
 	if id:
 		formatted_msg += rf' [italic gray11]\[{id}][/]'
 
-	console.print(rf'[dim]\[[magenta4]DBG[/]] {formatted_msg}[/]')
+	try:
+		console.print(rf'[dim]\[[magenta4]DBG[/]] {formatted_msg}[/]', highlight=False)
+	except Exception:
+		console.print(rf'[dim]\[[magenta4]DBG[/]] <MARKUP_DISABLED>{rich_escape(formatted_msg)}</MARKUP_DISABLED>[/]', highlight=False)  # noqa: E501
+		if 'rich' in DEBUG:
+			raise
 
 
 def escape_mongodb_url(url):
@@ -498,6 +427,7 @@ def extract_domain_info(input, domain_only=False):
 
 	Args:
 		input (str): An URL or FQDN.
+		domain_only (bool): Return only the registered domain name.
 
 	Returns:
 		tldextract.ExtractResult: Extracted info.
@@ -507,9 +437,9 @@ def extract_domain_info(input, domain_only=False):
 	if not result or not result.domain or not result.suffix:
 		return None
 	if domain_only:
-		if not validators.domain(result.registered_domain):
+		if not validators.domain(result.top_domain_under_public_suffix):
 			return None
-		return result.registered_domain
+		return result.top_domain_under_public_suffix
 	return result
 
 
@@ -787,15 +717,12 @@ def process_wordlist(val):
 	if template_wordlist:
 		val = template_wordlist
 
-	if Path(val).exists():
-		return val
-	else:
-		return download_file(
-			val,
-			target_folder=CONFIG.dirs.wordlists,
-			offline_mode=CONFIG.offline_mode,
-			type='wordlist'
-		)
+	return download_file(
+		val,
+		target_folder=CONFIG.dirs.wordlists,
+		offline_mode=CONFIG.offline_mode,
+		type='wordlist'
+	)
 
 
 def convert_functions_to_strings(data):
@@ -815,3 +742,84 @@ def convert_functions_to_strings(data):
 		return json.dumps(data.__name__)  # or use inspect.getsource(data) if you want the actual function code
 	else:
 		return data
+
+
+def headers_to_dict(header_opt):
+	headers = {}
+	for header in header_opt.split(';;'):
+		split = header.strip().split(':')
+		key = split[0].strip()
+		val = ':'.join(split[1:]).strip()
+		headers[key] = val
+	return headers
+
+
+def format_object(obj, color='magenta', skip_keys=[]):
+	if isinstance(obj, list) and obj:
+		return ' [' + ', '.join([f'[{color}]{rich_escape(item)}[/]' for item in obj]) + ']'
+	elif isinstance(obj, dict) and obj.keys():
+		obj = {k: v for k, v in obj.items() if k.lower().replace('-', '_') not in skip_keys}
+		if obj:
+			return ' [' + ', '.join([f'[bold {color}]{rich_escape(k)}[/]: [{color}]{rich_escape(v)}[/]' for k, v in obj.items()]) + ']'  # noqa: E501
+	return ''
+
+
+def autodetect_type(target):
+	"""Autodetect the type of a target.
+
+	Args:
+		target (str): The target to autodetect the type of.
+
+	Returns:
+		str: The type of the target.
+	"""
+	if validators.url(target, simple_host=True):
+		return URL
+	elif validate_cidr_range(target):
+		return CIDR_RANGE
+	elif validators.ipv4(target) or validators.ipv6(target) or target == 'localhost':
+		return IP
+	elif validators.domain(target):
+		return HOST
+	elif validators.domain(target.split(':')[0]):
+		return HOST_PORT
+	elif validators.mac_address(target):
+		return MAC_ADDRESS
+	elif validators.email(target):
+		return EMAIL
+	elif validators.iban(target):
+		return IBAN
+	elif validators.uuid(target):
+		return UUID
+	elif Path(target).exists():
+		return PATH
+	elif validators.slug(target):
+		return SLUG
+
+	return str(type(target).__name__).lower()
+
+
+def validate_cidr_range(target):
+	if '/' not in target:
+		return False
+	try:
+		ipaddress.ip_network(target, False)
+		return True
+	except ValueError:
+		return False
+
+
+def get_versions_from_string(string):
+	"""Get versions from a string.
+
+	Args:
+		string (str): String to get versions from.
+
+	Returns:
+		list[str]: List of versions.
+	"""
+	regex = r'v?[0-9]+\.[0-9]+\.?[0-9]*\.?[a-zA-Z]*'
+	matches = re.findall(regex, string)
+	if not matches:
+		return []
+	return matches