This repository has been archived on 2023-12-07. You can view files and clone it, but cannot push or open issues or pull requests.
drone-cleanup-agents/cleanupagents/Cli.py

151 lines
6.0 KiB
Python

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