#!/usr/bin/env bash
# memory_acquisition.sh - Cross-platform memory acquisition script
#
# Linux: Uses LiME, AVML, or /proc/kcore
# Windows (WSL/Cygwin guidance): Uses winpmem or DumpIt
#
# Usage:
#   ./memory_acquisition.sh [output_path] [method]
#
# Examples:
#   ./memory_acquisition.sh /tmp/memory.raw          # Auto-detect method
#   ./memory_acquisition.sh /tmp/memory.raw lime      # Force LiME
#   ./memory_acquisition.sh /tmp/memory.raw avml      # Force AVML
#   ./memory_acquisition.sh                           # Interactive mode

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

log_info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*"; }

# Default output path
OUTPUT_PATH="${1:-memory_$(date +%Y%m%d_%H%M%S).raw}"
METHOD="${2:-auto}"

detect_os() {
    case "$(uname -s 2>/dev/null || echo Unknown)" in
        Linux*)     echo "linux" ;;
        Darwin*)    echo "macos" ;;
        CYGWIN*|MINGW*|MSYS*) echo "windows" ;;
        *)          echo "unknown" ;;
    esac
}

check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        log_error "Memory acquisition requires root/sudo privileges."
        log_info "Re-run with: sudo $0 $*"
        exit 1
    fi
}

verify_dump() {
    local dump_path="$1"
    if [ ! -f "$dump_path" ]; then
        log_error "Memory dump file not created: $dump_path"
        return 1
    fi

    local size
    size=$(stat -f%z "$dump_path" 2>/dev/null || stat -c%s "$dump_path" 2>/dev/null || echo 0)

    if [ "$size" -lt 1048576 ]; then
        log_error "Memory dump is suspiciously small: $size bytes"
        return 1
    fi

    log_info "Memory dump created: $dump_path"
    log_info "Size: $(echo "scale=2; $size / 1073741824" | bc 2>/dev/null || echo "$size bytes") GB"

    # Calculate hashes for integrity verification
    log_info "Calculating verification hashes..."
    if command -v sha256sum &>/dev/null; then
        local hash
        hash=$(sha256sum "$dump_path" | awk '{print $1}')
        log_info "SHA-256: $hash"
        echo "$hash  $dump_path" > "${dump_path}.sha256"
        log_info "Hash saved to: ${dump_path}.sha256"
    elif command -v shasum &>/dev/null; then
        local hash
        hash=$(shasum -a 256 "$dump_path" | awk '{print $1}')
        log_info "SHA-256: $hash"
        echo "$hash  $dump_path" > "${dump_path}.sha256"
    fi

    return 0
}

# --- Linux acquisition methods ---

acquire_lime() {
    local output="$1"
    log_info "Attempting acquisition with LiME..."

    if ! lsmod 2>/dev/null | grep -q lime; then
        if [ -f /opt/lime/lime.ko ]; then
            log_info "Loading LiME kernel module..."
            insmod /opt/lime/lime.ko "path=$output format=lime"
        elif command -v apt-get &>/dev/null; then
            log_warn "LiME module not found. Install with:"
            log_warn "  sudo apt-get install linux-headers-\$(uname -r)"
            log_warn "  git clone https://github.com/504ensicslabs/LiME"
            log_warn "  cd LiME/src && make"
            return 1
        else
            log_error "LiME kernel module not found."
            return 1
        fi
    else
        log_info "LiME already loaded, using existing module."
        # Trigger dump via /proc interface
        echo "$output" > /proc/lime 2>/dev/null || {
            log_error "Failed to trigger LiME dump."
            return 1
        }
    fi

    # Wait for dump to complete
    sleep 2
    verify_dump "$output"
}

acquire_avml() {
    local output="$1"
    log_info "Attempting acquisition with AVML..."

    if ! command -v avml &>/dev/null; then
        log_warn "AVML not found in PATH."
        log_info "Download from: https://github.com/microsoft/avml/releases"
        log_info "  wget https://github.com/microsoft/avml/releases/latest/download/avml"
        log_info "  chmod +x avml && sudo mv avml /usr/local/bin/"
        return 1
    fi

    avml "$output"
    verify_dump "$output"
}

acquire_proc_kcore() {
    local output="$1"
    log_info "Attempting acquisition from /proc/kcore..."
    log_warn "WARNING: /proc/kcore acquisition may not capture all memory regions."
    log_warn "Prefer LiME or AVML for forensically sound acquisition."

    if [ ! -r /proc/kcore ]; then
        log_error "/proc/kcore not readable."
        return 1
    fi

    # Get physical memory size
    local mem_size
    mem_size=$(awk '/MemTotal/ {print $2 * 1024}' /proc/meminfo)
    log_info "Physical memory size: $((mem_size / 1073741824)) GB"

    dd if=/proc/kcore of="$output" bs=1M count="$((mem_size / 1048576 + 1))" 2>/dev/null || {
        log_warn "dd of /proc/kcore may have partial results."
    }

    verify_dump "$output"
}

acquire_proc_mem() {
    local output="$1"
    log_info "Attempting acquisition from /proc/mem via fmem..."

    if [ ! -c /dev/fmem ]; then
        log_error "/dev/fmem not available. Install fmem module first."
        return 1
    fi

    local mem_size
    mem_size=$(awk '/MemTotal/ {print $2 * 1024}' /proc/meminfo)

    dd if=/dev/fmem of="$output" bs=1M count="$((mem_size / 1048576 + 1))" 2>/dev/null
    verify_dump "$output"
}

acquire_linux() {
    local output="$1"
    local method="$2"

    check_root

    case "$method" in
        lime)
            acquire_lime "$output"
            ;;
        avml)
            acquire_avml "$output"
            ;;
        kcore)
            acquire_proc_kcore "$output"
            ;;
        fmem)
            acquire_proc_mem "$output"
            ;;
        auto)
            log_info "Auto-detecting best acquisition method for Linux..."
            if command -v avml &>/dev/null; then
                acquire_avml "$output"
            elif lsmod 2>/dev/null | grep -q lime || [ -f /opt/lime/lime.ko ]; then
                acquire_lime "$output"
            elif [ -c /dev/fmem ]; then
                acquire_proc_mem "$output"
            elif [ -r /proc/kcore ]; then
                acquire_proc_kcore "$output"
            else
                log_error "No acquisition method available."
                log_info "Install one of: AVML, LiME, or fmem"
                log_info "  AVML:  https://github.com/microsoft/avml/releases"
                log_info "  LiME:  https://github.com/504ensicslabs/LiME"
                exit 1
            fi
            ;;
        *)
            log_error "Unknown Linux method: $method"
            log_info "Available: auto, lime, avml, kcore, fmem"
            exit 1
            ;;
    esac
}

# --- macOS acquisition ---

acquire_macos() {
    local output="$1"
    check_root

    log_warn "macOS memory acquisition is limited due to SIP."
    log_info "Options:"
    log_info "  1. Use osxpmem (requires SIP disabled or kext approval)"
    log_info "     sudo osxpmem -o $output"
    log_info "  2. Create a VM snapshot and extract memory"
    log_info "  3. Use macOS built-in crash dumps:"
    log_info "     sudo nvram boot-args=\"debug=0x144\""
    log_info "     Then trigger NMI with: sudo dtrace -w -n 'BEGIN{panic()}'"

    if command -v osxpmem &>/dev/null; then
        log_info "Found osxpmem, attempting acquisition..."
        osxpmem -o "$output"
        verify_dump "$output"
    else
        log_error "osxpmem not found. Download from: https://github.com/Velocidex/c-aff4"
        exit 1
    fi
}

# --- Windows guidance ---

acquire_windows_guidance() {
    local output="$1"

    log_info "=== Windows Memory Acquisition Guide ==="
    echo ""
    log_info "Run one of these tools on the Windows system:"
    echo ""
    echo "  Option 1: winpmem (recommended)"
    echo "    Download: https://github.com/Velocidex/WinPmem/releases"
    echo "    Run:      winpmem_mini_x64.exe $output"
    echo ""
    echo "  Option 2: DumpIt"
    echo "    Download: https://www.toolwar.com/dumpit"
    echo "    Run:      DumpIt.exe /quiet /output $output"
    echo ""
    echo "  Option 3: FTK Imager"
    echo "    Download: https://www.exterro.com/ftk-imager"
    echo "    Use GUI:  File > Capture Memory..."
    echo ""
    echo "  Option 4: Magnet RAM Capture"
    echo "    Download: https://www.magnetforensics.com/resources/magnet-ram-capture/"
    echo ""
    echo "  Option 5: VM Snapshot"
    echo "    VMware:     .vmem file in VM directory"
    echo "    Hyper-V:    Export VM > .bin file"
    echo "    VirtualBox: Use 'VBoxManage debugvm <name> dumpvmcore --filename $output'"
    echo ""
    log_info "After acquisition, verify with:"
    echo "    certutil -hashfile $output SHA256"
}

# --- Main ---

main() {
    log_info "Memory Acquisition Script"
    log_info "Output: $OUTPUT_PATH"
    log_info "Method: $METHOD"
    echo ""

    # Check output directory exists
    local output_dir
    output_dir=$(dirname "$OUTPUT_PATH")
    if [ ! -d "$output_dir" ]; then
        log_error "Output directory does not exist: $output_dir"
        exit 1
    fi

    # Check available disk space
    local avail_kb
    avail_kb=$(df -k "$output_dir" | awk 'NR==2 {print $4}')
    local mem_kb
    mem_kb=$(awk '/MemTotal/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    if [ "$mem_kb" -gt 0 ] && [ "$avail_kb" -lt "$((mem_kb * 2))" ]; then
        log_warn "Low disk space. Need ~$((mem_kb * 2 / 1048576)) GB, have ~$((avail_kb / 1048576)) GB"
    fi

    local os_type
    os_type=$(detect_os)
    log_info "Detected platform: $os_type"

    case "$os_type" in
        linux)
            acquire_linux "$OUTPUT_PATH" "$METHOD"
            ;;
        macos)
            acquire_macos "$OUTPUT_PATH"
            ;;
        windows)
            acquire_windows_guidance "$OUTPUT_PATH"
            ;;
        *)
            log_error "Unsupported platform: $os_type"
            exit 1
            ;;
    esac
}

main "$@"