#!/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/

tmp_data="/tmp_data_scanning"
cache="${tmp_data}/cache"
log_dir="${tmp_data}/logs"
log_file="${log_dir}/testing.log"
error_log="${log_dir}/error.log"
scan_finished_flag="${log_dir}/scan_finished"
original_err_log="${log_dir}/original_err.log"

# Function to handle cleanup and logging
cleanup() {
    local exit_code=$?
    local error_message="Script exited with code $exit_code"
    
    case $exit_code in
        0) error_message="Script completed successfully" ;;
        1) error_message="General error occurred" ;;
        2) error_message="Misuse of shell builtins" ;;
        124) error_message="Timeout occurred" ;;
        126) error_message="Command invoked cannot execute" ;;
        127) error_message="Command not found" ;;
        137) error_message="Script received SIGKILL" ;;
        139) error_message="Segmentation fault" ;;
        143) error_message="Script received SIGTERM" ;;
        200) error_message="Trying to create file in /root is not allowed" ;;
        201) error_message="Space attack" ;;
        202) error_message="nofile limit exceeded" ;;
        203) error_message="nproc limit exceeded" ;;
        204) error_message="Memory limit exceeded" ;;
        205) error_message="Suspicious action detected" ;;
        254) error_message="Ctrl+C pressed" ;;
        *) error_message="Unknown error occurred with exit code $exit_code" ;;
    esac
    
    if [ $exit_code -eq 0 -o $exit_code -eq 124 ]; then
        return
    fi
    echo "$(date): $error_message" > "${error_log}"
    
    # Kill any remaining background processes
    kill $(jobs -p) 2>/dev/null

    # Additional cleanup actions can be added here
    # For example, removing temporary files:
    rm -f "${scan_finished_flag}"
}

# Function to handle errors
handle_error() {
    local exit_code=$?
    echo "Error occurred with exit code: $exit_code"
    exit $exit_code
}

# Set up the traps
trap cleanup EXIT
trap handle_error ERR

check_point="${MODEL_TO_SCAN}" # This one should be a local call

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

cd "${MYLOC}"

rm -rf "${tmp_data}"/*
mkdir -p "${log_dir}" "${cache}"
# Create a symbolic link to the cache directory if it doesn't exist
# [ ! -e "${HOME}/.cache" ] && ln -s ${cache} "${HOME}/.cache"

# Testing storage space attack
# dd if=/dev/zero of=${tmp_data}/aaa bs=1M count=128

if ! timeout {{TIMEOUT_SECONDS}} strace -f -o "${log_file}" -s 512 -e trace=connect,openat,bind,write,unlink,unlinkat,execve,clone,clone3,fork,vfork,access /opt/conda/envs/myenv/bin/python testcase.py "${check_point}" "${log_file}" 2>"${original_err_log}"; then
  exit_code=$?
  if [ $exit_code -ne 124 -a $exit_code -ne 0 ]; then
    echo "#################### Error detected ########################"
    cat "${original_err_log}"
    echo "############################################################"
    exit $exit_code
  fi
fi

# Remove all null characters from the log file and save the output to a temporary file
tr -d '\000' < "${log_file}" > "${log_file}.clean"
rm -f "${log_file}"; mv "${log_file}.clean" "${log_file}"

# # Remove some known IPs
# echo
# HUGGINGFACE_IPS=`dig huggingface.co | grep ';; ANSWER SECTION:' -A 5 | grep 'huggingface.co' | awk '{print $5}' | tr '\n' '|' | sed s'/|$//g'`
HUGGINGFACE_IPS="127.0.0.1"
# echo "HuggingFace IPs: ${HUGGINGFACE_IPS}"
CURR_DIR="$(pwd)"
echo

# Accessing those resources are normal
WHITELIST="(${HUGGINGFACE_IPS}|/opt/conda/lib/python3.10/site-packages/|/lib/|/proc/|/share/|/share\"|\.so|, \"[^/]|/opt/conda/|\"/sys/|\"/dev|\"/etc/localtime|\"/root/.cache|${CURR_DIR})"

# Try to detect storage space attack
remaining_space="$(df -h "${tmp_data}" | tail -n 1 | awk '{print $4}')"
if [ "${remaining_space}" == "0" ]; then
  echo "#################### Storage space attack ###################"
  echo "No space left on device for ${tmp_data}"
  echo "############################################################"
  exit 201
fi

touch "${scan_finished_flag}"

# cat "${log_file}" | egrep '^[0-9]+[ ]+(connect.*AF_INET|openat|bind\(.*sa_family=AF_INET.*sin_port.*sin_addr=inet_addr)' | egrep -v "sin_port=htons\(53\)" | egrep -v "${WHITELIST}" > ${error_log}
# cnt=`wc -l ${error_log} | awk '{print $1}'`
# if [ ${cnt} -gt 0 ]; then
# 	echo "#################### Suspicious behavior ###################"
# 	cat ${error_log}
# 	echo "############################################################"
#   exit 205
# else
#   echo "################ No suspicious behavior found ##############"
#   touch ${scan_finished_flag}
# fi
