shepherd-agent/shepherd/core.py

"""
Core shepherd module, tying together main service functionality.
"""


import os
from pathlib import Path
from datetime import datetime
import toml
import logging
import click
from copy import deepcopy

from . import scheduler
from . import config
from . import plugin
from . import control


# Future implementations of checking config differences should be done on
# the hash of the nested conf dict, so comments shouldn't affect this.

# save old config to somewhere in the shepherd root dir - probably need to
# implement a TOML writer in the config module.

# later on, there's going to be an issue with a new config being applied
# remotely, then the system restarting, and an old edit in /boot being
# applied over the top...
# Fix this by saving the working config to /boot when new config applied
# remotely.


# Relative pathnames here are all relative to "root_dir"
def define_core_config(confdef):
    """
    Defines the config definition by populating the ConfigDefinition passed in ``confdef`` - the same pattern plugins use
    """
    confdef.add_def("name", config.StringDef(
        helptext="Identifying name for this device"))

    confdef.add_def("hostname",
                    config.StringDef(default="", optional=True,
                                     helptext="If set, changes the system hostname"))
    confdef.add_def("plugin_dir",
                    config.StringDef(default="~/shepherd-plugins", optional=True,
                                     helptext="Optional directory for Shepherd to look for plugins in."))
    confdef.add_def("root_dir",
                    config.StringDef(default="~/shepherd", optional=True,
                                     helptext="Operating directory for shepherd to place working files."))
    confdef.add_def("custom_config_path",
                    config.StringDef(optional=True,
                                     helptext="Path to custom config layer TOML file."))
    confdef.add_def("generated_config_path",
                    config.StringDef(default="shepherd-generated.toml", optional=True,
                                     helptext="Path to custom file Shepherd will generate to show compiled config that was used and any errors in validation."))

    confdef.add_def("control_server", config.StringDef())
    confdef.add_def("control_api_key", config.StringDef())


def resolve_core_conf_paths(core_conf):
    """
    Set the cwd to ``root_dir`` and resolve other core config paths relative to that.
    Also expands out any "~" user characters.
    """
    core_conf["root_dir"] = str(Path(core_conf["root_dir"]).expanduser().resolve())
    os.chdir(core_conf["root_dir"])
    core_conf["plugin_dir"] = str(Path(core_conf["plugin_dir"]).expanduser().resolve())
    core_conf["custom_config_path"] = str(
        Path(core_conf["custom_config_path"]).expanduser().resolve())
    core_conf["generated_config_path"] = str(
        Path(core_conf["generated_config_path"]).expanduser().resolve())


def load_config_layer(confman, config_source):
    """
    Load a config layer, find the necessary plugin classes, then validate it.
    """
    # Load in config layer
    confman.load_overlay(config_source)

    # Get the core config so we can find the plugin directory
    core_config = confman.get_config_bundle("shepherd")
    plugin_dir = core_config["plugin_dir"]

    # List other table names to get plugins we need to load
    plugin_names = confman.get_bundle_names()
    plugin_names.remove("shepherd")

    # Load plugins to get their conf defs
    plugin_classes = plugin.find_plugins(plugin_names, plugin_dir)
    for plugin_name, plugin_class in plugin_classes.items():
        new_conf_def = config.ConfDefinition()
        plugin_class.define_config(new_conf_def)
        confman.add_confdef(plugin_name, new_conf_def)

    # Get plugin configs
    plugin_configs = confman.get_config_bundles(plugin_classes.keys())
    return (core_config, plugin_classes, plugin_configs)


def compile_config(default_config_path):
    """
    Run through the process of assembling the various config layers, falling back to working
    ones where necessary. Also gathers needed plugin classes in the process.
    """

    # Create core confdef and populate it
    core_confdef = config.ConfDefinition()
    define_core_config(core_confdef)

    confman = config.ConfigManager()
    confman.add_confdef("shepherd", core_confdef)

    # Default config. This must validate to continue.
    try:
        core_conf, plugin_classes, plugin_configs = load_config_layer(
            confman, Path(default_config_path).expanduser())
        logging.info(F"Loaded default config from {default_config_path}")
    except:
        logging.error(F"Failed to load default config from {default_config_path}")
        raise

    # Resolve and freeze local install paths that shouldn't be changed or affect loading custom config
    confman.freeze_value("shepherd", "root_dir")
    confman.freeze_value("shepherd", "plugin_dir")
    confman.freeze_value("shepherd", "custom_config_path")
    confman.freeze_value("shepherd", "generated_config_path")
    resolve_core_conf_paths(core_conf)

    # Pull out custom config path and save current good root_config
    custom_config_path = core_conf["custom_config_path"]
    saved_root_config = deepcopy(confman.root_config)

    # Custom config layer. If this fails, maintain default config but continue on to Control layer
    try:
        core_conf, plugin_classes, plugin_configs = load_config_layer(
            confman, custom_config_path)
        logging.info(F"Loaded custom config from {custom_config_path}")
    except Exception as e:
        logging.error(
            F"Failed to load custom config from {custom_config_path}. Falling back to default config.", exc_info=e)
        confman.root_config = saved_root_config

    # Freeze Shepherd Control related config.
    confman.freeze_value("shepherd", "control_server")
    confman.freeze_value("shepherd", "control_api_key")
    resolve_core_conf_paths(core_conf)

    # Save current good root_config
    saved_root_config = deepcopy(confman.root_config)

    # Shepherd Control config layer. If this fails, maintain current local config.
    try:
        control_config = control.get_config(core_conf["root_dir"])
        try:
            core_conf, plugin_classes, plugin_configs = load_config_layer(
                confman, control_config)
            logging.info(F"Loaded cached Shepherd Control config")
        except Exception as e:
            logging.error(
                F"Failed to load cached Shepherd Control config. Falling back to local config.", exc_info=e)
            confman.root_config = saved_root_config
    except:
        logging.warning("No cached Shepherd Control config available.")

    confman.dump_to_file(core_conf["generated_config_path"])

    return core_conf, plugin_classes, plugin_configs


@click.group(invoke_without_command = True)
#help="Path to default config TOML file"

@click.argument('default_config', default="shepherd-default.toml", type=click.Path())
@click.pass_context
def cli(ctx, default_config):
    """
    Core service. Expects the default config to be set as an argument.
    """
    #argparser = argparse.ArgumentParser(description="Keep track of a mob "
    #                                                "of roaming Pis")
    #argparser.add_argument("configfile", nargs='?', metavar="configfile",
    #                       help="Path to configfile", default="shepherd.toml")
    #argparser.add_argument(
    #    '-e', '--noedit', help="Disable the editable config temporarily", action="store_true", default=False)
    #argparser.add_argument("-t", "--test", help="Test and interface function of the from 'plugin:function'",
    #                       default=None)

    #args = argparser.parse_args()

    core_conf, plugin_classes, plugin_configs = compile_config(default_config)

    if args.test is None:
        control.init_control(core_conf, plugin_configs)

    scheduler.init_scheduler(core_conf)
    plugin.init_plugins(plugin_classes, plugin_configs, core_conf)
    scheduler.restore_jobs()

    print(str(datetime.now()))

    if ctx.invoked_subcommand is not None:
        return

    print('Press Ctrl+{0} to exit'.format('Break' if os.name == 'nt' else 'C'))
    try:
        scheduler.start()
    except (KeyboardInterrupt, SystemExit):
        pass

@click.argument('plugin_function')
@cli.command()
def test():
    if args.test is not None:
        (test_plugin, test_func) = args.test.split(':')
        func = getattr(shepherd.plugin.plugin_functions[test_plugin], test_func)
        print(func())
        return