diff --git a/sausage b/sausage index 3a79d68..8f0e755 100755 --- a/sausage +++ b/sausage @@ -1,164 +1,170 @@ #!/usr/bin/python3 +# © All rights reserved. ECOLE POLYTECHNIQUE FEDERALE DE LAUSANNE, +# Switzerland +# SCITAS - Scientific IT and Application Support, 2021 +# See the LICENSE.txt file for more details. + + import configparser import requests import json import argparse import getpass from datetime import date from datetime import datetime def valid_date(date): try: validate = datetime.strptime(date, "%Y-%m-%d") return validate except ValueError: msg = "Not a valid date: '{0}', YYYY-MM-DD expected.".format(date) raise argparse.ArgumentTypeError(msg) class AppArgs(object): def __init__(self): self.response = {} self.parser = argparse.ArgumentParser(prog = 'Sausage', description = 'SCITAS Account Usage.', formatter_class=argparse.ArgumentDefaultsHelpFormatter) self.add_args() def add_args(self): self.parser.add_argument('-u','--user', help='If not provided whoami is considered') self.parser.add_argument('-a','--account', help='all users from an account are printed') self.parser.add_argument('-s','--start', help='Start date - format YYYY-MM-DD', type=valid_date) self.parser.add_argument('-e','--end', help='End date - format YYYY-MM-DD', type=valid_date) args = self.parser.parse_args() if args.start and args.end is None: self.parser.error("range requires both dates (--start and --end)") if args.end: if args.start is None: self.parser.error("range requires both dates (--start and --end)") if args.end < args.start: self.parser.error("start date must be earlier than end date") self.response = { "user": args.user, "account": args.account, "start": args.start, "end": args.end } class GetData(object): def __init__(self,conf,args): self.cfg_parser = configparser.ConfigParser() self.cfg_parser.read(conf) self.server = self.cfg_parser.get("server", "url") + ":" + self.cfg_parser.get("server", "port") self.message = [] self.vseparator = " | " self.hseparator = "---------------------------------------------------------" if args["start"]: self.firstday = str(args["start"].date()) else: self.firstday = str(date.today().replace(day=1)) if args["end"]: self.lastday = str(args["end"].date()) else: self.lastday = str(date.today()) self.request(args) def printbox(self): if self.format == 1: head_b = "" title = "ACCOUNT: " elif self.format == 2: head_b = "Account" + self.vseparator title = "USER: " head_a = "Cluster" + self.vseparator head_c = "Core-hours" + self.vseparator head_d = "CHF" data = self.response.json() carbon = 0 element = data["name"] self.message.append(self.hseparator) self.message.append(title + element) self.message.append("Global usage from " + self.firstday + " to " + self.lastday) self.message.append(self.hseparator) self.message.append(f"{head_a : >12}{head_b : >20}{head_c : >13}{head_d : >10}") self.message.append(self.hseparator) for key, value in data.items(): if isinstance(value, dict): if self.format == 1: head_a = key + self.vseparator chf = "{:.2f}".format(value['chf']) time = "{:.2f}".format(value['time']) carbon += value['co2'] head_c = str(time) + self.vseparator head_d = str(chf) self.message.append(f"{head_a : >12}{head_b : >20}{head_c : >13}{head_d : >10}") elif self.format == 2: head_b = key + self.vseparator for k,v in value.items(): head_a = k + self.vseparator if v['chf'] > 0.00999: chf = "{:.2f}".format(v['chf']) else: chf = "-" if v['time'] > 0.00999: time = "{:.2f}".format(v['time']) else: time = "-" carbon += v['co2'] head_c = str(time) + self.vseparator head_d = str(chf) self.message.append(f"{head_a : >12}{head_b : >20}{head_c : >13}{head_d : >10}") carbon = "{:.2f}".format(carbon / 1000) self.message.append(self.hseparator) self.message.append("Estimated carbon footprint: " + str(carbon) + " kg/co2 equivalent") def request(self,args): # First case : without arguments if all(v == None for k,v in args.items()): self.response = requests.get(self.server + '/user/' + getpass.getuser()) self.format = 2 # Second case : only with user elif all(v == None for v in [args["account"], args["start"], args["end"]]) and args["user"]: self.response = requests.get(self.server + '/user/' + args["user"]) self.format = 2 # Third case : only with account elif all(v == None for v in [args["user"], args["start"], args["end"]]) and args["account"]: self.response = requests.get(self.server + '/account/' + args["account"]) self.format = 1 # Fourth case : with user and account (without range) elif all(v != None for v in [args["user"], args["account"]]) and all(v == None for v in [args["start"], args["end"]]): self.response = requests.get(self.server + '/account/' + args["account"] + '/' + args["user"]) self.format = 2 # Fifth case : with range elif all(v != None for v in [args["start"], args["end"]]) and (args["user"] or args["account"]): if args["account"]: self.response = requests.get(self.server + '/range/account/' + self.firstday + '/' + self.lastday + '/' + args["account"]) self.format = 1 elif args["user"]: self.response = requests.get(self.server + '/range/user/' + self.firstday + '/' + self.lastday + '/' + args["user"]) self.format = 2 if self.response.status_code == 200: self.printbox() conf_file = "/etc/sausage/sausage.cfg" options = AppArgs() getdata = GetData(conf_file,options.response) for item in getdata.message: print("#" + "{0:^57}".format(item) + "#")