#!/bin/bash
###########################################################################
#
#          FILE:  check_btrfs_quota
# 
#         USAGE:  check_btrfs_quota
#                   -m <mountpount>  : set the btrfs Mountpoint
#                   -s <subvolname>  : set the subvolume name 
#                                      multiple subvols can be comma seperated
#                   -e               : Check excl
#                   -r               : Check rfer
#                   -f <path>        : Check recursive
#                   -w <warn>        : percentage used for warning
#                   -c <crit>        : percentage used for critical
#                   -S               : force a fs sync before calculation
#                 debug              Create a tracefile for remote debugging
# 
#        AUTHOR:  Tobias D. Oestreicher
#
#     COPYRIGHT:  B1 Systems GmbH, Vohburg, Germany
#      HOMEPAGE:  http://www.b1-systems.de/
#       LICENSE:  GPLv3 <http://www.gnu.org/licenses/gpl.txt>
#       VERSION:  0.1
#       CREATED:  29.09.2015
#
###########################################################################
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the 
# Free Software Foundation; either version 3 of the License.
#
# This program is distributed 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, see <http://www.gnu.org/licenses/>.
#
###########################################################################
#     CHANGELOG:  (19.10.2015) TDOe - First Publishing 
###########################################################################

# Color Settings
RED='\e[0;31m'
GREEN='\e[0;32m'
BOLD='\e[1m'
RESCOL='\e[0m'

BTRFSPROG=$( which btrfs )

NAGIOS_OK=0
NAGIOS_WARNING=1
NAGIOS_CRITICAL=2
NAGIOS_UNKNOWN=3

###########################################################################
#          D O   N O T   E D I T   B E L O W   T H I S   L I N E
###########################################################################

#CUR=""
#MAX=""
FORCESYNC=0;

# Handle user abort
USE_TRAP=true
if [ ${USE_TRAP} == "true" ];then
    trap 'trap_func' TERM INT
fi


# Handle debug output
DBGSTR='debug'
DBG_LOG_DIR='/tmp'
CUR_UNIX_DATE=$(date +%s)
TRACE_FILE="${DBG_LOG_DIR}/$(basename ${0}).${CUR_UNIX_DATE}.trc"
OPTS=`echo "$@"|grep -c "$DBGSTR"` 
if [ $OPTS -gt 0 ];then
echo '###########################################################################'
echo '#                          DEBUGGING IS ENABLED                           #'
echo '###########################################################################'
    export PS4='+${BASH_SOURCE}:${LINENO}:${FUNCNAME[0]}: '
    exec 2>${TRACE_FILE}
	set -xv
fi


###########################################################################
help_msg()                           # NULL
###########################################################################
	{
		echo "check_btrfs_quota v. 0.1   - BTRFS Quota Check -   (c) Tobias D. Oestreicher"
		echo "-----------------------------------------------------------------------------"
		echo "Usage:"
		echo "  `basename $0` [options]"
		echo ""
		echo "Options:"
		echo "  -m <mountpount>         : set the btrfs Mountpoint"
		echo "  -s <subvolname>         : set subvolume name or comma seperated list"
                echo "  -q <qgroup>             : set qgroup or comma seperated list"
		echo "  -e                      : Check excl"
		echo "  -r                      : Check rfer"
		echo "  -f <path>               : Check recursive"
		echo "  -w <warn>               : percentage used for warning"
		echo "  -c <crit>               : percentage used for critical"
		echo "  -S                      : force a fs sync before calculation"
		echo "  debug                   : creating tracefile for debug"
		echo ""
		echo "Example:"
		echo "  `basename $0` -m /mnt -s subvol1,subvol2,subvol3 -r -w 75 -c 90"
		echo "  `basename $0` -m /mnt -q 0/15,0/256,1/100 -e -w 75 -c 90"
		echo "  `basename $0` -m /mnt -f /mnt/subvol3 -r -w 75 -c 90 -S"
		echo "-----------------------------------------------------------------------------"
		echo ""
		exit

	}

###########################################################################
trap_func()
###########################################################################
    {
        echo "Script aborted by user..."
        echo -n "cleaning up "
        echo -n ". "
        echo -n ". "
        echo -n ". "
        echo "done"
        exit 1
    }



###########################################################################
get_subvol_id()
###########################################################################
    {
        SUBVOLID=$($BTRFSPROG subvol list $MOUNTPOINT|grep -e" ${1}$"| awk '{print $2}')
        echo $SUBVOLID;
    }



###########################################################################
get_subvol_stats()
###########################################################################
    {
        SUBVOLID=$(get_subvol_id ${1})
        if [ "${SUBVOLID}" == "" ];then
            echo "Subvol ${sv} not existing"
            exit ${NAGIOS_UNKNOWN}
        fi
        SUBVOLDATA=$($BTRFSPROG qgroup show -r -e $MOUNTPOINT --sort=qgroupid,rfer,excl,max_rfer,max_excl --kbytes|grep -e "^.*/${SUBVOLID}")
        SUBVOLDATA=$(echo $SUBVOLDATA|awk '{print $2,$3,$4,$5}')
        
    }

###########################################################################
get_qgroup_stats()
###########################################################################
    {
        QGROUPDATA=$($BTRFSPROG qgroup show -r -e $MOUNTPOINT --sort=qgroupid,rfer,excl,max_rfer,max_excl --kbytes|grep -e "^${1}[ 	]")
        if [ $(echo ${QGROUPDATA}|wc -c) -le 1 ]; then
            echo "qGroup ${qg} not existing"
            exit ${NAGIOS_UNKNOWN}
        fi
        QGROUPDATA=$(echo $QGROUPDATA|awk '{print $2,$3,$4,$5}')
    }


###########################################################################
force_fs_sync()
###########################################################################
    {
        SYNC=$($BTRFSPROG filesystem sync $1 2>/dev/null)
    }

###########################################################################
get_rfer_percentage()
###########################################################################
    {
        CUR=$(echo $SUBVOLDATA|awk '{print $1}'|awk -F '.' '{print $1}')
        MAX=$(echo $SUBVOLDATA|awk '{print $3}'|awk -F '.' '{print $1}')
        if [ "$MAX" == "none" ];then MAX=0;fi 
        if [ $MAX -eq 0 ]; then
            RFER_USED_PCT=0
        else
            RFER_USED_PCT=$(echo "$CUR / ($MAX/100)"|bc)
        fi
        echo $RFER_USED_PCT $CUR $MAX
    }

###########################################################################
get_excl_percentage()
###########################################################################
    {
        CUR=$(echo $SUBVOLDATA|awk '{print $2}'|awk -F '.' '{print $1}')
        MAX=$(echo $SUBVOLDATA|awk '{print $4}'|awk -F '.' '{print $1}')
        if [ "$MAX" == "none" ];then MAX=0;fi 
        if [ $MAX -eq 0 ]; then
            EXCL_USED_PCT=0
        else
            EXCL_USED_PCT=$(echo "$CUR / ($MAX/100)"|bc)
        fi
        echo $EXCL_USED_PCT $CUR $MAX
    }


###########################################################################
options()                           # OPT1 = $@
###########################################################################
	{
		COUNT=$#
		if [ "$COUNT" == "0" ]; then help_msg; fi
	
		while [ -n "$1" ];do
			case $1 in
				-m)
					shift
					MOUNTPOINT=$1
					;;
				-s)
                                        shift
                                        SUBVOLNAME=$1
					;;
                                -e)
					METH="excl"
                                        ;;
                                -r)
                                        METH="rfer"
                                        ;;
                                -q)
                                        shift
                                        QGROUPID=$1
                                        ;;
                                -f)
                                        shift
                                        PATHRECURSIVE=$1
                                        ;;
                                -w)
                                        shift
                                        WARN=$1
                                        ;;
                                -c)
                                        shift
                                        CRIT=$1
                                        ;;
                                -S)
                                        FORCESYNC=1;
                                        ;;
				-h|--help)
					help_msg;
					;;
			esac
			shift
		done	

            if [ -z "${METH}" ] || [ -z "${WARN}" ] || [ -z "${CRIT}" ] || [ -z "${MOUNTPOINT}" ] || ( [ -z "${SUBVOLNAME}" ] && [ -z "${QGROUPID}" ] && [ -z "${PATHRECURSIVE}" ] ); then
		help_msg;
                exit ${NAGIOS_UNKNOWN}
            fi
            if [ ${FORCESYNC} -eq 1 ]; then
                force_fs_sync ${MOUNTPOINT}
            fi


            EXITCODE=${NAGIOS_OK}
            PERFDATA=""
            PLUGINLONGOUTPUT=""
            if [ -n "${SUBVOLNAME}" ];then
            # multisubvol start
                SHORTTXTOUT=${SUBVOLNAME}
                for sv in $(echo ${SUBVOLNAME}|tr ',' ' ');do
                    get_subvol_stats $sv


                    if [ "${METH}" == "excl" ];then
                        RES=$(get_excl_percentage)
                    else
                        RES=$(get_rfer_percentage)
                    fi
                    RESPCT=$(echo $RES|awk '{print $1}')

                    if [ "${PLUGINLONGOUTPUT}" == '' ]; then
                        PLUGINLONGOUTPUT="${sv}\t${RESPCT}%\t${SUBVOLDATA}"
                    else
                        PLUGINLONGOUTPUT="${PLUGINLONGOUTPUT}\n${sv}\t${RESPCT}%\t${SUBVOLDATA}"
                    fi

                    if [ "${PERFDATA}" == "" ];then
                        PERFDATA="${sv}_${METH}_pct=$RESPCT" 
                    else
                        PERFDATA="${PERFDATA} ${sv}_${METH}_pct=$RESPCT" 
                    fi
                    CUR=$(echo $RES|awk '{print $2}')
                    MAX=$(echo $RES|awk '{print $3}')
    
                    if [ ${RESPCT} -ge ${CRIT} ] && [ "${EXITCODE}" != "${NAGIOS_CRITICAL}" ]; then
                        EXITCODE=${NAGIOS_CRITICAL}
                        #echo "set EXIT CODE to NAGIOS_CRIT"
                    fi

                    if [ ${RESPCT} -ge ${WARN} ] && [ ${RESPCT} -lt ${CRIT} ] && [ "${EXITCODE}" != "${NAGIOS_CRITICAL}" ]; then
                        EXITCODE=${NAGIOS_WARNING}
                        #echo "set EXIT CODE to NAGIOS_WARN"
                    fi
    
                done
            # multisubvol end
            elif [ -n "${QGROUPID}" ];then
            # multiqgroup start
                SHORTTXTOUT=${QGROUPID}
                for qg in $(echo ${QGROUPID}|tr ',' ' ');do
                    get_qgroup_stats ${qg}

                    SUBVOLDATA=$QGROUPDATA
                    if [ "${METH}" == "excl" ];then
                        RES=$(get_excl_percentage)
                    else
                        RES=$(get_rfer_percentage)
                    fi
                    RESPCT=$(echo $RES|awk '{print $1}')


                    if [ "${PLUGINLONGOUTPUT}" == '' ]; then
                        PLUGINLONGOUTPUT="${qg}\t${RESPCT}%\t${QGROUPDATA}"
                    else
                        PLUGINLONGOUTPUT="${PLUGINLONGOUTPUT}\n${qg}\t${RESPCT}%\t${QGROUPDATA}"
                    fi




                    if [ "${PERFDATA}" == "" ];then
                        PERFDATA="${qg}_${METH}_pct=$RESPCT" 
                    else
                        PERFDATA="${PERFDATA} ${qg}_${METH}_pct=$RESPCT" 
                    fi
                    CUR=$(echo $RES|awk '{print $2}')
                    MAX=$(echo $RES|awk '{print $3}')
    
                    if [ ${RESPCT} -ge ${CRIT} ] && [ "${EXITCODE}" != "${NAGIOS_CRITICAL}" ]; then
                        EXITCODE=${NAGIOS_CRITICAL}
                        #echo "set EXIT CODE to NAGIOS_CRIT"
                    fi
                    if [ ${RESPCT} -ge ${WARN} ] && [ ${RESPCT} -lt ${CRIT} ] && [ "${EXITCODE}" != "${NAGIOS_CRITICAL}" ]; then
                        EXITCODE=${NAGIOS_WARNING}
                        #echo "set EXIT CODE to NAGIOS_WARN"
                    fi
    
                done
            # multiqgroup end
            elif [ -n ${PATHRECURSIVE} ];then
            # pathrecursive start
                ANCH_QGRPS=$(btrfs qg s ${PATHRECURSIVE} -F --sort=qgroupid,rfer,excl,max_rfer,max_excl --kbytes 2>/dev/null|awk '{print $1 }'|grep  "^[0-9]")
                if [ -z "${ANCH_QGRPS}" ];then
                    echo "Unknown: Path ${PATHRECURSIVE} not existing or not a btrfs"
                    exit ${NAGIOS_UNKNOWN}
                else
                    if [ $(echo -e $ANCH_QGRPS|wc -w) -gt 1 ]; then
                        for anchester in $ANCH_QGRPS;do
                            get_qgroup_stats ${anchester}

                            SUBVOLDATA=$QGROUPDATA
                            if [ "${METH}" == "excl" ];then
                                RES=$(get_excl_percentage)
                            else
                                RES=$(get_rfer_percentage)
                            fi
                            RESPCT=$(echo $RES|awk '{print $1}')
                            CUR=$(echo $RES|awk '{print $2}')
                            MAX=$(echo $RES|awk '{print $3}')
                            if [ "${PLUGINLONGOUTPUT}" == '' ]; then
                                PLUGINLONGOUTPUT="${anchester}\t${RESPCT}%\t${QGROUPDATA}"
                            else
                                PLUGINLONGOUTPUT="${PLUGINLONGOUTPUT}\n${anchester}\t${RESPCT}%\t${QGROUPDATA}"
                            fi
                            if [ "${PERFDATA}" == "" ];then
                                PERFDATA="${anchester}_${METH}_pct=$RESPCT"
                            else
                                PERFDATA="${PERFDATA} ${anchester}_${METH}_pct=$RESPCT"
                            fi


                            if [ ${RESPCT} -ge ${CRIT} ] && [ "${EXITCODE}" != "${NAGIOS_CRITICAL}" ]; then
                                EXITCODE=${NAGIOS_CRITICAL}
                                #echo "set EXIT CODE to NAGIOS_CRIT"
                            fi
                            if [ ${RESPCT} -ge ${WARN} ] && [ ${RESPCT} -lt ${CRIT} ] && [ "${EXITCODE}" != "${NAGIOS_CRITICAL}" ]; then
                                EXITCODE=${NAGIOS_WARNING}
                                #echo "set EXIT CODE to NAGIOS_WARN"
                            fi



                            #echo "${anchester} $RESPCT $CUR $MAX"


                        done
                    else
                        echo "Unknown: Only own qgroup, no ancesters. Please use -s or -q to check"
                        exit ${NAGIOS_UNKNOWN}

                    fi 
                fi

            # pathrecursive end
            fi


            case $EXITCODE in

                ${NAGIOS_OK})
                               OUTTXT="OK:"
                               ;;
                ${NAGIOS_WARNING})
                               OUTTXT="WARNING:"
                               ;;
                ${NAGIOS_CRITICAL})
                               OUTTXT="CRITICAL:"
                               ;;
                *)
                               OUTTXT="UNKNOWN:"
                               ;;
            esac
            PERFDATA=$(echo ${PERFDATA}|tr '/' '-')
            if [ -z "${PATHRECURSIVE}" ];then
                echo -e "${OUTTXT} Checked for ${METH}-quota on ${SHORTTXTOUT}|${PERFDATA}\n${PLUGINLONGOUTPUT}"
            else
                echo -e "${OUTTXT} Checked recursive for ${METH}-quota on ${PATHRECURSIVE}|${PERFDATA}\n${PLUGINLONGOUTPUT}"
            fi
            exit $EXITCODE

	}



###########################################################################
###########################################################################
###                               M A I N                               ###
###########################################################################
###########################################################################
options $@
