Merge pull request #11 from scallister/separate-cli-logic

Separate cli logic

README.md

@@ -82,11 +82,14 @@ Several ways to use your own plugins with Fossor
 3. Build and wrap Fossor
    - Fossor is both a CLI and a Python library. Import the Fossor engine into your own tool, and add any local plugins you've made by passing a module or a filesystem path to add_plugins().
    ```python
-   from fossor.engine import Fossor
-   from fossor.cli import main
-   f = Fossor()
-   f.add_plugins('/opt/fossor')
-   main()
+    from fossor.engine import Fossor
+    from fossor.cli import main
+
+    @fossor_cli_flags
+    def main(context, **kwargs):
+        f = Fossor()
+
+        f.run(**kwargs)
    ```
 
 ### Example Plugins

fossor/cli.py

@@ -11,87 +11,169 @@
 import logging
 import setproctitle
 import parsedatetime as pdt
+
+from functools import wraps
 from datetime import datetime, timedelta
+
 from fossor.engine import Fossor
 
 default_plugin_dir = '/opt/fossor'
 
 
-@click.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True, help_option_names=['-h', '--help']))  # noqa: C901
-@click.pass_context
-@click.option('-p', '--pid', help='Pid to investigate.')
-@click.option('--product', help='Product to investigate.')
-@click.option('--plugin-dir', default=default_plugin_dir, help=f'Import all plugins from this directory. Default dir: {default_plugin_dir}')
-@click.option('--hours', default=1, help='Sets start-time to X hours ago, defaults to 24 hours. Plugins may optionally implement and use this.')
-@click.option('-r', '--report', default='StdOut', help='Report Plugin to run.')
-@click.option('--start-time', default='', help='Plugins may optionally implement and use this.')
-@click.option('--end-time', default='', help='Plugins may optionally implement and use this. Defaults to now.')
-@click.option('-t', '--time-out', default=600, help='Default timeout for plugins.')
-@click.option('-d', '--debug', is_flag=True)
-@click.option('-v', '--verbose', is_flag=True)
-@click.option('--no-truncate', is_flag=True)
-def main(context, pid, product, plugin_dir, hours, report, start_time, end_time, time_out, debug, verbose, no_truncate):
+def setup_logging(ctx, param, value):
+    level = logging.CRITICAL
+    if value:
+        level = logging.DEBUG
+    logging.getLogger("requests").setLevel(logging.WARNING)
+    logging.basicConfig(stream=sys.stderr, level=level)
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level)
+    return value
+
+
+def get_timestamp(value):
+    '''Return timestamp from a human readable date'''
+
+    if not value:
+        return
+
+    # If already a timestamp return a float
+    try:
+        timestamp = float(value)
+        return timestamp
+    except ValueError:
+        pass
+
+    # Convert human readable string to timestamp
+    cal = pdt.Calendar()
+    timestamp = (cal.parseDT(value)[0]).timestamp()
+    return timestamp
+
+
+def set_start_time(ctx, param, value):
+    if 'start_time' in ctx.params and ctx.params['start_time'] is not None:
+        return ctx.params['start_time']
+    return get_timestamp(value)
+
+
+def set_end_time(ctx, param, value):
+    return get_timestamp(value)
+
+
+def set_relative_start_time(ctx, param, value):
+    '''Overrides start time relative to number of {value} hours'''
+    hours = value
+
+    start_time = (datetime.now() - timedelta(hours=hours)).timestamp()
+    ctx.params['start_time'] = start_time
+
+
+class CsvList(click.ParamType):
+    name = 'Csv-List'
+
+    def convert(self, value, param, ctx):
+        result = value.split(',')
+        return result
+
+
+def fossor_cli_flags(f):
+    '''Add default Fossor CLI flags'''
+    # Flags will appear in reverse order of how they  are listed here:
+    f = add_dynamic_args(f)  # Must be applied after all other click options since this requires click's context object to be passed
+
+    # Add normal flags
+    csv_list = CsvList()
+    f = click.option('--black-list', 'blacklist', type=csv_list, help='Do not run these plugins.')(f)
+    f = click.option('--white-list', 'whitelist', type=csv_list, help='Only run these plugins.')(f)
+    f = click.option('--truncate/--no-truncate', 'truncate', show_default=True, default=True, is_flag=True)(f)
+    f = click.option('-v', '--verbose', is_flag=True)(f)
+    f = click.option('-d', '--debug', is_flag=True, callback=setup_logging)(f)
+    f = click.option('-t', '--time-out', 'timeout', show_default=True, default=600, help='Default timeout for plugins.')(f)
+    f = click.option('--end-time', callback=set_end_time, help='Plugins may optionally implement and use this. Defaults to now.')(f)
+    f = click.option('--start-time', callback=set_start_time, help='Plugins may optionally implement and use this.')(f)
+    f = click.option('-r', '--report', type=click.STRING, show_default=True, default='StdOut', help='Report Plugin to run.')(f)
+    f = click.option('--hours', type=click.INT, default=24, show_default=True, callback=set_relative_start_time,
+                     help='Sets start-time to X hours ago. Plugins may optionally implement and use this.')(f)
+    f = click.option('--plugin-dir', default=default_plugin_dir, show_default=True, help=f'Import all plugins from this directory.')(f)
+    f = click.option('-p', '--pid', type=click.INT, help='Pid to investigate.')(f)
+
+    # Required for parsing dynamics arguments
+    f = click.pass_context(f)
+    f = click.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True, help_option_names=['-h', '--help']))(f)
+    return f
+
+
+def add_dynamic_args(function):
+    @wraps(function)
+    def wrapper(*args, **kwargs):
+        context = args[0]
+        dynamic_kwargs = parse_dynamic_args(context.args)
+        kwargs = {**kwargs, **dynamic_kwargs}  # Merge these dict objects
+        function(*args, **kwargs)
+    return wrapper
+
+
+def set_process_title(arg=None):
+    '''
+    Decorator for setting the process title
+    Can be applied in two ways:
+    @set_process_title
+    @set_process_title('new_process_title')
+
+    Defaults to 'fossor'
+
+    '''
+    title = 'fossor'
+    if type(arg) == str:
+        title = arg
+
+    def real_decorator(function):
+        @wraps(function)
+        def wrapper(*args, **kwargs):
+            setproctitle.setproctitle(title + ' ' + ' '.join(sys.argv[1:]))
+            function(*args, **kwargs)
+        return wrapper
+    # If called as @set_process_title
+    if callable(arg):
+        return real_decorator(arg)
+    # If called as @set_process_title('new_title')
+    return real_decorator
+
+
+@fossor_cli_flags
+@set_process_title
+def main(context, **kwargs):
     """Fossor is a plugin oriented tool for automating the investigation of broken hosts and services.
 
     \b
     Advanced:
     Multiple additional internal variables may be passed on the command line in this format: name="value".
     This is intended for testing or automation.
     """  # \b makes click library honor paragraphs
-    setproctitle.setproctitle('fossor ' + ' '.join(sys.argv[1:]))
-    if debug:
-        log_level = logging.DEBUG
-        verbose = True
-    else:
-        log_level = logging.CRITICAL
-        logging.getLogger("requests").setLevel(logging.WARNING)
-    logging.basicConfig(stream=sys.stdout, level=log_level)
-    log = logging.getLogger(__name__)
+
     f = Fossor()
+    f.run(**kwargs)
+
 
-    # Add pre-defined args
-    f.add_variable('timeout', time_out)
-    if product:
-        f.add_variable('product', product)
-    if pid:
-        f.add_variable('pid', pid)
-    if no_truncate:
-        f.add_variable('truncate', False)
-    f.add_variable('verbose', verbose)
-
-    # Add dynamic args
-    log.debug(f"Remaining Arguments that will be used for dynamic args now that hardcoded cli flags have been removed: {context.args}")
+def parse_dynamic_args(context_args):
+    '''Return a dict of arguments with their values. Intended to parse leftover arguments from click's context.arg list'''
+    log = logging.getLogger(__name__)
+    log.debug(f"Remaining Arguments that will be used for dynamic args now that hardcoded cli flags have been removed: {context_args}")
     format_help_message = "Please use --help to see valid flags. Or use the name=value method for setting internal variables."
-    for arg in context.args:
+
+    kwargs = {}
+    for arg in context_args:
         a = arg.strip()
         if a.startswith('-'):
             raise ValueError(f"Argument: {arg} is invalid {format_help_message}.")
         try:
             name, value = a.split('=')
-            f.add_variable(name, value)
+            kwargs[name] = value
             log.debug(f"Using dynamic variable {name}={value}")
         except Exception as e:
             log.exception(f"Failed to add argument: \"{arg}\". {format_help_message} Exception was: {e}")
             raise e
-
-    # Get start/end times
-    cal = pdt.Calendar()
-    if start_time:
-        start_time = (cal.parseDT(start_time)[0]).timestamp()
-    else:
-        start_time = (datetime.now() - timedelta(hours=hours)).timestamp()
-    f.add_variable('start_time', str(start_time))
-    if end_time:
-        end_time = (cal.parseDT(end_time)[0]).timestamp()
-    else:
-        end_time = datetime.now().timestamp()
-    f.add_variable('end_time', str(end_time))
-
-    # Import plugin dir if it exists
-    f.add_plugins(plugin_dir)
-
-    log.debug("Starting fossor")
-    return f.run(report=report)
+    return kwargs
 
 
 if __name__ == '__main__':

fossor/engine.py

@@ -102,12 +102,16 @@ def _import_submodules_by_path(self, path: str) -> set:
     def add_plugins(self, source=fossor):
         '''
         Recursively return a dict of plugins (classes) that inherit from the given parent_class) from within that source.
-        source accepts a path as a string or a module.
+        source accepts either a python module or a filesystem path as a string.
         '''
+        if source is None:
+            return
+
         if type(source) == str:
             modules = self._import_submodules_by_path(path=source)
         else:
             modules = self._import_submodules_by_module(source)
+
         for module in modules:
             for obj_name, obj in module.__dict__.items():
                 # Get objects from each module that look like plugins for the Check abstract class
@@ -282,6 +286,11 @@ def _convert_simple_type(self, value):
 
         return value
 
+    def add_variables(self, **kwargs):
+        for name, value in kwargs.items():
+            self.add_variable(name=name, value=value)
+        return True
+
     def add_variable(self, name, value):
         '''Adds variable, converts numbers and booleans to their types'''
         if name in self.variables:
@@ -295,21 +304,35 @@ def add_variable(self, name, value):
         self.log.debug("Added variable {name} with value of {value} (type={type})".format(name=name, value=value, type=type(value)))
         return True
 
-    def run(self, report='StdOut', check_whitelist=[], check_blacklist=[], variable_plugin_whitelist=[], variable_plugin_blacklist=[]):
+    def _process_whitelist(self, whitelist):
+        if whitelist:
+            whitelist = [name.casefold() for name in whitelist]
+            self.variable_plugins = {plugin for plugin in self.variable_plugins if plugin.get_name().casefold() in whitelist}
+            self.check_plugins = {plugin for plugin in self.check_plugins if plugin.get_name().casefold() in whitelist}
+
+    def _process_blacklist(self, blacklist):
+        if blacklist:
+            blacklist = [name.casefold() for name in blacklist]
+            self.variable_plugins = {plugin for plugin in self.variable_plugins if plugin.get_name().casefold() not in blacklist}
+            self.check_plugins = {plugin for plugin in self.check_plugins if plugin.get_name().casefold() not in blacklist}
+
+    def run(self, report='StdOut', **kwargs):
         '''Runs Fossor with the given report. Method returns a string of the report output.'''
+        self.log.debug("Starting Fossor")
+
+        self._process_whitelist(kwargs.get('whitelist'))
+        self._process_blacklist(kwargs.get('blacklist'))
+
+        # Add kwargs as variables
+        self.add_variables(**kwargs)
+
+        # Add directory plugins
+        self.add_plugins(kwargs.get('plugin_dir'))
 
         # Gather Variables
-        if variable_plugin_whitelist:
-            self.variable_plugins = {plugin for plugin in self.variable_plugins if plugin.get_name() in variable_plugin_whitelist}
-        for variable_plugin in variable_plugin_blacklist:
-            self.variable_plugins = {plugin for plugin in self.variable_plugins if plugin.get_name() not in variable_plugin_blacklist}
         self.get_variables()
 
-        # Run Checks
-        if check_whitelist:
-            self.check_plugins = {plugin for plugin in self.check_plugins if plugin.get_name() in check_whitelist}
-        if check_blacklist:
-            self.check_plugins = {plugin for plugin in self.check_plugins if plugin.get_name() not in check_blacklist}
+        # Run checks
         output_queue = self._run_plugins_parallel(self.check_plugins)
 
         # Run Report

test/test_fossor.py

@@ -33,41 +33,25 @@ def test_convert_simple_type():
     assert type(f._convert_simple_type('foo')) == str
 
 
-def test_variable_plugin_whitelist():
+def test_plugin_whitelist():
     f = Fossor()
     f.add_variable('timeout', 1)
     assert len(f.variable_plugins) > 2
-    whitelist = ['Hostname', 'PidExe']
-    f.run(variable_plugin_whitelist=whitelist)
+    assert len(f.check_plugins) > 2
+    whitelist = ['BuddyInfo', 'LoadAvg', 'hostname', 'pidexe']  # Should be case insensitive
+    f.run(whitelist=whitelist)
     assert len(f.variable_plugins) == 2
+    assert len(f.check_plugins) == 2
 
 
-def test_variable_plugin_blacklist():
+def test_plugin_blacklist():
     f = Fossor()
     f.add_variable('timeout', 1)
     assert fossor.variables.hostname.Hostname in f.variable_plugins
-    blacklist = ['Hostname']
-    f.run(variable_plugin_blacklist=blacklist)
+    blacklist = ['hostname', 'BuddyInfo', 'LoadAvg']  # Should be case insensitive
+    f.run(blacklist=blacklist)
     assert 'Hostname' not in f.variable_plugins
+    assert 'BuddyInfo' not in f.check_plugins
+    assert 'LoadAvg' not in f.check_plugins
     assert len(f.variable_plugins) > 0
-
-
-def test_check_plugin_whitelist():
-    f = Fossor()
-    f.add_variable('timeout', 1)
-    assert len(f.check_plugins) > 2
-    whitelist = ['BuddyInfo', 'LoadAvg']
-    f.run(check_whitelist=whitelist)
-    assert len(f.check_plugins) == 2
-    assert fossor.checks.buddyinfo.BuddyInfo in f.check_plugins
-
-
-def test_check_plugin_blacklist():
-    f = Fossor()
-    f.add_variable('timeout', 1)
-    assert len(f.check_plugins) > 2
-    assert fossor.checks.buddyinfo.BuddyInfo in f.check_plugins
-    blacklist = ['BuddyInfo', 'LoadAvg']
-    f.run(check_blacklist=blacklist)
-    assert fossor.checks.buddyinfo.BuddyInfo not in f.check_plugins
-    assert len(f.check_plugins) > 1
+    assert len(f.check_plugins) > 0