diff --git a/scripts/stories/STORY.md b/scripts/stories/STORY.md index 3d0101de..356d744d 100644 --- a/scripts/stories/STORY.md +++ b/scripts/stories/STORY.md @@ -60,9 +60,9 @@ secator z default mydomain.com --worker **Callbacks (library mode):** ```py from secator.runners import Workflow -from secator.config import ConfigLoader +from secator.template import TemplateLoader -config = ConfigLoader(name='workflows/host_recon') +config = TemplateLoader(name='workflows/host_recon') hooks = { Task: { 'on_item': { diff --git a/secator/__init__.py b/secator/__init__.py index 17a66e0f..e69de29b 100644 --- a/secator/__init__.py +++ b/secator/__init__.py @@ -1,537 +0,0 @@ -import os -from pathlib import Path -from subprocess import call, DEVNULL -from typing import Dict, List -from typing_extensions import Annotated, Self - -import requests -import yaml -from dotmap import DotMap -from pydantic import AfterValidator, BaseModel, model_validator, ValidationError - -from secator.rich import console, console_stdout - -Directory = Annotated[Path, AfterValidator(lambda v: v.expanduser())] -StrExpandHome = Annotated[str, AfterValidator(lambda v: v.replace('~', str(Path.home())))] - -ROOT_FOLDER = Path(__file__).parent.parent -LIB_FOLDER = ROOT_FOLDER / 'secator' -CONFIGS_FOLDER = LIB_FOLDER / 'configs' - - -class StrictModel(BaseModel, extra='forbid'): - pass - - -class Directories(StrictModel): - bin: Directory = Path.home() / '.local' / 'bin' - data: Directory = Path.home() / '.secator' - templates: Directory = '' - reports: Directory = '' - wordlists: Directory = '' - cves: Directory = '' - payloads: Directory = '' - revshells: Directory = '' - celery: Directory = '' - celery_data: Directory = '' - celery_results: Directory = '' - - @model_validator(mode='after') - def set_default_folders(self) -> Self: - """Set folders to be relative to the data folders if they are unspecified in config.""" - for folder in ['templates', 'reports', 'wordlists', 'cves', 'payloads', 'revshells', 'celery', 'celery_data', 'celery_results']: # noqa: E501 - rel_target = '/'.join(folder.split('_')) - val = getattr(self, folder) or self.data / rel_target - setattr(self, folder, val) - return self - - -class Debug(StrictModel): - level: int = 0 - component: str = '' - - -class Celery(StrictModel): - broker_url: str = 'filesystem://' - broker_pool_limit: int = 10 - broker_connection_timeout: float = 4.0 - broker_visibility_timeout: int = 3600 - override_default_logging: bool = True - result_backend: StrExpandHome = '' - - -class Cli(StrictModel): - github_token: str = '' - record: bool = False - stdin_timeout: int = 1000 - - -class Runners(StrictModel): - input_chunk_size: int = 1000 - progress_update_frequency: int = 60 - skip_cve_search: bool = False - - -class HTTP(StrictModel): - socks5_proxy: str = 'socks5://127.0.0.1:9050' - http_proxy: str = 'https://127.0.0.1:9080' - store_responses: bool = False - proxychains_command: str = 'proxychains' - freeproxy_timeout: int = 1 - - -class Tasks(StrictModel): - exporters: List[str] = ['json', 'csv'] - - -class Workflows(StrictModel): - exporters: List[str] = ['json', 'csv'] - - -class Scans(StrictModel): - exporters: List[str] = ['json', 'csv'] - - -class Payloads(StrictModel): - templates: Dict[str, str] = { - 'lse': 'https://github.com/diego-treitos/linux-smart-enumeration/releases/latest/download/lse.sh', - 'linpeas': 'https://github.com/carlospolop/PEASS-ng/releases/latest/download/linpeas.sh', - 'sudo_killer': 'https://github.com/TH3xACE/SUDO_KILLER/archive/refs/heads/V3.zip' - } - - -class Wordlists(StrictModel): - defaults: Dict[str, str] = {'http': 'bo0m_fuzz', 'dns': 'combined_subdomains'} - templates: Dict[str, str] = { - 'bo0m_fuzz': 'https://raw.githubusercontent.com/Bo0oM/fuzz.txt/master/fuzz.txt', - 'combined_subdomains': 'https://raw.githubusercontent.com/danielmiessler/SecLists/master/Discovery/DNS/combined_subdomains.txt' # noqa: E501 - } - lists: Dict[str, List[str]] = {} - - -class GoogleAddon(StrictModel): - enabled: bool = False - drive_parent_folder_id: str = '' - credentials_path: str = '' - - -class WorkerAddon(StrictModel): - enabled: bool = False - - -class MongodbAddon(StrictModel): - enabled: bool = False - url: str = 'mongodb://localhost' - update_frequency: int = 60 - - -class Addons(StrictModel): - google: GoogleAddon = GoogleAddon() - worker: WorkerAddon = WorkerAddon() - mongodb: MongodbAddon = MongodbAddon() - - -class SecatorConfig(StrictModel): - dirs: Directories = Directories() - debug: Debug = Debug() - celery: Celery = Celery() - cli: Cli = Cli() - runners: Runners = Runners() - http: HTTP = HTTP() - tasks: Tasks = Tasks() - workflows: Workflows = Workflows() - scans: Scans = Scans() - payloads: Payloads = Payloads() - wordlists: Wordlists = Wordlists() - addons: Addons = Addons() - offline_mode: bool = False - - -class Config(DotMap): - """Config class. - - Examples: - >>> config = Config.parse() # get default config. - >>> config = Config.parse({'dirs': {'data': '/tmp/'}) # get custom config (from dict). - >>> config = Config.parse(path='/path/to/config.yml') # get custom config (from YAML file). - >>> config.print() # print config without defaults. - >>> config.print(partial=False) # print full config. - >>> config.set('addons.google.enabled', False) # set value in config. - >>> config.save() # save config back to disk. - """ - - _error = False - - def get(self, key=None, print=True): - """Retrieve a value from the configuration using a dotted path. - - Args: - key (str | None): Dotted key path. - print (bool): Print the resulting value. - - Returns: - Any: value at key. - """ - value = self - if key: - for part in key.split('.'): - value = value[part] - if value is None: - console.print(f'[bold red]Key {key} does not exist.[/]') - return None - if print: - if key: - yaml_str = Config.dump(DotMap({key: value}), partial=False) - else: - yaml_str = Config.dump(self, partial=False) - Config.print_yaml(yaml_str) - return value - - def set(self, key, value, set_partial=True): - """Set a value in the configuration using a dotted path. - - Args: - key (str | None): Dotted key path. - value (Any): Value. - partial (bool): Also set value in partial config (written to disk). - - Returns: - bool: Success boolean. - """ - # Get existing value - existing_value = self.get(key, print=False) - - # Convert dotted key path to the corresponding uppercase key used in _keymap - map_key = key.upper().replace('.', '_') - success = False - if map_key in self._keymap: - # Traverse to the second last key to handle the setting correctly - target = self - partial = self._partial - for part in self._keymap[map_key][:-1]: - target = target[part] - if set_partial: - partial = partial[part] - - # Set the value on the final part of the path - final_key = self._keymap[map_key][-1] - - # Convert the value to the correct type based on the current value type - try: - if isinstance(existing_value, bool): - if isinstance(value, str): - value = value.lower() in ("true", "1", "t") - elif isinstance(value, (int, float)): - value = True if value == 1 else False - elif isinstance(existing_value, int): - value = int(value) - elif isinstance(existing_value, float): - value = float(value) - if existing_value != value: - target[final_key] = value - if set_partial: - partial[final_key] = value - success = True - except ValueError: - success = False - # console.print(f'[bold red]{key}: cannot cast value "{value}" to {type(existing_value).__name__}') - else: - console.print(f'[bold red]Key "{key}" not found in config keymap[/].') - return success - - def save(self, target_path: Path = None, partial=True): - """Save config as YAML on disk. - - Args: - target_path (Path | None): If passed, saves the config to this path. - partial (bool): Save partial config. - """ - if not target_path: - if not self._path: - return - target_path = self._path - with target_path.open('w') as f: - f.write(Config.dump(self, partial=partial)) - self._path = target_path - - def print(self, partial=True): - """Print config. - - Args: - partial (bool): Print partial config only. - """ - yaml_str = self.dump(self, partial=partial) - yaml_str = f'# {self._path}\n\n{yaml_str}' if self._path and partial else yaml_str - Config.print_yaml(yaml_str) - - @staticmethod - def parse(data: dict = {}, path: Path = None, env_overrides: bool = False): - """Parse config. - - Args: - data (dict): Config data. - path (Path | None): Path to YAML config. - env_overrides (bool): Apply env overrides. - - Returns: - Config: instance of Config object. - None: if the config was not loaded properly or there are validation errors. - """ - if path: - data = Config.read_yaml(path) - - # Load data - try: - config = Config.load(SecatorConfig, data) - config._valid = True - - # HACK: set default result_backend if unset - if not config.celery.result_backend: - config.celery.result_backend = f'file://{config.dirs.celery_results}' - - except ValidationError as e: - error_str = str(e).replace('\n', '\n ') - if path: - error_str.replace('SecatorConfig', f'SecatorConfig ({path})') - console.print(f'[bold red]:x: {error_str}') - # console.print('[bold green]Using default config.[/]') - config = Config.parse() - config._valid = False - - # Set hidden attributes - keymap = Config.build_key_map(config) - partial = Config(data) - config._partial = partial - config._path = path - config._keymap = keymap - - # Override config values with environment variables - if env_overrides: - config.apply_env_overrides() - data = {k: v for k, v in config.toDict().items() if not k.startswith('_')} - config = Config.parse(data, env_overrides=False) # re-validate config - config._partial = partial - config._path = path - - return config - - @staticmethod - def load(schema, data: dict = {}): - """Validate a config using Pydantic. - - Args: - data (dict): Config dict. - - Returns: - Config: instance of Config object. - """ - return Config(schema(**data).model_dump()) - - @staticmethod - def read_yaml(yaml_path): - """Read YAML from path. - - Args: - yaml_path (Path): path to yaml config. - - Returns: - dict: Loaded data. - """ - with yaml_path.open('r') as f: - data = yaml.load(f.read(), Loader=yaml.Loader) - return data or {} - - @staticmethod - def print_yaml(string): - """Print YAML string using rich. - - Args: - string (str): YAML string. - """ - from rich.syntax import Syntax - data = Syntax(string, 'yaml', theme='ansi-dark', padding=0, background_color='default') - console_stdout.print(data) - - @staticmethod - def dump(config, partial=True): - """Safe dump config as yaml: - - `Path`, `PosixPath` and `WindowsPath` objects are translated to strings. - - Home directory in paths is replaced with the tilde '~'. - - Returns: - str: YAML dump. - """ - import yaml - from pathlib import Path, PosixPath, WindowsPath - - # Get home dir - home = str(Path.home()) - - # Custom dumper to add line breaks between items and a path representer to translate paths to strings - class LineBreakDumper(yaml.SafeDumper): - def write_line_break(self, data=None): - super().write_line_break(data) - if len(self.indents) == 1: - super().write_line_break() - - def posix_path_representer(dumper, data): - path = str(data) - if path.startswith(home): - path = path.replace(home, '~') - return dumper.represent_scalar('tag:yaml.org,2002:str', path) - - LineBreakDumper.add_representer(Path, posix_path_representer) - LineBreakDumper.add_representer(PosixPath, posix_path_representer) - LineBreakDumper.add_representer(WindowsPath, posix_path_representer) - - # Get data dict - data = config.toDict() - - # HACK: Replace home dir in result_backend - if isinstance(config, Config): - data['celery']['result_backend'] = data['celery']['result_backend'].replace(home, '~') - del data['_path'] - if partial: - data = data['_partial'] - else: - del data['_partial'] - - data = {k: v for k, v in data.items() if not k.startswith('_')} - return yaml.dump(data, Dumper=LineBreakDumper, sort_keys=False) - - @staticmethod - def build_key_map(config, base_path=[]): - key_map = {} - for key, value in config.items(): - if key.startswith('_'): # ignore - continue - current_path = base_path + [key] - if isinstance(value, dict): - key_map.update(Config.build_key_map(value, current_path)) - else: - key_map['_'.join(current_path).upper()] = current_path - return key_map - - def apply_env_overrides(self): - """Override config values from environment variables.""" - # Build a map of keys from the config - key_map = Config.build_key_map(self) - - # Prefix for environment variables to target - prefix = "SECATOR_" - - # Loop through environment variables - for var in os.environ: - if var.startswith(prefix): - # Remove prefix and get the path from the key map - key = var[len(prefix):] - if key in key_map: - path = '.'.join(k.lower() for k in key_map[key]) - value = os.environ[var] - - # Set the new value recursively - success = self.set(path, value, set_partial=False) - if success: - console.print(f'[bold green4]{var} (override success)[/]') - else: - console.print(f'[bold red]{var} (override failed: cannot update value)[/]') - else: - console.print(f'[bold red]{var} (override failed: key not found in config)[/]') - - -def download_files(data: dict, target_folder: Path, offline_mode: bool, type: str): - """Download remote files to target folder, clone git repos, or symlink local files. - - Args: - data (dict): Dict of name to url or local path prefixed with 'git+' for Git repos. - target_folder (Path): Target folder for storing files or repos. - type (str): Type of files to handle. - offline_mode (bool): Offline mode. - """ - for name, url_or_path in data.items(): - if url_or_path.startswith('git+'): - # Clone Git repository - git_url = url_or_path[4:] # remove 'git+' prefix - repo_name = git_url.split('/')[-1] - if repo_name.endswith('.git'): - repo_name = repo_name[:-4] - target_path = target_folder / repo_name - if not target_path.exists(): - console.print(f'[bold turquoise4]Cloning git {type} [bold magenta]{repo_name}[/] ...[/] ', end='') - if offline_mode: - console.print('[bold orange1]skipped [dim][offline[/].[/]') - continue - try: - call(['git', 'clone', git_url, str(target_path)], stderr=DEVNULL, stdout=DEVNULL) - console.print('[bold green]ok.[/]') - except Exception as e: - console.print(f'[bold red]failed ({str(e)}).[/]') - data[name] = target_path.resolve() - elif Path(url_or_path).exists(): - # Create a symbolic link for a local file - local_path = Path(url_or_path) - target_path = target_folder / local_path.name - if not target_path.exists(): - console.print(f'[bold turquoise4]Symlinking {type} [bold magenta]{name}[/] ...[/] ', end='') - try: - target_path.symlink_to(local_path) - console.print('[bold green]ok.[/]') - except Exception as e: - console.print(f'[bold red]failed ({str(e)}).[/]') - data[name] = target_path.resolve() - else: - # Download file from URL - ext = url_or_path.split('.')[-1] - filename = f'{name}.{ext}' - target_path = target_folder / filename - if not target_path.exists(): - try: - console.print(f'[bold turquoise4]Downloading {type} [bold magenta]{filename}[/] ...[/] ', end='') - if offline_mode: - console.print('[bold orange1]skipped [dim](offline)[/].[/]') - continue - resp = requests.get(url_or_path, timeout=3) - resp.raise_for_status() - with open(target_path, 'wb') as f: - f.write(resp.content) - console.print('[bold green]ok.[/]') - except requests.RequestException as e: - console.print(f'[bold red]failed ({str(e)}).[/]') - continue - data[name] = target_path.resolve() - - -# Load configs -default_config = Config.parse() -data_root = default_config.dirs.data -config_path = data_root / 'config.yml' -if not config_path.exists(): - if not data_root.exists(): - console.print(f'[bold turquoise4]Creating directory [bold magenta]{data_root}[/] ... [/]', end='') - data_root.mkdir(parents=False) - console.print('[bold green]ok.[/]') - console.print( - f'[bold turquoise4]Creating user conf [bold magenta]{config_path}[/]... [/]', end='') - config_path.touch() - console.print('[bold green]ok.[/]') -CONFIG = Config.parse(path=config_path, env_overrides=True) - -# Create directories if they don't exist already -for name, dir in CONFIG.dirs.items(): - if not dir.exists(): - console.print(f'[bold turquoise4]Creating directory [bold magenta]{dir}[/] ... [/]', end='') - dir.mkdir(parents=False) - console.print('[bold green]ok.[/]') - -# Download wordlists and set defaults -download_files(CONFIG.wordlists.templates, CONFIG.dirs.wordlists, CONFIG.offline_mode, 'wordlist') -for category, name in CONFIG.wordlists.defaults.items(): - if name in CONFIG.wordlists.templates.keys(): - CONFIG.wordlists.defaults[category] = str(CONFIG.wordlists.templates[name]) - -# Download payloads -download_files(CONFIG.payloads.templates, CONFIG.dirs.payloads, CONFIG.offline_mode, 'payload') - -# Print config -if CONFIG.debug.component == 'config': - CONFIG.print() diff --git a/secator/celery.py b/secator/celery.py index 01c397c3..6f3f0e95 100644 --- a/secator/celery.py +++ b/secator/celery.py @@ -9,7 +9,7 @@ # from pyinstrument import Profiler # TODO: make pyinstrument optional from rich.logging import RichHandler -from secator import CONFIG +from secator.config import CONFIG from secator.rich import console from secator.runners import Scan, Task, Workflow from secator.runners._helpers import run_extractors diff --git a/secator/cli.py b/secator/cli.py index 1765cc12..9ca4338a 100644 --- a/secator/cli.py +++ b/secator/cli.py @@ -14,8 +14,8 @@ from rich.markdown import Markdown from rich.rule import Rule -from secator import CONFIG, ROOT_FOLDER, Config, default_config, config_path -from secator.config import ConfigLoader +from secator.config import CONFIG, ROOT_FOLDER, Config, default_config, config_path +from secator.template import TemplateLoader from secator.decorators import OrderedGroup, register_runner from secator.definitions import ADDONS_ENABLED, ASCII, DEV_PACKAGE, OPT_NOT_SUPPORTED, VERSION from secator.installer import ToolInstaller, fmt_health_table_row, get_health_table, get_version_info @@ -27,7 +27,7 @@ click.rich_click.USE_RICH_MARKUP = True ALL_TASKS = discover_tasks() -ALL_CONFIGS = ConfigLoader.load_all() +ALL_CONFIGS = TemplateLoader.load_all() ALL_WORKFLOWS = ALL_CONFIGS.workflow ALL_SCANS = ALL_CONFIGS.scan diff --git a/secator/config.py b/secator/config.py index e61b6eec..17a66e0f 100644 --- a/secator/config.py +++ b/secator/config.py @@ -1,137 +1,537 @@ -import glob +import os from pathlib import Path +from subprocess import call, DEVNULL +from typing import Dict, List +from typing_extensions import Annotated, Self +import requests import yaml from dotmap import DotMap +from pydantic import AfterValidator, BaseModel, model_validator, ValidationError -from secator.rich import console -from secator import CONFIG, CONFIGS_FOLDER +from secator.rich import console, console_stdout -CONFIGS_DIR_KEYS = ['workflow', 'scan', 'profile'] +Directory = Annotated[Path, AfterValidator(lambda v: v.expanduser())] +StrExpandHome = Annotated[str, AfterValidator(lambda v: v.replace('~', str(Path.home())))] +ROOT_FOLDER = Path(__file__).parent.parent +LIB_FOLDER = ROOT_FOLDER / 'secator' +CONFIGS_FOLDER = LIB_FOLDER / 'configs' -def load_config(name): - """Load a config by name. - Args: - name: Name of the config, for instances profiles/aggressive or workflows/domain_scan. +class StrictModel(BaseModel, extra='forbid'): + pass + + +class Directories(StrictModel): + bin: Directory = Path.home() / '.local' / 'bin' + data: Directory = Path.home() / '.secator' + templates: Directory = '' + reports: Directory = '' + wordlists: Directory = '' + cves: Directory = '' + payloads: Directory = '' + revshells: Directory = '' + celery: Directory = '' + celery_data: Directory = '' + celery_results: Directory = '' + + @model_validator(mode='after') + def set_default_folders(self) -> Self: + """Set folders to be relative to the data folders if they are unspecified in config.""" + for folder in ['templates', 'reports', 'wordlists', 'cves', 'payloads', 'revshells', 'celery', 'celery_data', 'celery_results']: # noqa: E501 + rel_target = '/'.join(folder.split('_')) + val = getattr(self, folder) or self.data / rel_target + setattr(self, folder, val) + return self + + +class Debug(StrictModel): + level: int = 0 + component: str = '' + + +class Celery(StrictModel): + broker_url: str = 'filesystem://' + broker_pool_limit: int = 10 + broker_connection_timeout: float = 4.0 + broker_visibility_timeout: int = 3600 + override_default_logging: bool = True + result_backend: StrExpandHome = '' + + +class Cli(StrictModel): + github_token: str = '' + record: bool = False + stdin_timeout: int = 1000 + + +class Runners(StrictModel): + input_chunk_size: int = 1000 + progress_update_frequency: int = 60 + skip_cve_search: bool = False + + +class HTTP(StrictModel): + socks5_proxy: str = 'socks5://127.0.0.1:9050' + http_proxy: str = 'https://127.0.0.1:9080' + store_responses: bool = False + proxychains_command: str = 'proxychains' + freeproxy_timeout: int = 1 + + +class Tasks(StrictModel): + exporters: List[str] = ['json', 'csv'] + + +class Workflows(StrictModel): + exporters: List[str] = ['json', 'csv'] + + +class Scans(StrictModel): + exporters: List[str] = ['json', 'csv'] + + +class Payloads(StrictModel): + templates: Dict[str, str] = { + 'lse': 'https://github.com/diego-treitos/linux-smart-enumeration/releases/latest/download/lse.sh', + 'linpeas': 'https://github.com/carlospolop/PEASS-ng/releases/latest/download/linpeas.sh', + 'sudo_killer': 'https://github.com/TH3xACE/SUDO_KILLER/archive/refs/heads/V3.zip' + } + + +class Wordlists(StrictModel): + defaults: Dict[str, str] = {'http': 'bo0m_fuzz', 'dns': 'combined_subdomains'} + templates: Dict[str, str] = { + 'bo0m_fuzz': 'https://raw.githubusercontent.com/Bo0oM/fuzz.txt/master/fuzz.txt', + 'combined_subdomains': 'https://raw.githubusercontent.com/danielmiessler/SecLists/master/Discovery/DNS/combined_subdomains.txt' # noqa: E501 + } + lists: Dict[str, List[str]] = {} - Returns: - dict: Loaded config. + +class GoogleAddon(StrictModel): + enabled: bool = False + drive_parent_folder_id: str = '' + credentials_path: str = '' + + +class WorkerAddon(StrictModel): + enabled: bool = False + + +class MongodbAddon(StrictModel): + enabled: bool = False + url: str = 'mongodb://localhost' + update_frequency: int = 60 + + +class Addons(StrictModel): + google: GoogleAddon = GoogleAddon() + worker: WorkerAddon = WorkerAddon() + mongodb: MongodbAddon = MongodbAddon() + + +class SecatorConfig(StrictModel): + dirs: Directories = Directories() + debug: Debug = Debug() + celery: Celery = Celery() + cli: Cli = Cli() + runners: Runners = Runners() + http: HTTP = HTTP() + tasks: Tasks = Tasks() + workflows: Workflows = Workflows() + scans: Scans = Scans() + payloads: Payloads = Payloads() + wordlists: Wordlists = Wordlists() + addons: Addons = Addons() + offline_mode: bool = False + + +class Config(DotMap): + """Config class. + + Examples: + >>> config = Config.parse() # get default config. + >>> config = Config.parse({'dirs': {'data': '/tmp/'}) # get custom config (from dict). + >>> config = Config.parse(path='/path/to/config.yml') # get custom config (from YAML file). + >>> config.print() # print config without defaults. + >>> config.print(partial=False) # print full config. + >>> config.set('addons.google.enabled', False) # set value in config. + >>> config.save() # save config back to disk. """ - path = CONFIGS_FOLDER / f'{name}.yaml' - if not path.exists(): - console.log(f'Config "{name}" could not be loaded.') - return - with path.open('r') as f: - return yaml.load(f.read(), Loader=yaml.Loader) - - -def find_configs(): - results = {'scan': [], 'workflow': [], 'profile': []} - dirs_type = [CONFIGS_FOLDER] - if CONFIG.dirs.templates: - dirs_type.append(CONFIG.dirs.templates) - paths = [] - for dir in dirs_type: - dir_paths = [ - Path(path) - for path in glob.glob(str(dir).rstrip('/') + '/**/*.y*ml', recursive=True) - ] - paths.extend(dir_paths) - for path in paths: - with path.open('r') as f: + + _error = False + + def get(self, key=None, print=True): + """Retrieve a value from the configuration using a dotted path. + + Args: + key (str | None): Dotted key path. + print (bool): Print the resulting value. + + Returns: + Any: value at key. + """ + value = self + if key: + for part in key.split('.'): + value = value[part] + if value is None: + console.print(f'[bold red]Key {key} does not exist.[/]') + return None + if print: + if key: + yaml_str = Config.dump(DotMap({key: value}), partial=False) + else: + yaml_str = Config.dump(self, partial=False) + Config.print_yaml(yaml_str) + return value + + def set(self, key, value, set_partial=True): + """Set a value in the configuration using a dotted path. + + Args: + key (str | None): Dotted key path. + value (Any): Value. + partial (bool): Also set value in partial config (written to disk). + + Returns: + bool: Success boolean. + """ + # Get existing value + existing_value = self.get(key, print=False) + + # Convert dotted key path to the corresponding uppercase key used in _keymap + map_key = key.upper().replace('.', '_') + success = False + if map_key in self._keymap: + # Traverse to the second last key to handle the setting correctly + target = self + partial = self._partial + for part in self._keymap[map_key][:-1]: + target = target[part] + if set_partial: + partial = partial[part] + + # Set the value on the final part of the path + final_key = self._keymap[map_key][-1] + + # Convert the value to the correct type based on the current value type try: - config = yaml.load(f.read(), yaml.Loader) - type = config.get('type') - if type: - results[type].append(path) - except yaml.YAMLError as exc: - console.log(f'Unable to load config at {path}') - console.log(str(exc)) - return results - - -class ConfigLoader(DotMap): - - def __init__(self, input={}, name=None, **kwargs): - if name: - name = name.replace('-', '_') # so that workflows have a nice '-' in CLI - config = self._load_from_name(name) - elif isinstance(input, str) or isinstance(input, Path): - config = self._load_from_file(input) + if isinstance(existing_value, bool): + if isinstance(value, str): + value = value.lower() in ("true", "1", "t") + elif isinstance(value, (int, float)): + value = True if value == 1 else False + elif isinstance(existing_value, int): + value = int(value) + elif isinstance(existing_value, float): + value = float(value) + if existing_value != value: + target[final_key] = value + if set_partial: + partial[final_key] = value + success = True + except ValueError: + success = False + # console.print(f'[bold red]{key}: cannot cast value "{value}" to {type(existing_value).__name__}') else: - config = input - super().__init__(config) - - def _load_from_file(self, path): - if isinstance(path, str): - path = Path(path) - if not path.exists(): - console.log(f'Config path {path} does not exists', style='bold red') - return - with path.open('r') as f: - return yaml.load(f.read(), Loader=yaml.Loader) - - def _load_from_name(self, name): - return load_config(name) - - @classmethod - def load_all(cls): - configs = find_configs() - return ConfigLoader({ - key: [ConfigLoader(path) for path in configs[key]] - for key in CONFIGS_DIR_KEYS - }) - - def get_tasks_class(self): - from secator.runners import Task - tasks = [] - for name, conf in self.tasks.items(): - if name == '_group': - group_conf = ConfigLoader(input={'tasks': conf}) - tasks.extend(group_conf.get_tasks_class()) + console.print(f'[bold red]Key "{key}" not found in config keymap[/].') + return success + + def save(self, target_path: Path = None, partial=True): + """Save config as YAML on disk. + + Args: + target_path (Path | None): If passed, saves the config to this path. + partial (bool): Save partial config. + """ + if not target_path: + if not self._path: + return + target_path = self._path + with target_path.open('w') as f: + f.write(Config.dump(self, partial=partial)) + self._path = target_path + + def print(self, partial=True): + """Print config. + + Args: + partial (bool): Print partial config only. + """ + yaml_str = self.dump(self, partial=partial) + yaml_str = f'# {self._path}\n\n{yaml_str}' if self._path and partial else yaml_str + Config.print_yaml(yaml_str) + + @staticmethod + def parse(data: dict = {}, path: Path = None, env_overrides: bool = False): + """Parse config. + + Args: + data (dict): Config data. + path (Path | None): Path to YAML config. + env_overrides (bool): Apply env overrides. + + Returns: + Config: instance of Config object. + None: if the config was not loaded properly or there are validation errors. + """ + if path: + data = Config.read_yaml(path) + + # Load data + try: + config = Config.load(SecatorConfig, data) + config._valid = True + + # HACK: set default result_backend if unset + if not config.celery.result_backend: + config.celery.result_backend = f'file://{config.dirs.celery_results}' + + except ValidationError as e: + error_str = str(e).replace('\n', '\n ') + if path: + error_str.replace('SecatorConfig', f'SecatorConfig ({path})') + console.print(f'[bold red]:x: {error_str}') + # console.print('[bold green]Using default config.[/]') + config = Config.parse() + config._valid = False + + # Set hidden attributes + keymap = Config.build_key_map(config) + partial = Config(data) + config._partial = partial + config._path = path + config._keymap = keymap + + # Override config values with environment variables + if env_overrides: + config.apply_env_overrides() + data = {k: v for k, v in config.toDict().items() if not k.startswith('_')} + config = Config.parse(data, env_overrides=False) # re-validate config + config._partial = partial + config._path = path + + return config + + @staticmethod + def load(schema, data: dict = {}): + """Validate a config using Pydantic. + + Args: + data (dict): Config dict. + + Returns: + Config: instance of Config object. + """ + return Config(schema(**data).model_dump()) + + @staticmethod + def read_yaml(yaml_path): + """Read YAML from path. + + Args: + yaml_path (Path): path to yaml config. + + Returns: + dict: Loaded data. + """ + with yaml_path.open('r') as f: + data = yaml.load(f.read(), Loader=yaml.Loader) + return data or {} + + @staticmethod + def print_yaml(string): + """Print YAML string using rich. + + Args: + string (str): YAML string. + """ + from rich.syntax import Syntax + data = Syntax(string, 'yaml', theme='ansi-dark', padding=0, background_color='default') + console_stdout.print(data) + + @staticmethod + def dump(config, partial=True): + """Safe dump config as yaml: + - `Path`, `PosixPath` and `WindowsPath` objects are translated to strings. + - Home directory in paths is replaced with the tilde '~'. + + Returns: + str: YAML dump. + """ + import yaml + from pathlib import Path, PosixPath, WindowsPath + + # Get home dir + home = str(Path.home()) + + # Custom dumper to add line breaks between items and a path representer to translate paths to strings + class LineBreakDumper(yaml.SafeDumper): + def write_line_break(self, data=None): + super().write_line_break(data) + if len(self.indents) == 1: + super().write_line_break() + + def posix_path_representer(dumper, data): + path = str(data) + if path.startswith(home): + path = path.replace(home, '~') + return dumper.represent_scalar('tag:yaml.org,2002:str', path) + + LineBreakDumper.add_representer(Path, posix_path_representer) + LineBreakDumper.add_representer(PosixPath, posix_path_representer) + LineBreakDumper.add_representer(WindowsPath, posix_path_representer) + + # Get data dict + data = config.toDict() + + # HACK: Replace home dir in result_backend + if isinstance(config, Config): + data['celery']['result_backend'] = data['celery']['result_backend'].replace(home, '~') + del data['_path'] + if partial: + data = data['_partial'] + else: + del data['_partial'] + + data = {k: v for k, v in data.items() if not k.startswith('_')} + return yaml.dump(data, Dumper=LineBreakDumper, sort_keys=False) + + @staticmethod + def build_key_map(config, base_path=[]): + key_map = {} + for key, value in config.items(): + if key.startswith('_'): # ignore + continue + current_path = base_path + [key] + if isinstance(value, dict): + key_map.update(Config.build_key_map(value, current_path)) else: - tasks.append(Task.get_task_class(name)) - return tasks - - def get_workflows(self): - return [ConfigLoader(name=f'workflows/{name}') for name, _ in self.workflows.items()] - - def get_workflow_supported_opts(self): - opts = {} - tasks = self.get_tasks_class() - for task_cls in tasks: - task_opts = task_cls.get_supported_opts() - for name, conf in task_opts.items(): - supported = opts.get(name, {}).get('supported', False) - opts[name] = conf - opts[name]['supported'] = conf['supported'] or supported - return opts - - def get_scan_supported_opts(self): - opts = {} - workflows = self.get_workflows() - for workflow in workflows: - workflow_opts = workflow.get_workflow_supported_opts() - for name, conf in workflow_opts.items(): - supported = opts.get(name, {}).get('supported', False) - opts[name] = conf - opts[name]['supported'] = conf['supported'] or supported - return opts - - @property - def supported_opts(self): - return self.get_supported_opts() - - def get_supported_opts(self): - opts = {} - if self.type == 'workflow': - opts = self.get_workflow_supported_opts() - elif self.type == 'scan': - opts = self.get_scan_supported_opts() - elif self.type == 'task': - tasks = self.get_tasks_class() - if tasks: - opts = tasks[0].get_supported_opts() - return dict(sorted(opts.items())) + key_map['_'.join(current_path).upper()] = current_path + return key_map + + def apply_env_overrides(self): + """Override config values from environment variables.""" + # Build a map of keys from the config + key_map = Config.build_key_map(self) + + # Prefix for environment variables to target + prefix = "SECATOR_" + + # Loop through environment variables + for var in os.environ: + if var.startswith(prefix): + # Remove prefix and get the path from the key map + key = var[len(prefix):] + if key in key_map: + path = '.'.join(k.lower() for k in key_map[key]) + value = os.environ[var] + + # Set the new value recursively + success = self.set(path, value, set_partial=False) + if success: + console.print(f'[bold green4]{var} (override success)[/]') + else: + console.print(f'[bold red]{var} (override failed: cannot update value)[/]') + else: + console.print(f'[bold red]{var} (override failed: key not found in config)[/]') + + +def download_files(data: dict, target_folder: Path, offline_mode: bool, type: str): + """Download remote files to target folder, clone git repos, or symlink local files. + + Args: + data (dict): Dict of name to url or local path prefixed with 'git+' for Git repos. + target_folder (Path): Target folder for storing files or repos. + type (str): Type of files to handle. + offline_mode (bool): Offline mode. + """ + for name, url_or_path in data.items(): + if url_or_path.startswith('git+'): + # Clone Git repository + git_url = url_or_path[4:] # remove 'git+' prefix + repo_name = git_url.split('/')[-1] + if repo_name.endswith('.git'): + repo_name = repo_name[:-4] + target_path = target_folder / repo_name + if not target_path.exists(): + console.print(f'[bold turquoise4]Cloning git {type} [bold magenta]{repo_name}[/] ...[/] ', end='') + if offline_mode: + console.print('[bold orange1]skipped [dim][offline[/].[/]') + continue + try: + call(['git', 'clone', git_url, str(target_path)], stderr=DEVNULL, stdout=DEVNULL) + console.print('[bold green]ok.[/]') + except Exception as e: + console.print(f'[bold red]failed ({str(e)}).[/]') + data[name] = target_path.resolve() + elif Path(url_or_path).exists(): + # Create a symbolic link for a local file + local_path = Path(url_or_path) + target_path = target_folder / local_path.name + if not target_path.exists(): + console.print(f'[bold turquoise4]Symlinking {type} [bold magenta]{name}[/] ...[/] ', end='') + try: + target_path.symlink_to(local_path) + console.print('[bold green]ok.[/]') + except Exception as e: + console.print(f'[bold red]failed ({str(e)}).[/]') + data[name] = target_path.resolve() + else: + # Download file from URL + ext = url_or_path.split('.')[-1] + filename = f'{name}.{ext}' + target_path = target_folder / filename + if not target_path.exists(): + try: + console.print(f'[bold turquoise4]Downloading {type} [bold magenta]{filename}[/] ...[/] ', end='') + if offline_mode: + console.print('[bold orange1]skipped [dim](offline)[/].[/]') + continue + resp = requests.get(url_or_path, timeout=3) + resp.raise_for_status() + with open(target_path, 'wb') as f: + f.write(resp.content) + console.print('[bold green]ok.[/]') + except requests.RequestException as e: + console.print(f'[bold red]failed ({str(e)}).[/]') + continue + data[name] = target_path.resolve() + + +# Load configs +default_config = Config.parse() +data_root = default_config.dirs.data +config_path = data_root / 'config.yml' +if not config_path.exists(): + if not data_root.exists(): + console.print(f'[bold turquoise4]Creating directory [bold magenta]{data_root}[/] ... [/]', end='') + data_root.mkdir(parents=False) + console.print('[bold green]ok.[/]') + console.print( + f'[bold turquoise4]Creating user conf [bold magenta]{config_path}[/]... [/]', end='') + config_path.touch() + console.print('[bold green]ok.[/]') +CONFIG = Config.parse(path=config_path, env_overrides=True) + +# Create directories if they don't exist already +for name, dir in CONFIG.dirs.items(): + if not dir.exists(): + console.print(f'[bold turquoise4]Creating directory [bold magenta]{dir}[/] ... [/]', end='') + dir.mkdir(parents=False) + console.print('[bold green]ok.[/]') + +# Download wordlists and set defaults +download_files(CONFIG.wordlists.templates, CONFIG.dirs.wordlists, CONFIG.offline_mode, 'wordlist') +for category, name in CONFIG.wordlists.defaults.items(): + if name in CONFIG.wordlists.templates.keys(): + CONFIG.wordlists.defaults[category] = str(CONFIG.wordlists.templates[name]) + +# Download payloads +download_files(CONFIG.payloads.templates, CONFIG.dirs.payloads, CONFIG.offline_mode, 'payload') + +# Print config +if CONFIG.debug.component == 'config': + CONFIG.print() diff --git a/secator/decorators.py b/secator/decorators.py index 6bd91e5f..114dad0f 100644 --- a/secator/decorators.py +++ b/secator/decorators.py @@ -6,7 +6,7 @@ from rich_click.rich_group import RichGroup from secator.definitions import ADDONS_ENABLED, OPT_NOT_SUPPORTED -from secator import CONFIG +from secator.config import CONFIG from secator.runners import Scan, Task, Workflow from secator.utils import (deduplicate, expand_input, get_command_category, get_command_cls) diff --git a/secator/definitions.py b/secator/definitions.py index 51f5e41b..370de910 100644 --- a/secator/definitions.py +++ b/secator/definitions.py @@ -5,7 +5,7 @@ from dotenv import find_dotenv, load_dotenv from importlib.metadata import version -from secator import CONFIG, ROOT_FOLDER +from secator.config import CONFIG, ROOT_FOLDER load_dotenv(find_dotenv(usecwd=True), override=False) diff --git a/secator/exporters/gdrive.py b/secator/exporters/gdrive.py index d56e988b..dca897a5 100644 --- a/secator/exporters/gdrive.py +++ b/secator/exporters/gdrive.py @@ -2,7 +2,7 @@ import csv import yaml -from secator import CONFIG +from secator.config import CONFIG from secator.exporters._base import Exporter from secator.rich import console from secator.utils import pluralize diff --git a/secator/hooks/mongodb.py b/secator/hooks/mongodb.py index 1053252b..df557252 100644 --- a/secator/hooks/mongodb.py +++ b/secator/hooks/mongodb.py @@ -5,7 +5,7 @@ from bson.objectid import ObjectId from celery import shared_task -from secator import CONFIG +from secator.config import CONFIG from secator.output_types import OUTPUT_TYPES from secator.runners import Scan, Task, Workflow from secator.utils import debug, escape_mongodb_url diff --git a/secator/installer.py b/secator/installer.py index cafa1d19..ec81afac 100644 --- a/secator/installer.py +++ b/secator/installer.py @@ -11,7 +11,7 @@ from secator.rich import console from secator.runners import Command -from secator import CONFIG +from secator.config import CONFIG class ToolInstaller: diff --git a/secator/runners/_base.py b/secator/runners/_base.py index ee88ca0a..7adb6767 100644 --- a/secator/runners/_base.py +++ b/secator/runners/_base.py @@ -15,7 +15,7 @@ from rich.progress import SpinnerColumn, TextColumn, TimeElapsedColumn from secator.definitions import DEBUG -from secator import CONFIG +from secator.config import CONFIG from secator.output_types import OUTPUT_TYPES, OutputType, Progress from secator.report import Report from secator.rich import console, console_stdout @@ -49,7 +49,7 @@ class Runner: """Runner class. Args: - config (secator.config.ConfigLoader): Loaded config. + config (secator.config.TemplateLoader): Loaded config. targets (list): List of targets to run task on. results (list): List of existing results to re-use. workspace_name (str): Workspace name. diff --git a/secator/runners/command.py b/secator/runners/command.py index 9f884fe7..12dc5e89 100644 --- a/secator/runners/command.py +++ b/secator/runners/command.py @@ -10,9 +10,9 @@ from fp.fp import FreeProxy -from secator.config import ConfigLoader +from secator.template import TemplateLoader from secator.definitions import OPT_NOT_SUPPORTED, OPT_PIPE_INPUT -from secator import CONFIG +from secator.config import CONFIG from secator.runners import Runner from secator.serializers import JSONSerializer from secator.utils import debug @@ -104,7 +104,7 @@ class Command(Runner): def __init__(self, input=None, **run_opts): # Build runnerconfig on-the-fly - config = ConfigLoader(input={ + config = TemplateLoader(input={ 'name': self.__class__.__name__, 'type': 'task', 'description': run_opts.get('description', None) diff --git a/secator/runners/scan.py b/secator/runners/scan.py index 1b020698..aea793f0 100644 --- a/secator/runners/scan.py +++ b/secator/runners/scan.py @@ -1,7 +1,7 @@ import logging -from secator.config import ConfigLoader -from secator import CONFIG +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 @@ -52,7 +52,7 @@ def yielder(self): # Run workflow workflow = Workflow( - ConfigLoader(name=f'workflows/{name}'), + TemplateLoader(name=f'workflows/{name}'), targets, results=[], run_opts=run_opts, diff --git a/secator/runners/task.py b/secator/runners/task.py index af6e7853..b02cc3d5 100644 --- a/secator/runners/task.py +++ b/secator/runners/task.py @@ -1,6 +1,6 @@ from secator.definitions import DEBUG from secator.output_types import Target -from secator import CONFIG +from secator.config import CONFIG from secator.runners import Runner from secator.utils import discover_tasks diff --git a/secator/runners/workflow.py b/secator/runners/workflow.py index e5e9c5c9..47f68b29 100644 --- a/secator/runners/workflow.py +++ b/secator/runners/workflow.py @@ -1,6 +1,6 @@ from secator.definitions import DEBUG from secator.output_types import Target -from secator import CONFIG +from secator.config import CONFIG from secator.runners._base import Runner from secator.runners.task import Task from secator.utils import merge_opts @@ -81,7 +81,7 @@ def get_tasks(self, obj, targets, workflow_opts, run_opts): """Get tasks recursively as Celery chains / chords. Args: - obj (secator.config.ConfigLoader): Config. + obj (secator.config.TemplateLoader): Config. targets (list): List of targets. workflow_opts (dict): Workflow options. run_opts (dict): Run options. diff --git a/secator/tasks/_categories.py b/secator/tasks/_categories.py index dfc427cc..1def1b10 100644 --- a/secator/tasks/_categories.py +++ b/secator/tasks/_categories.py @@ -11,7 +11,7 @@ RATE_LIMIT, REFERENCES, RETRIES, SEVERITY, TAGS, THREADS, TIMEOUT, URL, USER_AGENT, USERNAME, WORDLIST) from secator.output_types import Ip, Port, Subdomain, Tag, Url, UserAccount, Vulnerability -from secator import CONFIG +from secator.config import CONFIG from secator.runners import Command from secator.utils import debug diff --git a/secator/tasks/dnsxbrute.py b/secator/tasks/dnsxbrute.py index dc52d8cf..817d1582 100644 --- a/secator/tasks/dnsxbrute.py +++ b/secator/tasks/dnsxbrute.py @@ -1,6 +1,6 @@ from secator.decorators import task from secator.definitions import (DOMAIN, HOST, RATE_LIMIT, RETRIES, THREADS, WORDLIST, EXTRA_DATA) -from secator import CONFIG +from secator.config import CONFIG from secator.output_types import Subdomain from secator.tasks._categories import ReconDns diff --git a/secator/tasks/httpx.py b/secator/tasks/httpx.py index ea482b24..04081b8f 100644 --- a/secator/tasks/httpx.py +++ b/secator/tasks/httpx.py @@ -8,7 +8,7 @@ MATCH_WORDS, METHOD, OPT_NOT_SUPPORTED, PROXY, RATE_LIMIT, RETRIES, THREADS, TIMEOUT, URL, USER_AGENT) -from secator import CONFIG +from secator.config import CONFIG from secator.tasks._categories import Http from secator.utils import sanitize_url diff --git a/secator/tasks/katana.py b/secator/tasks/katana.py index ff6f6f5c..ff13f785 100644 --- a/secator/tasks/katana.py +++ b/secator/tasks/katana.py @@ -12,7 +12,7 @@ RATE_LIMIT, RETRIES, STATUS_CODE, STORED_RESPONSE_PATH, TECH, THREADS, TIME, TIMEOUT, URL, USER_AGENT, WEBSERVER, CONTENT_LENGTH) -from secator import CONFIG +from secator.config import CONFIG from secator.output_types import Url, Tag from secator.tasks._categories import HttpCrawler diff --git a/secator/template.py b/secator/template.py new file mode 100644 index 00000000..3c41732d --- /dev/null +++ b/secator/template.py @@ -0,0 +1,137 @@ +import glob +from pathlib import Path + +import yaml +from dotmap import DotMap + +from secator.rich import console +from secator.config import CONFIG, CONFIGS_FOLDER + +TEMPLATES_DIR_KEYS = ['workflow', 'scan', 'profile'] + + +def load_template(name): + """Load a config by name. + + Args: + name: Name of the config, for instances profiles/aggressive or workflows/domain_scan. + + Returns: + dict: Loaded config. + """ + path = CONFIGS_FOLDER / f'{name}.yaml' + if not path.exists(): + console.log(f'Config "{name}" could not be loaded.') + return + with path.open('r') as f: + return yaml.load(f.read(), Loader=yaml.Loader) + + +def find_templates(): + results = {'scan': [], 'workflow': [], 'profile': []} + dirs_type = [CONFIGS_FOLDER] + if CONFIG.dirs.templates: + dirs_type.append(CONFIG.dirs.templates) + paths = [] + for dir in dirs_type: + dir_paths = [ + Path(path) + for path in glob.glob(str(dir).rstrip('/') + '/**/*.y*ml', recursive=True) + ] + paths.extend(dir_paths) + for path in paths: + with path.open('r') as f: + try: + config = yaml.load(f.read(), yaml.Loader) + type = config.get('type') + if type: + results[type].append(path) + except yaml.YAMLError as exc: + console.log(f'Unable to load config at {path}') + console.log(str(exc)) + return results + + +class TemplateLoader(DotMap): + + def __init__(self, input={}, name=None, **kwargs): + if name: + name = name.replace('-', '_') # so that workflows have a nice '-' in CLI + config = self._load_from_name(name) + elif isinstance(input, str) or isinstance(input, Path): + config = self._load_from_file(input) + else: + config = input + super().__init__(config) + + def _load_from_file(self, path): + if isinstance(path, str): + path = Path(path) + if not path.exists(): + console.log(f'Config path {path} does not exists', style='bold red') + return + with path.open('r') as f: + return yaml.load(f.read(), Loader=yaml.Loader) + + def _load_from_name(self, name): + return load_template(name) + + @classmethod + def load_all(cls): + configs = find_templates() + return TemplateLoader({ + key: [TemplateLoader(path) for path in configs[key]] + for key in TEMPLATES_DIR_KEYS + }) + + def get_tasks_class(self): + from secator.runners import Task + tasks = [] + for name, conf in self.tasks.items(): + if name == '_group': + group_conf = TemplateLoader(input={'tasks': conf}) + tasks.extend(group_conf.get_tasks_class()) + else: + tasks.append(Task.get_task_class(name)) + return tasks + + def get_workflows(self): + return [TemplateLoader(name=f'workflows/{name}') for name, _ in self.workflows.items()] + + def get_workflow_supported_opts(self): + opts = {} + tasks = self.get_tasks_class() + for task_cls in tasks: + task_opts = task_cls.get_supported_opts() + for name, conf in task_opts.items(): + supported = opts.get(name, {}).get('supported', False) + opts[name] = conf + opts[name]['supported'] = conf['supported'] or supported + return opts + + def get_scan_supported_opts(self): + opts = {} + workflows = self.get_workflows() + for workflow in workflows: + workflow_opts = workflow.get_workflow_supported_opts() + for name, conf in workflow_opts.items(): + supported = opts.get(name, {}).get('supported', False) + opts[name] = conf + opts[name]['supported'] = conf['supported'] or supported + return opts + + @property + def supported_opts(self): + return self.get_supported_opts() + + def get_supported_opts(self): + opts = {} + if self.type == 'workflow': + opts = self.get_workflow_supported_opts() + elif self.type == 'scan': + opts = self.get_scan_supported_opts() + elif self.type == 'task': + tasks = self.get_tasks_class() + if tasks: + opts = tasks[0].get_supported_opts() + return dict(sorted(opts.items())) diff --git a/secator/utils.py b/secator/utils.py index a2cd45fd..b12a1c9c 100644 --- a/secator/utils.py +++ b/secator/utils.py @@ -20,7 +20,7 @@ from rich.markdown import Markdown from secator.definitions import (DEBUG, DEBUG_COMPONENT, VERSION, DEV_PACKAGE) -from secator import CONFIG, ROOT_FOLDER, LIB_FOLDER +from secator.config import CONFIG, ROOT_FOLDER, LIB_FOLDER from secator.rich import console logger = logging.getLogger(__name__) diff --git a/tests/integration/test_workflows.py b/tests/integration/test_workflows.py index 0643aeeb..32fc67cc 100644 --- a/tests/integration/test_workflows.py +++ b/tests/integration/test_workflows.py @@ -4,7 +4,7 @@ import warnings from time import sleep -from secator.config import ConfigLoader +from secator.template import TemplateLoader from secator.runners import Task from secator.output_types import Target, Port, Url from secator.definitions import DEBUG @@ -128,7 +128,7 @@ def test_adhoc_workflow(self): } } } - config = ConfigLoader(conf) + config = TemplateLoader(conf) workflow = Workflow( config, targets=['localhost'], diff --git a/tests/performance/loadtester.py b/tests/performance/loadtester.py index 0d8471eb..e65007d5 100644 --- a/tests/performance/loadtester.py +++ b/tests/performance/loadtester.py @@ -1,7 +1,7 @@ import eventlet eventlet.monkey_patch() from secator.runners import Workflow -from secator.config import ConfigLoader +from secator.template import TemplateLoader from secator.rich import console from secator.celery import * import os @@ -29,7 +29,7 @@ def create_runner(index): register('json', json.dumps, json.loads, content_type='application/json', content_encoding='utf-8') - config = ConfigLoader(name='workflows/subdomain_recon') + config = TemplateLoader(name='workflows/subdomain_recon') run_opts = { 'sync': False, # 'print_start': True, diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index ff04a0ea..f83169bc 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,5 +1,5 @@ import unittest -from secator import Config, CONFIG +from secator.config import Config, CONFIG class TestConfig(unittest.TestCase): diff --git a/tests/unit/test_offline.py b/tests/unit/test_offline.py index 50f01860..e77103cf 100644 --- a/tests/unit/test_offline.py +++ b/tests/unit/test_offline.py @@ -20,7 +20,7 @@ def test_offline_cve_lookup(self): self.assertEqual(result, None) def test_offline_downloads(self): - from secator import download_files, CONFIG + from secator.config import download_files, CONFIG download_files( {'pyproject.toml': 'https://raw.githubusercontent.com/freelabz/secator/main/pyproject.toml'}, CONFIG.dirs.data,