#!/usr/bin/env bash
# fetch.sh — Compile CTF challenge binaries for fold binary validation
#
# Targets:
#   1. picoCTF 2022 — "buffer overflow 0" (CWE-121, stack-based overflow)
#   2. picoCTF 2023 — "VNE" (ARM64 binary, reverse-engineering challenge)
#   3. vuln-demo       — custom CWE-121/CWE-134/CWE-416
#
# Sources live in src/ (hand-authored or reconstructed from public writeups).
# This script compiles them for ARM64 Mach-O, x86_64 Mach-O, and x86_64 ELF.
#
# Usage:
#   ./fetch.sh              # compile all
#   ./fetch.sh --clean      # remove all compiled artifacts
#   ./fetch.sh --list       # list what is available

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
BIN_DIR="$SCRIPT_DIR/bin"
SRC_DIR="$SCRIPT_DIR/src"
mkdir -p "$BIN_DIR"

OS_NATIVE="$(uname -s)"

# CFLAGS separated by platform because macOS ld rejects ELF-only flags.
# macOS: Mach-O linker doesn't support -z, -no-pie, or -static the ELF way.
CFLAGS_MACOS="-O0 -g -fno-stack-protector -Wno-implicit-function-declaration"
CFLAGS_ELF="-O0 -g -fno-stack-protector -z execstack -no-pie -static"

# Colour helpers (safe for non-TTY)
BOLD=""; RED=""; GREEN=""; YELLOW=""; NC=""
[ -t 2 ] && BOLD="\033[1m" RED="\033[31m" GREEN="\033[32m" YELLOW="\033[33m" NC="\033[0m"

say()  { printf "${GREEN}==>${NC} %s\n" "$*"; }
warn() { printf "${YELLOW}WARN:${NC} %s\n" "$*" >&2; }
err()  { printf "${RED}ERROR:${NC} %s\n" "$*" >&2; exit 1; }

# ——————————————————————————————————————————————
# Compile a single source for macOS (ARM64 + x86_64 Mach-O)
# ——————————————————————————————————————————————
compile_macos() {
    local src="$1" out_base="$2"
    say "  -> compiling ARM64 Mach-O"
    clang $CFLAGS_MACOS -arch arm64 "$src" -o "${out_base}-arm64-macho" || warn "ARM64 compile failed"
    chmod +x "${out_base}-arm64-macho" 2>/dev/null || true
    rm -rf "${out_base}-arm64-macho.dSYM" 2>/dev/null || true

    say "  -> compiling x86_64 Mach-O"
    clang $CFLAGS_MACOS -arch x86_64 "$src" -o "${out_base}-x86_64-macho" || warn "x86_64 Mach-O compile failed"
    chmod +x "${out_base}-x86_64-macho" 2>/dev/null || true
    rm -rf "${out_base}-x86_64-macho.dSYM" 2>/dev/null || true
}

# ——————————————————————————————————————————————
# Compile a single source for x86_64 ELF via Docker
# ——————————————————————————————————————————————
compile_elf_docker() {
    local src="$1" out="$2"
    if ! command -v docker &>/dev/null || ! docker info &>/dev/null 2>&1; then
        warn "Docker not available — skipping x86_64 ELF build"
        return 0
    fi
    say "  -> compiling x86_64 ELF via Docker"
    local src_name
    src_name="$(basename "$src")"
    docker run --rm --platform linux/amd64 \
        -v "$src:/src/${src_name}:ro" \
        -v "$BIN_DIR:/out" \
        alpine:latest \
        sh -c "apk add --no-cache gcc musl-dev >/dev/null 2>&1 && gcc $CFLAGS_ELF /src/${src_name} -o /out/$(basename "$out")" \
        || warn "x86_64 ELF cross-compile failed"
    chmod +x "$out" 2>/dev/null || true
}

# ——————————————————————————————————————————————
# Compile one target across all available platforms
# ——————————————————————————————————————————————
compile_target() {
    local name="$1"
    local src="$SRC_DIR/${name}.c"
    local out_base="$BIN_DIR/$name"

    if [ ! -f "$src" ]; then
        warn "source not found: $src — skipping $name"
        return 0
    fi

    say "Compiling $name"

    if [ "$OS_NATIVE" = "Darwin" ]; then
        compile_macos "$src" "$out_base"
    fi

    compile_elf_docker "$src" "${out_base}-x86_64-elf"
}

# ——————————————————————————————————————————————
# Main
# ——————————————————————————————————————————————

case "${1:-}" in
    --clean|clean)
        say "Removing all compiled binaries"
        rm -rf "$BIN_DIR"
        mkdir -p "$BIN_DIR"
        ;;
    --list|list)
        say "Available targets:"
        echo "  pico2022-buffer-overflow-0    — picoCTF 2022 CWE-121"
        echo "  pico2023-vne-arm64            — picoCTF 2023 ARM64 RE"
        echo "  vuln-demo                     — custom CWE-121/CWE-134/CWE-416"
        echo ""
        echo "Binaries in $BIN_DIR:"
        ls -lh "$BIN_DIR" 2>/dev/null || echo "  (none yet — run ./fetch.sh)"
        ;;
    *)
        echo "=== Fold CTF Challenge Binary Compiler ==="
        echo "Platform: $(uname -m) / $(uname -s)"
        echo ""
        compile_target "pico2022-buffer-overflow-0"
        compile_target "pico2023-vne-arm64"
        compile_target "vuln-demo"
        echo ""
        echo "=== Done ==="
        echo "Binaries in $BIN_DIR:"
        ls -lh "$BIN_DIR" 2>/dev/null || echo "  (none)"
        ;;
esac