#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# Full GSP firmware analysis pipeline
# Usage: ./analyze_gsp_firmware.sh <firmware_bin> [output_dir]

set -euo pipefail

FIRMWARE_BIN="${1:-~/Git/nvidia-linux-firmware/nvidia/tu102/gsp/gsp-570.144.bin}"
OUTPUT_DIR="${2:-/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP}"
PROJECT_DIR="/home/rmcore/Git/integdev_gpu_drv/open-iov/tools/ghidra/gsp-tu104-570.144"
PROJECT_NAME="gsp-tu104-570.144"

FIRMWARE_BIN=$(eval echo $FIRMWARE_BIN)

echo "=== GSP Firmware Analysis Pipeline ==="
echo "Firmware: $FIRMWARE_BIN"
echo "Output: $OUTPUT_DIR"
echo "Project: $PROJECT_DIR/$PROJECT_NAME"
echo ""

# Step 1: Extract .fwimage section
echo "Step 1: Extracting .fwimage section..."
FWIMAGE_BIN="$OUTPUT_DIR/gsp-570.144.fwimage.bin"
objcopy -O binary -j .fwimage "$FIRMWARE_BIN" "$FWIMAGE_BIN"
echo "  Saved to $FWIMAGE_BIN ($(stat -c%s "$FWIMAGE_BIN") bytes)"

# Step 2: Extract version
echo "Step 2: Extracting version..."
VERSION=$(readelf -x .fwversion "$FIRMWARE_BIN" | grep -A1 "Hex dump" | tail -1 | awk '{print $2$3$4$5}' | xxd -r -p)
echo "  Version: $VERSION"
echo "$VERSION" > "$OUTPUT_DIR/gsp_version.txt"

# Step 3: Extract signatures
echo "Step 3: Extracting signatures..."
for i in 0 1 2; do
    objcopy -O binary -j ".fwsignature$i" "$FIRMWARE_BIN" "$OUTPUT_DIR/gsp-570.144.fwsignature$i.bin" 2>/dev/null || true
done

# Step 4: Create Ghidra project for ELF container
echo "Step 4: Creating Ghidra project for ELF container..."
/snap/ghidra/47/ghidra/support/analyzeHeadless "$PROJECT_DIR" "$PROJECT_NAME" \
    -import "$FIRMWARE_BIN" \
    -processor "RISCV:LE:64:default" \
    -cspec "gcc" \
    -overwrite \
    -scriptPath "$PROJECT_DIR/scripts" \
    -postScript "extract_fwimage.py"

# Step 5: Create Ghidra project for .fwimage binary
echo "Step 5: Creating Ghidra project for .fwimage binary..."
FWIMAGE_PROJECT_DIR="/home/rmcore/Git/integdev_gpu_drv/open-iov/tools/ghidra/gsp-fwimage-570.144"
FWIMAGE_PROJECT_NAME="gsp-fwimage-570.144"
mkdir -p "$FWIMAGE_PROJECT_DIR/scripts"
cp "$PROJECT_DIR/scripts/analyze_gsp_fwimage.py" "$FWIMAGE_PROJECT_DIR/scripts/"

/snap/ghidra/47/ghidra/support/analyzeHeadless "$FWIMAGE_PROJECT_DIR" "$FWIMAGE_PROJECT_NAME" \
    -import "$FWIMAGE_BIN" \
    -processor "RISCV:LE:64:default" \
    -cspec "gcc" \
    -overwrite \
    -scriptPath "$FWIMAGE_PROJECT_DIR/scripts" \
    -postScript "analyze_gsp_fwimage.py"

# Step 6: Also analyze bootloader and booter
echo "Step 6: Analyzing bootloader and booter..."
BOOTLOADER_BIN="$OUTPUT_DIR/bootloader-570.144.bin"
BOOTER_LOAD_BIN="$OUTPUT_DIR/booter_load-570.144.bin"
BOOTER_UNLOAD_BIN="$OUTPUT_DIR/booter_unload-570.144.bin"

cp ~/Git/nvidia-linux-firmware/nvidia/tu102/gsp/bootloader-570.144.bin "$BOOTLOADER_BIN"
cp ~/Git/nvidia-linux-firmware/nvidia/tu102/gsp/booter_load-570.144.bin "$BOOTER_LOAD_BIN"
cp ~/Git/nvidia-linux-firmware/nvidia/tu102/gsp/booter_unload-570.144.bin "$BOOTER_UNLOAD_BIN"

for name in bootloader booter_load booter_unload; do
    bin_file="$OUTPUT_DIR/${name}-570.144.bin"
    proj_dir="/home/rmcore/Git/integdev_gpu_drv/open-iov/tools/ghidra/gsp-${name}-570.144"
    proj_name="gsp-${name}-570.144"
    mkdir -p "$proj_dir/scripts"
    cp "$PROJECT_DIR/scripts/analyze_gsp_fwimage.py" "$proj_dir/scripts/analyze_${name}.py"
    
    # Modify script for raw binary
    sed -i 's/program = getCurrentProgram()/program = getCurrentProgram()\n# Raw binary - set base address to 0/' "$proj_dir/scripts/analyze_${name}.py"
    
    /snap/ghidra/47/ghidra/support/analyzeHeadless "$proj_dir" "$proj_name" \
        -import "$bin_file" \
        -processor "RISCV:LE:64:default" \
        -cspec "gcc" \
        -overwrite \
        -scriptPath "$proj_dir/scripts" \
        -postScript "analyze_${name}.py" 2>&1 | tail -20
done

echo ""
echo "=== Analysis Complete ==="
echo "Results in: $OUTPUT_DIR"
echo ""
echo "Key files:"
ls -la "$OUTPUT_DIR"/*.json "$OUTPUT_DIR"/*.txt "$OUTPUT_DIR"/*.c 2>/dev/null | head -20