#!/bin/bash

# A tool to install and uninstall an Intel Pin wrapper for any given binary.
#
# MODES:
#   -i, --install /path/to/binary   : Moves the original binary to a safe location
#                                     and replaces it with a wrapper script.
#
#   -u, --uninstall /path/to/binary : Restores the original binary and cleans up.

# --- Configuration ---
PIN_ROOT=${PIN_ROOT:-} # Let the script check for its existence.

# --- Directories ---
# Base directory to securely store original binaries
SAFE_BIN_DIR="/var/coverage/bin"
# Directory for Pin output logs
LOG_DIR="/var/coverage/data"
# Directory to search for the Pin Tool
PIN_TOOL_SEARCH_DIR="/var/coverage"

# A unique comment to identify our wrapper scripts
WRAPPER_ID_COMMENT="# Pin Wrapper generated by wrap.sh"

# --- Script Logic ---
# Exit immediately if a command exits with a non-zero status.
set -e

# --- Initial Sanity Checks ---
# Check if PIN_ROOT is set and valid before doing anything else.
if [ -z "$PIN_ROOT" ] || [ ! -d "$PIN_ROOT" ]; then
    echo "Error: PIN_ROOT environment variable is not set or does not point to a valid directory." >&2
    echo "Please set PIN_ROOT to the root of your Intel Pin installation." >&2
    exit 1
fi

# Automatically find the Pin Tool library.
echo "--> Searching for Pin Tool 'FuncTracer.so'..."
PIN_TOOL=$(find "$PIN_TOOL_SEARCH_DIR" -name "FuncTracer.so" -type f -print -quit)

if [ -z "$PIN_TOOL" ]; then
    echo "Error: Could not automatically find 'FuncTracer.so' within '$PIN_TOOL_SEARCH_DIR'." >&2
    echo "Please ensure the Pin tool is compiled and present in a subdirectory." >&2
    exit 1
fi
echo "    Pin Tool found at: $PIN_TOOL"

# --- HELP FUNCTION ---
show_help() {
    echo "Usage: $0 [MODE] /path/to/binary"
    echo ""
    echo "Modes:"
    echo "  -i, --install   Install the instrumentation wrapper for a binary."
    echo "  -u, --uninstall Restore the original binary and remove the wrapper."
    echo ""
    echo "Example:"
    echo "  sudo ./wrap.sh -i /usr/bin/gzip"
    echo "  sudo ./wrap.sh -u /usr/bin/gzip"
}

# --- INSTALL FUNCTION ---
handle_install() {
    local target_binary="$1"

    if [ -z "$target_binary" ]; then
        echo "Error: --install requires a target binary path." >&2
        show_help
        exit 1
    fi

    if [ ! -f "$target_binary" ] || [ -L "$target_binary" ]; then
        echo "Error: Target '$target_binary' is not a regular file." >&2
        exit 1
    fi
    
    # Check if it's already a wrapper
    if grep -q "$WRAPPER_ID_COMMENT" "$target_binary"; then
        echo "Error: '$target_binary' is already a wrapper. Use --uninstall first." >&2
        exit 1
    fi
    
    # Create the base directory for our backups
    mkdir -p "$SAFE_BIN_DIR"
    
    # Create a unique, secure subdirectory for the original binary
    local random_dir
    random_dir=$(mktemp -d -p "$SAFE_BIN_DIR" "XXXXXX")
    local binary_name=$(basename "$target_binary")
    local moved_binary_path="$random_dir/$binary_name"

    echo "--> Moving original binary to '$moved_binary_path'..."
    mv "$target_binary" "$moved_binary_path"
    
    echo "--> Creating new wrapper script at '$target_binary'..."
    
    # Create the wrapper script using a heredoc.
    # The PIN_TOOL variable now contains the auto-detected path.
    cat <<WRAPPER_EOF > "$target_binary"
#!/bin/bash
${WRAPPER_ID_COMMENT} on $(date)
# Original Binary: ${moved_binary_path}

# Configuration for the wrapper
PIN_ROOT=\${PIN_ROOT:-${PIN_ROOT}}
PIN_TOOL="${PIN_TOOL}"
LOG_DIR="${LOG_DIR}"
ORIGINAL_BINARY="${moved_binary_path}"

# Ensure log directory exists
mkdir -p "\$LOG_DIR"

# Generate a unique log filename
binary_name=\$(basename "\$0")
timestamp=\$(date "+%Y%m%d-%H%M%S")
nano_seconds=\$(date "+%N")
log_file="\$LOG_DIR/\${binary_name}_\${timestamp}_\${nano_seconds}.log"

# Execute the original binary with Pin
exec "\$PIN_ROOT/pin" -t "\$PIN_TOOL" -logfile "\$log_file" -- "\$ORIGINAL_BINARY" "\$@"
WRAPPER_EOF

    chmod +x "$target_binary"
    
    echo "Installation complete for '$target_binary'."
}

# --- UNINSTALL FUNCTION ---
handle_uninstall() {
    local target_binary="$1"
    
    if [ -z "$target_binary" ]; then
        echo "Error: --uninstall requires a target binary path." >&2
        show_help
        exit 1
    fi

    if [ ! -f "$target_binary" ] || ! grep -q "$WRAPPER_ID_COMMENT" "$target_binary"; then
        echo "Error: '$target_binary' is not a valid wrapper script. Nothing to uninstall." >&2
        exit 1
    fi

    # Extract the original binary path from the wrapper script
    local original_binary_path
    original_binary_path=$(grep '# Original Binary:' "$target_binary" | head -n1 | cut -d' ' -f4)

    if [ -z "$original_binary_path" ] || [ ! -f "$original_binary_path" ]; then
        echo "Error: Could not find original binary path in wrapper or path is invalid." >&2
        exit 1
    fi
    
    local safe_dir
    safe_dir=$(dirname "$original_binary_path")

    echo "--> Restoring original binary for '$target_binary'..."
    mv "$original_binary_path" "$target_binary"
    echo "    Original binary restored."

    echo "--> Cleaning up secure storage directory '$safe_dir'..."
    rmdir "$safe_dir"
    echo "    Done."
    echo "Uninstall complete for '$target_binary'."
}

# --- SCRIPT ENTRY POINT ---
case "$1" in
    -i|--install)
        handle_install "$2"
        ;;
    -u|--uninstall)
        handle_uninstall "$2"
        ;;
    -h|--help|*)
        show_help
        exit 1
        ;;
esac

exit 0
