#!/usr/bin/env bash
# Cross-platform UPX unpacking wrapper with verification.
#
# Usage:
#   bash upx_unpack.sh <input_file> [output_file]
#
# Features:
#   - Detects UPX availability on both Linux and Windows
#   - Verifies the file is UPX-packed before attempting
#   - Creates backup of original file
#   - Validates unpacked binary integrity
#   - Reports entropy change and string count improvement

set -euo pipefail

# --- Configuration ---
MIN_ENTROPY_DROP=0.5  # Minimum entropy decrease to consider unpack successful

# --- Color output (when terminal supports it) ---
if [ -t 1 ] && command -v tput &>/dev/null; then
    RED=$(tput setaf 1 2>/dev/null || echo "")
    GREEN=$(tput setaf 2 2>/dev/null || echo "")
    YELLOW=$(tput setaf 3 2>/dev/null || echo "")
    RESET=$(tput sgr0 2>/dev/null || echo "")
else
    RED=""
    GREEN=""
    YELLOW=""
    RESET=""
fi

info()  { echo "${GREEN}[+]${RESET} $*"; }
warn()  { echo "${YELLOW}[!]${RESET} $*"; }
error() { echo "${RED}[-]${RESET} $*" >&2; }

# --- Usage ---
usage() {
    echo "Usage: $0 <input_file> [output_file]"
    echo ""
    echo "Unpack a UPX-packed executable with verification."
    echo ""
    echo "Arguments:"
    echo "  input_file    Path to the packed executable"
    echo "  output_file   Path for the unpacked output (default: <input>_unpacked.<ext>)"
    echo ""
    echo "Examples:"
    echo "  $0 malware.exe"
    echo "  $0 malware.exe unpacked_malware.exe"
    echo "  $0 packed_elf unpacked_elf"
    exit 1
}

# --- Argument parsing ---
if [ $# -lt 1 ]; then
    usage
fi

INPUT_FILE="$1"

if [ ! -f "$INPUT_FILE" ]; then
    error "File not found: $INPUT_FILE"
    exit 1
fi

# Generate output filename if not specified
if [ $# -ge 2 ]; then
    OUTPUT_FILE="$2"
else
    BASENAME="${INPUT_FILE%.*}"
    EXTENSION="${INPUT_FILE##*.}"
    if [ "$BASENAME" = "$INPUT_FILE" ]; then
        # No extension
        OUTPUT_FILE="${INPUT_FILE}_unpacked"
    else
        OUTPUT_FILE="${BASENAME}_unpacked.${EXTENSION}"
    fi
fi

# --- Detect UPX ---
find_upx() {
    # Check standard locations
    if command -v upx &>/dev/null; then
        echo "upx"
        return 0
    fi

    # Windows-specific paths
    local win_paths=(
        "/c/Tools/upx/upx.exe"
        "/c/Program Files/upx/upx.exe"
        "/c/ProgramData/chocolatey/bin/upx.exe"
        "$HOME/tools/upx/upx.exe"
    )

    for path in "${win_paths[@]}"; do
        if [ -f "$path" ]; then
            echo "$path"
            return 0
        fi
    done

    # Linux-specific paths
    local linux_paths=(
        "/usr/bin/upx"
        "/usr/local/bin/upx"
        "$HOME/.local/bin/upx"
        "$HOME/tools/upx"
    )

    for path in "${linux_paths[@]}"; do
        if [ -f "$path" ] && [ -x "$path" ]; then
            echo "$path"
            return 0
        fi
    done

    return 1
}

UPX_BIN=$(find_upx) || {
    error "UPX not found. Install it:"
    error "  Linux:   sudo apt install upx-ucl  OR  sudo pacman -S upx"
    error "  macOS:   brew install upx"
    error "  Windows: choco install upx  OR  download from https://github.com/upx/upx/releases"
    exit 1
}

info "Using UPX: $UPX_BIN"

# --- Check if file is UPX-packed ---
check_upx_packed() {
    local file="$1"

    # Method 1: Check for UPX magic string
    if grep -qao "UPX!" "$file" 2>/dev/null; then
        return 0
    fi

    # Method 2: Check section names
    if command -v strings &>/dev/null; then
        if strings "$file" 2>/dev/null | grep -q "^UPX[0-2]$"; then
            return 0
        fi
    fi

    # Method 3: Try upx -t (test)
    if "$UPX_BIN" -t "$file" &>/dev/null; then
        return 0
    fi

    return 1
}

info "Checking if file is UPX-packed..."

if ! check_upx_packed "$INPUT_FILE"; then
    warn "File does not appear to be UPX-packed."
    warn "The file may use a modified UPX header or a different packer."
    read -p "Attempt unpacking anyway? (y/N) " -n 1 -r 2>/dev/null || REPLY="y"
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        info "Aborted."
        exit 0
    fi
fi

# --- Calculate entropy (for verification) ---
calculate_entropy() {
    local file="$1"
    python3 -c "
import math, sys
data = open(sys.argv[1], 'rb').read()
if not data:
    print('0.0')
    sys.exit()
counts = [0]*256
for b in data:
    counts[b] += 1
length = len(data)
entropy = -sum((c/length)*math.log2(c/length) for c in counts if c > 0)
print(f'{entropy:.4f}')
" "$file" 2>/dev/null || echo "N/A"
}

count_strings() {
    local file="$1"
    if command -v strings &>/dev/null; then
        strings "$file" 2>/dev/null | wc -l | tr -d ' '
    else
        echo "N/A"
    fi
}

# --- Pre-unpack metrics ---
info "Gathering pre-unpack metrics..."
PRE_ENTROPY=$(calculate_entropy "$INPUT_FILE")
PRE_STRINGS=$(count_strings "$INPUT_FILE")
PRE_SIZE=$(stat -c%s "$INPUT_FILE" 2>/dev/null || stat -f%z "$INPUT_FILE" 2>/dev/null || echo "N/A")

info "Pre-unpack: entropy=$PRE_ENTROPY, strings=$PRE_STRINGS, size=$PRE_SIZE bytes"

# --- Create backup ---
BACKUP_FILE="${INPUT_FILE}.bak"
if [ ! -f "$BACKUP_FILE" ]; then
    cp "$INPUT_FILE" "$BACKUP_FILE"
    info "Backup created: $BACKUP_FILE"
fi

# --- Attempt unpacking ---
info "Unpacking with UPX..."

# Copy input to output location first, then unpack in-place
# This preserves the original file
cp "$INPUT_FILE" "$OUTPUT_FILE"

if "$UPX_BIN" -d "$OUTPUT_FILE" 2>&1; then
    info "UPX unpacking command completed."
else
    UPX_EXIT=$?
    error "UPX unpacking failed (exit code: $UPX_EXIT)"

    # Try with -f (force) flag
    warn "Retrying with force flag..."
    cp "$INPUT_FILE" "$OUTPUT_FILE"
    if "$UPX_BIN" -d -f "$OUTPUT_FILE" 2>&1; then
        info "UPX unpacking succeeded with force flag."
    else
        error "UPX unpacking failed even with force flag."
        rm -f "$OUTPUT_FILE"

        warn "Possible causes:"
        warn "  - Modified UPX header (try fixing UPX! magic bytes)"
        warn "  - Custom UPX build with modified decompressor"
        warn "  - File is not actually UPX-packed"
        warn ""
        warn "Manual fixes to try:"
        warn "  1. Restore UPX! magic: find offset and patch bytes"
        warn "  2. Use older UPX version matching the packer version"
        warn "  3. Manual unpacking via debugger (set BP on tail jump)"
        exit 1
    fi
fi

# --- Post-unpack verification ---
if [ ! -f "$OUTPUT_FILE" ]; then
    error "Output file was not created."
    exit 1
fi

info "Verifying unpacked binary..."

POST_ENTROPY=$(calculate_entropy "$OUTPUT_FILE")
POST_STRINGS=$(count_strings "$OUTPUT_FILE")
POST_SIZE=$(stat -c%s "$OUTPUT_FILE" 2>/dev/null || stat -f%z "$OUTPUT_FILE" 2>/dev/null || echo "N/A")

info "Post-unpack: entropy=$POST_ENTROPY, strings=$POST_STRINGS, size=$POST_SIZE bytes"

# --- Report ---
echo ""
echo "========================================="
echo "  UPX Unpacking Results"
echo "========================================="
echo "  Input:    $INPUT_FILE"
echo "  Output:   $OUTPUT_FILE"
echo "  Backup:   $BACKUP_FILE"
echo "-----------------------------------------"
echo "  Metric        Before      After"
echo "  Entropy       $PRE_ENTROPY        $POST_ENTROPY"
echo "  Strings       $PRE_STRINGS        $POST_STRINGS"
echo "  Size (bytes)  $PRE_SIZE      $POST_SIZE"
echo "========================================="

# Verify entropy decreased
if [ "$PRE_ENTROPY" != "N/A" ] && [ "$POST_ENTROPY" != "N/A" ]; then
    python3 -c "
pre = float('$PRE_ENTROPY')
post = float('$POST_ENTROPY')
drop = pre - post
if drop >= $MIN_ENTROPY_DROP:
    print('[+] Entropy decreased by {:.4f} - unpacking likely successful'.format(drop))
elif drop > 0:
    print('[!] Entropy decreased by only {:.4f} - verify manually'.format(drop))
else:
    print('[-] Entropy did not decrease - unpacking may have failed')
" 2>/dev/null || true
fi

# Check if the unpacked file is a valid executable
if command -v file &>/dev/null; then
    FILE_TYPE=$(file "$OUTPUT_FILE" 2>/dev/null || echo "unknown")
    info "File type: $FILE_TYPE"
fi

info "Done. Unpacked file: $OUTPUT_FILE"