secator 0.15.1__py3-none-any.whl → 0.16.1__py3-none-any.whl

secator/template.py

@@ -1,30 +1,28 @@
-import glob
+import yaml
 
 from collections import OrderedDict
-from pathlib import Path
-
-import yaml
 from dotmap import DotMap
+from pathlib import Path
 
-from secator.config import CONFIG, CONFIGS_FOLDER
-from secator.rich import console
-from secator.utils import convert_functions_to_strings, debug
 from secator.output_types import Error
-
-TEMPLATES = []
+from secator.rich import console
 
 
 class TemplateLoader(DotMap):
 
 	def __init__(self, input={}, name=None, **kwargs):
 		if name:
-			if '/' not in name:
+			split = name.split('/')
+			if len(split) != 2:
 				console.print(Error(message=f'Cannot load {name}: you should specify a type for the template when loading by name (e.g. workflow/<workflow_name>)'))  # noqa: E501
 				return
-			_type, name = name.split('/')
-			config = next((p for p in TEMPLATES if p['type'] == _type and p['name'] == name in str(p)), None)
+			_type, _name = tuple(split)
+			if _type.endswith('s'):
+				_type = _type[:-1]
+			from secator.loader import find_templates
+			config = next((p for p in find_templates() if p['type'] == _type and p['name'] == _name), None)
 			if not config:
-				console.print(Error(message=f'Template {name} not found in loaded templates'))
+				console.print(Error(message=f'Template {_type}/{_name} not found in loaded templates'))
 				config = {}
 		elif isinstance(input, dict):
 			config = input
@@ -35,9 +33,6 @@ class TemplateLoader(DotMap):
 			config = self._load(input)
 		super().__init__(config, **kwargs)
 
-	def add_to_templates(self):
-		TEMPLATES.append(self)
-
 	def _load_from_path(self, path):
 		if not path.exists():
 			console.print(Error(message=f'Config path {path} does not exists'))
@@ -48,16 +43,6 @@ class TemplateLoader(DotMap):
 	def _load(self, input):
 		return yaml.load(input, Loader=yaml.Loader)
 
-	@property
-	def supported_opts(self):
-		"""Property to access supported options easily."""
-		return self._collect_supported_opts()
-
-	@property
-	def flat_tasks(self):
-		"""Property to access tasks easily."""
-		return self._extract_tasks()
-
 	def print(self):
 		"""Print config as highlighted yaml."""
 		config = self.toDict()
@@ -69,81 +54,210 @@ class TemplateLoader(DotMap):
 		yaml_highlight = Syntax(yaml_str, 'yaml', line_numbers=True)
 		console.print(yaml_highlight)
 
-	def _collect_supported_opts(self):
-		"""Collect supported options from the tasks extracted from the config."""
-		tasks = self._extract_tasks()
-		opts = {}
-		for _, task_info in tasks.items():
-			task_class = task_info['class']
-			if task_class:
-				task_opts = task_class.get_supported_opts()
-				for name, conf in task_opts.items():
-					if name not in opts or not opts[name].get('supported', False):
-						opts[name] = convert_functions_to_strings(conf)
-		return opts
-
-	def _extract_tasks(self):
-		"""Extract tasks from any workflow or scan config.
-
-		Returns:
-			dict: A dict of task full name to task configuration containing the keyts keys ['name', 'class', 'opts']).
-		"""
-		from secator.runners import Task
-		tasks = OrderedDict()
-
-		def parse_config(config, prefix=''):
-			for key, value in config.items():
-				if key.startswith('_group'):
-					parse_config(value, prefix)
-				elif value:
-					task_name = f'{prefix}/{key}' if prefix else key
-					name = key.split('/')[0]
-					if task_name not in tasks:
-						tasks[task_name] = {'name': name, 'class': Task.get_task_class(name), 'opts': {}}
-					tasks[task_name]['opts'] = value.toDict()
-
-		if not self.type:
-			return tasks
-
-		elif self.type == 'task':
-			tasks[self.name] = {'name': self.name, 'class': Task.get_task_class(self.name)}
-
-		elif self.type == 'scan':
-			# For each workflow in the scan, load it and incorporate it with a unique prefix
-			for wf_name, _ in self.workflows.items():
-				name = wf_name.split('/')[0]
-				config = TemplateLoader(name=f'workflow/{name}')
-				wf_tasks = config.flat_tasks
-				# Prefix tasks from this workflow with its name to prevent collision
-				for task_key, task_val in wf_tasks.items():
-					unique_task_key = f"{wf_name}/{task_key}"  # Append workflow name to task key
-					tasks[unique_task_key] = task_val
-
-		elif self.type == 'workflow':
-			# Normal parsing of a workflow
-			parse_config(self.tasks)
-
-		return dict(tasks)
-
-
-def find_templates():
-	results = []
-	dirs = [CONFIGS_FOLDER]
-	if CONFIG.dirs.templates:
-		dirs.append(CONFIG.dirs.templates)
-	paths = []
-	for dir in dirs:
-		config_paths = [
-			Path(path)
-			for path in glob.glob(str(dir).rstrip('/') + '/**/*.y*ml', recursive=True)
-		]
-		debug(f'Found {len(config_paths)} templates in {dir}', sub='template')
-		paths.extend(config_paths)
-	for path in paths:
-		config = TemplateLoader(input=path)
-		debug(f'Loaded template from {path}', sub='template')
-		results.append(config)
-	return results
-
-
-TEMPLATES = find_templates()
+
+def get_short_id(id_str, config_name):
+	"""Remove config name prefix from ID string if present.
+
+	Args:
+		id_str: The ID string to process
+		config_name: The config name prefix to remove
+
+	Returns:
+		str: ID string with prefix removed, or original string if no prefix found
+	"""
+	if id_str.startswith(config_name):
+		return id_str.replace(config_name + '.', '')
+	return id_str
+
+
+def get_config_options(config, exec_opts=None, output_opts=None, type_mapping=None):
+	"""Extract and normalize command-line options from configuration.
+
+	Args:
+		config: Configuration object (task, workflow, or scan)
+		exec_opts: Execution options dictionary (optional)
+		output_opts: Output options dictionary (optional)
+		type_mapping: Type mapping for option types (optional)
+
+	Returns:
+		OrderedDict: Normalized options with metadata
+	"""
+	from secator.tree import build_runner_tree, walk_runner_tree, get_flat_node_list
+	from secator.utils import debug
+	from secator.runners.task import Task
+
+	# Task config created on-the-fly
+	if config.type == 'task':
+		config = TemplateLoader({
+			'name': config.name,
+			'type': 'workflow',
+			'tasks': {config.name: {}}
+		})
+
+	# Get main info
+	tree = build_runner_tree(config)
+	nodes = get_flat_node_list(tree)
+	exec_opts = exec_opts or {}
+	output_opts = output_opts or {}
+	type_mapping = type_mapping or {}
+	all_opts = OrderedDict({})
+
+	# Log current config and tree
+	debug(f'[magenta]{config.name}[/]', sub=f'cli.{config.name}')
+	debug(f'{tree.render_tree()}', sub=f'cli.{config.name}')
+
+	# Process global execution options
+	for opt in exec_opts:
+		opt_conf = exec_opts[opt].copy()
+		opt_conf['prefix'] = 'Execution'
+		all_opts[opt] = opt_conf
+
+	# Process global output options
+	for opt in output_opts:
+		opt_conf = output_opts[opt].copy()
+		opt_conf['prefix'] = 'Output'
+		all_opts[opt] = opt_conf
+
+	# Process config options
+	# a.k.a:
+	# - default YAML config options, defined in default_options: key in the runner YAML config
+	# - new options defined in options: key in the runner YAML config
+	config_opts_defaults = config.default_options.toDict()
+	config_opts = config.options.toDict()
+	for k, v in config_opts.items():
+		all_opts[k] = v
+		all_opts[k]['prefix'] = f'{config.type}'
+
+	def find_same_opts(node, nodes, opt_name, check_class_opts=False):
+		"""Find options with the same name that are defined in other nodes of the same type."""
+		same_opts = []
+		for _ in nodes:
+			if _.id == node.id or _.type != node.type:
+				continue
+			node_task = None
+			if check_class_opts:
+				node_task = Task.get_task_class(_.name)
+				if opt_name not in node_task.opts:
+					continue
+				opts_value = node_task.opts[opt_name]
+			else:
+				if opt_name not in _.opts:
+					continue
+				opts_value = _.opts[opt_name]
+			name_str = 'nodes' if not check_class_opts else 'tasks'
+			debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] -> [bold green]{opt_name}[/] found in other {name_str} [bold blue]{_.id}[/]', sub=f'cli.{config.name}.same', verbose=True)  # noqa: E501
+			same_opts.append({
+				'id': _.id,
+				'task_name': node_task.__name__ if node_task else None,
+				'name': _.name,
+				'value': opts_value,
+			})
+		if same_opts:
+			other_tasks = ", ".join([f'[bold yellow]{_["id"]}[/]' for _ in same_opts])
+			debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] -> [bold green]{opt_name}[/] found in {len(same_opts)} other {name_str}: {other_tasks}', sub=f'cli.{config.name}.same', verbose=True)  # noqa: E501
+		return same_opts
+
+	def process_node(node):
+		debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] ({node.type})', sub=f'cli.{config.name}')
+
+		if node.type not in ['task', 'workflow']:
+			return
+
+		# Process workflow options
+		# a.k.a the new options defined in options: key in the workflow YAML config;
+		if node.type == 'workflow':
+			for k, v in node.opts.items():
+				same_opts = find_same_opts(node, nodes, k)
+				conf = v.copy()
+				opt_name = k
+				conf['prefix'] = f'{node.type.capitalize()} {node.name}'
+				if len(same_opts) > 0:  # opt name conflict, change opt name
+					opt_name = f'{node.name}.{k}'
+					debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] -> [bold green]{k}[/] renamed to [bold green]{opt_name}[/] [dim red](duplicated)[/]', sub=f'cli.{config.name}')  # noqa: E501
+				all_opts[opt_name] = conf
+			return
+
+		# Process task options
+		# a.k.a task options defined in their respective task classes
+		cls = Task.get_task_class(node.name)
+		task_opts = cls.opts.copy()
+		task_opts_meta = cls.meta_opts.copy()
+		task_opts_all = {**task_opts, **task_opts_meta}
+		node_opts = node.opts or {}
+		ancestor_opts_defaults = node.ancestor.default_opts or {}
+		node_id_str = get_short_id(node.id, config.name)
+
+		for k, v in task_opts_all.items():
+			conf = v.copy()
+			conf['prefix'] = f'Task {node.name}'
+			default_from_config = node_opts.get(k) or ancestor_opts_defaults.get(k) or config_opts_defaults.get(k)
+			opt_name = k
+			same_opts = find_same_opts(node, nodes, k)
+
+			# Found a default in YAML config, either in task options, or workflow options, or config options
+			if default_from_config:
+				conf['required'] = False
+				conf['default'] = default_from_config
+				conf['default_from'] = node_id_str
+				if node_opts.get(k):
+					conf['default_from'] = node_id_str
+					conf['prefix'] = 'Config'
+				elif ancestor_opts_defaults.get(k):
+					conf['default_from'] = get_short_id(node.ancestor.id, config.name)
+					conf['prefix'] = f'{node.ancestor.type.capitalize()} {node.ancestor.name}'
+				elif config_opts_defaults.get(k):
+					conf['default_from'] = config.name
+					conf['prefix'] = 'Config'
+				mapped_value = cls.opt_value_map.get(opt_name)
+				if mapped_value:
+					if callable(mapped_value):
+						default_from_config = mapped_value(default_from_config)
+					else:
+						default_from_config = mapped_value
+				conf['default'] = default_from_config
+				if len(same_opts) > 0:  # change opt name to avoid conflict
+					conf['prefix'] = 'Config'
+					opt_name = f'{conf["default_from"]}.{k}'
+					debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] -> [bold green]{k}[/] renamed to [bold green]{opt_name}[/] [dim red](default set in config)[/]', sub=f'cli.{config.name}')  # noqa: E501
+
+			# Standard meta options like rate_limit, delay, proxy, etc...
+			elif k in task_opts_meta:
+				conf['prefix'] = 'Meta'
+				debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] -> [bold green]{k}[/] changed prefix to [bold cyan]Meta[/]', sub=f'cli.{config.name}')  # noqa: E501
+
+			# Task-specific options
+			elif k in task_opts:
+				same_opts = find_same_opts(node, nodes, k, check_class_opts=True)
+				if len(same_opts) > 0:
+					applies_to = set([node.name] + [_['name'] for _ in same_opts])
+					conf['applies_to'] = applies_to
+					conf['prefix'] = 'Shared task'
+					debug(f'[bold]{config.name}[/] -> [bold blue]{node.id}[/] -> [bold green]{k}[/] changed prefix to [bold cyan]Common[/] [dim red](duplicated {len(same_opts)} times)[/]', sub=f'cli.{config.name}')  # noqa: E501
+			else:
+				raise ValueError(f'Unknown option {k} for task {node.id}')
+			all_opts[opt_name] = conf
+
+	walk_runner_tree(tree, process_node)
+
+	# Normalize all options
+	debug('[bold yellow3]All opts processed. Showing defaults:[/]', sub=f'cli.{config.name}')
+	normalized_opts = OrderedDict({})
+	for k, v in all_opts.items():
+		v['reverse'] = False
+		v['show_default'] = True
+		default_from = v.get('default_from')
+		default = v.get('default', False)
+		if isinstance(default, bool) and default is True:
+			v['reverse'] = True
+		if type_mapping and 'type' in v:
+			v['type'] = type_mapping.get(v['type'], str)
+		short = v.get('short')
+		k = k.replace('.', '-').replace('_', '-').replace('/', '-')
+		from_str = default_from.replace('.', '-').replace('_', '-').replace('/', '-') if default_from else None
+		if not default_from or from_str not in k:
+			v['short'] = short if short else None
+		else:
+			v['short'] = f'{from_str}-{short}' if short else None
+		debug(f'\t[bold]{k}[/] -> [bold green]{v.get("default", "N/A")}[/] [dim red](default from {v.get("default_from", "N/A")})[/]', sub=f'cli.{config.name}')  # noqa: E501
+		normalized_opts[k] = v
+	return normalized_opts

secator/thread.py

@@ -4,21 +4,21 @@ from secator.output_types import Error
 
 
 class Thread(threading.Thread):
-    """A thread that returns errors in their join() method as secator.output_types.Error."""
+	"""A thread that returns errors in their join() method as secator.output_types.Error."""
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.error = None
+	def __init__(self, *args, **kwargs):
+		super().__init__(*args, **kwargs)
+		self.error = None
 
-    def run(self):
-        try:
-            if hasattr(self, '_target'):
-                self._target(*self._args, **self._kwargs)
-        except Exception as e:
-            self.error = Error.from_exception(e)
+	def run(self):
+		try:
+			if hasattr(self, '_target'):
+				self._target(*self._args, **self._kwargs)
+		except Exception as e:
+			self.error = Error.from_exception(e)
 
-    def join(self, *args, **kwargs):
-        super().join(*args, **kwargs)
-        if self.error:
-            return self.error
-        return None
+	def join(self, *args, **kwargs):
+		super().join(*args, **kwargs)
+		if self.error:
+			return self.error
+		return None

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