#!/bin/bash

# Define the usage message
usage() {
  echo "Usage: $0 -u <url> [-w <warning>] [-c <critical>] [-C <contains>] [-L] [-I] [-H <headers>] [-p] [-R] [-s]"
  echo "[Arguments]:
  -u Specify the URL to check (required).
  -w Set the warn level for response time (default: 1 second).
  -c Set the critical level for response time (default: 2 seconds).
  -C If the body does not contain this string, return CRITICAL.
  -L Follow redirects.
  -I Insecure mode (--insecure).
  -H Specify headers. Formatted like \"Header1: value,Header2: value\"
  -p Print the curl command and exit
  -R Set curl --resolve option.
  -s Ignore the response status code."
  exit 3
}

# Parse the command-line arguments
while getopts "u:w:c:C:H:R:LhIps" opt; do
  case $opt in
  u)
    URL=$OPTARG
    ;;
  w)
    WARNING_LEVEL=$OPTARG
    ;;
  c)
    CRITICAL_LEVEL=$OPTARG
    ;;
  L)
    FOLLOW_REDIRECTS="-L"
    ;;
  C)
    CONTAINS=$OPTARG
    ;;
  I)
    INSECURE="--insecure"
    ;;
  H)
    HEADERS=$OPTARG
    ;;
  p)
    PRINT_ONLY=true
    ;;
  R)
    RESOLVE="--resolve $OPTARG"
    ;;
  s)
    IGNORE_STATUS_CODE=true
    ;;
  h)
    usage
    ;;
  *)
    usage
    ;;
  esac
done

WARNING_LEVEL=${WARNING_LEVEL:-1}
CRITICAL_LEVEL=${CRITICAL_LEVEL:-2}
#FOLLOW_REDIRECTS=${FOLLOW_REDIRECTS:-""}
PRINT_ONLY=${PRINT_ONLY:-false}
IGNORE_STATUS_CODE=${IGNORE_STATUS_CODE:-false}

if [ -z "$URL" ]; then
  usage
fi

HEADER_ARGS=""
IFS=',' read -ra values <<<"$HEADERS"
for value in "${values[@]}"; do
  HEADER_ARGS+=' -H "'$value'" '
done

TMP_ERROR_LOG=$(mktemp)
TMP_RESPONSE=$(mktemp)

if $PRINT_ONLY; then
  printf "%s" "curl --output \"$TMP_RESPONSE\" -s -w \"%{http_code}\n%{time_total}\" $(echo "${HEADER_ARGS[@]}" | tr -s ' ') $FOLLOW_REDIRECTS $INSECURE $RESOLVE $URL\n"
  exit 3
fi

RESPONSE=$(curl --output "$TMP_RESPONSE" -w "%{http_code}\n%{time_total}" $(echo "${HEADER_ARGS[@]}" | tr -s ' ') $FOLLOW_REDIRECTS $INSECURE $RESOLVE $URL 2>"$TMP_ERROR_LOG")

# shellcheck disable=SC2181
status=$?
if [ $status -ne 0 ]; then
  case $status in
  1)
    msg="CRITICAL: Unsupported protocol"
    ;;
  3)
    msg="CRITICAL: Malformed URL $URL"
    ;;
    #    5)
    #      msg="CRITICAL: Couldn't resolve proxy $proxy"
    #      ;;
  6)
    msg="CRITICAL: Couldn't resolve host $URL"
    ;;
    #  7)
    #    msg="CRITICAL: Couldn't connect to proxy $proxy"
    #    ;;
  22)
    msg="CRITICAL: Server returned http code >= 400"
    ;;
  52)
    msg="CRITICAL: Server returned empty response (52)"
    ;;
  56)
    msg="CRITICAL: Failure recieving network data (56)"
    ;;
  60)
    msg="CRITICAL: SSL/TLS connection problem (60)"
    ;;
  *)
    echo "UNKNOWN: $status - $URL"
    exit 3
    ;;
  esac

  echo -e "$msg"
  #  echo "Error log:"
  #  cat "$TMP_ERROR_LOG"
  rm -rf "$TMP_ERROR_LOG"
  exit 2
fi
rm -rf "$TMP_ERROR_LOG"

RESPONSE_CODE=$(echo "$RESPONSE" | head -n 1)
RESPONSE_TIME=$(printf "%.3f" $(echo "$RESPONSE" | tail -n 1))

OUTPUT_MSG=""
OUTPUT_CODE=0

if ([ $RESPONSE_CODE -eq 200 ] || $IGNORE_STATUS_CODE) && [ "$(echo "$RESPONSE_TIME < $CRITICAL_LEVEL" | bc -l)" -eq 1 ]; then
  OUTPUT_MSG="OK: ${RESPONSE_TIME}s - $URL"
  OUTPUT_CODE=0
elif ([ $RESPONSE_CODE -eq 200 ] || $IGNORE_STATUS_CODE) && [ "$(echo "$RESPONSE_TIME < $WARNING_LEVEL" | bc -l)" -eq 1 ]; then
  OUTPUT_MSG="WARNING: response time is slow ($RESPONSE_TIME seconds)."
  OUTPUT_CODE=1
elif ([ $RESPONSE_CODE -eq 200 ] || $IGNORE_STATUS_CODE); then
  OUTPUT_MSG"CRITICAL: response time is very slow ($RESPONSE_TIME seconds)."
  OUTPUT_CODE=2
else
  OUTPUT_MSG="CRITICAL: website did not return 200, response was $RESPONSE_CODE code."
  OUTPUT_CODE=2
fi

if [[ -n ${CONTAINS+x} ]]; then
  if ! grep -q "$CONTAINS" "$TMP_RESPONSE"; then
    OUTPUT_MSG+="\nCRITICAL: response did not contain required string!\nFound: $(cat "$TMP_RESPONSE")"
    OUTPUT_CODE=2
    #  else
    #    OUTPUT_MSG+="\nOK: response contained required string."
  fi
fi
rm -rf "$TMP_RESPONSE"

if $IGNORE_STATUS_CODE && [[ $RESPONSE_CODE -ne 200 ]]; then
  OUTPUT_MSG+="\nResponse code was $RESPONSE_CODE."
fi

OUTPUT_MSG+=" | response_time=${RESPONSE_TIME}s;$WARNING_LEVEL;$CRITICAL_LEVEL;0"

case $OUTPUT_CODE in
#0)
#  echo "OK: $URL"
#  ;;
1)
  echo "WARNING: $URL"
  ;;
2)
  echo "CRITICAL: $URL"
  ;;
3)
  echo "UNKNOWN: $URL"
  ;;
esac

echo -e "$OUTPUT_MSG"
exit $OUTPUT_CODE