#!/usr/bin/env bash
# stage-pe32-modules.sh — pull every PE32/PE32+ module out of a UEFIExtract
# dump tree and stage it into a flat folder with a real, readable name,
# ready for Ghidra's batch importer (analyzeHeadless -import <dir>).
#
# UEFIExtract's own layout buries each module several GUID-named directories
# deep (e.g. ".../8C8CE578-8A3D-4F1C-9935-896185C32DD3/48 PchS3Peim/1 PE32
# image section/body.bin"), which is unusable for a batch import. This script
# renames each one after its actual module folder (e.g. "PchS3Peim.efi"),
# de-duplicating collisions with a numeric suffix.
#
# v2: fixes a real bug in v1. Most modules sit under an extra "Compressed
# section" (or "GUIDed section") wrapper directory between the real module
# name and the "PE32 image section" folder — v1 only looked one level up,
# so it named these "Compressed_section_NNN" instead of the real module
# name. This version walks UP past any known structural/wrapper directory
# name (and bare GUIDs) until it finds the real one.
#
# Usage:
#   1. Extract a BIOS/UEFI image with UEFIExtract first:
#        UEFIExtract firmware.bin
#      This produces "firmware.bin.dump/" alongside the original file.
#   2. Run this script against that dump directory:
#        ./stage-pe32-modules.sh firmware.bin.dump ./staged_modules
#   3. Batch-import everything into a Ghidra project:
#        analyzeHeadless <project_dir> <project_name> \
#          -import ./staged_modules -log import.log

set -euo pipefail

DUMP_DIR="${1:?Usage: $0 <uefiextract_dump_dir> <output_dir>}"
OUT_DIR="${2:?Usage: $0 <uefiextract_dump_dir> <output_dir>}"

mkdir -p "$OUT_DIR"

# Directory names that are UEFIExtract structure, never a real module name.
# Matched case-insensitively after stripping the leading "NNN " index.
is_generic() {
  local name_lc
  name_lc=$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')
  case "$name_lc" in
    "pe32 image section"|"pe32+ image section"|"compressed section"| \
    "guided section"|"raw section"|"freeform subtype guid section"| \
    "dxe dependency section"|"pei dependency section"|"ui section"| \
    "version section"|"firmware volume image section"|"free space"| \
    "bios region"|"me region"|"descriptor region")
      return 0 ;;
  esac
  # bare GUID (8-4-4-4-12 hex), also never a real module name on its own
  if [[ "$1" =~ ^[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}$ ]]; then
    return 0
  fi
  return 1
}

find "$DUMP_DIR" -type d -iname "*PE32 image section*" 2>/dev/null | while read -r d; do
  src="$d/body.bin"
  [ -f "$src" ] || continue

  # Walk up from the PE32 section's parent, skipping generic wrapper dirs,
  # until we find the real module name (or run out of path).
  cur="$(dirname "$d")"
  modname=""
  while [ "$cur" != "$DUMP_DIR" ] && [ "$cur" != "." ] && [ "$cur" != "/" ]; do
    base=$(basename "$cur" | sed -E 's/^[0-9]+ //')
    if ! is_generic "$base"; then
      modname="$base"
      break
    fi
    cur="$(dirname "$cur")"
  done
  [ -n "$modname" ] || modname="unnamed"

  # sanitize for use as a filename
  modname=$(printf '%s' "$modname" | tr -c 'A-Za-z0-9._-' '_')

  target="$OUT_DIR/$modname.efi"
  i=2
  while [ -e "$target" ]; do
    target="$OUT_DIR/${modname}_${i}.efi"
    i=$((i + 1))
  done
  cp "$src" "$target"
done

echo "Staged $(ls "$OUT_DIR" | wc -l) modules into $OUT_DIR"