151 lines
6.0 KiB
Python
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()
|