#!/bin/bash

AIDE_DB_DIR="/var/lib/aide"
AIDE_DB_PATH="${AIDE_DB_DIR}/aide.db.gz"
AIDE_NEW_DB_PATH="${AIDE_DB_DIR}/aide.db.new.gz"
CRON_FILE_PATH="/etc/crontab"
LOG_PATH="/var/log/aide/aide_service.log"
CRON_EXEC_CMD="/usr/sbin/aide_cron.sh"
CRON_CMD="*/30 * * * * ${CRON_EXEC_CMD}"
CRON_MATE_RULE=" \/usr\/sbin\/aide_cron\.sh/d"
TMP_CRONTAB_PATH="/tmp/aide_crontab.tmp"

function log() {
  DATE=$(date +"%Y-%m-%d %H:%M:%S")
  echo "$DATE - $1" >> ${LOG_PATH}
}

# Define a function to set a timer
set_timer() {
  interval=$1
  unit=$2
   
  # Parse the command line arguments and set the timer
  case "$unit" in
    "minute")
      cron_interval="*/$interval * * * *"
      ;;
    "hour")
      cron_interval="0 */$interval * * *"
      ;;
    "day")
      cron_interval="0 0 */$interval * *"
      ;;
    *)
      echo "Invalid unit. Please use 'minute', 'hour', or 'day'." >&2
      exit 1
      ;;
  esac
  
  CRON_CMD="$cron_interval ${CRON_EXEC_CMD}"
  
  echo "Scheduled task has been set to run every $interval $unit."
}

modify_timer() {
  interval=$1
  unit=$2
  set_timer $interval $unit
  /usr/bin/crontab -l > ${TMP_CRONTAB_PATH}
  /usr/bin/sed -i ${CRON_MATE_RULE} ${TMP_CRONTAB_PATH} &>>${LOG_PATH}
  echo "${CRON_CMD}" >> ${TMP_CRONTAB_PATH}
  /usr/bin/crontab ${TMP_CRONTAB_PATH}
  rm ${TMP_CRONTAB_PATH}
  echo "Scheduled task has been modified to run every $interval $unit."
}

function aide_start() {
  # Run AIDE initialization if database file doesn't exist
  echo "Start to initialize AIDE and enable scheduled checks..."
  /usr/bin/rm -f ${AIDE_NEW_DB_PATH} &>>${LOG_PATH}
  if [ ! -f ${AIDE_DB_PATH} ]; then
      echo "Start to init aide db..."
      /usr/sbin/aide --init &>>${LOG_PATH} 
      mv ${AIDE_NEW_DB_PATH} ${AIDE_DB_PATH} &>>${LOG_PATH}
      echo "End to init aide db."
  fi
  echo "Start to create a scheduled monitoring task..."
  /usr/bin/crontab -l > ${TMP_CRONTAB_PATH}
  /usr/bin/sed -i ${CRON_MATE_RULE} ${TMP_CRONTAB_PATH} &>>${LOG_PATH}
  echo "${CRON_CMD}" >> ${TMP_CRONTAB_PATH}
  /usr/bin/crontab ${TMP_CRONTAB_PATH}
  rm ${TMP_CRONTAB_PATH}
  echo "End to create a scheduled monitoring task."
  echo "End to initialize AIDE and enable scheduled checks."
}

function cron_task() { 
  CRON_CMD=""
  matching_job=$(/usr/bin/crontab -l | grep "${CRON_EXEC_CMD}")
  if [[ -n "$matching_job" ]]; then
    CRON_CMD=$matching_job
    echo "The current scheduled tasks is ${CRON_CMD}"
    return 1
  fi
  return 0
}

function aide_stop() {
  echo "Start to disable scheduled checks..."
  cron_task
  if [ $? -eq 1 ]; then
    /usr/bin/crontab -l | grep -v "${CRON_EXEC_CMD}" | crontab -
  fi
  echo "End to disable scheduled checks."
}

function aide_update() {
  echo "Start to update aide..."
  aide_stop
  echo "Start to update aide db..."
  /usr/sbin/aide --update &>>${LOG_PATH}
  rm -f ${AIDE_DB_DIR}/aide.db_*_bk.gz 
  current_time=$(date +"%Y_%m_%d_%H_%M_%S")
  mv ${AIDE_DB_PATH} "${AIDE_DB_DIR}/aide.db_${current_time}_bk.gz" &>>${LOG_PATH}
  mv ${AIDE_NEW_DB_PATH} ${AIDE_DB_PATH} &>>${LOG_PATH}
  echo "End to update aide db."
  if [ -n "${CRON_CMD}" ]; then
    echo "Start to create a scheduled monitoring task..."
    echo "${CRON_CMD}" > ${TMP_CRONTAB_PATH}
    /usr/bin/crontab ${TMP_CRONTAB_PATH}
    rm ${TMP_CRONTAB_PATH}
    echo "End to create a scheduled monitoring task."
  fi
  echo "End to update aide."
}

# If no arguments are provided, display the help information
if [ -z "$1" ]; then
  echo "Usage: script_name [options]"
  echo "Options:"
  echo "   -s, --start <interval> <unit>       Initialize AIDE and enable scheduled checks. interval stands for"
  echo "	                               time interval for scheduled checks and unit represents the unit "      
  echo "                                       for the time interval (minute/hour/day).for example: -m 5 minute."
  echo "   -t, --stop                          Disable scheduled checks."
  echo "   -u, --update                        Update AIDE DB and enable scheduled checks."
  echo "   -m, --modify <interval> <unit>      Modify the time interval for scheduled checks. interval stands for"
  echo "                                       time interval for scheduled checks and unit represents the unit for "
  echo "                                       the time interval (minute/hour/day).for example: -m 5 minute."
  echo "   -q, --query                         View current scheduled tasks."
  exit 1
fi

# Handle different options
case "$1" in
  -s|--start)
    if [[ $2 == "" ]] || [[ $3 == "" ]]; then
      echo "Method usage failed. Please try again."
      exit 1  
    fi
    if ! [[ $2 =~ ^[1-9][0-9]*$ ]]; then
      echo "Invalid interval. Please enter a positive integer." >&2
      exit 1
    fi
    interval=$2
    unit=$3
    set_timer $interval $unit
    aide_start
    ;;
  -t|--stop)
    aide_stop
    ;;
  -u|--update)
    aide_update
    ;;
  -m|--modify)
    if [[ $2 == "" ]] || [[ $3 == "" ]]; then
      echo "Method usage failed. Please try again."
      exit 1
    fi
    if ! [[ $2 =~ ^[1-9][0-9]*$ ]]; then
      echo "Invalid interval. Please enter a positive integer." >&2
      exit 1
    fi
    interval=$2
    unit=$3
    modify_timer $interval $unit
    ;;
  -q|--query)
    cron_task
    ;;
  -h|--help)
    echo "Usage: script_name [options]"
    echo "Options:"
    echo "   -s, --start <interval> <unit>       Initialize AIDE and enable scheduled checks. interval stands for"
    echo "                                       time interval for scheduled checks and unit represents the unit "
    echo "                                       for the time interval (minute/hour/day).for example: -m 5 minute."
    echo "   -t, --stop                          Disable scheduled checks."
    echo "   -u, --update                        Update AIDE DB and enable scheduled checks."
    echo "   -m, --modify <interval> <unit>      Modify the time interval for scheduled checks. interval stands for"
    echo "                                       time interval for scheduled checks and unit represents the unit for "
    echo "                                       the time interval (minute/hour/day).for example: -m 5 minute."
    echo "   -q, --query                         View current scheduled tasks."
    ;;
  *)
    echo "Invalid argument. Please enter a valid option."
    exit 1
    ;;
esac
exit 0
