#!/usr/bin/python2.7
#
# Univention Quota
#  read default quota-settings and write into quota-table
#
# Copyright 2004-2014 Univention GmbH
#
# http://www.univention.de/
#
# All rights reserved.
#
# The source code of this program is made available
# under the terms of the GNU Affero General Public License version 3
# (GNU AGPL V3) as published by the Free Software Foundation.
#
# Binary versions of this program provided by Univention to you as
# well as other copyrighted, protected or trademarked materials like
# Logos, graphics, fonts, specific documentations and configurations,
# cryptographic keys etc. are subject to a license agreement between
# you and Univention and not subject to the GNU AGPL V3.
#
# In the case you use this program under the terms of the GNU AGPL V3,
# the program is provided in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License with the Debian GNU/Linux or Univention distribution in file
# /usr/share/common-licenses/AGPL-3; if not, see
# <http://www.gnu.org/licenses/>.


import sys
import os
import subprocess
import pwd
import time
import univention.config_registry
import univention.uldap
import univention.lib.policy_result


def get_mountpoint_from_share_path(sharepath):
	sharepath = os.path.abspath(sharepath)
	while not os.path.ismount(sharepath):
		sharepath = os.path.dirname(sharepath)
	return sharepath


def test_for_true(value):
	# Using Syntax definition TrueFalseUp: Possible values are 'TRUE' and 'FALSE'
	return value == 'TRUE'


def get_blocks(size):
	size = size.upper()

	if size.endswith('KB'):
		return int(size.strip('KB'))
	if size.endswith('MB'):
		return int(size.strip('MB')) * 1024
	if size.endswith('GB'):
		return int(size.strip('GB')) * 1024 * 1024
	if size.endswith('B'):
		return int(size.strip('B')) / 1024
	if int(size[-1]) in range(0, 10):
		return int(size) / 1024


def get_shares():
	lo = univention.uldap.getMachineConnection(ldap_master=False)
	shares = {}
	for dn, attrs in lo.search(filter='(&(objectClass=univentionShare)(univentionShareHost=%s.%s))' % (ucr.get('hostname', ''), ucr.get('domainname', '')), attr=['univentionSharePath']):
		policy_result = univention.lib.policy_result.policy_result(dn)[0]
		shares[attrs['univentionSharePath'][0]] = {
				'inodeSoftLimit': int(policy_result.get('univentionQuotaSoftLimitInodes', ['0'])[0]),
				'inodeHardLimit': int(policy_result.get('univentionQuotaHardLimitInodes', ['0'])[0]),
				'spaceSoftLimit': get_blocks(policy_result.get('univentionQuotaSoftLimitSpace', ['0'])[0]),
				'spaceHardLimit': get_blocks(policy_result.get('univentionQuotaHardLimitSpace', ['0'])[0]),
				'reapplyQuota': test_for_true(policy_result.get('univentionQuotaReapplyEveryLogin', ['FALSE'])[0]),
				'mountpoint': get_mountpoint_from_share_path(attrs['univentionSharePath'][0])
			}

	return shares


def get_mountpoints_with_quota(user, shares):
	_mountpoints = {}
	p = subprocess.Popen(['quota', '-vwpu', user], stdout=subprocess.PIPE)
	# quota_stdout has the parseable quota output
	# format:
	# Dateisystemquotas fuer user univention2 (uid 2006):
	# Dateisystem Bloecke   Quota   LimitGnadenfrist Dateien   Quota   LimitGnadenfrist
	# /dev/mapper/vg_ucs-rootfs      24  307200  409600       0       8       0       0       0
	# /dev/vda2       0       0       0       0       0       0       0       0
	quota_stdout, stderr_empty = p.communicate()
	quota_stdout = quota_stdout.splitlines()

	with open('/etc/mtab', 'r') as etc_mtab:
		mtab = etc_mtab.readlines()

	# Create dict of all mountpoints who have usrquota activated.
	# Set updatequota flag to True if setquota has to be called.
	for share, attrs in shares.iteritems():
		if not attrs['mountpoint'] in _mountpoints:
			# We do not know this MP, so get its quota settings
			# check what the current userquota is (0 0 0 0 is unlimited/undefined)
			for line in mtab:
				current_mountpoint = line.split()[1]  # shares are defined on mountpoints
				current_device = line.split()[0]  # quotas are defined on devices
				if current_mountpoint == attrs['mountpoint'] and 'usrquota' in line:
					for quotaline in quota_stdout:
						quota_values = quotaline.split()
						if quota_values[0] == current_device:
							initialValue = False
							# If quota is unset, set flag that the quota should be set.
							if quota_values[2] == '0' and quota_values[3] == '0' and quota_values[6] == '0' and quota_values[7] == '0':
								initialValue = True
							# Now the current quota values are known
							_mountpoints[current_mountpoint] = {
									'spaceSoftLimit': quota_values[2],
									'spaceHardLimit': quota_values[3],
									'inodeSoftLimit': quota_values[6],
									'inodeHardLimit': quota_values[7],
									'updatequota': initialValue
											}
		
		# The mountpoint will not be in _mountpoints if no quota is activated for it
		# Continue with next share definition
		if not attrs['mountpoint'] in _mountpoints:
			continue

		# If the policy flag for updating the quota is set, set flag that the quota should be set.
		if attrs['reapplyQuota']:
			_mountpoints[attrs['mountpoint']]['updatequota'] = True

		# Compare the share policy against the current mp quota. Always set share quota if current mp value = 0
		for policy_value in ['inodeSoftLimit', 'inodeHardLimit', 'spaceSoftLimit', 'spaceHardLimit']:
			if _mountpoints[attrs['mountpoint']][policy_value] == '0' or int(attrs[policy_value]) < int(_mountpoints[attrs['mountpoint']][policy_value]):
				_mountpoints[attrs['mountpoint']][policy_value] = str(attrs[policy_value])

	return _mountpoints


if __name__ == '__main__':
	if '-h' in sys.argv or '--help' in sys.argv:
		print '%s - write user quota from shares into quota-table' % sys.argv[0]
		print 'Usage: %s [username]' % sys.argv[0]
		print '    username defaults to $USER\n'
		print 'Example:'
		print '  %s someuser' % sys.argv[0]
		print '     set up quota for someuser'
		print '  %s' % sys.argv[0]
		print '     set up quota for $USER'
		sys.exit(0)

	ucr = univention.config_registry.ConfigRegistry()
	ucr.load()
	log = open(ucr.get('quota/logfile', '/dev/null'), 'a+')
	TIMEFORMAT = '%a %b %d %H:%M:%S %Z %Y'
	user = sys.argv[1] if len(sys.argv) > 1 else os.environ['USER']
	uid = pwd.getpwnam(user)[2]

	log.write('-------------------------------------------------------- start\n')
	log.write('%s\n' % time.strftime(TIMEFORMAT))
	if uid < 1000:
		log.write('Won\'t set quota for system-users (%s has UID %i)\n' % (user, uid))
		sys.exit(1)
	log.write('Set quota for user %s\n' % user)

	shares = get_shares()
	mountpoints = get_mountpoints_with_quota(user, shares)
	if not shares:
		log.write('No shares found on this host\n')
	elif not mountpoints:
		log.write('No mountpoints available on this host to configure any quota\n')
	else:
		log.write('Found %i shares and %i mountpoints\n' % (len(shares), len(mountpoints)))
		for mountpoint_path, attrs in mountpoints.iteritems():
			if attrs['updatequota']:
				log.write('Set quota on mountpoint "%s": /usr/sbin/setquota -u %s %s %s %s %s %s\n' % (mountpoint_path, user, attrs['spaceSoftLimit'], attrs['spaceHardLimit'], attrs['inodeSoftLimit'], attrs['inodeHardLimit'], mountpoint_path))
				set_quota = subprocess.Popen(['/usr/sbin/setquota', '-u', user, str(attrs['spaceSoftLimit']), str(attrs['spaceHardLimit']), str(attrs['inodeSoftLimit']), str(attrs['inodeHardLimit']), mountpoint_path])
				set_quota.wait()
			else:
				log.write('Do not set/update quota for mountpoint "%s"\n' % mountpoint_path)

	log.write('-------------------------------------------------------- exit\n')
