#!/usr/bin/python3
#
# Univention Quota
#  read default quota-settings and write into quota-table
#
# SPDX-FileCopyrightText: 2004-2025 Univention GmbH
# SPDX-License-Identifier: AGPL-3.0-only


import os
import pickle  # noqa: S403
import pwd
import re
import subprocess
import sys
import time

import univention.config_registry
import univention.lib.fstab


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[:-2])
    if size.endswith('MB'):
        return int(size[:-2]) * 1024
    if size.endswith('GB'):
        return int(size[:-2]) * 1024 * 1024
    if size.endswith('TB'):
        return int(size[:-2]) * 1024 * 1024 * 1024
    if size.endswith('B'):
        return int(size[:-1]) // 1024
    return int(size) // 1024


def get_shares():
    SHARE_CACHE_DIR = '/var/cache/univention-quota/'

    shares = {}
    for f in os.listdir(SHARE_CACHE_DIR):
        filename = os.path.join(SHARE_CACHE_DIR, f)
        if os.path.isdir(filename):
            continue
        with open(filename, 'rb') as fd:
            _dn, attrs, policy_result = pickle.load(fd, encoding='bytes')
        sharepath = attrs['univentionSharePath'][0].decode('UTF-8')
        shares[sharepath] = {
            '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(sharepath),
        }

    return shares


class Mountpoints:

    def __init__(self, user, shares):
        self.mountpoints_with_quota = {}
        self.user = user
        self.shares = shares
        # Create dict of all mountpoints who have usrquota activated.
        # Set updatequota flag to True if setquota has to be called.
        # check what the current userquota is (0 0 0 0 is unlimited/undefined)
        for mount in self.mtab():
            if not mount['quotaenabled']:
                continue

            quotas = self.get_quota()
            for quota in quotas.values():
                if quota['device'] == mount['device']:
                    # If quota is unset, set flag that the quota should be set:
                    initialValue = quota['bsoft'] == '0' and quota['bhard'] == '0' and quota['fsoft'] == '0' and quota['fhard'] == '0'
                    # Now the current quota values are known
                    self.mountpoints_with_quota[mount['mountpoint']] = {
                        'spaceSoftLimit': quota['bsoft'],
                        'spaceHardLimit': quota['bhard'],
                        'inodeSoftLimit': quota['fsoft'],
                        'inodeHardLimit': quota['fhard'],
                        'reapplyQuota': initialValue,
                    }

        # add shares whose quotas have never been set before
        for share in shares.values():
            if share['mountpoint'] not in self.mountpoints_with_quota.keys():
                self.mountpoints_with_quota[share['mountpoint']] = share

        for mountpoint in self.mountpoints_with_quota.keys():
            shares_on_mount = self.get_shares_with_mountpoint(mountpoint)
            if not shares_on_mount:
                self.mountpoints_with_quota[mountpoint]['reapplyQuota'] = False
                continue

            for share_on_mount in shares_on_mount:
                if share_on_mount['reapplyQuota']:
                    self.mountpoints_with_quota[mountpoint]['reapplyQuota'] = True

            for policy_value in ['inodeSoftLimit', 'inodeHardLimit', 'spaceSoftLimit', 'spaceHardLimit']:
                smallest_quota = self.calculate_smallest_quota(shares_on_mount, policy_value)
                self.mountpoints_with_quota[mountpoint][policy_value] = smallest_quota

    def get_shares_with_mountpoint(self, mountpoint):
        _shares = []
        for _share in self.shares.values():
            if _share['mountpoint'] == mountpoint:
                _shares.append(_share)
        return _shares

    def calculate_smallest_quota(self, shares_on_mount, policy_value):
        # from one policy value of different shares on the same mount, calculate the smallest one.
        # ignore 0 for the comparison, except every policy value is 0
        filtered_zero = []
        for _share in shares_on_mount:
            if _share[policy_value] > 0:
                filtered_zero.append(_share[policy_value])
        if not filtered_zero:
            return 0
        smallest = min(filtered_zero)
        return smallest

    def get_quota(self):
        # parses the output of 'quota -vwpu <user>'
        devices = {}
        p = subprocess.Popen(['quota', '-vwpu', '--always-resolve', self.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, _ = p.communicate()
        regex = re.compile('^ *(?P<device>[^ ]*) *(?P<bused>[0-9]*). *(?P<bsoft>[0-9]*) *(?P<bhard>[0-9]*) *(?P<btime>[0-9]*) *(?P<fused>[0-9]*). *(?P<fsoft>[0-9]*) *(?P<fhard>[0-9]*) *(?P<ftime>[0-9]*)$')
        for quotaline in quota_stdout.decode('UTF-8').split('\n'):
            matches = regex.match(quotaline)
            if not matches:
                continue
            grp = matches.groupdict()
            devices[grp['device']] = grp
        return devices

    def mtab(self):
        return [{
                'mountpoint': mt.mount_point,
                'device': mt.spec,
                'quotaenabled': mt.hasopt('usrquota') or (mt.hasopt('usrjquota=') and mt.hasopt('jqfmt=vfsv')),
                } for mt in univention.lib.fstab.File('/etc/mtab')]


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 = Mountpoints(user, shares).mountpoints_with_quota
    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.items():
            if attrs['reapplyQuota']:
                log.write('Set quota on mountpoint "%s": /usr/sbin/setquota --always-resolve -u %s %s %s %s %s %s\n' % (mountpoint_path, user, attrs['spaceSoftLimit'], attrs['spaceHardLimit'], attrs['inodeSoftLimit'], attrs['inodeHardLimit'], mountpoint_path))
                command = ['/usr/sbin/setquota', '--always-resolve', '-u', user, str(attrs['spaceSoftLimit']), str(attrs['spaceHardLimit']), str(attrs['inodeSoftLimit']), str(attrs['inodeHardLimit']), mountpoint_path]
                set_quota = subprocess.Popen(command, stderr=log)
                set_quota.wait()
            else:
                log.write('Do not set/update quota for mountpoint "%s"\n' % mountpoint_path)

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