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 their 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_home_quota() {
	local user="$1"

	local response=$(curl -sk -w "%{http_code}" "https://coldfront.arc.vt.edu/api/arc/quota/home/" \
		-H "Authorization: Token $(cat "/home/$user/$CACHE_DIR/token")" 2>/dev/null)
	local http_code="${response: -3}"
	if [ "$http_code" == "404" ]; then
		echo -e "${YELLOW}No /home quota on record for $user.${NC}"
		return 0
	fi
	if [ "$http_code" != "200" ]; then
		local detail="HTTP $http_code"
		[ -z "$response" ] && detail="no response from server"
		echo -e "${RED}Error: could not retrieve /home usage for $user from ColdFront ($detail).${NC}" >&2
		return 1
	fi
	local json="${response%???}"

	local home=$(echo "$json" | python3 <(cat <<'EOF'
import sys, json
from datetime import datetime

quota = json.load(sys.stdin)
usage = quota.get("capacity_usage") or 0
limit = quota.get("limit") or 0
modified = quota.get("modified")
last_updated = ""
if modified:
    last_updated = datetime.fromisoformat(modified.replace("Z", "+00:00")).strftime("%b %-d %H:%M")
print(f"{usage}|{limit}|{last_updated}")
EOF
) 2>/dev/null)
	if [ -z "$home" ]; then
		echo -e "${RED}Error: could not parse the /home usage returned for $user by ColdFront.${NC}" >&2
		return 1
	fi

	local usage_bytes=$(echo "$home" | awk -F\| '{print $1}')
	local limit_bytes=$(echo "$home" | awk -F\| '{print $2}')
	local last_updated=$(echo "$home" | awk -F\| '{print $3}')

	local usage_GB=$(bytes_to_GB "$usage_bytes")
	local limit_GB=$(bytes_to_GB "$limit_bytes")

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

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_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}"
	if [ "$http_code" != "200" ]; then
		local detail="HTTP $http_code"
		[ -z "$response" ] && detail="no response from server"
		echo -e "${RED}Error: could not retrieve project storage and compute usage for $user from ColdFront ($detail).${NC}" >&2
		return 1
	fi
	local json="${response%???}"
	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;
	local rc=0
	for user in $user_list; do
		print_home_quota "$user" || rc=1
		echo
		print_rest_api_quota "$user" || rc=1
	done

	return $rc
}

main "$@"