#!/usr/bin/env bash

generic_version_regex="^[0-9]{1,2}\.[0-9]{1,2}.[0-9]{1,2}$"
version_regex="^[3-9]\.[0-9]{1,2}.[0-9]{1,2}$"
pycti_default_version="6.8.8"

# little trick for exact matching in arrays
declare -A env_arguments=(["prod"]=1 ["test"]=1 ["ci"]=1)
declare -A test_mode=(["test"]=1 ["ci"]=1)
declare -A cmd_arguments=(["build"]=1 ["up"]=1 ["start"]=1 ["restart"]=1 ["down"]=1 ["stop"]=1 ["kill"]=1 ["logs"]=1 ["ps"]=1)

declare -A path_mapping=(["default"]="docker/default.yml" ["postgres"]="docker/postgres.override.yml" ["rabbitmq"]="docker/rabbitmq.override.yml" ["test"]="docker/test.override.yml" ["ci"]="docker/ci.override.yml" ["custom"]="docker/custom.override.yml" ["traefik"]="docker/traefik.yml" ["traefik_prod"]="docker/traefik_prod.yml" ["traefik_local"]="docker/traefik_local.yml" ["multi_queue"]="docker/multi-queue.override.yml" ["test_multi_queue"]="docker/test.multi-queue.override.yml" ["flower"]="docker/flower.override.yml" ["test_flower"]="docker/test.flower.override.yml" ["elastic"]="docker/elasticsearch.override.yml" ["https"]="docker/https.override.yml" ["nfs"]="docker/nfs.override.yml" ["redis"]="docker/redis.override.yml" ["nginx_default"]="docker/nginx.override.yml" ["ollama"]="docker/ollama.override.yml" ["test_ollama"]="docker/test.ollama.override.yml" ["ci_ollama"]="docker/ci.ollama.override.yml" ["misp"]="docker/misp.override.yml")
print_synopsis () {
  echo "SYNOPSIS"
  echo -e "    start <env> <command> [OPTIONS]"
  echo -e "    start -h|--help\n"
}

print_help () {
  print_synopsis
  echo "ARGUMENTS"
  echo "    <env>    Environment for execution"
  echo -e "    <command>    Command for docker compose\n"
  echo "OPTIONS"
  echo "    --project_name <name>       Specify project name."
  echo "    --version <value>           Choose the version you would like to install"
  echo "                                (>=3.0.0). Works only in 'prod' mode. Default"
  echo "                                version is the most recently released."
  echo "    --all_analyzers             Uses every integration."
  echo "    --tor_analyzers             Uses the integrations/tor_analyzers/compose.yml"
  echo "                                file."
  echo "    --malware_tools_analyzers   Uses the integrations/malware_tools_analyzers/"
  echo "                                compose.yml file."
  echo "    --thug                      Uses the integrations/thug/compose.yml"
  echo "                                file."
  echo "    --cyberchef                 Uses the integrations/cyberchef/compose.yml"
  echo "                                file."
  echo "    --pcap_analyzers            Uses the integrations/pcap_analyzers/compose.yml"
  echo "                                file."
  echo "    --bbot                      Uses the integrations/bbot/compose.yml"
  echo "    --phunter                   Uses the integrations/phunter/compose.yml"
  echo "    --multi_queue               Uses the multiqueue.override.yml compose file."
  echo "    --nfs                       Uses the nfs.override.yml compose file."
  echo "    --traefik_prod              Uses the traefik.yml and traefik_prod.yml compose file."
  echo "    --traefik_local             Uses the traefik.yml and traefik_local.yml compose file."
  echo "    --use-external-database     Do NOT use postgres.override.yml compose file."
  echo "    --use-external-redis        Do NOT use redis.override.yml compose file."
  echo "    --rabbitmq                  Uses the rabbitmq.override.yml compose file."
  echo "    --flower                    Uses the flower.override.yml compose file."
  echo "    --ollama                    Uses the ollama.override.yml compose file"
  echo "                                to enable the local LLM chatbot service."
  echo "    --misp                      Uses the misp.override.yml compose file"
  echo "                                to spin up a local MISP instance."
  echo "    --custom                    Uses custom.override.yml to leverage your"
  echo "                                customized configuration."
  echo "    --debug-build               See more verbose output from the build."
  echo "    --elastic                   This spins up Elasticsearch on your"
  echo "                                machine (might need >=16GB of RAM)."
  echo "    --https                     This leverage the https.override.yml file that"
  echo "                                can be used to host IntelOwl with HTTPS and your"
  echo "                                own certificate."
  echo "    --pycti-version <value>     The pycti version to choose. This must match the"
  echo "                                OpenCTI server version you are connecting to."
  echo "                                Default is ${pycti_default_version}."
}

check_parameters () {
  if [[ $# == 1 && ( $1 == "--help" || $1 == "-h" ) ]]; then
    print_help
    exit 0
  fi
  if [[ $# -lt 2 ]]; then
    echo "Error! Provide at least the environment and a command." >&2
    print_synopsis
    exit 2
  fi

  if ! [[ ${env_arguments["${1}"]} ]]; then
    echo "Error! Illegal environment specified." >&2
    exit 2
  fi
  if ! [[ ${cmd_arguments["${2}"]} ]]; then
    echo "Error! Illegal command specified." >&2
    exit 2
  fi

  env_argument=$1
  cmd_argument=$2
}

set_defaults_values () {
  project_name="intel_owl"
  version=$current_version
  export PYCTI_VERSION=$pycti_default_version
}

if ! docker compose version > /dev/null 2>&1; then
  echo "Run ./initialize.sh to install Docker Compose 2"
fi

check_parameters "$@" && shift 2

# shellcheck source=docker/.env disable=SC1091
. docker/.env

current_version=${REACT_APP_INTELOWL_VERSION/"v"/""}

docker_analyzers=("pcap_analyzers" "tor_analyzers" "malware_tools_analyzers" "thug" "cyberchef" "phoneinfoga" "phishing_analyzers" "nuclei_analyzer" "bbot" "phunter")


for value in "${docker_analyzers[@]}"; do
  path_mapping["${value}"]+="integrations/${value}/compose.yml"
  path_mapping["${value}.test"]+="integrations/${value}/compose-tests.yml"
  # cannot use a list as value of associative array. we have to use a string and convert dynamically
  path_mapping["all_analyzers"]+="${path_mapping[${value}]} "
  path_mapping["all_analyzers.test"]+="${path_mapping["${value}.test"]} "
done

if [[ ${test_mode["${env_argument}"]} ]]; then
  is_test=true
  test_appendix=".test"

# shellcheck source=docker/.env.start.test disable=SC1091
  . docker/.env.start.test

else
  is_test=false
  test_appendix=""
fi

# parse arguments and put them in an associative array to better use them later
declare -A params
declare -A analyzers
project_name=""
set_defaults_values
while [[ $# -gt 0 ]]; do
  case $1 in
  --project_name)
    project_name=$2
    shift 2
  ;;
  -v | --version)
    if ! [[ $2 =~ $version_regex ]]; then
      echo "Error! Wrong version format." >&2
      exit 1
    fi
    version=$2
    shift 2
  ;;
  --all_analyzers)
    analyzers["all_analyzers"]=true
    shift 1
  ;;
  --tor_analyzers)
    analyzers["tor_analyzers"]=true
    shift 1
    ;;
  --nuclei_analyzer)
    analyzers["nuclei_analyzer"]=true
    shift 1
    ;;
  --malware_tools_analyzers)
    analyzers["malware_tools_analyzers"]=true
    shift 1
    ;;
  --thug)
    analyzers["thug"]=true
    shift 1
    ;;
  --cyberchef)
    analyzers["cyberchef"]=true
    shift 1
    ;;
  --pcap_analyzers)
    analyzers["pcap_analyzers"]=true
    shift 1
    ;;
  --phoneinfoga)
    analyzers["phoneinfoga"]=true
    shift 1
    ;;
  --phishing_analyzers)
    analyzers["phishing_analyzers"]=true
    shift 1
    ;;
  --bbot)
    analyzers["bbot"]=true
    shift 1
    ;;
  --phunter)
    analyzers["phunter"]=true
    shift 1
    ;;
  --multi_queue)
    params["multi_queue"]=true
    shift 1
  ;;
  --nfs)
    params["nfs"]=true
    shift 1
  ;;
  --use-external-database)
    params["use_external_database"]=true
    shift 1
  ;;
  --use-external-redis)
    params["use_external_redis"]=true
    shift 1
  ;;
  --rabbitmq)
    params["rabbitmq"]=true
    shift 1
  ;;
  --sqs)
    params["sqs"]=true
    shift 1
  ;;
  --flower)
    params["flower"]=true
    shift 1
  ;;
  --ollama)
    params["ollama"]=true
    shift 1
  ;;
  --misp)
    params["misp"]=true
    shift 1
  ;;
  --custom)
    params["custom"]=true
    shift 1
  ;;
  --debug-build)
    debug_build=true
    shift 1
  ;;
  --elastic)
    params["elastic"]=true
    ./create_elastic_certs
    shift 1
  ;;
  --pycti-version)
    if ! [[ $2 =~ $generic_version_regex ]]; then
      echo "Error! Wrong pycti version format." >&2
      exit 1
    fi
    export PYCTI_VERSION=$2
    shift 2
  ;;
  --https)
    params["https"]=true
    shift 1
  ;;
  --traefik_prod)
    params["traefik_prod"]=true
    shift 1
  ;;
  --traefik_local)
    params["traefik_local"]=true
    shift 1
  ;;
  -h | --help)
    print_help
    exit 0
    ;;
  --) shift 1; break;; # explicit end of arguments
  *)
    echo "Error! Invalid option $1."
    exit 1
    ;;
  esac
done

# here all variables should be parsed and ready for use
cmd_py_version=("up" "build")
if [[ ( ! $env_argument == "test" || ! ${cmd_py_version[*]} =~ $cmd_argument ) && \
  ( $PYCTI_VERSION != "$pycti_default_version" ) ]]; then
    echo "pycti_version options are valid only while running in" >&2
    echo "'test' mode and while building a new image. This is because they can change" >&2
    echo "the version of those library only during the build of a new Docker Image." >&2
    exit 11
fi

# check if all_analyzers and other flags have been used
for value in "${docker_analyzers[@]}"; do
  if [ "${analyzers["$value"]}" ]; then
    docker_flags+=("${value}")
  fi
done
if [ "${analyzers["all_analyzers"]}" ] && [ ${#docker_flags[@]} -ne 0 ]; then
  echo "It is not possible to select both --all_analyzers and another docker container."
  exit 1
fi

# default file
compose_files=("${path_mapping["default"]}")
if ! [ "${params["use_external_database"]}" ]; then
  compose_files+=("${path_mapping["postgres"]}")
fi
if ! [ "${params["use_external_redis"]}" ]; then
  compose_files+=("${path_mapping["redis"]}")
fi

if [ "${params["rabbitmq"]}" ]; then
  compose_files+=("${path_mapping["rabbitmq"]}")
elif [ "${params["sqs"]}" ]; then
  compose_files+=("${path_mapping["sqs"]}")
fi

if [ "$is_test" = true ]; then
  compose_files+=("${path_mapping["$env_argument"]}")
fi

# Check for the traefik_prod or traefik_local argument and include traefik base compose

traefik_enabled=false
if [ "${params["traefik_prod"]}" ] || [ "${params["traefik_local"]}" ]; then
  compose_files+=("${path_mapping["traefik"]}")
  traefik_enabled=true
fi

# Add the default nginx configuration if traefik is not used
if [ "$traefik_enabled" = false ]; then
  compose_files+=("${path_mapping["nginx_default"]}")
fi

# add all the other ones
for value in "${!params[@]}"; do
  if [ "${params[$value]}" ]; then
    compose_files+=("${path_mapping["$value"]}")
  fi
done

# add all the test files
if [[ $env_argument == "test" ]]; then
  test_values=("multi_queue" "flower" "ollama")
  for value in "${test_values[@]}"; do
    if [ "${params["$value"]}" ]; then
      compose_files+=("${path_mapping["test_$value"]}")
    fi
  done
elif [[ $env_argument == "ci" ]]; then
  # In ci mode --ollama pulls in ollama.override.yml (via the params loop above), whose chatbot
  # worker is pinned to the pulled release image; repoint it to the built :ci image.
  if [ "${params["ollama"]}" ]; then
    compose_files+=("${path_mapping["ci_ollama"]}")
  fi
fi

# add and parse analyzers
if [ "${analyzers["all_analyzers"]}" ]; then
  IFS=', ' read -r -a all_analyzers_array <<< "${path_mapping["all_analyzers"]}"
  for analyzer in "${all_analyzers_array[@]}" ; do
    compose_files+=("${analyzer}")
  done
  if [ "$is_test" = true ]; then
    IFS=', ' read -r -a all_analyzers_array_test <<< "${path_mapping["all_analyzers${test_appendix}"]}"
    for analyzer in "${all_analyzers_array_test[@]}" ; do
      compose_files+=("${analyzer}")
    done
  fi
else
  # add the single analyzers compose files
  if [ ${#docker_flags[@]} -ne 0 ]; then
    for analyzer in "${docker_flags[@]}"; do
      compose_files+=("${path_mapping["$analyzer"]}")
      if [ "$is_test" = true ]; then
        compose_files+=("${path_mapping["$analyzer$test_appendix"]}")
      fi
    done
  fi
fi

if [[ $env_argument == "prod" && $version != "${current_version}" ]]; then
  echo "Requested version ${version} is different from current version ${current_version}"
  directory=$(git config --global --get safe.directory)
  if ! [ $? ] || ! [[ "${directory}" == $(pwd) ]]; then
    git config --global --add safe.directory "$(pwd)"
  fi
  git checkout "v${version}"
fi

if [ "$debug_build" ]; then
  export BUILDKIT_PROGRESS="plain"
fi
export DOCKER_BUILDKIT=1

for value in "${compose_files[@]}" ; do
  if [ -n "${value}" ]; then
    to_run+=("-f" "$value")
  fi
done

if grep "docker" <<< "$(groups)" > /dev/null 2>&1; then
  docker compose --project-directory docker "${to_run[@]}" -p "$project_name" "$cmd_argument" "$@"
else
  sudo docker compose --project-directory docker "${to_run[@]}" -p "$project_name" "$cmd_argument" "$@"
fi
