#!/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_scanning="/tmp_data_scanning"
cache="${tmp_data_scanning}/cache"
log_dir="${tmp_data_scanning}/logs"
log_file="${log_dir}/testing.log"
error_log="${log_dir}/error.log"
suc_log="${log_dir}/scan_finished"
original_err_log="${log_dir}/original_err.log"

tmp_data_parsing="/tmp_data_parsing"
parsing_log_dir="${tmp_data_parsing}/logs"
parsing_log_file="${parsing_log_dir}/parsing.log"
parsing_original_testing_log="${parsing_log_dir}/original_testing.log"
parsing_error_log="${parsing_log_dir}/error.log"
parsing_suspicious_log="${parsing_log_dir}/suspicious.log"
tmp_suspicious_log="${parsing_log_dir}/tmp_suspicious.log"
parsing_suc_log="${parsing_log_dir}/parsing_finished"
parsing_original_err_log="${parsing_log_dir}/original_err.log"
mkdir -p "${parsing_log_dir}"

RET=0

# If the suc_log doesn't exist, it means the scan is not finished successfully. We should output the error_log
cnt=0
MAX_TM={{TIMEOUT_SECONDS}}
while [ "${cnt}" -lt "${MAX_TM}" ]; do
  if [ -f "${suc_log}" ]; then
    break
  fi
  sleep 10
  cnt=$((cnt+10))
done

if [ ! -f "${suc_log}" ]; then
  echo -n "The scan is not finished successfully." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  if [ -f "${error_log}" ]; then
    echo " Please check the error log." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    echo "------------------------------------------------------------" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    cat "${error_log}" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    echo "------------------------------------------------------------" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  else
    echo " No error log found however. The scan might be abnormally terminated and logging can't even be written out." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  fi
  RET=201
  ls -lh "${tmp_data_parsing}"
fi

# # Remove some known IPs
ALLOWED_IP_ADDRESS="(127.0.0.1|18.238.192\..*)"
# echo "HuggingFace IPs: ${HUGGINGFACE_IPS}"
SCANN_RUNNING_DIR="/share" # `pwd`

# Accessing those resources are normal
WHITELIST0="open(at)?\([^, ]*, \"(?!(\.\.|/))" # This is to filter "remain" paths starting with "/" or ".."
WHITELIST="(((connect|bind)\(.*${ALLOWED_IP_ADDRESS})|openat\([^,]*, \"((/usr/lib|/usr/local/lib|/lib|/proc|/opt/conda|/sys|/dev|/usr/share).*\".*O_RDONLY|(/etc/ld.so.cache|/opt/conda/envs/myenv/.*/__pycache__/)|/home/sandboxuser/.cache|${SCANN_RUNNING_DIR}|(/dev/null|(/usr|/var)?/tmp/))|unlink(at)?\(\"/tmp/)|(open(at)?|unlink(at)?)\(([^,]*, )?\"/dev/shm/sem" #; echo "WHITELIST: ${WHITELIST}"
OPENSSL_WHITELIST="open(at)?\([^, ]*, \"(/etc/ssl/openssl.cnf|/etc/pki/tls/openssl.d|/etc/crypto-policies/back-ends/opensslcnf.config)\".*"

# CUDA white list
CUDA_OPENAT_WHITELIST="open(at)?\([^, ]*, \"(/dev/.*(nvidia|cuda)|/usr/local/cuda/lib64|${HOME}/.nv/nvidia|/proc/|/etc/nvidia/).*\".*"
CUDA_UNLINK_WHITELIST="unlink(at)?\(\d+,\s*\"(.*\.pyc|.*_remote_module_non_scriptable.py|__pycache__)\""
CUDA_WRITE_WHITELIST="write\(\d+,\s*\"(cuda-EvtHandlr|cuda\w*)\""

# Try to detect storage space attack
remaining_space="$(df -h "${tmp_data_scanning}" | tail -n 1 | awk '{print $4}')"
if [ "${remaining_space}" == "0" ]; then
  echo -e "\n\n########## Potential Storage space attack ###################" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "No space left on device for ${tmp_data_scanning} - the folder for testing logs." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "This is abnormal as the scan log should not be that large." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The scan log should contain the system calls related to such abnormal file writing." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "############################################################\n\n" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  RET=202
fi

# Get the suspicious logs
# cat "${log_file}" | grep -P '^[0-9]+\s+(connect\(.*AF_INET|openat\(|bind\(.*sa_family=AF_INET.*sin_port.*sin_addr=inet_addr|write\(([^12]|\d\d+),|exec(ve)?\(|unlink(at)?\(|(v)?fork\(|access\()' | grep -Pv "${WHITELIST0}" | grep -Pv "${WHITELIST}" > ${tmp_suspicious_log}
cat "${log_file}" | grep -P '^[0-9]+\s+(connect\(.*AF_INET|openat\(|bind\(.*sa_family=AF_INET.*sin_port.*sin_addr=inet_addr|write\(([^12]|\d\d+),|exec(ve)?\(|unlink(at)?\(|(v)?fork\()' \
  | grep -Pv "${WHITELIST0}" | grep -Pv "${WHITELIST}" \
  | grep -Pv "${OPENSSL_WHITELIST}" \
  | grep -Pv "${CUDA_OPENAT_WHITELIST}" | grep -Pv "${CUDA_UNLINK_WHITELIST}" | grep -Pv "${CUDA_WRITE_WHITELIST}" \
  > "${tmp_suspicious_log}"

# Ignore read/write calls without corresponding openat calls
bash /share/cleanup_irrelevant_io.sh "${tmp_suspicious_log}" > "${tmp_suspicious_log}_clean"
cat "${tmp_suspicious_log}_clean" > "${tmp_suspicious_log}"; rm -f "${tmp_suspicious_log}_clean"

# Try to detect OOM - by checking the SIGKILL signal
main_pid="$(cat "${log_file}" | awk '{print $1}' | sort | uniq -c | sort -rn | head -1 | awk '{print $2}')"
sed 's/[+\-]\+\s*//g' -i "${log_file}" # Remove the +/- signs

kill_log_msg="$(cat "${log_file}" | grep -P "^${main_pid}\s+.*killed.*\s+SIGK[A-Z]+")"
if [ -n "${kill_log_msg}" ]; then
  sig="$(echo "${kill_log_msg}" | grep -Po 'SIG[A-Z]+')"
  echo -e "\n\n############## Potential Kill signal #####################" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The main process ${main_pid} is terminated by a signal - ${sig}" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "It seems like it was killed by OOM killer due to memory exhaustion." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "This is a suspicious behavior unless you've set a too low \"memory\" value in the \"config.ini\" file" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The scan log should contain the system calls related to such abnormal signal." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "############################################################\n\n" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "Unexpected ${kill_log_msg} signal detected." >> "${tmp_suspicious_log}"
  RET=203
fi

# Try to detect openning too many files - by checking "(Too many open files)" in the strace log
too_many_files_signal="$(cat "${log_file}" | grep -P "(Too many open files)" | sort -u | wc -l | awk '{print $1}')"
if [ "${too_many_files_signal}" -eq 1 ]; then
  echo -e "\n\n############## Unexpected too many open files ################" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The log contains the signal of \"Too many open files\"." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "This is abnormal as the ML model handling should not result in so many open files." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The scan log should contain the system calls related to such abnormal signal. You can also check your setting of \`nofile_limit\` in the \`config.ini\` file." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "############################################################\n\n" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "Unexpected \"Too many open files\" - (probably caused by exhuasting the \"nofile\" setting) suspicious behavior in ML model handling" >> "${tmp_suspicious_log}"
  RET=204
fi

# Try to detect fork bomb - exceed the `nproc` setting by checking the number of fork/vfork/clone system calls
fork_bomb_signal="$(cat "${log_file}" | grep -P '^[0-9]+\s+(fork\(|vfork\(|clone\d*\().*\(Resource temporarily unavailable\)\s*$' | wc -l | awk '{print $1}')"
if [ "${fork_bomb_signal}" -gt 0 ]; then
  echo -e "\n\n############## Potential Fork bomb #####################" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The log contains the signal of \"Resource temporarily unavailable\"." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "This is abnormal as the ML model handling should not exhaust the OS allowed maximum number of processes, unless the \`nproc_limit\` is set to a too low value." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "The scan log should contain the system calls related to such abnormal signal." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "############################################################\n\n" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
  echo -e "Unexpected \"Resource temporarily unavailable\" (probably caused by exhausting the \`nproc\` setting) - suspicious behavior in ML model handling" >> "${tmp_suspicious_log}"
  RET=205
fi

# Try to detect error stack thrown out by Python inside the original error log
if [ -f "${original_err_log}" ]; then
  error_stack="$(cat "${original_err_log}" | grep -P 'Traceback \(most recent call last\):')"
  if [ -n "${error_stack}" ]; then
    echo -e "\n\n############## Potential Python error stack #####################" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    echo -e "The original error log contains the Python error stack." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    echo -e "Please carefully read it to ensure that the custom test case is valid and the libraries in the requirements.txt is sufficient." >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    echo -e "############################################################\n\n" >> "${parsing_log_file}" 2>>"${parsing_error_log}"
    echo -e "Unexpected Python error stack detected. Please ensure the test case can start successfully." >> "${tmp_suspicious_log}"
    RET=206
  fi
fi

# Check if there are any suspicious logs
cnt="$(wc -l "${tmp_suspicious_log}" | awk '{print $1}')"
if [ "${cnt}" -gt 0 ]; then
  /opt/conda/envs/myenv/bin/python /share/explain_sys_calls.py "${tmp_suspicious_log}" >> "${parsing_suspicious_log}" 2>>"${parsing_error_log}"
  RET=210
fi

cp "${original_err_log}" "${parsing_original_err_log}"
cp "${log_file}" "${parsing_original_testing_log}"

rm -f "${tmp_suspicious_log}"

# tail -f /dev/null

exit "${RET}"
