#!/bin/bash
set -eu

# Define exit codes
EXIT_DOCKER_INSTALL_FAILED=20
EXIT_ACR_AUTH_FAILED=30

###########
# FUNCTIONS
###########

# Validates the provided URL to ensure it follows the HTTP/HTTPS format.

logdate() {
   echo "[$(date +'%F %T')] $1"
}

validate_url() {
  if [[ ! $1 =~ ^https?://.+ ]]; then
    logdate "Invalid URL. Please enter a valid http/https URL."
    exit 1
  fi
}

# Validates the provided IP address to ensure it conforms to IPv4 format.
validate_ip() {
  if [[ ! $1 =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
    logdate "Invalid IP address format. Please enter a valid IP."
    exit 1
  fi
}

# Validates the security token to ensure it is a 32-character lowercase alphanumeric string.
validate_token() {
  if [[ ! $1 =~ ^[a-z0-9]{32}$ ]]; then
    logdate "Invalid token. Please enter a 32-character alphanumeric lowercase string."
    exit 1
  fi
}

# Validates the provided vSOC version to ensure it matches the required numeric format with four segments.
validate_version() {
  if [[ ! $1 =~ ^[0-9]+\.[0-9]+\.[0-9]+[\+\-_a-zA-Z]+$ ]]; then
    logdate "Invalid version. Please enter a version like 16.xx.xx.xx"
    exit 1
  fi
}

# Checks for available ports, starting at port 9090.
find_available_port() {
    local port=9090
    local max_attempts=100
    local attempt=1

    while ss -tuln | grep -q ":$port\b"; do
        if ((attempt >= max_attempts)); then
            logdate "No available ports found after $max_attempts attempts."
            exit 1
        fi
        ((port++))
        ((attempt++))
    done
    echo "$port"
}

# Checks the operational status of specified Docker containers and reports issues.
check_container_status() {
    # Checks the agent container status
    container_agent_status=$(sudo docker ps --filter "name=$1" --format "{{.Names}} {{.Status}}" | grep -q "$1.*Up" && echo "running" || echo "not running")
    container_executor_status=$(sudo docker ps --filter "name=$2" --format "{{.Names}} {{.Status}}" | grep -q "$2.*Up" && echo "running" || echo "not running")

    # Sets the agent and executor status variables.
    agent_status=$container_agent_status
    executor_status=$container_executor_status

    # Checks whether the containers are in a "Restarting" state. If either container is in this state, the script outputs an error message and terminates.
    if sudo docker ps --filter "name=$1" --format "{{.Names}} {{.Status}}" | grep -q "$1.*Restarting"; then
        logdate "installation failed, please check $1's log"
        exit 1
    fi
    if sudo docker ps --filter "name=$2" --format "{{.Names}} {{.Status}}" | grep -q "$2.*Restarting"; then
        logdate "installation failed, please check $2's log"
        exit 1
    fi

    # Notifies the user of any container that is not in a running state.
    if [[ "$agent_status" == "not running" ]]; then
        logdate "$1 container is not running"
        return 1
    fi
    if [[ "$executor_status" == "not running" ]]; then
        logdate "$2 container is not running"
        return 2
    fi
    return 0
}

######################################
# INSTALL DOCKER + AUTHENTICATE TO ACR
######################################

# Ensures Docker is installed and attempts automatic installation if missing.
function install_docker() {
  if ! command -v docker >/dev/null 2>&1; then
      logdate "Docker is not installed. Installing Docker..."

      # Installs Docker based on the helper script
      logdate "This script will attempt to install Docker automatically by downloading and executing this helper script provided by Docker Inc: https://get.docker.com."
      read -p "Would you like to download and execute this script to install Docker automatically? (Y/N): " user_response

      if [[ $user_response == "y" || $user_response == "Y" ]]; then
          curl -fsSL https://get.docker.com -o get-docker.sh
          sudo sh ./get-docker.sh
          sudo systemctl enable docker --now

          # Verifies whether Docker was installed successfully.
          if ! command -v docker >/dev/null 2>&1; then
              logdate "Docker installation failed."
              exit $EXIT_DOCKER_INSTALL_FAILED
          fi
      else
          logdate "Docker not installed. Please install manually and re-run this script."
          exit $EXIT_DOCKER_INSTALL_FAILED
      fi
  fi
}

function auth_docker_registry() {
  local docker_registry=$1
  if [[ ${docker_registry} != "ACR" ]];then
    logdate "docker registry not support, exit 1"
    exit 1
  fi

  SERVICE_ACCOUNT_NAME="d3acrpull"
  SERVICE_ACCOUNT_FILE="service_account_d3soar.key"
  SERVICE_ENDPOINT="d3soar.azurecr.io"

  # Uses the service account file, if found, to log in to ACR. If not found or if login fails, exit with an error code and display an appropriate message.
  if [ -f "$SERVICE_ACCOUNT_FILE" ]; then
      logdate "Logging in to Azure Container Registry using service account key..."
      LOGIN_OUTPUT=$(cat "$SERVICE_ACCOUNT_FILE" | sudo docker login -u ${SERVICE_ACCOUNT_NAME} --password-stdin ${SERVICE_ENDPOINT} 2>&1)
      logdate "$LOGIN_OUTPUT"

      # Checks for "Login Succeeded" in the output.
      if logdate "$LOGIN_OUTPUT" | grep -q "Login Succeeded"; then
          logdate "Generating Docker Compose file..."
      else
          logdate "Docker login failed."
          exit $EXIT_ACR_AUTH_FAILED
      fi
  else
      logdate "Service account file $SERVICE_ACCOUNT_FILE not found. Please contact your D3 representative to obtain it."
      exit $EXIT_ACR_AUTH_FAILED
  fi
}


#################################
# GENERATE docker-compose-d3.yaml
#################################
function generate_docker_compose_yaml() {
  # Loads the vSOC URL and security token from vars.conf configuration file.
  if [ -f "./vars.conf" ]; then
      source vars.conf
  else
      logdate "'vars.conf' doesn't exist. Create vars.conf with two line-separated key=value pairs: vsoc_url and security_token."
      exit 1
  fi

  # Validates that required variables are set.
  if [[ -z "${vsoc_url:-}" || -z "${security_token:-}" ]]; then
    logdate "Error: Missing required variables (vsoc_url, security_token). Verify vars.conf has these keys set."
    exit 1
  fi

  # Obtains the host machine's IP address.
  host_ip=$(ip -o route get "8.8.8.8" 2>/dev/null | awk '{print $7}')

  # Generates a unique GUID for the remote Python agent.
  python_remote_guid=$(uuidgen | tr -d '-' | tr '[:lower:]' '[:upper:]')

  # Finds an available Docker port.
  docker_port=$(find_available_port)

  # Sets the agent display name.
  agent_display_name="D3Agent - $host_ip"

  # Retrieves the Docker GID.
  docker_gid=$(getent group docker | cut -d: -f3)

  # Ensures the vSOC URL is properly formatted with a trailing slash.
  vsoc_url="${vsoc_url%/}/"

  # Validates the loaded configuration values.
  validate_ip "$host_ip"
  validate_token "$security_token"
  validate_url "$vsoc_url"

  # Constructs the full URL.
  url="${vsoc_url}api/REST/GetD3Version"

  # Fetches the content of the URL specified in $url.
  html=$(curl -s -k "$url")

  # Extracts the vSOC version number.
  vsoc_version=$(echo "$html" | grep -oP '[0-9]+\.[0-9]+\.[0-9]+[\+-_0-9a-zA-Z]+' |sed 's/+/_/g')

  # Verifies whether the vsoc_version value matches the expected format.
  validate_version "$vsoc_version"

  # Checks whether vsoc_version is empty. If so, prints an error and exits.
  if [[ -z "$vsoc_version" ]]; then
          logdate "Version information not found."
      exit 1
  fi

  # Generate unique names for the containers.
  container_agent=d3agent_$RANDOM
  container_executor=d3executor_$RANDOM

  DOCKER_REGISTRY="d3soar.azurecr.io"

  # Generates the Docker Compose YAML
  cat <<EOF >docker-compose-d3.yml
services:
  d3agent:
    container_name: $container_agent
    image: "${DOCKER_REGISTRY}/d3prod/d3agent:${vsoc_version}"
    restart: always
    environment:
      - REMOTE_SERVER_URL=${vsoc_url}
      - PROXY_IDENTITY=${security_token}
      - PYTHON_REMOTE_URL=http://${host_ip}:${docker_port}/
      - PYTHON_REMOTE_GUID=${python_remote_guid}
      - SERVICE_DISPLAY_NAME=${agent_display_name}
      - DOCKER_GID=${docker_gid}
      - DOCKER_PORT=${docker_port}
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock

  d3executor:
    container_name: $container_executor
    privileged: true
    image: "${DOCKER_REGISTRY}/d3prod/d3executor:${vsoc_version}"
    ports:
      - "${docker_port}:9090"
    restart: always
    environment:
      - EXECUTOR_GUID=${python_remote_guid}
EOF

  logdate "A Docker Compose file for D3 Proxy Agent has been generated successfully: docker-compose-d3.yml"

}

function start_container() {
  local docker_compose_cmd=$1

  # Prompts the user to start the containers.
  read -p "Would you like to start the containers? (Y/N): " user_response

  if [[ $user_response == "y" || $user_response == "Y" ]]; then
      # Attempts to bring up the Docker containers.
      sudo ${docker_compose_cmd} -f docker-compose-d3.yml up -d --remove-orphans || logdate "Failed to bring up Docker containers. Verify the contents of 'docker-compose-d3.yml' and run 'sudo docker compose -f docker-compose-d3.yml up -d' to manually start the containers."

      logdate "Checking container status..."
      counter=1

      # Performs up to 5 health checks on the containers. If any container fails, the loop exits early.
      while [ $counter -le 3 ]; do
          # Checks container statuses.
          check_container_status $container_agent $container_executor

          # Checks whether the agent container is running.
          if [[ $? -eq 1 ]]; then
              logdate "Unhealthy Agent!"
              break
          fi

          # Checks whether the executor container is running.
          if [[ $? -eq 2 ]]; then
              logdate "Unhealthy Executor!"
              break
          fi

          logdate "Health Check #$counter/3 passed"
          counter=$((counter + 1))

          sleep 5
      done

      # Checks final container statuses and print a success or failure message with troubleshooting suggestions.
      if [[ "$agent_status" != "running" ]] || [[ "$executor_status" != "running" ]]; then
        logdate "Installation failed. You can check the container statuses with sudo ${docker_compose_cmd} -f docker-compose-d3.yml ps' or the container logs with sudo ${docker_compose_cmd} -f docker-compose-d3.yml logs"
      else
        logdate "Installation succesful! Log in to your VSOC instance at $vsoc_url to check the Agent Status from the UI."
      fi

  else
      logdate "Verify the contents of 'docker-compose-d3.yml' and run sudo ${docker_compose_cmd} -f docker-compose-d3.yml up -d to manually start the containers."
  fi
}

function adapter_podman() {
  if [[ $(sudo docker version 2>&1 |grep -Eo podman) != 'podman' ]];then
      return
  fi

  local to_install=false
  podman-compose -v || to_install=true
  if [ "${to_install}" = "true" ];then
      logdate "install podman-compose ... "
      sudo yum install -y podman-compose
  else
      logdate "podman-compose has been installed"
  fi

  # if not installed, will exit directly
  podman-compose -v

  id docker || sudo useradd docker
  SERVICE_CONF='/usr/lib/systemd/system/podman.service'
  SOCK_FILE='/var/run/podman/podman.sock'
  sudo cp ${SERVICE_CONF} ${SERVICE_CONF}.bak
  cat <<EOF >./podman.service
[Unit]
Description=Podman API Service
Requires=podman.socket
After=podman.socket
Documentation=man:podman-system-service(1)
StartLimitIntervalSec=0

[Service]
Delegate=true
Type=exec
KillMode=process
Environment=LOGGING="--log-level=info"
ExecStart=/usr/bin/podman \$LOGGING system service --time=0
ExecStopPost=/bin/chown root:docker /run/podman/podman.sock && ln -s /var/run/docker.sock

[Install]
WantedBy=default.target
EOF
  sudo cp ./podman.service ${SERVICE_CONF}
  sudo systemctl daemon-reload
  sudo systemctl restart podman
  sudo systemctl status podman
  if [[ "$(pgrep podman)"x != x ]];then
      logdate "podman is running"
  fi
  logdate "display sock file"
  sudo ls -al ${SOCK_FILE} /var/run/docker.sock
}

function main() {
  D3_DOCKER_REGISTRY="ACR"
  # install docker
  install_docker
  # adapt podman
  adapter_podman
  # auth to acr
  auth_docker_registry ${D3_DOCKER_REGISTRY}
  logdate 'auth done'
  # generate docker-compose config
  generate_docker_compose_yaml

  # start d3agent/d3executor container
  DOCKER_COMPOSE_CMD="docker compose"
  if [[ $(sudo docker version 2>&1 |grep -Eo podman) = 'podman' ]];then
      DOCKER_COMPOSE_CMD="podman-compose"
      logdate "use podman-compose to start container ..."
  fi
  start_container "${DOCKER_COMPOSE_CMD}"
}

main

logdate ">>>>>> DONE <<<<<<"