#!/bin/bash
#
# Copyright (c) 2017-2018, SyLabs, Inc. All rights reserved.
# Copyright (c) 2017, SingularityWare, LLC. All rights reserved.
#
# See the COPYRIGHT.md file at the top-level directory of this distribution and at
# https://github.com/singularityware/singularity/blob/master/COPYRIGHT.md.
#
# This file is part of the Singularity Linux container project. It is subject to the license
# terms in the LICENSE.md file found in the top-level directory of this distribution and
# at https://github.com/singularityware/singularity/blob/master/LICENSE.md. No part
# of Singularity, including this file, may be copied, modified, propagated, or distributed
# except according to the terms contained in the LICENSE.md file.
#
# This file also contains content that is covered under the LBNL/DOE/UC modified
# 3-clause BSD license and is subject to the license terms in the LICENSE-LBNL.md
# file found in the top-level directory of this distribution and at
# https://github.com/singularityware/singularity/blob/master/LICENSE-LBNL.md.


set -u

if [ -z "${SINGULARITY_libexecdir:-}" ]; then
    echo "ERROR: SINGULARITY_libexecdir not defined in environment"
    exit 2
fi

if [ -z "${SINGULARITY_MESSAGELEVEL:-}" ]; then
    echo "Warning: SINGULARITY_MESSAGELEVEL is undefined, temporarily setting to '5' (all messages)"
    SINGULARITY_MESSAGELEVEL=5
fi

if [ -z "${USER:-}" ]; then
    USER=`id -un`
    export USER
fi
if [ -z "${HOME:-}" ]; then
    HOME=`getent passwd "$USER" | cut -d : -f 6`
    export HOME
fi


message() {
    LEVEL="${1:-}"
    MESSAGE="${2:-}"
    if [ -z "$MESSAGE" ]; then
        return 0
    fi
    shift
    shift

    # if level is symbolic, get a numeric value
    case "$LEVEL" in
	INFO) LEVNUM=1 ;;
	VERBOSE) LEVNUM=4 ;;
	DEBUG) LEVNUM=5 ;;
        [1-5]) LEVNUM="$LEVEL" ;;
        *) LEVNUM=1 ;; # default to INFO
    esac

    case "$LEVEL" in
        e|error|E|ERROR)
            tput -Txterm setaf 1 2>/dev/null
            printf "ERROR: $MESSAGE" "$@" 1>&2
            tput -Txterm sgr0 2>/dev/null
        ;;
        a|ABORT)
            tput -Txterm setaf 1 2>/dev/null
            printf "ABORT: $MESSAGE" "$@" 1>&2
            tput -Txterm sgr0 2>/dev/null
        ;;
        w|warn|warning|W|WARN|WARNING)
            tput -Txterm setaf 3 2>/dev/null
            printf "WARNING: $MESSAGE" "$@" 1>&2
            tput -Txterm sgr0 2>/dev/null
        ;;
        1|INFO)
            if [ "$LEVNUM" -le "$SINGULARITY_MESSAGELEVEL" ]; then
                printf "$MESSAGE" "$@"
            fi
        ;;
        p|P|CHECKPASS|PASS)
            tput -Txterm setaf 2 2>/dev/null
            printf "PASS: $MESSAGE" "$@" 1>&2
            tput -Txterm sgr0 2>/dev/null
        ;;
        NOTIFY|notify)
            tput -Txterm setaf 6 2>/dev/null
            printf "$MESSAGE" "$@" 1>&2
            tput -Txterm sgr0 2>/dev/null
        ;;
        f|F|CHECKFAIL|FAIL)
            tput -Txterm setaf 1 2>/dev/null
            printf "FAIL: $MESSAGE\n" "$@" 1>&2
            tput -Txterm sgr0 2>/dev/null
        ;;
        [2-5]|VERBOSE|DEBUG)
            if [ "$LEVNUM" -le "$SINGULARITY_MESSAGELEVEL" ]; then
                printf "$MESSAGE" "$@" 1>&2
            fi
        ;;
    esac

    return 0
}


singularity_key_get() {
    KEY="${1:-}"
    FILE="${2:-}"
    if OUT=`egrep -i "^$KEY:" $FILE`; then
        echo "$OUT" | head -n 1 | sed -e "s@^$KEY:\s*@@i" | sed -e "s@\s*#.*@@"
        return 0
    fi
    return 1
}


singularity_keys_get() {
    KEY="${1:-}"
    FILE="${2:-}"
    egrep -i "^$KEY:" "$FILE" | while read i; do
        echo "$i" | sed -e "s@^$KEY:\s*@@i" | sed -e "s@\s*#.*@@"
    done | tr '\n' ' '
    echo

    return 0
}

singularity_daemon_glob() {
    if ! USERID=`id -ru`; then
        message ERROR "Could not ascertain user ID\n"
        exit 255
    fi

    if ! HOST=`hostname`; then
        message ERROR "Could not ascertain hostname\n"
        ABORT 255
    fi

    if ! HOME=`getent passwd ${USERID} | cut -d: -f6`; then
        message ERROR "Could not discover user's home directory\n"
        ABORT 255
    fi

    for arg in "$@"; do
        for i in ${HOME}/.singularity/daemon/${HOST}/${arg}; do
            echo "$i"
        done
    done
    return 0
}

singularity_daemon_file() {
    SINGULARITY_DAEMON_NAME="${1:-}"

    if [ -z "${SINGULARITY_DAEMON_NAME:-}" ]; then
        message ERROR "singularity_daemon_file() called with no process name\n"
        ABORT 255
    fi

    if ! USERID=`id -ru`; then
        message ERROR "Could not ascertain user ID\n"
        ABORT 255
    fi
    
    if ! HOST=`hostname`; then
        message ERROR "Could not ascertain hostname\n"
        ABORT 255
    fi

    if ! HOME=`getent passwd ${USERID} | cut -d: -f6`; then
        message ERROR "Could not discover user's home directory\n"
        ABORT 255
    fi

    # This is a configurable option, needs to be set accordingly!
    SINGULARITY_DAEMON_FILE="${HOME}/.singularity/daemon/${HOST}/${SINGULARITY_DAEMON_NAME}"
    export SINGULARITY_DAEMON_FILE SINGULARITY_DAEMON_NAME

    message 2 "Using Daemon configuration file: ${SINGULARITY_DAEMON_FILE}\n"

    return 0
}

singularity_calculate_size() {
    FOLDER="${1:-}"
    if [ ! -d "${FOLDER}" ]; then
        message ERROR "Folder not found ($FOLDER)\n"
    fi

    eval du -sm ${FOLDER} | cut -f1 # MB

    return 0
}

ABORT() {
    RETVAL="${1:-}"
    if [ -z "$RETVAL" ]; then
        RETVAL=1
    fi
    message ABORT "Aborting with RETVAL=$RETVAL\n"
    exit $RETVAL
}

check_pattern() {
    STRING="${1:-}"
    PATTERN="${2:-}"
    case "$PATTERN" in
        $STRING)
            true
        ;;
        *)
            return 1
        ;;
    esac
    return 0
}

cmd() {
    message 2 " + %-68.68s" "$*"
    "$@" >/dev/null 2>&1
    RETVAL=$?
    if [ $RETVAL -eq 0 ]; then
        message 2 "OK\n"
    else
        message 2 "ERROR\n"
    fi
    return $RETVAL
}

eval_abort() {
    eval "$@"
    RETVAL=$?
    if [ $RETVAL -ne 0 ]; then
        exit $RETVAL
    fi
    return 0
}

stest() {
    ERROR="${1:-}"
    TMPFILE=`mktemp`
    shift
    message 2 " + %-80.80s " "$*"
    "$@" >$TMPFILE 2>&1
    CODE="$?"
    if [ "$ERROR" = "0" -a "$CODE" != "0" ]; then
        message 2 "%13s ERROR\n" "(retval=$CODE)"
        tail "$TMPFILE"
        echo "Full output in: $TMPFILE"
        exit 1
    elif [ "$ERROR" != "0" -a "$CODE" = "0" ]; then
        message 2 "%13s ERROR\n" "(retval=$CODE)"
        tail "$TMPFILE"
        echo "Full output in: $TMPFILE"
        exit 1
    else
        message 2 "%13s OK\n" "(retval=$CODE)"
    fi
    rm -f "$TMPFILE"
}


singularity_import() {
    MOD="${1:-}"
    if [ -z "$SINGULARITY_libexecdir" ]; then
        message ERROR "libexecdir not defined, are you running this from within Singularity?\n"
        exit 1
    fi
    if [ -f "$SINGULARITY_libexecdir/singularity/mods/$MOD.smod" ]; then
        . "$SINGULARITY_libexecdir/singularity/mods/$MOD.smod"
    else
        message ERROR "Could not load Singularity module: $MOD\n"
        exit 255
    fi
    return 0
}


# Different versions of which respond differently (print aliases, or take
# different arguments)
singularity_which() {
    i="${1:-}"
    # Avoid non-pathnames with . not in path, and directories
    case $i in
        .* | /*)
            if [ -f "$i" -a -x "$i" ]; then
                echo "$i"
                return 0
            fi
    esac
    for p in `echo $PATH | sed -e 's/:/ /g'`; do
        if [ -f "$p/$i" -a -x "$p/$i" ]; then
            echo "$p/$i"
            return 0
        fi
    done
    return 1
}


parse_opts() {
    NEWOPTS=""
    while [ -n "${1:-}" ]; do
        case "${1:-}" in
            -*=*)
                ARG1=`echo -n "${1:-}" | cut -d = -f 1`
                ARG2=`echo -n "${1:-}" | cut -d = -f 2`
                NEWOPTS="$NEWOPTS \"$ARG1\" \"$ARG2\""
                shift
                continue
            ;;
            --*)
                NEWOPTS="$NEWOPTS \"${1:-}\""
                shift
                continue
            ;;
            -*)
                for o in `echo "${1:-}"| sed 's/^-//' | sed 's/./-& /g'`; do
                    NEWOPTS="$NEWOPTS \"$o\""
                done
                shift
                continue
            ;;
            *)
                NEWOPTS="$NEWOPTS $@"
                break
            ;;
        esac
    done
    # Eww, this is bad I know... Would be better just to pass the variable
    # around without making ad-hoc modifications... Got a better idea, let
    # me know! (gmk)
    echo "$NEWOPTS" | sed -e 's/\\/\\\\/g'
}


compare_envs() {
    # during bootstrap, if we are bootstrapping from a local image, compare the
    # environment in the new image to that in the old image and alert the user
    # of any changes
    REMOVED=$(comm -23 ${SINGULARITY_STARTING_ENVIRONMENT} ${SINGULARITY_ENDING_ENVIRONMENT})
    ADDED=$(comm -13 ${SINGULARITY_STARTING_ENVIRONMENT} ${SINGULARITY_ENDING_ENVIRONMENT})
     
    if [ -n "${REMOVED:-}" -o -n "${ADDED:-}" ]; then
        message 1 "Environment variables were added, removed, and/or changed during bootstrap.\n"
    elif [ "${SINGULARITY_STARTING_ENVSHA1:-}" != "${SINGULARITY_ENDING_ENVSHA1:-}" ]; then
        message 1 "The environment of ${SINGULARITY_IMAGE} may differ from the environment of ${FROM}\n"
        message 1 "${FROM} sha1 sum of environment is ${SINGULARITY_STARTING_ENVSHA1}\n"
        message 1 "${SINGULARITY_IMAGE} sha1 sum of environment is ${SINGULARITY_ENDING_ENVSHA1}\n"
    fi
    if [ -n "${REMOVED:-}" ]; then
        message 1 "Variables unique to original image (${FROM})\n"
        message 1 "$REMOVED\n"
    fi
    if [ -n "${ADDED:-}" ]; then
        message 1 "Variables unique to new image (${SINGULARITY_IMAGE})\n"
        message 1 "$ADDED\n"
    fi
}


replace_string() {
   ORIGINAL="${1:-}"
   MATCH="${2:-}"
   REPLACE="${3:-}"
   echo $ORIGINAL | sed -e s/${MATCH}/${REPLACE}/g
}


zcat_compat() {
    if [ -z "${1:-}" ]; then
        message FAIL "zcat_compat: missing argument\n"
        exit 1
    fi

    FILE=$1
    HEAD_CHUNK=`head -c 3 $FILE | od -A n -t x1 -w3 | tr -d ' '`

    if [ "$HEAD_CHUNK" == "1f8b08" ]; then
        zcat "$1"
    else
        cat "$1"
    fi
}

nonroot_build_warning() {
    USERID=`id -ru`
    if [ "$USERID" != "0" ]; then
        message WARNING "Building container as an unprivileged user. If you run this container as root\n";
        message WARNING "it may be missing some functionality.\n"
    fi
}

is_deffile() {
# check if a file looks like, walks like, and quacks like a def file
    FILE2CHECK=$1
    
    # first make a guess based on naming convention
    case ${FILE2CHECK:-} in
        *.img)
            return 1
            ;;
        *.tar*)
            return 1
            ;;
        *.def)
            return 0
            ;;
        Singularity) 
            return 0
            ;;
        *)
            ;;
        esac

    # name is ambiguous.  is it a text file with def file words in it?
    if ! `file $FILE2CHECK | grep -q "ASCII\|Unicode"`; then
        return 1
    fi
        
    grep -qiE "^bootstrap:|^%setup|^%post|^%test|^%labels|^%files|^%runscript|^%environment" "$FILE2CHECK"
    RETVAL=$?

    return $RETVAL
}


is_image() {
# check if a file looks like, walks like, and quacks like an image file

    if [ -z "${1:-}" ]; then
        return 1
    fi

    eval "$SINGULARITY_libexecdir/singularity/bin/image-type" "$1" >/dev/null 2>&1
    RETVAL=$?

    return $RETVAL
}


is_tar() {
# check if a file looks like, walks like,... oh you get the idea
    FILE2CHECK=$1
    TARMAGIC="7573746172202000" # "ustar  \x00"
    PAX_TARMAGIC="7573746172003030" # "ustar.00"; dot denotes \x00 (NULL)

    zcat_compat "$FILE2CHECK" | head -c 320 | od -A n -t x1 -w320 | tr -d ' ' | grep -q -e "$TARMAGIC" -e "$PAX_TARMAGIC"
    RETVAL=$?
    return $RETVAL
}


##############################################################################
# CHECKS
##############################################################################


# Given a script and an assigned level and tag, determine if it should be
# executed to perform the check
# [LEVEL] [SCRIPT] [TAGS]

exec_check() {
    CHECKSUCCESS="${1:-}"
    CHECKLEVEL="${2:-}"
    CHECKSCRIPT="${3:-}"
    shift
    shift
    shift
    CHECKTAGS="${@}"

    # The default case is that we don't run
    RETVAL=0  

    # Don't perform check if tag not indicated
    PERFORM_CHECK=1

    # if not "all," check if user wants check performed
    if [ "$SINGULARITY_CHECKTAGS" != "all" ]; then
        has_tag "$CHECKTAGS" "$SINGULARITY_CHECKTAGS" || PERFORM_CHECK=0
    fi

    if [ $PERFORM_CHECK -eq 0 ]; then
        return $RETVAL
    fi
 
    RUNCHECK_LEVEL=3        # Perform all checks
    case "$CHECKLEVEL" in
        l|low|L|LOW)
            RUNCHECK_LEVEL=3
        ;;
        m|med|M|MED)
            RUNCHECK_LEVEL=2
        ;;
        h|high|H|HIGH)
            RUNCHECK_LEVEL=1
        ;;
    esac

    if [ "$RUNCHECK_LEVEL" -le "$SINGULARITY_CHECKLEVEL" ]; then
       message NOTIFY "START ${CHECKSCRIPT##*/} tags[${CHECKTAGS}] level[$CHECKLEVEL]\n"
       scheck 0 $CHECKSCRIPT # script must have interpreter specified
       printf " $CHECKSCRIPT \n"
    fi

    return $RETVAL
}


has_tag() {

    taglist="$1"
    tag="$2"
    if test "${taglist#*$tag}" != "$taglist"
    then
        return 0 # has tag, perform test
    else
        return 1 # no tag, skip test
    fi
}


scheck() {
    ERROR="${1:-}"
    shift
    message 2 " + %-80.80s " "$*"
    "$@"
    CODE="$?"
    if [ "$ERROR" = "0" -a "$CODE" != "0" ]; then
        message FAIL "(retval=$CODE)"
    elif [ "$ERROR" != "0" -a "$CODE" = "0" ]; then
        message FAIL "(retval=$CODE)"
    else
        message PASS "(retval=$CODE)"
    fi
    return $CODE
}

atexit() {
    trap "$1" EXIT
}

if [ -n "${SHELL_DEBUG:-}" ]; then
    set -x
fi
