#!/usr/bin/bash
# Copyright Atomicorp 2021-2026
# https://www.atomicorp.com
#
# AGPL 3.0
# This program is free software; you can redistribute it and/or modify
# it under the terms of the Affero GNU General Public License (AGPL)
#


# Globals
VERSION=1.1.1
OSSEC_HOME=/var/ossec
SERVER=updates.atomicorp.com
OSSEC_CRS_RULES_VERSION=0
OSSEC_CRS_THREAT_VERSION=0
OSSEC_CRS_RULES_VERSION_CURRENT=0
OSSEC_CRS_THREAT_VERSION_CURRENT=0
YES=
DEBUG=
INSECURE=


# Functions
prereq_check() {
	declare -a PREREQS=("curl" "tar" "gzip" "bc")
	IS_GOOD="1"

	for PREREQ in "${PREREQS[@]}"; do
	    if ! command -v "$PREREQ" > /dev/null 2>&1; then
	      IS_GOOD="0"
	      echo " ERROR: $PREREQ not found. $PREREQ must be installed."
	    fi
	done

	if [ ! -d "${OSSEC_HOME}" ]; then
	    IS_GOOD="0"
	    echo " ERROR: OSSEC package not found. OSSEC must be installed."
	fi

	if [ "$IS_GOOD" -ne "1" ]; then
	    echo
	    echo "Prerequisite check failed. Exiting!"
	    echo
	    exit 1
	fi
}

function ossec_check() {
        # OUM installs into etc/decoders.d and etc/rules.d. Legacy ossec.conf
        # files that still use per-file <include> directives will not load those
        # updates. Modern Atomicorp configs already set decoder_dir/rule_dir.
        local ossec_conf="${OSSEC_HOME}/etc/ossec.conf"

        if [ ! -f "${ossec_conf}" ]; then
                return 0
        fi

        if grep -q 'decoder_dir.*decoders\.d' "${ossec_conf}" \
                && grep -q 'rule_dir.*rules\.d' "${ossec_conf}"; then
                return 0
        fi

        echo
        echo "########################################################################"
        echo "  WARNING: ${ossec_conf} is not configured for rules.d/decoders.d"
        echo
        echo "  OUM updates are installed under etc/decoders.d and etc/rules.d."
        echo "  Legacy configs that only use <include> rule files will ignore them."
        echo "  Replace the <rules></rules> section with:"
        echo
        echo "  <rules>"
        echo "          <decoder_dir pattern=\".xml$\">etc/decoders.d</decoder_dir>"
        echo "          <rule_dir pattern=\".xml$\">etc/rules.d</rule_dir>"
        echo "          <list>etc/lists/audit-key</list>"
        echo "          <!-- Declare etc/lists/threat/threatN if using threatfeed -->"
        echo "  </rules>"
        echo "########################################################################"
        echo
}


function rawurlencode() {
    local string="${1}"
    local strlen=${#string}
    local encoded=""
    local pos c o

    for ((pos = 0; pos < strlen; pos++)); do
        c=${string:$pos:1}
        case "$c" in
        [-_.~a-zA-Z0-9]) o="${c}" ;;
        *) printf -v o '%%%02x' "'$c" ;;
        esac
        encoded+="${o}"
    done
    echo "${encoded}"
    REPLY="${encoded}"
}

# Parse a VERSION file without sourcing it (avoids remote code execution).
# Accepts only KEY=numeric[.numeric...] assignments for known keys.
load_version_file() {
  local file="$1"
  local scope="$2" # remote|current
  local line key val
  local rules=0
  local threat=0

  if [ ! -f "$file" ]; then
    if [ "$scope" = "current" ]; then
      OSSEC_CRS_RULES_VERSION_CURRENT=0
      OSSEC_CRS_THREAT_VERSION_CURRENT=0
    else
      OSSEC_CRS_RULES_VERSION=0
      OSSEC_CRS_THREAT_VERSION=0
    fi
    return 0
  fi

  while IFS= read -r line || [ -n "$line" ]; do
    [[ "$line" =~ ^[[:space:]]*# ]] && continue
    [[ "$line" =~ ^[[:space:]]*$ ]] && continue
    if [[ "$line" =~ ^([A-Z_]+)=([0-9]+([.][0-9]+)*)$ ]]; then
      key="${BASH_REMATCH[1]}"
      val="${BASH_REMATCH[2]}"
      case "$key" in
        OSSEC_CRS_RULES_VERSION) rules="$val" ;;
        OSSEC_CRS_THREAT_VERSION) threat="$val" ;;
      esac
    fi
  done < "$file"

  if [ "$scope" = "current" ]; then
    OSSEC_CRS_RULES_VERSION_CURRENT="$rules"
    OSSEC_CRS_THREAT_VERSION_CURRENT="$threat"
  else
    OSSEC_CRS_RULES_VERSION="$rules"
    OSSEC_CRS_THREAT_VERSION="$threat"
  fi
}

version_lt() {
  # Return 0 if $1 < $2 (numeric dotted versions via bc).
  local a="${1:-0}"
  local b="${2:-0}"
  (( $(echo "${a} < ${b}" | bc -l) ))
}


download() {
  local URL="$1"
  local FILENAME="$2"
  local OPTS=()
  local CURL_AUTH=()
  local RESPONSE RETVAL

  if [ -z "$DEBUG" ]; then
    OPTS+=(-s)
  fi

  if [ -n "$INSECURE" ]; then
    OPTS+=(--insecure)
  fi

  # Prefer curl --user so passwords with reserved URL characters work.
  if [ -n "${USERNAME}" ] && [ "${USERNAME}" != "USERNAME" ] \
      && [ -n "${PASSWORD}" ] && [ "${PASSWORD}" != "PASSWORD" ]; then
    CURL_AUTH=(--user "${USERNAME}:${PASSWORD}")
  fi

  RESPONSE=$(curl --write-out '%{http_code}' "${OPTS[@]}" "${CURL_AUTH[@]}" \
    "${URL}" -o "${FILENAME}")

  RETVAL=$?
  if [ "${RETVAL}" -ne 0 ]; then
    echo "ERROR: Download failed with ERROR (${RETVAL})"
    if [ "${RETVAL}" -eq 60 ]; then
	echo
      	echo "  ERROR: Peer certificate cannot be authenticated with known CA certificates."
	echo
    fi
    echo
    exit 1
  fi

  if [[ "${RESPONSE}" == 401 ]] ; then
	echo
	echo "  ERROR: request returned HTTP error code 401 [Username/Password Invalid]"
	echo
	exit 1
  elif [[ "${RESPONSE}" != 200 ]] ; then
	echo
	echo "  ERROR: request returned HTTP error code ${RESPONSE}"
	echo
	exit 1
  fi
}

show_help() {
  echo
  echo "OSSEC Updater Modified (OUM) $VERSION"
  echo
  echo "Usage: oum [options] COMMAND"
  echo
  echo " Options:"
  echo "  -y			Assume yes to prompts (non-interactive update)"
  echo "  -d			Debug output"
  echo "  -i			Allow insecure TLS (curl --insecure)"
  echo
  echo " Commands:"
  echo
  echo "  help			Display a helpful usage message"
  echo "  list			List pending updates"
  echo "  update		Update rules and threatfeed"
  echo "  upgrade		Alias for update"
  echo "  configure		Configure subscription credentials"
  echo "  install PACKAGE	Install PACKAGE from the atomic repo"
  echo "  version		Display version"
  echo
}


load_config() {
  if [ ! -f "${OSSEC_HOME}/etc/oum.conf" ]; then
    echo
    echo "ERROR: ${OSSEC_HOME}/etc/oum.conf not found"
    echo
    exit 1
  fi
  # shellcheck source=/dev/null
  source "${OSSEC_HOME}/etc/oum.conf"
}

# Safely rewrite USERNAME/PASSWORD in oum.conf (handles special characters).
write_oum_credentials() {
  local user="$1"
  local pass="$2"
  local conf="${OSSEC_HOME}/etc/oum.conf"
  local tmp

  tmp=$(mktemp)
  # Preserve non-credential settings from the existing config.
  if [ -f "$conf" ]; then
    grep -vE '^(USERNAME|PASSWORD)=' "$conf" > "$tmp" || true
  fi
  {
    printf 'USERNAME=%q\n' "$user"
    printf 'PASSWORD=%q\n' "$pass"
    cat "$tmp"
  } > "${conf}.new"
  mv -f "${conf}.new" "$conf"
  chmod 600 "$conf"
  rm -f "$tmp"
  USERNAME="$user"
  PASSWORD="$pass"
}

check_version() {
  if [ -f "${OSSEC_HOME}/tmp/VERSION" ]; then
    rm -f "${OSSEC_HOME}/tmp/VERSION"
  fi
  # VERSION channel is public; skip auth if credentials are unset.
  local saved_user="${USERNAME}"
  local saved_pass="${PASSWORD}"
  if [ "${USERNAME}" = "USERNAME" ] || [ -z "${USERNAME}" ]; then
    USERNAME=
    PASSWORD=
  fi
  download "https://${SERVER}/channels/rules/VERSION" "${OSSEC_HOME}/tmp/VERSION"
  USERNAME="${saved_user}"
  PASSWORD="${saved_pass}"

  load_version_file "${OSSEC_HOME}/tmp/VERSION" remote
  load_version_file "${OSSEC_HOME}/var/VERSION" current
}

update_version() {
  local release=$1
  local version=$2
  if [ ! -f "${OSSEC_HOME}/var/VERSION" ]; then
    echo "OSSEC_CRS_RULES_VERSION=0" > "${OSSEC_HOME}/var/VERSION"
    echo "OSSEC_CRS_THREAT_VERSION=0" >> "${OSSEC_HOME}/var/VERSION"
  fi

  sed -i "s/${release}.*/${release}=${version}/g" "${OSSEC_HOME}/var/VERSION"
}

show_updates() {
  local i
  local -a ARRAY1=()

  check_version

  if version_lt "$OSSEC_CRS_RULES_VERSION_CURRENT" "$OSSEC_CRS_RULES_VERSION"; then
    ARRAY1+=("OSSEC-CRS-Rules ${OSSEC_CRS_RULES_VERSION}")
  fi

  if version_lt "$OSSEC_CRS_THREAT_VERSION_CURRENT" "$OSSEC_CRS_THREAT_VERSION"; then
    ARRAY1+=("Atomicorp-Threatfeed ${OSSEC_CRS_THREAT_VERSION}")
  fi

  if [ ${#ARRAY1[@]} -ge 1 ]; then
    echo "Available Upgrades"
    for i in "${!ARRAY1[@]}"; do
      printf "%-32s\n" "${ARRAY1[i]}"
    done
  else
	echo "No updates available"
  fi
}

restore_rules_from_backup() {
  local dst_decoders="${OSSEC_HOME}/etc/decoders.d"
  local dst_rules="${OSSEC_HOME}/etc/rules.d"
  local src_decoders="${OSSEC_HOME}/var/backup/decoders.d"
  local src_rules="${OSSEC_HOME}/var/backup/rules.d"

  rm -f "${dst_decoders}/"*crs*xml 2>/dev/null || true
  rm -f "${dst_rules}/"*crs*xml 2>/dev/null || true

  mkdir -p "${dst_decoders}" "${dst_rules}"
  # Trailing /. copies contents into the destination (avoids nesting dirs).
  if [ -d "${src_decoders}" ]; then
    cp -a "${src_decoders}/." "${dst_decoders}/"
  fi
  if [ -d "${src_rules}" ]; then
    cp -a "${src_rules}/." "${dst_rules}/"
  fi
}

update_rules() {
  echo "  Downloading Rule update: ${OSSEC_CRS_RULES_VERSION}"
  download \
    "https://${SERVER}/channels/rules/ossec/ossec-crs-rules-${OSSEC_CRS_RULES_VERSION}.tar.gz" \
    "${OSSEC_HOME}/tmp/ossec-crs-rules-${OSSEC_CRS_RULES_VERSION}.tar.gz"

  pushd "${OSSEC_HOME}/tmp/" >/dev/null
    if [ -d ossec-rules ]; then
      rm -rf ossec-rules
    fi
    echo -n "  Extracting archive: "
    if ! tar xf "ossec-crs-rules-${OSSEC_CRS_RULES_VERSION}.tar.gz"; then
      echo "Failed"
      echo "  ERROR: archive could not be extracted"
      popd >/dev/null
      exit 1
    else
      echo "OK"
    fi

    # Back up rules
    echo -n "  Backup current rules: "
    if [ -d "${OSSEC_HOME}/var/backup" ]; then
      rm -rf "${OSSEC_HOME}/var/backup"
    fi
    mkdir -p "${OSSEC_HOME}/var/backup/decoders.d/"
    mkdir -p "${OSSEC_HOME}/var/backup/rules.d/"
    cp -a "${OSSEC_HOME}/etc/decoders.d/"*xml "${OSSEC_HOME}/var/backup/decoders.d/" 2>/dev/null || true
    cp -a "${OSSEC_HOME}/etc/rules.d/"*xml "${OSSEC_HOME}/var/backup/rules.d/" 2>/dev/null || true
    echo "OK"

    echo -n "  Applying base rule policy: "
    mkdir -p "${OSSEC_HOME}/etc/decoders.d/"
    mkdir -p "${OSSEC_HOME}/etc/rules.d/"
    rm -f "${OSSEC_HOME}/etc/decoders.d/"*crs*xml
    rm -f "${OSSEC_HOME}/etc/rules.d/"*crs*xml
    cp -a ossec-rules/decoders.d/*xml "${OSSEC_HOME}/etc/decoders.d/"
    cp -a ossec-rules/rules.d/*xml "${OSSEC_HOME}/etc/rules.d/"
    echo "OK"

    echo "  Excluding rulesets"
    for ruleset in $EXCLUDE_RULES; do
      if [ -f "${OSSEC_HOME}/etc/rules.d/${ruleset}" ]; then
        echo "    Disabling: $ruleset"
        rm -f "${OSSEC_HOME}/etc/rules.d/${ruleset}"
      fi
    done

    # Lint
    echo -n "  Verifying rules: "
    if ! "${OSSEC_HOME}/bin/ossec-analysisd" -t >/dev/null 2>&1; then
      echo "Failed"
      echo "  ERROR: Rule update failed lint"
      echo "  Reverting to last working copy"
      restore_rules_from_backup
      echo
      popd >/dev/null
      exit 1
    else
      echo "OK"
      update_version OSSEC_CRS_RULES_VERSION "${OSSEC_CRS_RULES_VERSION}"
    fi

    UPDATE_OUT+="  OSSEC CRS Rules ${OSSEC_CRS_RULES_VERSION}\n"
    rm -rf ossec-rules

  popd >/dev/null
}

update_threatfeed() {
  local shopt_was_nullglob=
  echo "  Downloading Atomicorp Threatfeed update: ${OSSEC_CRS_THREAT_VERSION}"
  download \
    "https://${SERVER}/channels/rules/ossec/atomicorp-threatfeed-${OSSEC_CRS_THREAT_VERSION}.tar.gz" \
    "${OSSEC_HOME}/tmp/atomicorp-threatfeed-${OSSEC_CRS_THREAT_VERSION}.tar.gz"

  pushd "${OSSEC_HOME}/tmp/" >/dev/null
    if [ -d atomicorp-threatfeed ]; then
      rm -rf atomicorp-threatfeed
    fi
    echo -n "  Extracting archive: "
    if ! tar xf "atomicorp-threatfeed-${OSSEC_CRS_THREAT_VERSION}.tar.gz"; then
      echo "Failed"
      echo "  ERROR: archive could not be extracted"
      popd >/dev/null
      exit 1
    else
      echo "OK"
    fi

    mkdir -p "${OSSEC_HOME}/etc/lists/threat/"
    rm -f "${OSSEC_HOME}/etc/lists/threat/"* >/dev/null 2>&1 || true
    cp -a atomicorp-threatfeed/* "${OSSEC_HOME}/etc/lists/threat/"

    chown root:ossec "${OSSEC_HOME}/etc/lists/threat"
    chmod 750 "${OSSEC_HOME}/etc/lists/threat"
    if shopt -q nullglob; then
      shopt_was_nullglob=1
    else
      shopt -s nullglob
    fi
    local threat_files=("${OSSEC_HOME}/etc/lists/threat/"*)
    if [ ${#threat_files[@]} -gt 0 ]; then
      chown root:ossec "${threat_files[@]}"
      chmod 640 "${threat_files[@]}"
    fi
    if [ -z "$shopt_was_nullglob" ]; then
      shopt -u nullglob
    fi

    if [ -x "${OSSEC_HOME}/bin/ossec-makelists" ]; then
      echo -n "  Building CDB lists: "
      if "${OSSEC_HOME}/bin/ossec-makelists" >/dev/null 2>&1; then
        echo "OK"
      else
        echo "Failed"
        echo "  WARNING: ossec-makelists reported an error; ensure threat lists"
        echo "  are declared in ${OSSEC_HOME}/etc/ossec.conf"
      fi
    fi

    update_version OSSEC_CRS_THREAT_VERSION "${OSSEC_CRS_THREAT_VERSION}"
    UPDATE_OUT+="  Atomicorp Threatfeed ${OSSEC_CRS_THREAT_VERSION}\n"
    rm -rf atomicorp-threatfeed

  popd >/dev/null
}

# Apply updates if they are available
update() {
  local i
  local -a ARRAY1=()
  local RESTART=

  if [[ "${USERNAME}" == "USERNAME" ]]; then
	echo
	echo "ERROR: oum credentials have not been configured. Run:"
	echo "  oum configure"
	echo
	exit 1
  fi

  check_version
  if [ -n "$DEBUG" ]; then
    echo "DEBUG: Compare ${OSSEC_CRS_RULES_VERSION_CURRENT} to ${OSSEC_CRS_RULES_VERSION}"
  fi

  UPDATE_OUT=""

  if version_lt "$OSSEC_CRS_RULES_VERSION_CURRENT" "$OSSEC_CRS_RULES_VERSION"; then
    ARRAY1+=("OSSEC-CRS-Rules ${OSSEC_CRS_RULES_VERSION}")
  fi
  if version_lt "$OSSEC_CRS_THREAT_VERSION_CURRENT" "$OSSEC_CRS_THREAT_VERSION"; then
    ARRAY1+=("Atomicorp-Threatfeed ${OSSEC_CRS_THREAT_VERSION}")
  fi

  if [[ ${#ARRAY1[@]} -ge 1 ]]; then

    echo "==============================================================================="
    echo " Component			Version"
    echo "==============================================================================="
    echo "Upgrading: "
    for i in "${!ARRAY1[@]}"; do
      printf "  %-28s\n" "${ARRAY1[i]}"
    done
    echo
    echo
    echo "==============================================================================="
    echo

    if [ -z "$YES" ]; then
      read -p "Is this ok [y/N]: " -r
      if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo "Operation aborted."
        exit 1
      fi
    fi

    for i in "${ARRAY1[@]}"; do
      if [[ "$i" =~ OSSEC-CRS-Rules ]]; then
        update_rules
      fi
      if [[ "$i" =~ Atomicorp-Threatfeed ]]; then
        update_threatfeed
      fi
    done

    RESTART=1
    echo
    echo "Upgraded:"
    echo -e "${UPDATE_OUT}"
    echo "Complete!"
  else
	echo "No packages marked for update"
  fi

  # Restart ossec
  if [[ -n ${RESTART} ]]; then
    echo -n "  Restarting OSSEC: "
    if [ -n "$DEBUG" ]; then
      echo
      "${OSSEC_HOME}/bin/ossec-control" restart
    else
      if "${OSSEC_HOME}/bin/ossec-control" restart >/dev/null 2>&1; then
        echo "OK"
      else
        echo "Failed"
      fi
    fi
    echo
  fi
}

configure() {
  local user_tmp PASSWORD PASSWORD2 PASSCONFIRMED failed
  local HTTP_CODE BODY cred_tmp

  echo
  echo "OSSEC Updater Modified (OUM) $VERSION"

  # Prompt for Username
  read -p "Please enter your subscription username: " user_tmp
  until [ "$user_tmp" != "" ]; do
    echo " ERROR: Username cannot be blank. "
    echo
    read -p "Please enter your subscription username:  " user_tmp
  done

  PASSCONFIRMED=0
  failed=0
  while [ $PASSCONFIRMED -lt 1 ]; do
    if [ $failed -gt 2 ]; then
        echo "Exiting: too many failed attempts."
        echo
        exit 1
    fi

    read -sp "Please enter your subscription password: " PASSWORD
    echo
    if [ "$PASSWORD" == "" ]; then
        echo "Exiting: Password is blank..."
        exit 1
    fi

    unset PASSWORD2
    read -sp "Please re-enter your subscription password: " PASSWORD2
    echo

    if [ "$PASSWORD" == "$PASSWORD2" ]; then
        PASSCONFIRMED=1
    else
        failed=$((failed + 1))
        echo
        echo "   Passwords do not match."
        echo
    fi
  done

  echo
  echo -n "Verifying account: "
  cred_tmp=$(mktemp)
  HTTP_CODE=$(curl -s -o "${cred_tmp}" -w '%{http_code}' \
    --user "${user_tmp}:${PASSWORD}" \
    "https://${SERVER}/channels/rules/ossec/README.md" || echo "000")
  BODY=$(cat "${cred_tmp}" 2>/dev/null || true)
  rm -f "${cred_tmp}"

  # Accept HTTP 200, or legacy body match used by older Atomicorp channel responses.
  if [[ "$HTTP_CODE" == "200" ]] || [[ "$BODY" == "Access to this repo requires registration" ]]; then
	  echo "Passed"
  else
	  echo "Failed"
	  echo
	  echo "  Username/Password credentials incorrect (HTTP ${HTTP_CODE})."
	  echo "  Please confirm the correct credentials were entered."
	  echo
	  exit 1
  fi

  # Persist only after successful verification (safe against special chars in password).
  write_oum_credentials "$user_tmp" "$PASSWORD"

  echo "Configuration Complete!"
  echo
  echo "To update the system, please run:"
  echo "    oum update "
  echo
}


# Installs a package via OUM
install_package() {
  local pkg="$1"
  local pkg_mgr=

  if [ -z "$pkg" ]; then
    echo " ERROR: No package specified."
    echo " Usage: oum install PACKAGE"
    exit 1
  fi

  echo
  echo "Installing Package: ${pkg}"
  echo

  if command -v dnf >/dev/null 2>&1; then
    pkg_mgr=dnf
  elif command -v yum >/dev/null 2>&1; then
    pkg_mgr=yum
  else
    echo " ERROR: Neither dnf nor yum was found."
    exit 1
  fi

  if ! "${pkg_mgr}" install -y --enablerepo=atomic "${pkg}"; then
    echo
    echo " ERROR: There was a problem installing ${pkg}!"
    echo
    exit 1
  else
    echo
    echo " ${pkg} successfully installed"
  fi
}


# Parse global options, then the subcommand.
while getopts ":ydi" opt; do
  case ${opt} in
    y)
      YES=1
      ;;
    d)
      DEBUG=1
      ;;
    i)
      INSECURE=1
      ;;
    \?)
      echo "Invalid Option: -$OPTARG" 1>&2
      exit 1
      ;;
  esac
done
shift $((OPTIND -1))

command="${1:-}"
if [ -n "$command" ]; then
  shift
fi

case "$command" in
  help|"")
    show_help
    ;;
  version)
    echo "OUM Version: $VERSION"
    ;;
  list)
    load_config
    prereq_check
    show_updates
    ;;
  configure)
    load_config
    prereq_check
    configure
    ;;
  update|upgrade)
    load_config
    prereq_check
    update
    ossec_check
    ;;
  install)
    install_package "$1"
    ;;
  *)
    show_help
    exit 1
    ;;
esac
