drone-cleanup-agents/cleanupagents/Cli.py

#!/usr/bin/env python3
"""Main program."""

import argparse
import json
import os
import pickle  # nosec
import shutil
from collections import defaultdict

from cleanupagents import __version__
from cleanupagents.Logging import SingleLog
from cleanupagents.Utils import humanize
from cleanupagents.Utils import normalize_path
from cleanupagents.Utils import run_command


class AgentCleanup:

    def __init__(self):
        self.args = self._cli_args()
        self.config = self._config()
        self.log = SingleLog(logfile=self.config["logfile"])
        self.logger = self.log.logger
        self.run()

    def _cli_args(self):
        parser = argparse.ArgumentParser(
            description=("Cleanup outdated and faulty drone agents"))
        parser.add_argument("-v", dest="log_level", action="append_const", const=-1,
                            help="increase log level")
        parser.add_argument("-d", "--dry-run", dest="dry_run", action="store_true", default=False,
                            help="dry run without modifications")
        parser.add_argument("-q", dest="log_level", action="append_const",
                            const=1, help="decrease log level")
        parser.add_argument("--version", action="version", version="%(prog)s {}".format(__version__))

        return parser.parse_args().__dict__

    def _config(self):
        config = defaultdict(dict)
        logfile = os.environ.get("DRONE_LOGFILE", "/var/log/drone-agents.log")
        autoscaler = os.environ.get("DRONE_SCALER")
        binary = normalize_path(os.environ.get("DRONE_BIN", shutil.which("drone")) or "./")

        config["required"] = ["drone_server", "drone_token", "drone_scaler"]
        config["logfile"] = normalize_path(logfile)
        config["drone_server"] = os.environ.get("DRONE_SERVER")
        config["drone_token"] = os.environ.get("DRONE_TOKEN")
        config["log_level"] = os.environ.get("DRONE_LOGLEVEL", "WARNING")
        config["statefile"] = "/tmp/droneclean.pkl"  # nosec
        config["pending"] = defaultdict(dict, {})

        if os.path.isfile(binary):
            config["drone_bin"] = binary

        if autoscaler:
            config["drone_scaler"] = [x.strip() for x in autoscaler.split(",")]

        if not os.path.exists(os.path.dirname(config["logfile"])):
            os.makedirs(os.path.dirname(config["logfile"]))

        config["dry_run"] = self.args["dry_run"]

        # Override correct log level from argparse
        levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
        log_level = levels.index(config["log_level"])
        if self.args["log_level"]:
            for adjustment in self.args["log_level"]:
                log_level = min(len(levels) - 1, max(log_level + adjustment, 0))
            config["log_level"] = levels[log_level]

        if os.path.isfile(config["statefile"]):
            with open(config["statefile"], "rb") as json_file:
                config["pending"] = pickle.load(json_file)  # nosec

        return config

    def _clean(self):
        pending = self.config["pending"]
        control = defaultdict(dict)
        dry = self.config["dry_run"]
        dryrun_msg = ""

        self.logger.debug("Pending object dump: {}".format(pending))
        if dry:
            dryrun_msg = "[DRYRUN] "

        for scaler in self.config["drone_scaler"]:
            error = []
            self.logger.info("{dryrun_msg}Cleanup agents for scaler '{scaler}'".format(
                scaler=scaler, dryrun_msg=dryrun_msg))

            res = run_command("{binfile} -s {server} -t {token} --autoscaler {scaler} server ls --format '{{{{ . | jsonify }}}}'".format(
                server=self.config["drone_server"], token=self.config["drone_token"], scaler=scaler, binfile=self.config["drone_bin"]))

            if res.returncode > 0:
                self.log.sysexit_with_message("Command error:\n{}".format(humanize(res.stdout)))

            for line in res.stdout.splitlines():
                try:
                    obj = json.loads(line)
                except json.decoder.JSONDecodeError:
                    self.logger.debug("Unable to parse line: {}".format(line))
                    continue

                if obj["state"] == "error":
                    error.append(obj["name"])

            for e in error:
                force = ""
                force_msg = ""

                if pending[e] and pending[e] < 2:
                    control[e] = pending[e] + 1
                elif pending[e] and pending[e] > 1:
                    force = "--force"
                    force_msg = "\nwill '--force' now"
                else:
                    control[e] = 1

                self.logger.info("Stopping '{agent}' ({triage}/3) {force}".format(
                    agent=e, triage=control[e] or 3, force=force_msg))

                if not dry:
                    res = run_command("{binfile} -s {server} -t {token} --autoscaler {scaler} server destroy {force} {agent}".format(
                        server=self.config["drone_server"], token=self.config["drone_token"], scaler=scaler, agent=e, force=force,
                        binfile=self.config["drone_bin"]))

                    if res.returncode > 0 and "client error 404" not in humanize(res.stdout):
                        self.log.sysexit_with_message("Command error:\n{}".format(humanize(res.stdout)))

        if not dry:
            with open(self.config["statefile"], "wb") as json_file:
                pickle.dump(control, json_file)  # nosec

    def run(self):
        try:
            self.log.set_level(self.config["log_level"])
        except ValueError as e:
            self.log.sysexit_with_message("Can not update log config.\n{}".format(str(e)))

        for c in self.config["required"]:
            if not self.config[c]:
                self.log.sysexit_with_message("Environment variable '{}' is required but empty or unset".format(c.upper()))

        if not self.config["drone_bin"]:
            self.log.sysexit_with_message("Drone binary not found in PATH or not defined by 'DRONE_BIN'")

        self._clean()