#!/bin/bash

# Dynamic Scan for Machine Learning Models version 0.1.
#
# Copyright (c) 2025 Oracle, Inc.
# Licensed under the Universal Permissive License v 1.0 as shown at
# https://oss.oracle.com/licenses/upl/

# Prerequisites check
# The user needs to have sudo privilege
# The docker needs to be installed, with docker compose support
function check_prerequisites() {
    if [ "$(id -u)" -ne 0 ]; then
        sudo -l | grep -qP '\(ALL.*\)\s*NOPASSWD:\s*ALL' 2>/dev/null
        if [ $? -ne 0 ]; then
          echo "Please run this script as a privileged user, eg: has no restriction on sudo without password."
          exit 1
        fi
    fi

    if ! command -v docker &> /dev/null; then
        echo "Docker is not installed"
        exit 1
    fi

    # Determine which Docker Compose command is available
    if command -v docker-compose &> /dev/null; then
        export DOCKER_COMPOSE_CMD="docker-compose"
    elif docker compose version &> /dev/null; then
        export DOCKER_COMPOSE_CMD="docker compose"
    else
        echo "Docker Compose is not installed or not available in PATH."
        exit 1
    fi
}

check_prerequisites

# Create a shell function to use the exported variable
docker_compose() {
    $DOCKER_COMPOSE_CMD $@
}

if [ $# -ne 1 ]; then
  echo "Usage: $0 <model_to_scan>"
  echo "Example: $0 \"bert-base-uncased_with_RCE\""
  echo "Note, please make sure the model_to_scan is in the \"share\" folder"
  exit 1
fi

MODEL_TO_SCAN="$1"

# Get a temp file
RW_FS_scanning=$(mktemp)
RW_FS_parsing=$(mktemp)
# Get a random name for the mount point
MOUNT_POINT_scanning="/mnt/loopback_$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 12 | head -n 1)"
MOUNT_POINT_parsing="/mnt/loopback_$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 12 | head -n 1)"

# For the cleanup logs
CLEANUP_LOG_FILE=$(mktemp)
touch "${CLEANUP_LOG_FILE}"

# Function to log messages
log_message() {
    echo "$(date): $1" >> "$CLEANUP_LOG_FILE"
}

# Add a trap and do necessary cleanup if exits abnormally
cleanup() {
    log_message "Starting cleanup..."

    # Stop the docker compose
    docker_compose down

    # Unmount filesystems
    if mountpoint -q "$MOUNT_POINT_scanning"; then
        sudo umount "$MOUNT_POINT_scanning" && log_message "Unmounted $MOUNT_POINT_scanning" || log_message "Failed to unmount $MOUNT_POINT_scanning"
    else
        log_message "$MOUNT_POINT_scanning is not mounted"
    fi

    if mountpoint -q "$MOUNT_POINT_parsing"; then
        sudo umount "$MOUNT_POINT_parsing" && log_message "Unmounted $MOUNT_POINT_parsing" || log_message "Failed to unmount $MOUNT_POINT_parsing"
    else
        log_message "$MOUNT_POINT_parsing is not mounted"
    fi

    # Remove files
    sudo rm -f "$RW_FS_scanning" && log_message "Removed $RW_FS_scanning" || log_message "Failed to remove $RW_FS_scanning"
    sudo rm -f "$RW_FS_parsing" && log_message "Removed $RW_FS_parsing" || log_message "Failed to remove $RW_FS_parsing"

    # Remove directories
    sudo rm -rf "$MOUNT_POINT_scanning" && log_message "Removed $MOUNT_POINT_scanning" || log_message "Failed to remove $MOUNT_POINT_scanning"
    sudo rm -rf "$MOUNT_POINT_parsing" && log_message "Removed $MOUNT_POINT_parsing" || log_message "Failed to remove $MOUNT_POINT_parsing"

    log_message "Cleanup completed."
    cat "$CLEANUP_LOG_FILE"
    rm -f "$CLEANUP_LOG_FILE"
}

# Set trap for various signals
trap cleanup EXIT SIGTERM SIGINT SIGHUP SIGQUIT SIGABRT

if [ ! -d "share/${MODEL_TO_SCAN}" ] || [ -z "$(ls -A "share/${MODEL_TO_SCAN}")" ]; then
    echo -e "\e[31mModel to scan (\"${MODEL_TO_SCAN}\") not found or is empty in the \"share\" folder. The model-to-be-scanned MUST be pre-downloaded and put to the \"share\" folder!\e[0m"
    exit 1
fi

MYLOC=`dirname "$0"`
if [ '/' != "`echo \"${MYLOC}\" | cut -c1`" ]; then
  if [ '.' = "${MYLOC}" ]; then
    MYLOC="$PWD"
  else
    MYLOC="$PWD/${MYLOC}"
  fi
fi

cd "${MYLOC}"

# Set the shared folder
SHARED_FOLDER="${MYLOC}/share"
SHARED_FOLDER_MOUNT_POINT_scanning=$(df -h "${SHARED_FOLDER}" | grep -v '^Filesystem' | awk '{print $1}') # Something like: /dev/sda1

# Get the executable path
PYTHON_EXEC="$(cat config.ini | grep '^PYTHON_PATH' | /usr/bin/awk -F'=' '{print $2}' | /usr/bin/awk -F'#' '{print $1}' | /usr/bin/sed 's/^[ \t]*//;s/[ \t]*$//' | /usr/bin/envsubst)"

# Prepare the env
"${PYTHON_EXEC}" "${MYLOC}/env_config.py"

# # Rebuild the image
./build.sh

if [ -f .env ]; then
    source .env 2>/dev/null
fi

# Set up two loopback filesystems
dd if=/dev/zero of="${RW_FS_scanning}" bs=1M count=128 > /dev/null 2>&1 # Creates a 128MB temp file - used for scanning logs
dd if=/dev/zero of="${RW_FS_parsing}" bs=1M count=128 > /dev/null 2>&1 # Creates a 128MB temp file - used for parsing logs
echo y | sudo mkfs.ext4 "${RW_FS_scanning}" > /dev/null 2>&1 # Formats the file as an ext4 file system
echo y | sudo mkfs.ext4 "${RW_FS_parsing}" > /dev/null 2>&1 # Formats the file as an ext4 file system
sudo mkdir -p "${MOUNT_POINT_scanning}" && sudo mount -o loop "${RW_FS_scanning}" "${MOUNT_POINT_scanning}"  && sudo chown -R "$USER":"$USER" "${MOUNT_POINT_scanning}" # Mounts the file system to a mount point - changing ownership is a MUST or the folder will be root owned and can't be written to in the container
sudo mkdir -p "${MOUNT_POINT_parsing}" && sudo mount -o loop "${RW_FS_parsing}" "${MOUNT_POINT_parsing}"  && sudo chown -R "$USER":"$USER" "${MOUNT_POINT_parsing}" # Mounts the file system to a mount point - changing ownership is a MUST or the folder will be root owned and can't be written to in the container

# Generate the docker-compose.yml file
"${PYTHON_EXEC}" "${MYLOC}/setup_env.py" --model_to_scan "${MODEL_TO_SCAN}"

# Modify the  "/mnt/loopback_scanning:/tmp_data_scanning" in the docker-compose.yml file to "${MOUNT_POINT_scanning}:/tmp_data_scanning"
sed -i "s|/mnt/.*:/\(tmp_data_scanning\)|${MOUNT_POINT_scanning}:/\1|g" docker-compose.yml
# Modify the  "/mnt/loopback_parsing:/tmp_data_parsing" in the docker-compose.yml file to "${MOUNT_POINT_parsing}:/tmp_data_parsing"
sed -i "s|/mnt/.*:/\(tmp_data_parsing\)|${MOUNT_POINT_parsing}:/\1|g" docker-compose.yml

# Ensure the compose is stopped
docker_compose down

# Run docker compose up and capture output
docker_compose up

echo "Please check logs at: ${MOUNT_POINT_parsing}/logs"

# If the "${MOUNT_POINT_parsing}/logs/suspicious.log" exists, and the length is greater than 0, then display the content in red color
echo -e "\n\n"
if [ -s "${MOUNT_POINT_parsing}/logs/suspicious.log" ]; then
   echo -e "\e[31mSuspicious behaviors detected:\e[0m"
   sed 's/^/\x1b[31m/; s/$/\x1b[0m/' "${MOUNT_POINT_parsing}/logs/suspicious.log"
   parsing_log="${MOUNT_POINT_parsing}/logs/parsing.log"
   if [ -s "${parsing_log}" ]; then
       echo -e "\nPotential reason:"
       cat "${parsing_log}"
   fi
   new_log_location="${SHARED_FOLDER}/suspicious_logs_`date '+%F_%H:%M:%S_%Z'`_${MODEL_TO_SCAN}"
   cp -R "${MOUNT_POINT_parsing}/logs" "${new_log_location}"
   echo -e "\n\nPlease refer to the logs at: ${new_log_location} for more details.\nYou can now press Ctrl+C to exit."
else
   # Display no issues found in green color
   echo -e "\e[32mNo system-affecting issue found.\e[0m\nYou can now press Ctrl+C to exit."
fi

tail -f /dev/null
