#!/ bin/ bash
set -e
# Check root privileges
[[ $EUID - ne θ ]] && { echo" Please run as root: sudo $θ"; exit 1; }
# Define software list and URLs
software_list=(" Wazuh" "ClamAV" " audit")
declare - A urls=(
  [wazuh]="https://pubapi.hmxpr.com/client/wazuh"
  [jq]="https://pubapi.hmxpr.com/client/jq"
  [wazuhdeb]="https://pubapi.hmxpr.com/client/wazuhdeb"
  [wazuhaudit]="https://pubapi.hmxpr.com/client/wazuhaudit"
  [clamavdeb]="https://pubapi.hmxpr.com/client/clamavdeb"
  [clamav]="https://pubapi.hmxpr.com/client/clamav"
  [clamconf]="https://pubapi.hmxpr.com/conf/clamav"
  [audit_rules]="https://pubapi.hmxpr.com/conf/audit"
  [audit]="https://pubapi.hmxpr.com/client/audit"
  [lib]="https://pubapi.hmxpr.com/client/lib"
  [tools]="https://pubapi.hmxpr.com/client/"
  )
# Initialize variables for command line arguments
WAZUH_MANAGER=""
WAZUH_AGENT_GROUP="default"
INSTALL_WAZUH=false
INSTALL_CLAMAV=false
INSTALL_AUDIT=false
# Function to display usage information
usage() {
    echo "Usage: $θ [-w] [-c] [-a] [-m MANAGER_ADDRESS] [-g AGENT_GROUP]"
    echo "-w: Install Wazuh"
    echo "-c: Install ClamAV"
    echo "-a: Install audit"
    echo "-m MANAGER_ADDRESS: Specify Wazuh manager address"
    echo "-g AGENT_GROUP: Specify Wazuh agent group"
    echo "-h: Display this help message"
    exit 1
    }
# Parse command line arguments
while getopts ":wcam:g:h" opt; do
  case ${opt} in
    w) INSTALL_WAZUH=true ;;
    c) INSTALL_CLAMAV=true ;;
    a) INSTALL_AUDIT=true ;;
    m) WAZUH_MANAGER=$OPTARG ;;
    g) WAZUH_AGENT_GROUP=$OPTARG ;;
    h) usage ;;
    \?) echo "Invalid option: $OPTARG" 1>&2; usage ;;
    :) echo "Invalid option: $OPTARG requires an argument" 1>&2; usage ;;
    esac
done
shift $((OPTIND - 1))
# Helper functions
download_file(){
  local url="$1"
  local output_file="$2"
  if command -v wget &> /dev/null; then
    wget -0 "$output_file" "$url" || return 1
  elif command -v curl &> /dev/null; then
    curl -k -o "$ output_file" "$url" || return 1
  elif command -v fetch &> /dev/null; then
    fetch -o "$ output_file" "$url" || return 1
  else
    echo "Error: No suitable download tool found. Please install wget, curl, or fetch. "
    exit 1
  fi
}
ensure_download_tool() {
  for tool in wget unzip; do
    if ! command -v "$tool" > /dev/null 2>&1; then
      echo "Installing $tool…"
      if [[ -f /etc/os-release ]]; then
        ./etc/os-release
        case "$ID" in
          ubuntu|debian)
            case "${VERSION_ID}" in
              20.04|22.04|24.04)
                download_file "${urls[tools]}${tool}${VERSION_ID/./}" "${tool}.deb"
                dpkg -i "${tool}.deb" || {echo "$tool install failed"; exit 1; }
                rm - fr "${tool}.deb"
            ;;
              *)
                echo "不支持的 Ubuntu 版本: $VERSION_ID"
                exit 1
            ;;
            esac
        ;;
          centos|rhel|fedora|amzn)
            download_file "${urls[tools]}${tool}c7" "${tool}.rpm"
            rpm -ivh ${tool}.rpm
        ;;
          rocky)
            download_file "${urls[tools]}${tool}r9" "${tool}.rpm"
            rpm -ivh ${tool}.rpm
        ;;
        esac
    rm - fr ${tool}.*
    else
      echo "Unable to determine OS version"
      exit 1
      fi
    fi
  done
}
check_software_installed() {
  if [[ -f /etc/os-release ]]; then
    ./etc/os-release
    case "$ID" in
      ubuntu|debian) dpkg -s "$1" &> /dev/null ;;
      centos|rocky|rhel|fedora|amzn) rpm -q "$1" &> /dev/null ;;
      *) echo "Unsupported OS: $ID"; exit 1 ;;
    esac
  else
    echo "Unable to determine OS type"
    exit 1
  fi
}
get_user_choice() {
  if [[ $INSTALL_WAZUH == false && $INSTALL_CLAMAV == false && $INSTALL_AUDIT == false ]]; then
    echo "Select software to install (comma-separated):"
    for i in "${!software_list[@]}"; do
      echo"$((i+1)). ${software_list[$i]}"
    done
    read -p "Enter numbers: " choices_raw
    IFS=',' read -ra choices_array <<< "$choices_raw"
    selected_software=()
    for choice in "${choices_array[@]}"; do
      selected_software+=("${software_list[$((choice-1))]}")
    done
    echo "Selected: ${selected_software[*]}"
  else
    selected_software=()
    $INSTALL_WAZUH && selected_software+=("Wazuh")
    $INSTALL_CLAMAV && selected_software+=("ClamAV")
    $INSTALL_AUDIT && selected_software+=("audit")
    echo "Selected from command line: ${selected_software[*]}"
  fi
}
install_wazuh() {
  local pkg="wazuh-agent"
  if check_software_installed "$pkg"; then
    read -p "$pkg already installed. Reinstall? (y/n): " confirm
    if [[ "$confirm" != "y" ]]; then
      echo "Skipping Wazuh installation."
      return
    fi
    echo "Removing existing Wazuh installation..."
    ./etc/os-release
    case"$ID" in
      ubuntu|debian)
        dpkg --purge wazuh-agent
      ;;
      centos|rocky|rhel|fedora|amzn)
        rpm -e wazuh-agent
      ;;
      *)
        echo "Unsupported OS for Wazuh removal: $ID"
        exit 1
      ;;
    esac
    echo "Existing Wazuh installation removed. "
  fi
  if [[ -z $WAZUH_MANAGER ]]; then
    read -p " Enter manager address: " WAZUH_MANAGER
    read -p " Enter agent group: " WAZUH_AGENT_GROUP
    read -p " Confirm manager: $WAZUH_MANAGER, group: $WAZUH_AGENT_GROUP (y/n): " confirm
    [[ "$confirm" != "y" ]] && return
  fi
  ./etc/os-release
  case "$ID" in
    ubuntu|debian)
      download_file "${urls[wazuhdeb]}" "wazuh-agent.deb"
      WAZUH_MANAGER=$WAZUH_MANAGER WAZUH_AGENT_GROUP=$WAZUH_AGENT_GROUP dpkg -i wazuh-agent.deb
    ;;
    centosIrocky|rhel|fedora|amzn)
      download_file "${urls[wazuh]}" " wazuh-agent.rpm"
      WAZUH_MANAGER=$WAZUH_MANAGER WAZUH_AGENT_GROUP=$WAZUH_AGENT_GROUP rpm -ivh wazuh-agent.rpm
    ;;
  esac
  echo -e"wazuh_command.remote_commands=1\nlogcollector.remote_commands=1" >> /var/ossec/etc/local_internal_options.conf
  download_file "${urls[wazuhaudit]}" "editrule.sh"
  yes|mv editrule.sh /var/ossec/active-response/bin/
  chown -R root:wazuh /var/ossec/active-response/bin/editrule.sh
  chmod 750 /var/ossec/active-response/bin/editrule.sh
  systemctl daemon-reload
  systemctl enable --now wazuh-agent
  if check_software_installed jq; then
    echo "jq is already installed."
  else
    ./etc/os-release
    case "$ID" in
      ubuntu|debian)
        download_file "${urls[jq]}${VERSION_ID/./}" "jj${VERSION_ID/./}.zip"
        case "${VERSION_ID}" in
          20.04)
            unzip jq2004.zip
            dpkg -i libonig5_6.9.4-1_amd64.deb
            dpkg -i libjq1_1.6-1_amd64.deb
            dpkg -i jq_1.6-1_amd64.deb
          ;;
          22.04)
            unzip jq2204.zip
            dpkg -i libonig5_6.9.7.1-2build1_amd64.deb
            dpkg -i libjq1_1.6-2.1ubuntu3_amd64.deb
            dpkg -i jq_1.6-2.1ubuntu3_amd64.deb
          ;;
          24.04)
            unzip jq2404.zip
            dpkg -i libonig5_6.9.9-1build1_amd64.deb
            dpkg -i libjq1_1.7.1-3build1_amd64.deb
            dpkg -i jq_1.7.1-3build1_amd64.deb
          ;;
        esac
      ;;
      centos|rhel|fedora|amzn)
        download_file "${urls[jq]}c7" "jq.zip"
        unzip jq.zip
        rpm -ivh oniguruma-6.8.2-2.el7.x86_64.rpm
        rpm -ivh jq-1.6-2.el7.x86_64.rpm
      ;;
      rocky)
        download_file "${urls[jq]}c9" "jq.rpm"
        rpm -ivh jq.rpm
      ;;
      *)
        echo "Unsupported OS for jq removal: $ID"
        exit 1
      ;;
    esac
  fi
  echo "Wazuh installed successfully"
  rm -fr wazuh-agent.*
  rm -fr oniguruma-6.8.2-2.el7.x86_64.rpm
  rm -fr lib*
  rm -fr jq*
}
install_clamav() {
  local pkg="clamav"
  if check_software_installed "$pkg"; then
  read -p "$pkg already installed. Reinstall? (y/n): " confirm
    if [[ "$confirm" != "y" ]]: then
      echo "Skipping $pkg installation. "
      return
    fi
    echo "Removing existing $pkg installation... "
    ./etc/os-release
    case "$ID" in
      ubuntu|debian)
        dpkg --purge $pkg
      ;;
      centos|rocky|rhel|fedora|amzn)
        rpm -e $pkg
      ;;
      *)
        echo "Unsupported OS for $pkg removal: $ID"
        exit 1
      ;;
    esac
    echo "Existing $pkg installation removed. "
  fi
  ./etc/os-release
  case "$ID" in
    ubuntu|debian)
      download_file "${urls[clamavdeb]}" "clamav.deb"
      dpkg -i clamav.deb
    ;;
    centos|rocky|rhel|fedora|amzn)
      download_file "${urls[clamav]}" "clamav.rpm"
      rpm -ivh clamav.rpm
    ;;
    *)
      echo "Error! Can't detect the OS Version!"
      exit 1
    ;;
  esac
  download_file "${urls[clamconf]}" "clamav.zip"
  unzip clamav.zip
  yes|mv clamd.conf /usr/local/etc/
  yes|mv freshclam.conf /usr/local/etc/
  yes|mv clamd.service /etc/systemd/system/
  getent group clamav > /dev/null || groupadd clamav
  id -u clamav > /dev/null 2>&1 || useradd -g clamav clamav
  if [[ ! -d /usr/local/clamav ]]; then
    mkdir -p /usr/local/clamav
    chown clamav:clamav /usr/local/clamav
  fi
  for log_file in /var/log/freshclam.log /var/log/clamav.log; do
    if [[ ! -f $log_file ]]; then
      touch $log_file
      chmod 644 $log_file
      chown clamav:clamav $log_file
    else
      chmod 644 $log_file
      chown clamav:clamav $log_file
    fi
  done
  /usr/local/bin/freshclam
  echo "pls wait for few seconds"
  sleep 10
  systemctl daemon-reload
  systemctl enable clamd
  systemctl restart clamd
  echo "$pkg installed and configured"
  CRON_JOB="0 12 * * * /usr/local/bin/freshclam --quiet"
  EXISTING_CRONTAB=$(crontab - l 2>/dev/null)
  if echo "$EXISTING_CRONTAB" |grep - Fq "$CRON_JOB"; then
    echo "cron 任务已存在, 无需添加。"
    exit θ
  fi
  (echo "$EXISTING_CRONTAB"; echo "$CRON_JOB") | crontab -
  if crontab -l |grep -Fq "$CRON_JOB"; then
    echo "cron 任务成功添加。"
  else
    echo "添加 cron 任务失败。" 1>&2
    exit 1
  fi
  rm -fr clamav.*
}
  install_audit() {
    download_file "${urls[audit_rules]}" "audit.rules"
    local pkg="auditd"
    if check_software_installed "$pkg"; then
      read -p "$pkg already installed. Reinstall? (y/n): " confirm
      if [[ "$confirm" != "y" ]]; then
        echo "Skipping audit installation. "
        return
      fi
      echo "Removing existing audit installation... "
      ./etc/os-release
      case "$ID" in
        ubuntu|debian)
          dpkg --purge $pkg
          echo "Existing $pkg installation removed. "
          download_file "${urls[audit]}${VERSION_ID/./}" "audit.deb"
          download_file "${urls[lib]}${VERSION_ID/./}" "libauparseθ.deb"
          dpkg -i libauparse0.deb audit.deb
          if [[ ! -d /etc/audit/rules.d ]]; then
            mkdir -p /etc/audit/rules.d
          fi
          yes|mv audit.rules/etc/audit/rules.d/
          sed -i 's/^#?RestartPreventExitStatus=.*/RestartPreventExitStatus=255/' /usr/lib/system/system/auditd.service
          systemctl daemon-reload
          systemctl enable --now auditd
        ;;
        centos|rocky|rhel|fedora|amzn)
          echo "Audit don't need to re-installed on CentOS/RHEL/Fedora. "
          output=$(systemctl status audit d 2>&1)
          if echo "$output" |grep -q "Warning: Journal has been rotated since unit was started. Log output isincomplete or unavailable."; then
            echo "检测到 auditd服务异常,尝试修复。。。。"
            sed -i 's/RefuseManualStop=yes/RefuseManualStop=no/' /usr/lib/systemd/system/auditd.service
            systemctl daemon-reload
            systemctl restart auditd
            sed -i 's/RefuseManualStop=no/RefuseManualStop=yes/' /usr/lib/systemd/system/auditd.service
            systemctl daemon-reload
            echo "auditd服务修复完成, 请检查日志"
          else
            echo "audit服务正常,继续下一步"
          fi
        ;;
      esac
    fi
    yes|mv audit.rules /etc/audit/rules.d/augenrules --load
    echo "Audit installed and configured"
    rm -fr audit.deb
    rm -fr libauparse0.deb
}
# Main program
ensure_download_tool
get_user_choice
for software in "${selected_software[@]}"; do
  case $software in
    "Wazuh")
      install_wazuh 
    ;;
    "ClamAV")
      install_clamav
    ;;
    "audit")
      install_audit
    ;;
  esac
done