#!/bin/bash

ENABLED=yes
[ "$ENABLED" != "yes" ] && exit 0
DEBUGMODE=no

unset ROOTFS; [[ -d /usr/lib/ublinux ]] || ROOTFS=.
SOURCE=${ROOTFS}/usr/lib/ublinux/functions; [[ -f ${SOURCE} ]] && . ${SOURCE} 2>/dev/null || exit 0
SOURCE=${ROOTFS}/usr/lib/ublinux/os-config; [[ -f ${SOURCE} ]] && . ${SOURCE} 2>/dev/null || exit 0
debug_mode "$0" "$@"

SYSCONF="${ROOTFS}/${SYSCONF}"
SOURCE=${SYSCONF}/config; [[ -f ${SOURCE} ]] && . ${SOURCE} 2>/dev/null
SOURCE=${SYSCONF}/security; [ -f ${SOURCE} ] && . ${SOURCE} 2>/dev/null

get_compat_unit(){
    if [[ ! ${ITEM_UNIT} =~ ".service"$|".slice"$|".scope"$ ]]; then
 	USER_ID=$(id -u ${ITEM_UNIT} 2>/dev/null)
 	[[ -n ${USER_ID} ]] && USER_SLICE="user-${USER_ID}.slice" || return 1
	COMPAT_UNIT=${USER_SLICE}
    else
	COMPAT_UNIT=${ITEM_UNIT}
    fi
}
# ARG1: CGROUP_QUOTA[superadmin]=MemoryHigh=500M,MemorySwapMax=100M,CPUQuota=400%
exec_cgroup_quota_set(){
## Квоты на ресурсы, через cgroup2. Механизм systemd или напрямую cgroup
    PARAM="$@"
    [[ -n $@ ]] && declare -A CGROUP_QUOTA && eval "${PARAM%%=*}='${PARAM#*=}'"
    if [[ -n ${CGROUP_QUOTA[@]} ]]; then
 	for ITEM_UNIT in "${!CGROUP_QUOTA[@]}"; do
 	    get_compat_unit
	    CGROUP_QUOTA[${ITEM_UNIT}]="${CGROUP_QUOTA[${ITEM_UNIT}]//,/\' }'"
	    CGROUP_QUOTA[${ITEM_UNIT}]="${CGROUP_QUOTA[${ITEM_UNIT}]//=/=\'}"
	    if readlink -fq ${ROOTFS}/usr/bin/init | grep -q "lib/systemd/systemd$"; then
		eval /usr/bin/systemctl set-property --runtime ${COMPAT_UNIT} ${CGROUP_QUOTA[${ITEM_UNIT}]}
		/usr/bin/systemctl daemon-reload
	    else
		true
	    fi
	done
    fi
}
# ARG1: CGROUP_QUOTA[superadmin]=
exec_cgroup_quota_remove(){
## Очистить квоты на ресурсы, через cgroup2. Механизм systemd или напрямую cgroup
    PARAM="$@"
    [[ -n ${PARAM} ]] && ITEM_UNIT=${PARAM#*[} && ITEM_UNIT=${ITEM_UNIT%%]*} || return 0
    [[ -n ${ITEM_UNIT} ]] || return 0
    get_compat_unit
    if readlink -fq ${ROOTFS}/usr/bin/init | grep -q "lib/systemd/systemd$"; then
        /usr/bin/systemctl set-property --runtime ${COMPAT_UNIT} $(cat /run/systemd/system.control/${COMPAT_UNIT}.d/* | grep -v "^[#\[]" | sed "s/=.*$/= /g" | tr -d \\n)
        #systemctl revert ${COMPAT_UNIT}
        /usr/bin/systemctl daemon-reload
    else
        true
    fi
}

################
##### MAIN #####
################

    if [[ -z $@ || $1 == set ]]; then
	shift
	exec_cgroup_quota_set $@
    elif [[ $1 == remove ]]; then
	shift
	exec_cgroup_quota_remove $@
    fi