#!/usr/bin/env bash
# =============================================================================
# setup.sh — Set up the OP-TEE 3.18.0 (CVE-2022-46152 vulnerable version)
#            environment based on QEMU v8
#
# Reference:
#   https://optee.readthedocs.io/en/latest/building/devices/qemu.html#qemu-v8
#
# Steps:
#   1. repo sync (if not already done)
#   2. Configure toolchain (reuse existing or download fresh)
#   3. Build all components (TF-A / OP-TEE OS / Linux / Buildroot)
#   4. Build PoC exploit and inject into rootfs
# =============================================================================
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OPTEE_DIR="${SCRIPT_DIR}/optee"
POC_DIR="${SCRIPT_DIR}/poc"

# Colored output helpers
info()  { echo -e "\033[1;34m[INFO]\033[0m  $*"; }
ok()    { echo -e "\033[1;32m[ OK ]\033[0m  $*"; }
warn()  { echo -e "\033[1;33m[WARN]\033[0m  $*"; }
die()   { echo -e "\033[1;31m[FAIL]\033[0m  $*" >&2; exit 1; }

# --------------------------------------------------------------------------- #
# 0. Prerequisite checks
# --------------------------------------------------------------------------- #
info "Checking prerequisites..."

REQUIRED_CMDS="repo git make python3 aarch64-linux-gnu-gcc"
for cmd in $REQUIRED_CMDS; do
    command -v "$cmd" >/dev/null 2>&1 || die "Missing dependency: $cmd"
done

[[ -d "${OPTEE_DIR}/.repo" ]] || die "${OPTEE_DIR}/.repo not found. Please run first:
  mkdir -p optee && cd optee
  repo init -u https://github.com/OP-TEE/manifest.git -m qemu_v8.xml -b 3.18.0"

ok "Prerequisites satisfied"

# --------------------------------------------------------------------------- #
# 1. repo sync (idempotent: skip if already synced)
# --------------------------------------------------------------------------- #
info "Running repo sync (fast no-op if already synced)..."
cd "${OPTEE_DIR}"
repo sync -j"$(nproc)" --no-clone-bundle 2>&1 || warn "repo sync had warnings, continuing..."

# --------------------------------------------------------------------------- #
# 2. Toolchain
# --------------------------------------------------------------------------- #
info "Configuring toolchain..."
cd "${OPTEE_DIR}/build"

# Skip download if aarch64/aarch32 toolchains already exist under optee/toolchains/
if [[ -d "${OPTEE_DIR}/toolchains/aarch64" && -d "${OPTEE_DIR}/toolchains/aarch32" ]]; then
    ok "Toolchain already present: ${OPTEE_DIR}/toolchains/"
else
    info "Downloading toolchain (first run may take ~10 minutes)..."
    make toolchains
    ok "Toolchain download complete"
fi

# --------------------------------------------------------------------------- #
# 3. Build all components
# --------------------------------------------------------------------------- #
info "Building all OP-TEE QEMU v8 components (first run may take 30-60 minutes)..."
info "  - Trusted Firmware-A v2.6"
info "  - OP-TEE OS 3.18.0  <- vulnerable version (CVE-2022-46152)"
info "  - Linux kernel (linaro-swg/optee-3.18.0)"
info "  - Buildroot 2021.11 rootfs"

cd "${OPTEE_DIR}/build"
make -j"$(nproc)" all 2>&1 | tee "${SCRIPT_DIR}/build.log"
ok "Build complete!"

# --------------------------------------------------------------------------- #
# 4. Build PoC exploit and inject into rootfs
# --------------------------------------------------------------------------- #
info "Building CVE-2022-46152 PoC..."

# Prepend toolchain bin directory to PATH
TC_AARCH64="${OPTEE_DIR}/toolchains/aarch64/bin"
export PATH="${TC_AARCH64}:${PATH}"

cd "${POC_DIR}"
make CROSS_COMPILE=aarch64-linux-gnu-
ok "PoC built: ${POC_DIR}/cve_2022_46152"

# Add PoC to rootfs overlay
OVERLAY_DIR="${OPTEE_DIR}/build/br-ext/board/qemu/overlay"
mkdir -p "${OVERLAY_DIR}/root"
cp "${POC_DIR}/cve_2022_46152" "${OVERLAY_DIR}/root/"

# Repack rootfs (rebuild only Buildroot, no full rebuild needed)
info "Injecting PoC into rootfs..."
cd "${OPTEE_DIR}/build"
make buildroot 2>&1 | tail -10
ok "rootfs repacked; PoC placed at /root/cve_2022_46152"

echo ""
ok "=========================================="
ok " Environment setup complete!"
ok "=========================================="
echo ""
echo "  Run QEMU:          ./run_qemu.sh"
echo "  Run inside QEMU:   /root/cve_2022_46152"
echo ""
echo "  Vulnerable version: OP-TEE OS 3.18.0 (CVE-2022-46152 unpatched)"
echo "  Fixed version:      OP-TEE OS 3.19.0 (commit 728616b)"