diff --git a/web/reNgine/common_func.py b/web/reNgine/common_func.py index c989f269..f7a98097 100644 --- a/web/reNgine/common_func.py +++ b/web/reNgine/common_func.py @@ -34,38 +34,38 @@ DISCORD_WEBHOOKS_CACHE = redis.Redis.from_url(CELERY_BROKER_URL) # EngineType utils # #------------------# def dump_custom_scan_engines(results_dir): - """Dump custom scan engines to YAML files. + """Dump custom scan engines to YAML files. - Args: - results_dir (str): Results directory (will be created if non-existent). - """ - custom_engines = EngineType.objects.filter(default_engine=False) - if not os.path.exists(results_dir): - os.makedirs(results_dir, exist_ok=True) - for engine in custom_engines: - with open(f'{results_dir}/{engine.engine_name}.yaml', 'w') as f: - config = yaml.safe_load(engine.yaml_configuration) - yaml.dump(config, f, indent=4) + Args: + results_dir (str): Results directory (will be created if non-existent). + """ + custom_engines = EngineType.objects.filter(default_engine=False) + if not os.path.exists(results_dir): + os.makedirs(results_dir, exist_ok=True) + for engine in custom_engines: + with open(os.path.join(results_dir, f"{engine.engine_name}.yaml"), 'w') as f: + f.write(engine.yaml_configuration) def load_custom_scan_engines(results_dir): - """Load custom scan engines from YAML files. The filename without .yaml will - be used as the engine name. + """Load custom scan engines from YAML files. The filename without .yaml will + be used as the engine name. - Args: - results_dir (str): Results directory containing engines configs. - """ - config_paths = [ - f for f in os.listdir(results_dir) - if os.path.isfile(os.path.join(results_dir, f)) - ] - for path in config_paths: - engine_name = path.replace('.yaml', '').split('/')[-1] - full_path = os.path.join(results_dir, path) - with open(full_path, 'r') as f: - yaml_configuration = yaml.safe_load(f) - engine, _ = EngineType.objects.get_or_create(engine_name=engine_name) - engine.yaml_configuration = yaml.dump(yaml_configuration) - engine.save() + Args: + results_dir (str): Results directory containing engines configs. + """ + config_paths = [ + f for f in os.listdir(results_dir) + if os.path.isfile(os.path.join(results_dir, f)) and f.endswith('.yaml') + ] + for path in config_paths: + engine_name = os.path.splitext(os.path.basename(path))[0] + full_path = os.path.join(results_dir, path) + with open(full_path, 'r') as f: + yaml_configuration = f.read() + + engine, _ = EngineType.objects.get_or_create(engine_name=engine_name) + engine.yaml_configuration = yaml_configuration + engine.save() #--------------------------------#