#! /usr/bin/python

# This script prints out a warning if the user is over quota.
# It is executed by /etc/bashrc whenever a user logs in.

import re, os, sys

sys.stdout.write("Checking your disk quota...")
sys.stdout.flush()

# SSH to the storage server to get local quotas since we don't want to try to get
# quotas on other mounted file systems.
quotaCommand = "/usr/bin/ssh -q -o PasswordAuthentication=no -x nas-0-0 \"/usr/bin/quota -l\""

quota = os.popen(quotaCommand, "r")
header = quota.readline()
fileSystem = quota.readline()
quotas = quota.readline()
quota.close()

# Check if the user is over-quota
match = re.search("\s+/dev/sda3\s+(\d+)\\*\s+(\d+)\s+(\d+)\s+(\d+|none)", quotas)
if match:
	blocks = match.group(1)
	quota = match.group(2)
	limit = match.group(3)
	grace = match.group(4)
	overMB = round((int(blocks)-int(quota))/(1024.0))
	warning = "\n\033[1;5;31mWARNING:\033[0;0;31m Your account is \033[1;31m" + str(int(overMB)) + "\033[0;31m MB over-quota.\n"
	warning = warning + "         Your remainng grace period is: \033[1;31m" + grace + " days.\033[0;m\n"
	sys.stdout.write(warning)
	
# If they are not over quota find out how much they are using
else:
	match = re.search("\s+/dev/sda3\s+(\d+)\s+(\d+)", quotas)
	if match:
		blocks = float(match.group(1))
		quota = float(match.group(2))
		usedGB = int(round(blocks/(1024.0*1024.0)))
		quotaGB = int(round(quota/(1024.0*1024.0)))
		ratio = blocks/quota
		
		# Based on the ratio display a different message
		if (ratio > 0.5):
			sys.stdout.write("\033[0;1myou are using " + str(int(ratio*100)) +"% of your quota.\033[0;m\n")
			sys.stdout.write("Used: " + str(usedGB) + "GB Limit: " + str(quotaGB) +"GB\n")
		elif (ratio > 0.1):
			sys.stdout.write("you are under quota. (Used: " + str(usedGB) +"GB)\n")
		else:
			sys.stdout.write("you are way under quota. (Used: " + str(usedGB) +"GB) \n")
			
		if (ratio > 0.75):
			sys.stdout.write("\n\n")
			sys.stdout.write("PLEASE KEEP IN MIND that we do not actually have enough space for\n")
			sys.stdout.write(" every user to use the full quota limit. If you do not actually need\n")
			sys.stdout.write(" as much space as youa re using it would be nice if you could delete\n")
			sys.stdout.write(" some of your files.\n\n")
	# If we did not get a quota response it's either because the user has no quota (root) or
	# they do not correctly have their keychain setup for passwordless authentication across
	# the cluster. If the former is the case alert them, if not warn them to fix their keychain.
	else:
		match = re.search(".*none.*", header)
		if match:
			sys.stdout.write("no quota enforced.\n")
		else:
			sys.stdout.write("\033[1;5;31mfailed to check disk quota.\033[0;m\n")
			sys.stdout.write("ERROR: you probably have a non-blank ssh keychain password.\n")
			sys.stdout.write("You should set your \033[0;1mkeychain\033[0;m password (not your ssh login password)\n")
			sys.stdout.write("to be blank to avoid having to enter your password when logging\n")
			sys.stdout.write("into different machines on bagels.\n")
			sys.stdout.write("To do this run: \033[0;1mssh-keygen -p\033[0;m and set the new password to be blank.\n")
			sys.stdout.write("For more info see \033[0;1mhttp://cva.stanford.edu/systems/cva_systems.html#Passwords\033[0;m.")