#! /bin/bash

###############################################################################
#
#   trim_disks - Provides support to enable and execute TRIM on disks which
#   require it.
#
#  Copyright (C) 2019 Cumulus Networks, Inc. All rights reserved
#
#  This software is subject to the Cumulus Networks End User License Agreement
#  available  at the following locations:
#
#  Internet: https://cumulusnetworks.com/downloads/eula/latest/view/
#  Cumulus Linux systems: /usr/share/cumulus/EULA.txt
#
#  Description: This script can be run to determine if there are any disks in
#      the system which require that TRIM be enabled. If so, then the discard
#      option is added to the /etc/fstab file for those disk's mounts and the
#      fstrim command is executed. But since that will only take affect when
#      those file systems are re-mounted, a cron job is added to run every six
#      hours executing fstrim on drives which require it but are not currently
#      mounted with the discard option.
#  Syntax:
#       trim_disks [-c] [-v...] [-f] [-t <hours>]
#
#       -c : Used when run as a cron job. Skips the processing which adds a
#            cron job.
#       -v... : Verbose level. Can be specified multiple times to increase the
#            amount of verbosity. Without this flag no output will typically
#            be produced.
#       -f : Force all disks to be classified as requiring TRIM. This is mainly
#            used for debugging.
#       -t <hours> : Specify the frequency the CRON job will run, in hours. The
#            default is 6.
#
###############################################################################

#------------------------------------------------------------------------------
#
#   Global variables and "constants"
#
#------------------------------------------------------------------------------

# Restrict non-path'd execuatbles to known good directories
PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin

# Command line switches
cron=0
verbose=0
force=0
hours=6

# Location of the cron file
cronfile="/etc/cron.d/trim_disks"

# Array which contains one entry for each disk device, i.e. /dev/sda (or
# partition, i.e. /dev/sda2) which requires TRIM support.
declare -a disk_devs

# Array which contains one entry for each mounted file system on a disk which
# requires TRIM.
declare -a trim_req_mounts

# Array which contains one entry for each mounted file system on a disk which
# requires TRIM but which is not currently mounted with the discard option.
declare -a no_discard_mounts

# Paths to executables used by this script
REALPATH="/usr/bin/realpath"
SUDO="/usr/bin/sudo"
ID="/usr/bin/id"
RM="/bin/rm"
CAT="/bin/cat"
TEE="/usr/bin/tee"
FSTRIM="/sbin/fstrim"
BASENAME="/usr/bin/basename"
SED="/bin/sed"
CP="/bin/cp"
AWK="/usr/bin/awk"
CMP="/usr/bin/cmp"

# The name of the script, used for printing output
base=$(${BASENAME} $0)


#------------------------------------------------------------------------------
#
#   Functions
#
#------------------------------------------------------------------------------

# Check to make sure that the tools we'll use are available
tool_check()
{
    TOOLS="${REALPATH} ${SUDO} ${ID} ${RM} ${CAT} ${TEE} ${FSTRIM}
           ${SED} ${CP} ${AWK} ${CMP}"
    for TOOL in ${TOOLS}
    do
        [ -x ${TOOL} ] || {
            echo "${base}:  The required ${TOOL} program does not exist or can't be executed."
            exit 1
        }
    done

    return 0
}


# Check to make sure that this script can be run with root privileges
root_check()
{
    MYID=$(${SUDO} -n ${ID} -u)
    [ -n "${MYID}" ] && [ ${MYID} -eq 0 ] || {
        echo "${base}:  You must have root privileges to run this command."
        exit 1
    }

    return 0
}

# Display some information about how to use this program
usage()
{
    /bin/cat << 'EOF'
   Syntax:
        trim_disks [-c] [-v...] [-f] [-t <hours>]

    This script can be run to determine if there are any disks in
    the system which require that TRIM be enabled. If so, then the discard
    option is added to the /etc/fstab file for those disk's mounts and the
    fstrim command is executed. But since that will only take affect when
    those file systems are re-mounted, a cron job is added to run every six
    hours executing fstrim on drives which require it but are not currently
    mounted with the discard option.

     -c : Used when run as a cron job. Skips the processing which adds a
          cron job.
     -v... : Verbose level. Can be specified multiple times to increase the
          amount of verbosity. Without this flag no output will typically
          be produced.
     -f : Force all disks to be classified as requiring TRIM. This is mainly
          used for debugging.
     -t <hours> : Specify the frequency the CRON job will run, in hours. The
          default is 6.
EOF
    exit 1
}

# Get the program options
get_options()
{
    while getopts "cvft:" o
    do
        case "${o}" in
            c)      cron=1
                    SUDO=""
                    ;;
            v)      verbose=$((verbose+1))
                    ;;
            f)      force=1
                    ;;
            t)      hours=$OPTARG
                    ;;
            *)      usage
                    ;;
        esac
    done
    shift $((OPTIND-1))
}

# Determine is something is an element of an array. The first parameter is the
# element to search for and the second parameter is the array in which to
# search.
isElement()
{
    local e match="$1"
    shift
    for e; do [[ "${e}" == "${match}" ]] && return 0; done
    return 1
}

# Populate the disk_devs array with a list of disk drives, or partitions thereof,
# which require TRIM support. This is done by examining the file names in the
# /dev/disk/by-id directory for disks which are known to require TRIM. Also, if
# the -f command line option was provided, all disks in that directory are
# treated as requiring TRIM. The number of such disks found is the return code.
getTrimRequiredDiskDevs()
{
    local file
    local dev

    # Does the file system have a directory with disks listed by ID?
    disk_devs=()
    if [ ! -d /dev/disk/by-id ]
    then
        [ ${verbose} -ge 1 ] && echo "The /dev/disk/by-id directory does not exist."
        return 0
    fi

    # Examine each file name in the directory
    for file in /dev/disk/by-id/*
    do
        # Does the file name contain one of the strings identifying TRIM-required disks?
        if [[ ${force} -eq 1 ]] ||
           [[ ${file} == *"_3ME3_"* ]] ||
           [[ ${file} == *"_3IE3_"* ]] ||
           [[ ${file} == *"_3IE4_"* ]]
        then
            # Get the name of the disk device, e.g. /dev/sda, and add it to the
            # disk_devs array if it is not already there.
            dev=$(${REALPATH} -e ${file})
            isElement ${dev} "${disk_devs[@]}" || disk_devs+=(${dev})
        fi
    done

    [ ${verbose} -ge 2 ] && echo "The TRIM required disks are: ${disk_devs[@]}"
    return ${#disk_devs[@]}
}

# Examine the /etc/mtab for mounts of disks which require TRIM but were not
# mounted with the discard option. All such mount points are added to the
# no_discard_mounts array and the number of elements in the array is returned.
getMountsWithoutRequiredDiscard()
{
    local line
    local words

    # Does the /etc/mtab file exist and is it readable?
    no_discard_mounts=()
    trim_req_mounts=()
    if [ ! -r /etc/mtab ]
    then
        [ ${verbose} -ge 1 ] && echo "The /etc/mtab file does not exist or cannot be read."
        return 0
    fi

    # Process each line the the /etc/mtab file. If the first token is a device
    # which requires TRIM and the discard option isn't in the options, then add
    # the mount point to the array.
    while read -r line
    do
        # Break the line into tokens and make sure there are enough of them
        words=( ${line} )
        [ ${#words[@]} -ge 4 ] || continue
        # Is this a device which requires TRIM?
        isElement ${words[0]} "${disk_devs[@]}" || continue
        # Add to list of mounts of trim required devices if not already there
        isElement ${words[1]} "${trim_req_mounts[@]}" ||
            trim_req_mounts+=(${words[1]})
        # Is the discard option set? If not, is this mount point already in the
        # array? If not, add it.
        [[ ${words[3]} == *"discard"* ]] ||
            isElement ${words[1]} "${no_discard_mounts[@]}" ||
                no_discard_mounts+=(${words[1]})
    done < /etc/mtab

    [ ${verbose} -ge 2 ] && echo "Mount points not running with TRIM: ${no_discard_mounts[@]}"
    return ${#no_discard_mounts[@]}
}

# Remove the cron file which does the periodic fstrim.
removeCronFile()
{
    if [ -e ${cronfile} ]
    then
        [ ${verbose} -ge 1 ] && echo "Removing cron file ${cronfile}."
        ${SUDO} ${RM} ${cronfile}
    fi
    return
}

# Add a system cron file which will run this script to execute fstrim, among
# other checks, every six hours.
addCronFile()
{
    local crontext

    # Has the cron file already been installed?
    [ ${force} -eq 1 ] || [ ! -e ${cronfile} ] || {
        [ ${verbose} -ge 1 ] && echo "Not adding cron file ${cronfile} because it already exists"
        return
    }

    # Does the trim_disks script exist?
    [ ${force} -eq 1 ] || [ -x "/usr/bin/trim_disks" ] || {
        [ ${verbose} -ge 1 ] && echo "Not adding cron file ${cronfile} because /usr/bin/trim_disks is not executable."
        return
    }

    read -r -d '' crontext <<EOF
#  Copyright (C) 2019 Cumulus Networks, Inc. All rights reserved
#
#  This software is subject to the Cumulus Networks End User License Agreement
#  available  at the following locations:
#
#  Internet: https://cumulusnetworks.com/downloads/eula/latest/view/
#  Cumulus Linux systems: /usr/share/cumulus/EULA.txt
#
#  Some systems have disks which require that the fstrim command be run
#  periodically. This cron file runs the trim_disks command to accomplish this.

SHELL=/bin/sh
PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin

# m h   dom mon dow user  command
  5 */${hours} *   *   *   root  /usr/bin/trim_disks -c
EOF
    [ ${verbose} -ge 1 ] && echo "Adding cron file ${cronfile}."
    ${SUDO} echo "${crontext}" | ${SUDO} ${TEE} ${cronfile} > /dev/null || {
        echo "An error was encountered installing the cron file."
    }
}

# Add the discard option to entries in /etc/fstab which require it. In order to
# correlate an entry in /etc/fstab with a disk which requires TRIM the /dev/disk
# directory structure is examined. The first token in each /etc/fstab line is
# the file system, which can be specified in a number of ways:
#    a device, i.e. /dev/sda3
#    a UUID, i.e. UUID=1685583e-6970-4769-a89c-c5c5b534d15a
#    a label, i.e. LABEL=CL-SYSTEM
#    a partition label, i.e. PARTLABEL=CL-SYSTEM
#    a partition UUID, i.e. PARTUUID=324767e4-1a5f-4f11-912a-5ff826cf17fc
# An array of file system specifications is created for all possible ways to
# specify filesystems of disks which require TRIM. If the first token in a line
# of /etc/fstab matches an entry in the array, then that line in /etc/fstab
# will have the discard option added, if it is not already there.
addDiscardOptionToFstab()
{
    local trim_req_fs
    local disk_id
    local prefix
    local file
    local dev
    local mount
    local sed_cmds

    # Is there an /etc/fstab file?
    if [ ! -e /etc/fstab ]
    then
        [ ${verbose} -ge 2 ] && echo "Not adding discard option because /etc/fstab does not exist"
        return
    fi

    # Build the trim_req_fs array. Start with the TRIM required device names.
    trim_req_fs=("${disk_devs[@]}")
    # Go through each of the directories in /dev/disk, except by-id
    for disk_id in label partlabel partuuid uuid
    do
        # Convert the prefix to uppercase then look at each file in the directory
        prefix=$(echo ${disk_id} | ${AWK} '{print toupper($0)}')
        for file in /dev/disk/by-${disk_id}/*
        do
            # Get the device to which the file points. Is it a device which
            # requires trim? If so, is the file system spec already in the
            # trim_req_fs array? If not, add it.
            dev=$(${REALPATH} -e ${file})
            mount=${prefix}"="$(${BASENAME} ${file})
            isElement ${dev} "${disk_devs[@]}" &&
                isElement ${mount} "${trim_req_fs[@]}" ||
                    trim_req_fs+=(${mount})
        done
    done
    [ ${verbose} -ge 2 ] && echo "Possible filesystem specifiers which require TRIM: ${trim_req_fs[@]}."

    # Create a list of sed commands which can be used to modify /etc/fstab to
    # add TRIM on the lines which require it. 
    sed_cmds=()
    for mount in "${trim_req_fs[@]}"
    do
        # If the discard option is not present, add it to the options if the first
        # token matches an entry in trim_req_fs.
        sed_cmds+=("-e /\(\S\+\s\+\)\{3\}\S*discard\S*\s\+/!s ^\(${mount}\(\s\+\S\+\)\{3\}\) \1,discard ")
    done
    [ ${verbose} -ge 2 ] && echo "sed commands to add discard option: ${sed_cmds[@]}."

    # To be safe, the /etc/fstab file is not modified directly. Instead, it is
    # copied to a temporary file /tmp/fstab.orig and the stream editor takes that
    # file as input, producing /tmp/fstab.new. The /tmp/fstab.new is then
    # copied back over /etc/fstab.
    ${SUDO} ${CAT} /etc/fstab > /tmp/fstab.orig || {
        [ ${verbose} -ge 1 ] && echo "Could not make a temporary copy of /etc/fstab"
        return
    }
    ${SED} "${sed_cmds[@]}" /tmp/fstab.orig > /tmp/fstab.new || {
        [ ${verbose} -ge 1 ] && echo "Could not add discard option to /tmp/fstab.new file"
        return
    }
    ${SUDO} ${CMP} -s /tmp/fstab.new /etc/fstab || {
        ${SUDO} ${CP} /tmp/fstab.new /etc/fstab || {
            [ ${verbose} -ge 1 ] && echo "Could not copy /tmp/fstab.new file to /etc/fstab"
            return
        }
        [ ${verbose} -ge 1 ] && echo "/etc/fstab has been modified to add discard option where required"
    }

    return
}

# Run the fstrim command on all file system which require TRIM but are not
# currently mounted with the discard option (mount points in no_discard_mounts
# array).
runFstrim()
{
    local mounts=("$@")
    local mount
    local fsoptions=""

    # If no paramter was supplied, use the no_discard_mounts array
    [ -z "${mounts}" ] && mounts=("${no_discard_mounts[@]}")

    # If in verbose mode, run fstrim with the --verbose flag
    [ ${verbose} -ge 1 ] && fsoptions+=' --verbose'

    # Run fstrim on each mount point
    for mount in "${mounts[@]}"
    do
        [ ${verbose} -ge 1 ] && {
            echo "Executing fstrim on ${mount}: ${SUDO} ${FSTRIM} ${fsoptions} ${mount}"
        }
        ${SUDO} ${FSTRIM} ${fsoptions} ${mount}
    done
    return
}

#------------------------------------------------------------------------------
#
#   Execution of script starts here.
#
#------------------------------------------------------------------------------

# Do some initial checks to make sure we can run this script. Are all the required
# tools present and accounted for, and are we either running as root or have
# passwordless sudo access?
get_options "$@"
tool_check
[ ${cron} -eq 0 ] && root_check

# Populate the disk_devs array with devices which requie TRIM. If there are none
# we can remove the cron file, if any, and stop.
getTrimRequiredDiskDevs && { 
    [ ${verbose} -ge 1 ] && echo "There are no disks in the system which require TRIM."
    removeCronFile
    exit 0
}

# Add the discard option to all entries in /etc/fstab which require it
addDiscardOptionToFstab

# Prepare to run fstrim by making a list of all mounts which require TRIM but
# are not currently mounted with the discard option. If there are none, then
# we no longer need to run the cron job, so remove it, run fstrim one last time
# (after disks were mounted with discard option) and exit.
getMountsWithoutRequiredDiscard && {
    [ ${verbose} -ge 1 ] && echo "There are no disks mounted without the required discard option."
    removeCronFile
    runFstrim "${trim_req_mounts[@]}"
    exit 0
}
runFstrim

# If we are running from the cron task, we can leave here.
[ ${cron} -ne 0 ] && {
    [ ${verbose} -ge 1 ] && echo "Execution of CRON script $@ complete."
    exit 0
}

# Add a system cron job so that fstrim is run periodically
addCronFile
exit 0
