#!/bin/bash

# A script to extract the memory contents of a process.
# It uses the following utils/ways to extract memory content for respective OS's
#     - Ubuntu:
#           the `scanmem` util
#     - Amazon Linux and RedHat Linux:
#           read data from /proc/<pid>/maps and parse it using the `dd` util
#                      to extract strings from memory regions
# This script requires sudo.
#
# Running:
# 1. Get the PID of the process you want to perform memory scraping on.
#      If you want to get the PID associated with a docker container you can try using
#      `docker top <container-id>`
#
# 2. Run this script providing the PID as the first argument and the name of the output file as the second argument.
#      `sudo ./scrape_sata.sh <pid> <output-filename>`
#
# 3. The above will create the following files:
#    Ubuntu:
#      a. <output_filename>           Contains all the binary data extracted from memory.
#      b. <output_filename>_strings   Contains strings of printable characters extracted from the file above.
#    Amazon Linux:
#      a. <output_filename>_strings   Contains strings of printable characters extracted from memory.
#    RedHat Linux:
#      a. <output_filename>_strings   Contains strings of printable characters extracted from memory.

if [ "$#" -ne 2 ]; then
    echo "Usage $0 <pid> <output-filename>"
    exit 1
fi

PID=$1
DUMP_FILE=$2
OS=$(grep ^NAME= </etc/os-release | cut -d "=" -f2 | xargs)

RED=$(tput setaf 1)
GREEN=$(tput setaf 2)
RESET=$(tput sgr0)

# Cleanup on listening EXIT trap
cleanup() {
    if [ "$OS" = "Ubuntu" ]; then
        rm -rf /tmp/eachdump /tmp/cmd /tmp/log
    elif [[ "$OS" = "Amazon Linux" || "$OS" = "Red Hat Enterprise Linux" ]]; then
        rm -f /tmp/*_mem_*.bin
    fi
}
trap 'cleanup' EXIT

# Verify that a process with <PID> exists
verify_pid() {
    if ps -p "$PID" &>/dev/null; then
        echo -e "Starting memory dump of process ${GREEN}${PID}${RESET}\n"
    else
        echo "${RED}Given pid: $PID is not valid.${RESET}"
        exit 1
    fi
}

# Install utility for scraping
install_util() {
    if ! command -v "${MEM_UTIL}" &>/dev/null; then
        echo "'${MEM_UTIL}' tool could not be found. Installing ..."

        # attempt to install, else exit
        sudo apt install ${MEM_UTIL}
        if ! command -v ${MEM_UTIL} &>/dev/null; then
            echo "${RED}'${MEM_UTIL}' util could not be installed.${RESET}"
            exit 1
        fi
        echo -e "${GREEN}'${MEM_UTIL}' successfully installed!${RESET}\n"
    fi
}

# Function to scrape memory for processes running on ubuntu
scrape_ubuntu() {
    install_util
    echo "Dumping memory regions from process using ${MEM_UTIL}"
    echo "lregions" >/tmp/lregion
    sudo ${MEM_UTIL} "$PID" </tmp/lregion >/tmp/maps_all 2>&1
    rm /tmp/lregion
    echo -e "Done dumping the map information using ${MEM_UTIL}\n"

    # Remove the line numbers from output
    sed -i 's/\[[^]]*\]//g' /tmp/maps_all
    # Keep the lines with string "bytes" in them.
    grep bytes /tmp/maps_all | awk '{print $1 " " $2}' | tr --delete , >/tmp/region_pairs
    rm /tmp/maps_all

    echo -n "Using ${MEM_UTIL} to read the memory regions"

    if [ -f "$DUMP_FILE" ]; then
        truncate -s 0 "$DUMP_FILE"
    else
        touch "$DUMP_FILE"
    fi

    while read LINE; do
        touch /tmp/eachdump
        echo "dump $LINE /tmp/eachdump" >/tmp/cmd
        sudo ${MEM_UTIL} "$PID" </tmp/cmd >/tmp/log 2>&1
        echo -n "."
        cat /tmp/eachdump >>"$DUMP_FILE"
    done </tmp/region_pairs

    echo -e "Done\n"
    echo -e "Dumped the data to file ${GREEN}${DUMP_FILE}${RESET}"

    # extract strings from mem dump
    if command -v strings &>/dev/null; then
        STRINGS_FILE="${DUMP_FILE}_strings"
        if [ -f "$STRINGS_FILE" ]; then
            truncate -s 0 "$STRINGS_FILE"
        else
            touch "$STRINGS_FILE"
        fi
        strings "$DUMP_FILE" >"$STRINGS_FILE"

        echo "Extracted printable strings to file ${GREEN}${DUMP_FILE}_strings${RESET}"
    fi
}

# Function to scrape memory for processes running on Amazon Linux/RHEL
scrape_amazon_or_rhel_linux() {
    echo "Dumping memory regions from process"
    # Read regions of contiguous memory from process
    sudo cat /proc/$1/maps | grep -Fv ".so" | grep " 0 " | awk '{print $1}' | ( IFS="-"
    # For each page in /proc/$1/mem get page size and convert it to bytes,
    # then read region of memory and dump to a file
    while read a b; do
        sudo dd if=/proc/$1/mem bs=$( getconf PAGESIZE ) iflag=skip_bytes,count_bytes \
        skip=$(( 0x$a )) count=$(( 0x$b - 0x$a )) of=/tmp/"$1_mem_$a.bin"  >/dev/null 2>&1
        # Read the String from the file 
        strings /tmp/"$1_mem_$a.bin" >> "$2"_strings 2>&1
        echo -n ".";
    done )
    echo -e "Done\n";
    echo -e "Dumped the data to file ${GREEN}${DUMP_FILE}_strings${RESET}";
}

verify_pid

# Scrape process memory pages according to OS.
if [ "$OS" = "Ubuntu" ]; then
    MEM_UTIL=scanmem
    scrape_ubuntu
elif [[ "$OS" = "Amazon Linux" || "$OS" = "Red Hat Enterprise Linux" ]]; then
    scrape_amazon_or_rhel_linux "$PID" "$DUMP_FILE"
else
    echo "Operating system '$OS' not supported."
fi
