#!/bin/bash
#
# Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.
#
set -euo pipefail

# Our options
declare -a DEVICES
declare STORAGE_DRIVER=
declare FORCE=
declare QUIET=

# Supported storage driver names here
declare -ar _SUPPORTED_DRIVERS=("btrfs" "overlay2")

# For cleanup
declare _CREATED_CONFIG=
declare _CREATED_MOUNT=
declare _MODIFIED_CONFIG=
declare _MODIFIED_FSTAB=
declare _SUCCESS=

function cleanup {
    # In the event that things go south, clean up various changes we might have made
    if [ -z ${_SUCCESS} ] ; then
        # Restore original configuration if we fail
        [ ${_CREATED_MOUNT} ] && sed -i "/added by $(basename $0)/d" /etc/fstab
        [ ${_MODIFIED_FSTAB} ] && mv /etc/fstab.save_$$ /etc/fstab
        [ ${_CREATED_CONFIG} ] && rm /etc/docker/daemon.json
        [ ${_MODIFIED_CONFIG} ] && restore_config
    else
        # Otherwise just clean up
        if [ ${_MODIFIED_CONFIG} ] ; then
            rm -f /etc/systemd/system/docker.service.d/*.save_$$
            rm -f /etc/sysconfig/docker-storage.save_$$
            rm -f /etc/docker/daemon.json.save_$$
        fi
        [ ${_MODIFIED_FSTAB} ] && rm /etc/fstab.save_$$
    fi
}

function log {
    # Implements the quiet option
    # XXX maybe log to fs if quiet?
    [ -z $QUIET ] && echo $@
}

function check_storage_driver_valid {
    # Ensure the storage driver is supported, see _SUPPORTED_DRIVERS above.
    for d in ${_SUPPORTED_DRIVERS[@]} ; do
        [ "$d" == ${STORAGE_DRIVER} ] && return 0
    done
    return 1
}

function create_btrfs {
    # Create btrfs file system on the devices passed.
    # XXX add raid-level and md/data options
    if ! yum list installed btrfs-progs >/dev/null 2>&1 ; then
        yum makecache --quiet fast && sudo yum install --quiet -y btrfs-progs
    fi

    log "Creating 'btrfs' file system on: ${DEVICES[@]}"

    # create btrfs on device list (--quiet still makes noise on stdout)
    mkfs.btrfs -q -f ${DEVICES[@]} 1> /dev/null
}

function create_overlay2 {
    # Create btrfs file system on the devices passed.
    # XXX add support for creating overlays on files
    # XXX support other file system types (e.g. ext4, btrfs)?
    if ! yum list installed xfsprogs >/dev/null 2>&1 ; then
        sudo yum makecache --quiet fast && sudo yum install --quiet -y xfsprogs
    fi

    log "Creating 'xfs' file system on: ${DEVICES[@]}"

    mkfs.xfs -q -f -n ftype=1 ${DEVICES[@]}
}

function setup_docker_root {
    # Set up the docker root directory and fstab entry, and then mount the file system
    if [ ! -d /var/lib/docker ] ; then
        mkdir -p /var/lib/docker
    fi

    tmp=$(blkid -o udev ${DEVICES[0]} | head -1)
    first_dev_uuid=${tmp#*=}
    [ ${STORAGE_DRIVER} == 'btrfs' ] && fstype=btrfs || fstype=xfs
    echo "UUID=${first_dev_uuid} /var/lib/docker ${fstype} defaults 0 0 # added by $(basename $0)" \
         >> /etc/fstab
    _CREATED_MOUNT=1

    if ! err=$(mount /var/lib/docker 2>&1) ; then
        log "Failed to mount file system: $err"
        return 1
    fi
}

function check_devices {
    # First, return non-zero if a device passed is not a block device, or if it is currently mounted.
    # Second, return non-zero if the device passed has a file system on it, unless force is set.
    if [ ${STORAGE_DRIVER} != 'btrfs' ] && [ ${#DEVICES[@]} -gt 1 ] ; then
        log "Only single device supported for overlay2 (xfs)"
        return 1
    fi

    declare ret=
    for device in ${DEVICES[@]} ; do
        if [ ! -b $device ] ; then
            log "Device '${device}' not a block device"
            ret=1
        fi
        if grep --quiet "^${device} " /proc/mounts ; then
            log "Device '${device}' is currently mounted"
            ret=1
        fi
    done

    declare fsfound=
    if [ -z $FORCE ] ; then
        for device in ${DEVICES[@]} ; do
            if ! file -s ${device} | grep -q ": data$" ; then
                log "File system found on '${device}'"
                fsfound=1
            fi
        done
        if [ $fsfound ] ;then
            log "Use different device(s) or force with -f."
            ret=1
        fi
    fi

    return $ret
}

function check_mounted_fs_compatible {
    # Called if no devices are passed from the user.
    # Check that the would-be docker root file system is compatible with the storage driver specified.
    for mnt in /var/lib/docker /var/lib /var ; do
        if [ -d ${mnt} ] ; then
            fstype=$(stat -f -c %T ${mnt})
            if [ ${STORAGE_DRIVER} == 'btrfs' ] ; then
                if [ ${fstype} != 'btrfs' ] ; then
                    log "Existing file system on ${mnt} not btrfs, provide device list."
                    return 1
                fi
            else    # overlay2
                if [ ${fstype} == 'xfs' ] ; then
                    if xfs_info ${mnt} | grep --quiet 'ftype=0' ; then
                        log "'xfs' on ${mnt} has unsupported ftype configuration," \
                            "provide device-list to create new file system."
                        return 1
                    fi
                else
                    log "'overlay2' supported on 'xfs' only, ${mnt} file system is '${fstype}.'"
                    return 1
                fi
            fi
            return 0
        fi
    done
}

function check_existing_root {
    # Check for an existing Docker root. If found, do not clobber and error out.
    # Check for entry in /etc/fstab for unmounted docker root, and clobber that if force is set.
    # XXX should --f flag clobber existing deployment / mount?
    if [[ -d /var/lib/docker && "$(ls -A /var/lib/docker)" ]] ; then
        log "Existing deployment found at /var/lib/docker, empty or remove directory and retry."
        return 1
    fi

    if grep " /var/lib/docker " /etc/fstab | grep --quiet -v "^#" ; then
        if [ $FORCE ] ; then
            sed -i.save_$$ "/ \/var\/lib\/docker /d" /etc/fstab
        else
            log "Mount already specified for /var/lib/docker in /etc/fstab, remove or force with -f."
            return 1
        fi
    fi
}

function check_and_set_config {
    # Check for existing configuration in /etc/docker, /etc/sysconfig and /etc/systemd/system.
    # If not forcing, trace where configuration is found and return non-zero.
    #
    # If we are forcing configuration, swap out any 'storage-driver' value for our new storage
    # driver in the daemon.json file (or create a new file), and make sure we blot out any conflicting
    # config elsewhere.
    ret=0
    if [ -f /etc/docker/daemon.json ] && \
            grep storage-driver /etc/docker/daemon.json | grep --quiet -v "^#" ; then
        if ! grep storage-driver /etc/docker/daemon.json | grep --quiet ${STORAGE_DRIVER} ; then
            if [ $FORCE ] ; then
                # Save a copy of the old file and swap in our new driver. Remove options.
                sed -i.save_$$ "s/\(\"storage-driver\": \).*/\1\"${STORAGE_DRIVER}\"/" \
                    /etc/docker/daemon.json
                sed -i '/storage-opt/d' /etc/docker/daemon.json # XXX issues here with multi-line opts
                log "Removed existing storage configuration from /etc/docker/daemon.json."
            else
                log "Found conflicting storage configuration in /etc/docker/daemon.json..."
                ret=1
            fi
        fi
    else
        # If there was no daemon.json, create one
        [ ! -d /etc/docker ] && mkdir -m 0700 /etc/docker && chown root:root /etc/docker
        echo "{ \"storage-driver\": \"${STORAGE_DRIVER}\" }" > /etc/docker/daemon.json
        cat << EOF > /etc/docker/daemon.json
{
    "storage-driver": "${STORAGE_DRIVER}"
}
EOF
        _CREATED_CONFIG=1
    fi

    if [ -d /etc/systemd/system/docker.service.d ] && \
            grep storage-driver /etc/systemd/system/docker.service.d/* | grep --quiet -v "^#" ; then
        if [ $FORCE ] ; then
            for f in $(grep -rl storage-driver /etc/systemd/system/docker.service.d) ; do
                if grep storage-driver $f |grep --quiet -v "^#" ; then
                    sed -i.save_$$ '/storage-driver/ s/^#*/#/' $f
                    sed -i '/storage-opts/ s/^#*/#/' $f
                fi
            done
            log "Removed existing storage configuration from /etc/systemd/system/docker.service.d."
        else
            log "Found existing storage configuration in /etc/systemd/system/docker.service.d..."
            ret=1
        fi
    fi

    if [ -f /etc/sysconfig/docker-storage ] && \
            grep --quiet "^DOCKER_STORAGE_OPTIONS" /etc/sysconfig/docker-storage ; then
        if [ $FORCE ] ; then
            sed -i.save_$$ '/^DOCKER_STORAGE_OPTIONS/ s/^#*/#/' /etc/sysconfig/docker-storage
            log "Removed existing storage configuration from /etc/sysconfig/docker-storage."
        else
            log "Found existing storage configuration in /etc/sysconfig/docker-storage..."
            ret=1
        fi
    fi

    [[ $FORCE && $ret ]] && _MODIFIED_CONFIG=1

    return $ret
}

function restore_config {
    # Called from cleanup, put things back in order before we fail
    [ -z $_MODIFIED_CONFIG ] && return
    [ -f /etc/docker/daemon.json.save_$$ ] && \
        mv /etc/docker/daemon.json.save_$$ /etc/docker/daemon.json
    [ -f /etc/sysconfig/docker-storage.save_$$ ] && \
        mv /etc/sysconfig/docker-storage.save_$$ /etc/sysconfig/docker-storage
    for f in /etc/systemd/system/docker.service.d/* ; do
        [ -f $f.save_$$ ] && \
            mv $f.save_$$ $f
    done
    if [ -d /etc/systemd/system/docker.service.d ] && \
            grep --quiet storage-driver /etc/systemd/system/docker.service.d/* ; then
        for f in $(grep -rl storage-driver /etc/systemd/system/docker.service.d) ; do
            [ -f $f.save_$$ ] && mv $f.save_$$ $f
        done
    fi
}

function usage {
    echo
    echo "$(basename $0) [-f] [-q] [-s (btrfs|overlay2)] [-d </dev/sdX> [-d </dev/sdY>] .. ]"
    echo "    -s  storage driver ('btrfs' or 'overlay2')"
    echo "    -d  device to create file system on, repeat option for multiple devices"
    echo "    -f  force option, clobber existing configuration and force file system creation"
    echo "    -q  quiet option, output only error messages"
    echo
    echo "  Create file system specified by driver on the provided device(s) and configure container "
    echo "  runtime storage options. Note selection of 'overlay2' driver specifies an 'xfs' file system."
}

# Parse options, then onto main script
_OPTIONS=':d:s:fqh'
while getopts $_OPTIONS OPT ; do
    case ${OPT} in
        d)
            if [[ ${OPTARG} == -* ]] ; then
                echo "Missing device list argument for -${OPT}"
                usage && exit 1
            fi
            DEVICES+=(${OPTARG})
            ;;
        s)
            if [[ ${OPTARG} == -* ]] ; then
                echo "Missing storage driver argument for -${OPT}"
                usage && exit 1
            fi
            STORAGE_DRIVER=${OPTARG}
            ;;
        f)
            FORCE=true
            ;;
        q)
            QUIET=true
            ;;
        h)
            usage && exit 0
            ;;
        \?)
            echo "Unknown option: -${OPTARG}"
            usage && exit 1
            ;;
        :)
            echo "Missing argument for -${OPTARG}"
            usage && exit 1
            ;;
        *)
            usage && exit 1
            ;;
    esac
done

# Make sure we don't have any errant arguments
shift $((OPTIND-1))
if [ $# -gt 0 ] ; then
    args=("$@")
    echo "Unsupported argument: ${args[0]}"
    usage && exit 1
fi

# Must be root to do many of the things this script does
if [ $EUID -ne 0 ]; then
    echo "This script must be run by a privileged user."
    exit 1
fi

# If user passed a device list, make sure they passed a storage driver as well
if [ -z ${STORAGE_DRIVER} ] ; then
    if [[ ${DEVICES[@]:-} ]] ; then
        log "Provided device list requires storage driver selection."
        exit 1
    fi
    exit 0
fi

# Ensure storage driver is on our list of supported drivers
if ! check_storage_driver_valid ; then
    log "Driver '${STORAGE_DRIVER}' unsupported, supported storage drivers: ${_SUPPORTED_DRIVERS[@]}"
    exit 1
fi

# Don't bounce Docker if it's already running
if systemctl -q is-active docker ; then
    log "Docker service is currently active, stop and retry."
    exit 1
fi

if [[ ${DEVICES[@]:-} ]] ; then
    # User passed us devices, make sure they are usable
    if ! check_devices ; then
        exit 1
    fi
else
    # With no devices passed, check that mounted fs type is compatible with the storage driver
    if ! check_mounted_fs_compatible ; then
        exit 1
    fi
fi

# See if there's an existing docker root, smash if force is set
if ! check_existing_root ; then
    exit 1
fi

# From here, we've potentially modified the system so set up our cleanup in case we fail
trap cleanup EXIT

# Check for existing config, if force is set clobber it with our new config
if ! check_and_set_config ; then
    if [ -z $FORCE ] ; then
        log "Remove existing configuration and retry or force with -f."
        exit 1
    fi
fi

# Finally, create the file system and set up the docker root
create_${STORAGE_DRIVER}
setup_docker_root
_SUCCESS=1
exit 0
