#!/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()