quota

#!/usr/bin/env bash

set -uo pipefail

CACHE_DIR=".arc_quota"

RED=$(tput setaf 1);
GREEN=$(tput setaf 2);
YELLOW=$(tput setaf 3);
NC=$(tput sgr0); # No color
HEADER_COLORS=($GREEN $GREEN $GREEN $GREEN $GREEN $GREEN $GREEN)

usage() {
 echo "
  SYNOPSIS
    $0 [-h] [user1 [user2 ... ]]

  DESCRIPTION
    Without any argument, quota displays filesystem quota information
    for the current user for the filesystems configured on the current
    system.

    A list of users may be optionally supplied and thier quota information
    will be displayed. Displaying quota information for any user other
    than the current user requires sudo or root privileges.

  OPTIONS
    -h      show this help text

  EXAMPLE
    $0
    $0 pid1 pid2 pid3
    $0 -h
  "
}

print_ess_gpfs_quota() {
    local user=$1
    local projects_cache="/home/$user/$CACHE_DIR/projects"
    if [ -s "$projects_cache" ]; then
        echo

        local last_updated="$(ls -l "$projects_cache" | awk '{print $6, $7, $8}')"
        json=$(cat "$projects_cache")
        while read line; do
            local colors=('' '' '' '' '' '' '')
            local path=$(echo $line | awk -F\| '{print $1}')
            local num_fields=$(awk -F\| '{print NF}' <<< "$line")

            # If account is expired
            if [ "$num_fields" -lt 3 ]; then
		local row=("$user" "$path" "0" "0" "EXPIRED" "EXPIRED" "$last_updated")
                print_row row[@] colors[@]
                continue
            fi

	    local block_usage_kib=$(echo $line | awk -F\| '{print $2}');
	    local block_quota_kib=$(echo $line | awk -F\| '{print $3}');

	    local block_usage_gib=$(KiB_to_GB "$block_usage_kib")
	    local block_quota_gib=$(KiB_to_GB "$block_quota_kib")

	    local files_usage=$(echo $line | awk -F\| '{print $4}');
	    local files_quota=$(echo $line | awk -F\| '{print $5}');

            local row=("$user" "$path" "$block_usage_gib" "$block_quota_gib" "$files_usage" "$files_quota" "$last_updated")
            print_row row[@] colors[@]

        done < <(cat "$projects_cache" | python3 <( echo '
import sys, json
for k, v in json.load(sys.stdin).items():
    path = k
    block_usage = v.get("blockUsage", 0)
    block_quota = v.get("blockQuota", 0)
    files_usage = v.get("filesUsage", 0)
    files_quota = v.get("filesQuota", 0)
    if block_usage == 0 and block_quota == 0 and files_usage == 0 and files_quota == 0:
        print(f"{path}|{path}")
    else:
        print(f"{path}|{block_usage}|{block_quota}|{files_usage}|{files_quota}")
'))
    fi
}


print_home_quota() {
	local user=$1

	local json="";
	local home_cache="/home/$user/$CACHE_DIR/home";
	local last_updated="";
	if [ -s "$home_cache" ]; then
		json=$(cat "$home_cache");
		last_updated="$(ls -l "$home_cache" | awk '{print $6, $7, $8}')"
	fi

	[ -z "$json" ] && return; # skip when response is empty

	local usabe_bytes=$(echo "$json" | grep -o 'capacity_usage": "[0-9]\+"' | grep -o "[0-9]\+")
	local limit_bytes=$(echo "$json" | grep -o 'limit": "[0-9]\+"' | grep -o "[0-9]\+")

	local usage_GiB=$(bytes_to_GB "$usabe_bytes")
	local limit_GiB=$(bytes_to_GB "$limit_bytes")

	local colors=('' '' '' '' '' '' '')
	local row=("$user" "/home" "$usage_GiB" "$limit_GiB" "-" "-" "$last_updated")
	print_row row[@] colors[@];
}

print_alloc_quotas() {
	local user=$1

	print_alloc_header;

	print_alloc_quota "$user" "slurm"
}

print_alloc_quota() {
	local user="$1";
	local alloc_type="$2";

	local json="";
	local alloc_cache="/home/$user/$CACHE_DIR/$alloc_type";
	[ -s "$alloc_cache" ] && json=$(cat "$alloc_cache");

	[ -z "$json" ] && return; # skip when response is empty

	while read line; do
		local pid=$(echo $line | awk -F\| '{print $1}');
		local alloc=$(echo $line | awk -F\| '{print $2}');
		local cluster=$(echo $line | awk -F\| '{print $3}');
		local quota=$(echo $line | awk -F\| '{print $4}');
		local left=$(echo $line | awk -F\| '{print $5}');
		local status=$(echo $line | awk -F\| '{print $6}');

		local colors=('' '' '' '' '' '' '')

		local note="";
		local now_epoch=0;
		if [ "$status" == "Revoked" ]; then
			continue;
		fi
		if [ "$status" == "Expired" ]; then
			note+="Allocation has expired. ";
			colors[1]=${YELLOW};
			colors[2]=${YELLOW};
			colors[5]=${YELLOW};
			colors[6]=${RED};
		fi
		if [ "$left" == "0" ]; then
			note+="Allocation has no funds. ";
			colors[1]=${YELLOW};
			colors[2]=${YELLOW};
			colors[4]=${YELLOW};
			colors[6]=${RED};
		fi

		printf "$(print_padding colors[@] 'header')" \
			"$pid" "$alloc" "$cluster" "$quota" "$left" "$status" "$note";

	done < <(cat "$alloc_cache" | python <( echo '
from __future__ import print_function;
import sys, json;
for (k, v) in json.load(sys.stdin).items(): 
	i=0;
	status = ""
	if "status" in v:
		status = v["status"];
	for f in v["funds"]: 
		alloc = f["allocated"];
		remain = f["remaining"];

		if i == 0: 
			print("'$user'"+"|"+k+"|"+f["cluster"]+"|"+alloc+"|"+remain+"|"+status+"|")
		else:
			print("||"+f["cluster"]+"|" + alloc+"|"+remain+"|"+status+"|");
		i += 1;
'))

}

parse_storage_quota() {
	local quota="$1"
	local fs="$2"
	[ -n "$quota" ] || return;

	local current_blocksize_col='3';
	local maximum_blocksize_col='4';
	local current_files_col='1';
	local maximum_files_col='2';

	local grace_col='7';

	if [ "$fs" == "/work" ]; then
		current_blocksize_col='4';
		maximum_blocksize_col='5';

		local grace_col='8';
	fi

	local stats=$(echo "$quota" | awk 'BEGIN { RS = "" ; FS = "\n" } { print $3 }') 
	local grace=$(echo $stats | awk -F\| '{print $1}' | cut -d' ' -f${grace_col}-)
	local current_blocksize_KiB=$(echo -e "$stats" | awk '{print $'$current_blocksize_col'}')
	local maximum_blocksize_KiB=$(echo -e "$stats" | awk '{print $'$maximum_blocksize_col'}')

	local current_blocksize_GiB=$(KiB_to_GB "$current_blocksize_KiB")
	local maximum_blocksize_GiB=$(KiB_to_GB "$maximum_blocksize_KiB")

	local current_files=$(echo "$stats" |  awk -F\| '{print $2}' | awk '{print $'$current_files_col'}')
	local maximum_files=$(echo "$stats" |  awk -F\| '{print $2}' | awk '{print $'$maximum_files_col'}')

	local note="";
	local colors=('' '' '' '' '' '' '')

	if [ "$current_blocksize_GiB" -gt "$maximum_blocksize_GiB" ]; then
		note="You have ${grace}to reduce your data below $maximum_blocksize_GiB GiB";
		colors[2]=${RED};
		colors[6]=${RED};
	fi

	local row=("$user" "$fs" "$current_blocksize_GiB" \
		"$maximum_blocksize_GiB" "$current_files" "$maximum_files" "$note")

  print_row row[@] colors[@]
}

print_storage_header() {
	printf "$(print_padding HEADER_COLORS[@] 'header')" \
		"USER" "FILESYS/SET" "DATA (GB)" "QUOTA (GB)" "FILES" "QUOTA" "LAST UPDATED"
}

print_alloc_header() {
	echo;
	printf "$(print_padding HEADER_COLORS[@] 'header')" \
		"USER" "ALLOCATION" "CLUSTER" "QUOTA (hrs)" "LEFT (hrs)" "STATUS" "NOTE";
}

print_row() {
	local row=("${!1}");
	local colors=("${!2-}");

	printf "$(print_padding colors[@])" "${row[@]}"
}

print_padding() {
	local colors=("${!1}");
	local pad_type=${2-};
	local pad_length="%-16 %-36 %-12 %-11 %-10 %-10 %-";
	local header_format=(s s s s s s s);
	local row_format=(s s .1f .0f s s s);
	local i=0;
	if [ "$pad_type" == "header" ]; then
		for p in $pad_length; do
			echo -n "${colors[$i]}${p}${header_format[$i]}${NC} "
			let i++;
		done
	else
		for p in $pad_length; do
			echo -n "${colors[$i]}${p}${row_format[$i]}${NC} "
			let i++;
		done
	fi
	echo "\n";
}

KiB_to_GB() {
	local size_in_KiB="$1"
	local KiB_in_GB=976563;
	printf '%.0f' $(echo "$size_in_KiB / $KiB_in_GB" | bc -l)
}

KiB_to_GiB() {
	local size_in_KiB="$1"
	local KiB_in_GiB=1048576;
	printf '%.0f' $(echo "$size_in_KiB / $KiB_in_GiB" | bc -l)
}

bytes_to_GiB() {
	local size_in_bytes="$1"
	local bytes_in_GiB=1073741824;
	echo "$size_in_bytes / $bytes_in_GiB" | bc -l
}

bytes_to_GB() {
	local size_in_bytes="$1"
	local bytes_in_GB=1000000000;
	echo "$size_in_bytes / $bytes_in_GB" | bc -l
}

print_rest_api_quota() {
	local user="$1"

	local response=$(curl -sk -w "%{http_code}" "https://coldfront.arc.vt.edu/api/arc/quota/" \
		-H "Authorization: Token $(cat "/home/$user/$CACHE_DIR/token")" 2>/dev/null)
	local http_code="${response: -3}"
	local json="${response:0:-3}"
	[ "$http_code" != "200" ] && return 1
	echo "$json" | python3 <(cat <<EOF
import sys, json
from datetime import datetime

RED    = "\033[31m"
GREEN  = "\033[32m"
YELLOW = "\033[33m"
NC     = "\033[0m"

WIDTHS = [16, 36, 12, 11, 10, 10, None]

def print_row(fields, colors=None, is_header=True):
    if colors is None:
        colors = [""] * 7
    parts = []
    for i, (val, width) in enumerate(zip(fields, WIDTHS)):
        color = colors[i] if i < len(colors) else ""
        if not is_header and i == 2:
            text = f"{float(val):.1f}" if val != "" else ""
        elif not is_header and i == 3:
            if isinstance(val, str):
                text = val
            else:
                text = "Infinity" if val < 0 else f"{float(val):.0f}"
        else:
            text = str(val)
        if width is not None:
            text = f"{text:<{width}}"
        parts.append(f"{color}{text}{NC}")
    print(" ".join(parts))

def get_attr(attributes, attr_type, field="value"):
    return [a[field] for a in attributes if a["allocation_attribute_type"] == attr_type and field in a]

pid = "$user"
allocations = json.load(sys.stdin)
compute = [a for a in allocations if a["resource_type"] == "Cluster"]
storage = [a for a in allocations if a["resource_type"] == "Storage"]

for alloc in storage:
    attrs = alloc["allocation_attribute"]
    storage_group = get_attr(attrs, "Storage_Group_Name")
    if not storage_group:
        continue
    usage = get_attr(attrs, "storage_usage")
    data_gb = float(usage[0].split('/')[0]) * 1000 if usage else 0.0
    quota = get_attr(attrs, "Storage Quota (GB)")
    quota_gb = float(quota[0]) if quota else 0.0
    files_usage = get_attr(attrs, "ess_files_usage")
    files_quota = get_attr(attrs, "ess_files_quota")
    last_updated = get_attr(attrs, "ess_files_usage", "modified")
    print_row([pid, f"/projects/{storage_group[0]}", data_gb, quota_gb,
               files_usage[0] if files_usage else "",
               files_quota[0] if files_quota else "",
               datetime.fromisoformat(last_updated[0]).strftime("%b %-d %H:%M") if last_updated else ""], is_header=False)

print()

print_row(["USER", "ALLOCATION", "CLUSTER", "QUOTA (hrs)", "LEFT (hrs)", "STATUS", "NOTE"],
          colors=[GREEN] * 7)
for alloc in compute:
    if alloc["status"] == "Revoked":
        continue
    attrs = alloc["allocation_attribute"]
    account = get_attr(attrs, "slurm_account_name")
    if not account:
        continue
    raw_billing = int(alloc.get("compute_billing", 0))
    unlimited = raw_billing < 0
    billing_hrs = raw_billing // 60
    status = alloc["status"]
    for count, entry in enumerate(alloc["compute_usage"]):
        cluster = entry["cluster"]
        used_secs = entry["rawusage"]
        left = billing_hrs - int(used_secs) // 3600
        colors = [""] * 7
        note = ""
        if status == "Expired":
            note += "Allocation has expired. "
            colors[1] = colors[2] = colors[5] = YELLOW
            colors[6] = RED
        if not unlimited and left == 0:
            note += "Allocation has no funds. "
            colors[1] = colors[2] = colors[4] = YELLOW
            colors[6] = RED
        quota_col = "Infinity" if unlimited else billing_hrs
        left_col = "Infinity" if unlimited else left
        row = [pid, account[0], cluster, quota_col, left_col, status, note.strip()] \
              if count == 0 else ["", "", cluster, quota_col, left_col, status, note.strip()]
        print_row(row, colors)
EOF
)
}


main() {
  while getopts "dh" opt; do
    case "$opt" in
      h) usage; exit 0
    ;;
    esac
  done

	if [ $# -gt 0 ]; then
			# If we have arguments then we need root.
			if [ 0 -ne $(id -u) ]; then
					echo sudo/root required to check quota of a user list
					exit 1
			fi

			user_list=$@
	else
			# If no arguments we are checking quota for the executing user.
			user_list=$USER
	fi

        echo -e "${YELLOW}Warning: quota values are not updated in real time and may take a few hours to refresh.${NC}"
	print_storage_header;
	for user in $user_list; do
		print_home_quota "$user"
		echo
		if ! print_rest_api_quota "$user"; then
			print_ess_gpfs_quota "$user"
			print_alloc_quotas "$user"
		fi
	done

}

main "$@"