shepherd-agent/shepherd/core.py

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


import os
from pathlib import Path
from datetime import datetime
import logging
import pkg_resources
import click

from configspec import *

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


@click.group(invoke_without_command=True)
#help="Path to default config TOML file"
@click.argument('default_config_path', required=False, type=click.Path())
@click.pass_context
def cli(ctx, default_config_path):
    """
    Core service. Expects the default config to be set as an argument.
    """
    version_text = pkg_resources.get_distribution("shepherd")
    logging.info(F"Initialising Shepherd Agent {version_text}")
    # 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()

    confman = ConfigManager()

    plugin_classes = compile_config_and_get_plugins(confman, default_config_path)

    control.init_control(core_conf, plugin_configs)

    scheduler.init_scheduler(core_conf)
    plugin.init_plugins(plugin_classes, plugin_configs, core_config)
    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


def compile_config_and_get_plugins(confman, 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 = ConfigSpecification()
    define_core_config(core_confdef)
    confman.add_confspec("shepherd", core_confdef)

    # ====Default Local Config Layer====
    # This must validate to continue.
    default_config_path = Path(default_config_path).expanduser()
    try:
        plugin_classes = load_config_layer_and_plugins(confman, default_config_path)
        logging.info(F"Loaded default config layer from {default_config_path}")
    except Exception:
        logging.error(F"Failed to load default config from {default_config_path}")
        raise

    # Resolve and freeze local install paths that shouldn't be changed from default config
    core_conf = confman.get_config_bundle("shepherd")
    resolve_core_conf_paths(core_conf)
    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")

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

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

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

    # Save current good config
    confman.save_fallback()

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

    logging.debug("Compiled config: %s", confman.root_config)
    confman.dump_to_file(core_conf["generated_config_path"])

    return plugin_classes


# Relative pathnames here are all relative to "root_dir"
def define_core_config(confspec):
    """
    Defines the config definition by populating the ConfigSpecification passed in
    ``confspec`` - the same pattern plugins use
    """
    confspec.add_specs([
        ("name", StringSpec(helptext="Identifying name for this device")),
        ("hostname", StringSpec(default="", optional=True, helptext="If set, changes the system"
                                " hostname")),
        ("plugin_dir", StringSpec(default="~/shepherd-plugins", optional=True,
                                  helptext="Optional directory for Shepherd to look for"
                                  " plugins in.")),
        ("root_dir", StringSpec(default="~/shepherd", optional=True,
                                helptext="Operating directory for shepherd to place"
                                " working files.")),
        ("custom_config_path", StringSpec(optional=True, helptext="Path to custom config"
                                          " layer TOML file.")),
        ("generated_config_path", StringSpec(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."))
    ])

    confspec.add_spec("control_server", StringSpec())
    confspec.add_spec("control_api_key", StringSpec())


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_and_plugins(confman, config_source):
    """
    Load a config layer, find the necessary plugin classes, then validate it.
    If this succeeds, the returned dict of plugin classes will directly match
    the bundle names in the config manager.
    """
    # 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 bundle names to get plugins we need to load
    plugin_names = confman.get_bundle_names()
    plugin_names.remove("shepherd")

    # Load plugins to get their config specifications
    plugin_classes = plugin.find_plugins(plugin_names, plugin_dir)
    for plugin_name, plugin_class in plugin_classes.items():
        new_conf_spec = ConfigSpecification()
        plugin_class.define_config(new_conf_spec)
        confman.add_confspec(plugin_name, new_conf_spec)

    # Validate all plugin configs
    confman.validate_bundles()

    return plugin_classes