#!/bin/bash
#
# Copyright 2013 Red Hat
# All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

set -e  # exit on the first non-zero status
set -u  # exit on unset variables
set -o pipefail

SCRIPT_NAME=$(basename $0)


function show_options {
    EXITVAL=${1:-1}
    echo "Usage: $SCRIPT_NAME [-h] [-w TIMEOUT] [-l LOOP_COUNT] [-f FAIL_MATCH] [-s SUCCESS_MATCH] --delay SLEEP_TIME -- COMMAND"
    echo
    echo "Waits for a command to fail, succeed, or timeout."
    echo
    echo "Options:"
    echo "      -h,--help                   -- this help"
    echo "      -w,--walltime TIMEOUT       -- Timeout after TIMEOUT seconds."
    echo "      -l,--looptimeout LOOP_COUNT -- Timeout after checking COMMAND LOOP_COUNT times."
    echo "      -d,--delay SLEEP_TIME       -- Seconds to sleep between checks of COMMAND."
    echo "      -s,--success-match          -- Output that indicates a success."
    echo "      -f,--fail-match             -- Output that indicates a short-circuit failure."
    echo
    echo "Execute the command in a loop until it succeeds, a timeout is reached, or"
    echo "a short-circuit failure occurs. Between each check of the command sleep for"
    echo "the number of seconds specified by SLEEP_TIME."
    echo
    echo "Examples:"
    echo "    wait_for -w 300 --delay 10 -- ping -c 1 192.0.2.2"
    echo "    wait_for -w 10 --delay 1 -- ls file_we_are_waiting_for"
    echo "    wait_for -w 30 --delay 3 -- date \| grep 8"
    echo "    wait_for -w 300 --delay 10 --fail-match CREATE_FAILED -- openstack stack show undercloud"
    echo "    wait_for -w 300 --delay 10 --success-match CREATE_COMPLETE -- openstack stack show undercloud"
    exit $EXITVAL
}

USE_WALLTIME=
TIMEOUT=
DELAY=

if [ -n "${SUCCESSFUL_MATCH_OUTPUT:-}" ]; then
    echo "DEPRECATION WARNING: Using env vars for specifying SUCCESSFUL_MATCH_OUTPUT is deprecated."
fi
SUCCESSFUL_MATCH_OUTPUT=${SUCCESSFUL_MATCH_OUTPUT:-""}
if [ -n "${FAIL_MATCH_OUTPUT:-}" ]; then
    echo "DEPRECATION WARNING: Using env vars for specifying FAIL_MATCH_OUTPUT is deprecated."
fi
FAIL_MATCH_OUTPUT=${FAIL_MATCH_OUTPUT:-""}

USE_ARGPARSE=0
# We have to support positional arguments for backwards compat
if [ -n "$1" -a "${1:0:1}" == "-" ]; then
    USE_ARGPARSE=1
else
    echo "DEPRECATION WARNING: Using positional arguments for wait_for is deprecated."
fi

if [ $USE_ARGPARSE -eq 1 ]; then
    set +e
    TEMP=$(getopt -o h,w:,l:,d:,s:,f: -l help,walltime:,looptimeout:,delay:,success-match:,fail-match: -n $SCRIPT_NAME -- "$@")
    if [ $? != 0 ]; then
        show_options;
    fi
    set -e

    # Note the quotes around `$TEMP': they are essential!
    eval set -- "$TEMP"

    while true ; do
        case "$1" in
            -h) show_options 0;;
            --help) show_options 0;;
            -w|--walltime) [ -n "$USE_WALLTIME" ] && show_options
                USE_WALLTIME=1
                TIMEOUT="$2"
                shift 2
                ;;
            -l|--looptimeout) [ -n "$USE_WALLTIME" ] && show_options
                USE_WALLTIME=0
                TIMEOUT="$2"
                shift 2
                ;;
            -d|--delay) DELAY="$2"; shift 2;;
            -s|--success-match) SUCCESSFUL_MATCH_OUTPUT="$2"; shift 2;;
            -f|--fail-match) FAIL_MATCH_OUTPUT="$2"; shift 2;;
            --) shift ; break ;;
        esac
    done
else
    TIMEOUT=${1:-""}
    DELAY=${2:-""}
    USE_WALLTIME=0
    shift 2 || true
fi

COMMAND="$@"

if [ -z "$TIMEOUT" -o -z "$DELAY" -o -z "$COMMAND" ]; then
    show_options
fi


ENDTIME=$(($(date +%s) + $TIMEOUT))
TIME_REMAINING=0
function update_time_remaining {
    CUR_TIME="$(date +%s)"
    TIME_REMAINING=$(($ENDTIME - $CUR_TIME))
}


OUTPUT=

function check_cmd {
    STATUS=0
    OUTPUT=$(eval $COMMAND 2>&1) || STATUS=$?
    if [[ -n "$SUCCESSFUL_MATCH_OUTPUT" ]] \
        && [[ $OUTPUT =~ $SUCCESSFUL_MATCH_OUTPUT ]]; then
        exit 0
    elif [[ -n "$FAIL_MATCH_OUTPUT" ]] \
        && [[ $OUTPUT =~ $FAIL_MATCH_OUTPUT ]]; then
        echo "Command output matched '$FAIL_MATCH_OUTPUT'. Here is the complete output failure:"
        echo "${OUTPUT}"
        echo "Exiting ..."
        exit 1
    elif [[ -z "$SUCCESSFUL_MATCH_OUTPUT" ]] && [[ $STATUS -eq 0 ]]; then
        # The command successfully completed and we aren't testing against
        # it's output so we have finished waiting.
        exit 0
    fi
}

i=0
while [ $USE_WALLTIME -eq 1 -o $i -lt $TIMEOUT ]; do
    if [ $USE_WALLTIME -eq 1 ]; then
        update_time_remaining
        if [ $TIME_REMAINING -le 0 ]; then
            break
        fi
    else
        i=$((i + 1))
    fi

    check_cmd

    if [ $USE_WALLTIME -eq 1 ]; then
        update_time_remaining
        if [ $TIME_REMAINING -lt $DELAY ]; then
            if [ $TIME_REMAINING -gt 0 ]; then
                sleep $TIME_REMAINING
                check_cmd
            fi
        else
            sleep $DELAY
        fi
    else
        sleep $DELAY
    fi
done
if [ $USE_WALLTIME -eq 1 ]; then
    SECONDS=$TIMEOUT
else
    SECONDS=$((TIMEOUT * DELAY))
fi
printf 'Timing out after %d seconds:\nCOMMAND=%s\nOUTPUT=%s\n' \
    "$SECONDS" "$COMMAND" "$OUTPUT"
exit 1