#!/usr/bin/ruby

# Simple script for setting user quotas.
# - Eric Butler <eric@extremeboredom.net>

# Configure quotas for users. 
# Example:
# {
# 	'groupname' => {
# 		'user1',  => 50,
# 		'user2',  => :unlimited,
# 		:default  => 100
# 	}
# }
# Quota sizes are in MB, the above example gives user1 a 50MB quota, user2 no
# quota, and everyone else a 100MB quota.


# Configure this:
QUOTA_INFO = {
	'users' => {
		'eric' 	    => :unlimited,
		'jazzfreak' => 200,
		:default    => 100
	}
}


def set_quotas(quota_info)
	user_quotas = {}

	quota_info.each do |group_name, user_info|
		puts "Reading group #{group_name}"
		if !user_info.include?(:default)
			raise "Must specify default quota for users in group #{group_name}."
		end
		group_users = users_in_group(group_name)
		group_users.each do |user|
			puts "Doing user #{user}"
			unless user_quotas.include?(user)
				if user_info.include?(user)
					user_quotas[user] = user_info[user]
				else
					user_quotas[user] = user_info[:default]
				end
			else
				puts "WARNING: Not setting quota for #{user} more than once."
			end
		end
	end

	user_quotas.each do |user, quota|
		if quota == :unlimited
			hard_quota = 0
			soft_quota = 0
		else
			hard_quota = quota.to_s + "M"
			soft_quota = (quota - 10).to_s + "M"
		end
		`quotatool -u #{user} -b -l #{hard_quota} /`
		`quotatool -u #{user} -b -q #{soft_quota} /`
	end
end

def users_in_group(groupname)
  `grep ^#{groupname} /etc/group | awk -F\: '{ print $4 }'`.chomp.split(',')
end

set_quotas(QUOTA_INFO)