def test_nonexistent_arg(self):
"""Makes sure that an exception is raised for unknown @ variables.
A recipe writer needs to define a parsing tuple for each @ variable used by
the recipe.
"""
recipe_args = {
'recipe_arg1': 'This should be replaced: @parameterone',
'recipe_arg2': 'This arg cannot be replaced @nonexistent',
}
parser.set_defaults(**config.Config.get_extra())
args = parser.parse_args([
'value_for_param_one',
'--optional_param',
'3',
'--spaces_param',
'So unique!',
'BOOM',
])
with self.assertRaises(ValueError):
imported = dftw_utils.import_args_from_dict(
recipe_args, vars(args), config.Config)
dftw_utils.check_placeholders(imported)
def test_import_args_from_cli(self):
"""Tries parsing the CLI arguments and updating a recipe dictionary."""
recipe_args = {
'recipe_arg1': 'This should remain intact',
'recipe_arg2': 'This should be replaced: @parameterone',
'recipe_arg3': 'This should be replaced by @optional_param',
'recipe_arg4': 'This includes spaces: @spaces_param',
'recipe_arg5': 'Characters after param: @explosion!',
}
expected_args = {
'recipe_arg1': 'This should remain intact',
'recipe_arg2': 'This should be replaced: value_for_param_one',
'recipe_arg3': 'This should be replaced by 3',
'recipe_arg4': 'This includes spaces: S P A C E',
'recipe_arg5': 'Characters after param: BOOM!',
}
parser.set_defaults(**config.Config.get_extra())
args = parser.parse_args([
'value_for_param_one',
'--optional_param',
'3',
'--spaces_param',
'S P A C E',
'BOOM',
])
imported_args = dftw_utils.import_args_from_dict(
recipe_args, vars(args), config.Config)
self.assertEqual(imported_args, expected_args)
def test_config_fills_missing_args(self):
"""Tests that a configuration file will fill-in arguments that are missing
from the CLI."""
provided_args = {'arg1': 'This should remain intact', 'arg2': '@config'}
expected_args = {
'arg1': 'This should remain intact',
'arg2': 'A config arg',
}
config.Config.load_extra_data('{"config": "A config arg"}')
imported_args = dftw_utils.import_args_from_dict(
provided_args, {}, config.Config)
self.assertEqual(imported_args, expected_args)
def _setup_module_thread(module_description):
"""Calls the module's setup() function and sets an Event object for it.
Args:
module_description (dict): Corresponding recipe module description.
"""
new_args = utils.import_args_from_dict(
module_description['args'], vars(args), self.config)
module = self._module_pool[module_description['name']]
try:
module.setup(**new_args)
except Exception as error: # pylint: disable=broad-except
self.add_error(
'An unknown error occurred: {0!s}'.format(error), critical=True)
self.events[module_description['name']] = threading.Event()
def main():
"""Main function for DFTimewolf."""
parser = argparse.ArgumentParser(
formatter_class=argparse.RawDescriptionHelpFormatter,
description=generate_help())
subparsers = parser.add_subparsers()
for registered_recipe in config.Config.get_registered_recipes():
recipe, recipe_args, documentation = registered_recipe
subparser = subparsers.add_parser(
recipe['name'],
formatter_class=utils.DFTimewolfFormatterClass,
description='{0:s}'.format(documentation))
subparser.set_defaults(recipe=recipe)
for switch, help_text, default in recipe_args:
subparser.add_argument(switch, help=help_text, default=default)
# Override recipe defaults with those specified in Config
# so that they can in turn be overridden in the commandline
subparser.set_defaults(**config.Config.get_extra())
args = parser.parse_args()
recipe = args.recipe
# Thread all collectors.
state = DFTimewolfState()
for module_description in recipe['modules']:
# Combine CLI args with args from the recipe description
new_args = utils.import_args_from_dict(module_description['args'],
vars(args), config.Config)
# Create the module object and start processing
module_name = module_description['name']
print('Running module {0:s}'.format(module_name))
module = config.Config.get_module(module_name)(state)
module.setup(**new_args)
state.check_errors()
try:
module.process()
except DFTimewolfError as error:
state.add_error(error.message, critical=True)
# Check for eventual errors and clean up after each round.
state.check_errors()
state.cleanup()
print('Recipe executed successfully.')
def test_cli_precedence_over_config(self):
"""Tests that the same argument provided via the CLI overrides the one
specified in the config file."""
provided_args = {
'arg1': 'I want whatever CLI says: @parameterone',
}
expected_args = {
'arg1': 'I want whatever CLI says: CLI WINS!',
}
config.Config.load_extra_data('{"parameterone": "CONFIG WINS!"}')
args = parser.parse_args(['CLI WINS!', 'BOOM'])
imported_args = dftw_utils.import_args_from_dict(
provided_args, vars(args), config.Config)
self.assertEqual(imported_args, expected_args)
def _setup_module_thread(module_description):
"""Calls the module's setup() function and sets an Event object for it.
Args:
module_description (dict): Corresponding recipe module description.
"""
new_args = utils.import_args_from_dict(module_description['args'],
vars(args), self.config)
module = self._module_pool[module_description['name']]
try:
module.setup(**new_args)
except Exception as error: # pylint: disable=broad-except
self.add_error(
'An unknown error occurred: {0!s}\nFull traceback:\n{1:s}'.
format(error, traceback.format_exc()),
critical=True)
self.events[module_description['name']] = threading.Event()
self.cleanup()
def register_recipe(cls, recipe, **kwargs):
"""Registers a dftimewolf recipe.
Registers a DFTimeWolf recipe with specified parameters. Parameters can be
specified in three ways, in order of precedence:
* Defined in config.json
* Passed as arguments to the register_recipe function call
* Passed as CLI args
Args:
recipe: imported python module representing the recipe.
**kwargs: parameters to be replaced in the recipe before checking the
CLI arguments.
"""
# Update kwargs with what we already loaded from config.json
kwargs.update(cls._extra_config)
recipe.contents = dftw_utils.import_args_from_dict(recipe.contents, kwargs)
recipe_name = recipe.contents['name']
cls._recipe_classes[recipe_name] = (recipe.contents, recipe.args)
def main():
"""Main function for DFTimewolf."""
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(
title=u'Available recipes',
description=u'List of currently loaded recipes',
help=u'Recipe-specific help')
for recipe, recipe_args in config.Config.get_registered_recipes():
subparser = subparsers.add_parser(recipe[u'name'],
description=u'{0:s}'.format(
recipe.__doc__))
subparser.set_defaults(recipe=recipe)
for switch, help_text in recipe_args:
subparser.add_argument(switch, help=help_text)
args = parser.parse_args()
recipe = args.recipe
console_out = dftw_utils.DFTimewolfConsoleOutput(sender=u'DFTimewolfCli',
verbose=True)
# COLLECTORS
# Thread collectors
console_out.StdOut(u'Collectors:')
for collector in recipe['collectors']:
console_out.StdOut(u' {0:s}'.format(collector[u'name']))
collector_objs = []
for collector in recipe[u'collectors']:
new_args = dftw_utils.import_args_from_dict(collector[u'args'],
vars(args))
collector_cls = config.Config.get_collector(collector[u'name'])
collector_objs.extend(collector_cls.launch_collector(**new_args))
# Wait for collectors to finish and collect output
collector_output = []
for collector_obj in collector_objs:
collector_obj.join()
collector_output.extend(collector_obj.results)
if recipe[u'processors']:
# PROCESSORS
# Thread processors
console_out.StdOut(u'Processors:')
for processor in recipe[u'processors']:
console_out.StdOut(u' {0:s}'.format(processor[u'name']))
processor_objs = []
for processor in recipe[u'processors']:
new_args = dftw_utils.import_args_from_dict(
processor[u'args'], vars(args))
new_args[u'collector_output'] = collector_output
processor_class = config.Config.get_processor(processor[u'name'])
processor_objs.extend(processor_class.launch_processor(**new_args))
# Wait for processors to finish and collect output
processor_output = []
for processor in processor_objs:
processor.join()
processor_output.extend(processor.output)
else:
processor_output = collector_output
# EXPORTERS
# Thread exporters
console_out.StdOut(u'Exporters:')
for exporter in recipe[u'exporters']:
console_out.StdOut(u' {0:s}'.format(exporter[u'name']))
exporter_objs = []
for exporter in recipe[u'exporters']:
new_args = dftw_utils.import_args_from_dict(exporter[u'args'],
vars(args))
new_args[u'processor_output'] = processor_output
exporter_class = config.Config.get_exporter(exporter[u'name'])
exporter_objs.extend(exporter_class.launch_exporter(**new_args))
# Wait for exporters to finish
exporter_output = []
for exporter in exporter_objs:
exporter.join()
exporter_output.extend(exporter.output)
console_out.StdOut(u'Recipe {0:s} executed successfully'.format(
recipe[u'name']))
def main():
"""Main function for DFTimewolf."""
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(
title='Available recipes',
description='List of currently loaded recipes',
help=
'Recipe-specific help. Run dftimewolf <RECIPE_NAME> -h for details.')
for registered_recipe in config.Config.get_registered_recipes():
recipe, recipe_args, documentation = registered_recipe
subparser = subparsers.add_parser(
recipe['name'], description='{0:s}'.format(documentation))
subparser.set_defaults(recipe=recipe)
for switch, help_text, default in recipe_args:
subparser.add_argument(switch, help=help_text, default=default)
args = parser.parse_args()
recipe = args.recipe
console_out = dftw_utils.DFTimewolfConsoleOutput(sender='DFTimewolfCli',
verbose=True)
# Thread all collectors.
console_out.StdOut('Collectors:')
for collector in recipe['collectors']:
console_out.StdOut(' {0:s}'.format(collector['name']))
collector_objects = []
for collector in recipe['collectors']:
new_args = dftw_utils.import_args_from_dict(collector['args'],
vars(args), config.Config)
collector_cls = config.Config.get_collector(collector['name'])
collector_objects.extend(collector_cls.launch_collector(**new_args))
# global_errors will contain any errors generated along the way by collectors,
# producers or exporters.
global_errors = []
# Wait for collectors to finish and collect output.
collector_output = []
for collector_obj in collector_objects:
collector_obj.join()
collector_output.extend(collector_obj.results)
if collector_obj.errors:
#TODO(tomchop): Add name attributes in module objects
error = (collector_obj.__class__.__name__,
", ".join(collector_obj.errors))
global_errors.append(error)
console_out.StdErr("ERROR:{0:s}:{1:s}\n".format(*error))
if recipe['processors']:
# Thread processors.
console_out.StdOut('Processors:')
for processor in recipe['processors']:
console_out.StdOut(' {0:s}'.format(processor['name']))
processor_objs = []
for processor in recipe['processors']:
new_args = dftw_utils.import_args_from_dict(
processor['args'], vars(args), config.Config)
new_args['collector_output'] = collector_output
processor_class = config.Config.get_processor(processor['name'])
processor_objs.extend(processor_class.launch_processor(**new_args))
# Wait for processors to finish and collect output
processor_output = []
for processor in processor_objs:
processor.join()
processor_output.extend(processor.output)
if processor.errors:
# Note: Should we fail if modules produce errors, or is warning the user
# enough?
# TODO(tomchop): Add name attributes in module objects.
error = (processor.__class__.__name__,
", ".join(processor.errors))
global_errors.append(error)
console_out.StdErr("ERROR:{0:s}:{1:s}\n".format(*error))
else:
processor_output = collector_output
# Thread all exporters.
if recipe['exporters']:
console_out.StdOut('Exporters:')
for exporter in recipe['exporters']:
console_out.StdOut(' {0:s}'.format(exporter['name']))
exporter_objs = []
for exporter in recipe['exporters']:
new_args = dftw_utils.import_args_from_dict(
exporter['args'], vars(args), config.Config)
new_args['processor_output'] = processor_output
exporter_class = config.Config.get_exporter(exporter['name'])
exporter_objs.extend(exporter_class.launch_exporter(**new_args))
# Wait for exporters to finish.
exporter_output = []
for exporter in exporter_objs:
exporter.join()
exporter_output.extend(exporter.output)
if exporter.errors:
#TODO(tomchop): Add name attributes in module objects
error = (exporter.__class__.__name__,
", ".join(exporter.errors))
global_errors.append(error)
console_out.StdErr("ERROR:{0:s}:{1:s}\n".format(*error))
else:
exporter_output = processor_output
if not global_errors:
console_out.StdOut('Recipe {0:s} executed successfully'.format(
recipe['name']))
else:
console_out.StdOut('Recipe {0:s} executed with {1:d} errors:'.format(
recipe['name'], len(global_errors)))
for error in global_errors:
console_out.StdOut(' {0:s}: {1:s}'.format(*error))