/**
 * uninformed research
 * http://www.uninformed.org
 *
 * Search/replace/dump for memory in a running process.
 *
 * Props to thief for the idea :)
 *
 * skape
 * 12/29/2002
 * mmiller@hick.org
 *
 * gcc -Wall -O3 memgrep.c -o memgrep
 *
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <linux/user.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <signal.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <ctype.h>

#include <elf.h>

#define MEMGREP_CMD_SEARCH        0x00000001  // Search memory
#define MEMGREP_CMD_REPLACE       0x00000002  // Replace memory
#define MEMGREP_CMD_SEARCHREPLACE 0x00000003  // Search and replace matches.
#define MEMGREP_CMD_DUMP          0x00000004  // Dump memory

#define MEMGREP_FLAG_VERBOSE   (1 << 0)  // Be verbose
#define MEMGREP_FLAG_PROMPT    (1 << 1)  // Prompt before replacing 
#define MEMGREP_FLAG_DUMPCLEAN (1 << 2)  // Dump data in a readable text, not hex.

typedef struct _mem_ctx {

	unsigned long flags;    // Flags

	int pid;                // Process id
	unsigned long *addrs;   // Start address(es)
	unsigned long numAddrs; // Number of start address(es)
	unsigned long length;   // The length to operate on

	struct {

		unsigned char *buf;   // In
		unsigned long length; // In

	} search;

	struct {

		unsigned char *buf;   // In
		unsigned long length; // In

	} replace;

	struct {

		unsigned char *buf;   // Out
		unsigned long length; // In/Out

		unsigned long padding;// The padding used when dumping
	} dump;

} MEM_CTX;

typedef struct _process_section_addrs {

	unsigned long rodata;    // read-only data, static
	unsigned long data;      // data, static
	unsigned long bss;       // bss, static

	unsigned long stack;     // stack, dynamic

} PROCESS_SECTION_ADDRS;

unsigned long memgrep(unsigned long cmd, MEM_CTX *ctx);
unsigned long memgrep_search(MEM_CTX *ctx);
unsigned long memgrep_replace(MEM_CTX *ctx);
unsigned long memgrep_searchreplace(MEM_CTX *ctx);
unsigned long memgrep_dump(MEM_CTX *ctx);

unsigned char attachProcess(int pid);
unsigned char detachProcess(int pid);
void safeCleanup(int signal);

unsigned char *getBufferFromAddressPid(int pid, unsigned long addr, unsigned long length);
unsigned long  putBufferInAddressPid(int pid, unsigned long addr, unsigned char *buf, unsigned long length);
unsigned long  getProcessSectionAddresses(int pid, PROCESS_SECTION_ADDRS *addrs);
void populateAddressesFromString(int pid, unsigned long **addrs, unsigned long *numAddrs, const char *string);
void translateToHex(const char *fullString, unsigned char **buf, unsigned long *bufLength);
void translateFormatToHex(const char *format, const char *string, unsigned char **buf, unsigned long *bufLength);

void displayHelp();

unsigned long _pid = 0;

int main(int argc, char **argv)
{
	unsigned long cmd = 0;
	extern char *optarg;
	char *addrs = NULL, *from = NULL, *to = NULL;
	MEM_CTX ctx;
	int c;

	memset(&ctx, 0, sizeof(ctx));

	while ((c = getopt(argc, argv, "p:dsra:l:f:t:b:cvh")) != EOF)
	{
		switch (c)
		{
			case 'p':
				ctx.pid = atoi(optarg) & 0xFFFF;
				break;
			case 'd':
				cmd = MEMGREP_CMD_DUMP;		
				break;
			case 's':
				if (cmd == MEMGREP_CMD_REPLACE)
					cmd = MEMGREP_CMD_SEARCHREPLACE;
				else
					cmd = MEMGREP_CMD_SEARCH;
				break;
			case 'r':
				if (cmd == MEMGREP_CMD_SEARCH)
					cmd = MEMGREP_CMD_SEARCHREPLACE;
				else
					cmd = MEMGREP_CMD_REPLACE;
				break;
			case 'a':
				addrs = optarg;
				break;
			case 'l':
				ctx.length = strtoul(optarg, NULL, 10);
				break;
			case 'f':
				from = optarg;
				break;
			case 't':
				to = optarg;
				break;
			case 'b':
				ctx.dump.padding = strtoul(optarg, NULL, 10);
				break;
			case 'c':
				ctx.flags |= MEMGREP_FLAG_DUMPCLEAN;
				break;
			case 'v':
				ctx.flags |= MEMGREP_FLAG_VERBOSE;
				break;
			case 'h':
				displayHelp();
				break;
		}
	}

	if (cmd == 0 || !addrs || !ctx.pid)
		displayHelp();

	_pid = ctx.pid;

	// Populate addresses
	populateAddressesFromString(ctx.pid, &ctx.addrs, &ctx.numAddrs, addrs);

	if (!ctx.numAddrs)
	{
		fprintf(stderr, "memgrep: No addresses specified.\n");
		return 0;
	}

	switch (cmd)
	{
		case MEMGREP_CMD_SEARCH:
			if (!from)
			{
				fprintf(stderr, "memgrep: Search specified but no criteria used.\n");
				return 0;
			}

			translateToHex(from, &ctx.search.buf, &ctx.search.length);
			break;
		case MEMGREP_CMD_REPLACE:
			if (!to)
			{
				fprintf(stderr, "memgrep: Replace specified but no data found (must specify -t when -r is used).\n");
				return 0;
			}
			
			translateToHex(to, &ctx.replace.buf, &ctx.replace.length);
			break;
		case MEMGREP_CMD_SEARCHREPLACE:
			if (!from || !to)
			{
				fprintf(stderr, "memgrep: Search/Replace specified bu either search or replace data was missing (must specify -f and -t when search/replacing).\n");
				return 0;
			}

			translateToHex(from, &ctx.search.buf, &ctx.search.length);
			translateToHex(to, &ctx.replace.buf, &ctx.replace.length);
			break;
	}

	memgrep(cmd, &ctx);

	return 0;
}

void displayHelp()
{
	fprintf(stdout, "memgrep -- Run-time memory searching, dumping and modifying utility. (12/30/2002)\n"
						 "Usage: ./memgrep [-p pid] [-d] [-r] [-s] [-a addr1,addr2,bss,addr3] [-l length] [-f fmt,search data] [-t fmt,replace data]\n"
						 "                 [-v] [-h]\n"
						 "\n"
						 "   -p         The process id to operate on.\n"
						 "   -d         Dump memory from the specified address(es) for the given length (-l).\n"
						 "   -r         Replace memory at the specified address(es).  If -s is also specified,\n"
						 "              only memory that matches the search criteria will be replaced\n"
						 "   -s         Search memory at the specified address(es).\n"
						 "   -a [addr]  The address(es) to operate on seperated by commas.  Addresses can be\n"
                   "              in the following format:\n"
						 "                 0x821c4ac\n"
						 "                 821c4ac\n"
						 "              Also, the following keywords can be used:\n"
						 "                 bss       -> Uses the VMA associated with the .bss section (commonly heap).\n"
						 "                 stack     -> Dynamically determines the current stack pointer.\n"
						 "                 rodata    -> Uses the VMA associated with the .rodata section (read-only data, ie, static text).\n"
						 "                 data      -> Uses the VMA associated with the .data section (data, ie, global variables).\n"
						 "   -l [len]   The length to use when searching or dumping.  A length of 0 means search\n"
						 "              till end-of-memory.\n"
						 "   -f [data]  This specifies the search criteria.  Multiple formats are accepted for ease\n"
						 "              of use.  Below are accepted formats and their examples:\n"
						 "                 s -> String format  (Ex: 's,Testing')\n"
						 "                 x -> Hex format     (Ex: 'x,00414100AB')\n"
						 "                 i -> Integer format (Ex: 'i,4724')\n"
						 "   -t [data]  This specifies the replace data.  The same formats used with the -f parameter\n"
						 "              are valid for the -t parameter.\n"
						 "   -b [pad]   Number of bytes of padding to use around dump addresses (default is 0).\n"
						 "   -c         Dump data in clear-text, not hex.\n"
						 "   -v         Verbose.\n"
						 "   -h         Help.\n"
			  			 "\n"
						 "   Example search (search for 'Jane' in .bss):\n\n"
						 "      ./memgrep -p 1335 -s -a bss -f s,Jane\n\n"
						 "   Example replace (replace memory at 0x8423143 and 0x8443147 with 0x00ff0041):\n\n"
						 "      ./memgrep -p 1335 -r -a 0x8423143,0x8443147 -t x,00ff0041\n\n"
						 "   Example search/replace (Replace 'Test' with 'Rest' in .bss and .rodata):\n\n"
						 "      ./memgrep -p 1335 -s -r -a bss,rodata -f s,Test -t s,Rest\n\n"
						 "   Example dump (Dump memory starting at 0x8422113 for 16 bytes):\n\n"
						 "      ./memgrep -p 1335 -d -a 0x8422113 -l 16\n\n"
						 "\n"
						 );

	exit(0);
}

void populateAddressesFromString(int pid, unsigned long **addrs, unsigned long *numAddrs, const char *string)
{
	const char *comma = NULL, *curr = string;
	unsigned char done = 0;
	unsigned long numElements = 0;
	PROCESS_SECTION_ADDRS sectionAddrs;

	memset(&sectionAddrs, 0, sizeof(sectionAddrs));

	while (curr && !done)
	{
		comma = strchr(curr, ',');

		if (comma)
			*(char *)comma = 0;
		else
			done = 1;

		// Invalid format
		if (strlen(curr) < 2)
			break;

		if (!*addrs)
			*addrs = (unsigned long *)malloc((++numElements) * sizeof(unsigned long));
		else
			*addrs = (unsigned long *)realloc(*addrs, (++numElements) * sizeof(unsigned long));

		if (!isdigit(*curr))
		{
			if (!sectionAddrs.bss)
				getProcessSectionAddresses(pid, &sectionAddrs);
				
			if (!strcmp(curr, "bss"))
				(*addrs)[numElements-1] = sectionAddrs.bss;
			else if (!strcmp(curr, "stack"))
				(*addrs)[numElements-1] = sectionAddrs.stack;
			else if (!strcmp(curr, "rodata"))
				(*addrs)[numElements-1] = sectionAddrs.rodata;
			else if (!strcmp(curr, "data"))
				(*addrs)[numElements-1] = sectionAddrs.data;	
		}
		else
		{
			if (curr[1] == 'x')
				(*addrs)[numElements-1] = strtoul(curr+2, NULL, 16);
			else
				(*addrs)[numElements-1] = strtoul(curr, NULL, 16);
		}

		if (comma)
			*(char *)comma = ',';

		curr = comma + 1;
	}

	*numAddrs = numElements;
}

unsigned long getProcessSectionAddresses(int pid, PROCESS_SECTION_ADDRS *addrs)
{
	Elf32_Ehdr    elfHeader;
	Elf32_Shdr    *sectionHeaders = NULL, *stringTableHeader = NULL;
	unsigned long index = 0;
	char          path[1024], *stringTable = NULL;
	int           fd = 0;

	path[sizeof(path) - 1] = path[0] = 0;

	snprintf(path, sizeof(path) - 1, "/proc/%d/exe", pid);

	if ((fd = open(path, O_RDONLY)) <= 0)
		return 0;

	do
	{
		if (lseek(fd, 0, SEEK_SET) < 0)
			break;
		
		if (read(fd, &elfHeader, sizeof(elfHeader)) < 0)
			break;

		if (!(sectionHeaders = (Elf32_Shdr *)malloc(elfHeader.e_shentsize * elfHeader.e_shnum)))
			break;
		
		if (lseek(fd, elfHeader.e_shoff, SEEK_SET) < 0)
			break;

		// Read the section headers

		if (read(fd, sectionHeaders, elfHeader.e_shentsize * elfHeader.e_shnum) < 0)
			break;

		// Read the string table

		if (!(stringTableHeader = &sectionHeaders[elfHeader.e_shstrndx]))
			break;

		if (lseek(fd, stringTableHeader->sh_offset, SEEK_SET) < 0)
			break;

		if (!(stringTable = (char *)malloc(stringTableHeader->sh_size)))
			break;

		if (read(fd, stringTable, stringTableHeader->sh_size) < 0)
			break;

		for (index = 0; index < elfHeader.e_shnum; index++)
		{
			const char *name = &stringTable[sectionHeaders[index].sh_name];

			if (!name)
				continue;

			if (!strcmp(name, ".bss"))
				addrs->bss = sectionHeaders[index].sh_addr;
			else if (!strcmp(name, ".rodata"))
				addrs->rodata = sectionHeaders[index].sh_addr;
			else if (!strcmp(name, ".data"))
				addrs->data = sectionHeaders[index].sh_addr;
		}

	} while (0);

	if (stringTable)
		free(stringTable);
	if (sectionHeaders)
		free(sectionHeaders);

	close(fd);

	// Get current stack address

	if (attachProcess(pid))
	{
		struct user_regs_struct regs;

		memset(&regs, 0, sizeof(regs));

		if (ptrace(PTRACE_GETREGS, pid, NULL, &regs) != -1)
			addrs->stack = regs.esp;		
		else
			perror("ptrace(GETREGS)");	  
		
		detachProcess(pid);
	}

	return 1;
}

unsigned char attachProcess(int pid)
{
	if (ptrace(PTRACE_ATTACH, pid, NULL, NULL) < 0)
	{
		perror("ptrace(ATTACH)");

		return 0;
	}

	// waitpid doesn't work for me, anyone else had that problem?
	while (ptrace(PTRACE_PEEKDATA, pid, 0, NULL) == -1 && errno == 3)
		usleep(100);
	
	signal(SIGINT, safeCleanup);

	return 1;
}

unsigned char detachProcess(int pid)
{
	if (ptrace(PTRACE_DETACH, pid, NULL, NULL) < 0)
		perror("ptrace(DETACH)");

	signal(SIGINT, SIG_DFL);

	return 1;
}

void safeCleanup(int signal)
{
	detachProcess(_pid);

	exit(0);
}

void translateToHex(const char *fullString, unsigned char **buf, unsigned long *bufLength)
{
	const char *comma = strchr(fullString, ',');

	if (!comma)
		return;

	*(char *)comma = 0;

	translateFormatToHex(fullString, comma + 1, buf, bufLength);

	*(char *)comma = ',';
}

void translateFormatToHex(const char *format, const char *string, unsigned char **buf, unsigned long *bufLength)
{
	switch (format[0])
	{
		case 's': // String
			{
				unsigned long x;
				*bufLength = strlen(string);
				*buf       = (unsigned char *)malloc(*bufLength);

				if (!*buf)
					return;

				for (x = 0; x < *bufLength; x++)
					(*buf)[x] = string[x];
			}
			break;
		case 'i': // Integer
			{
				*bufLength = sizeof(long);
				*buf       = (unsigned char *)malloc(*bufLength);

				*(*((long **)buf)) = strtol(string, NULL, 10);
			}
			break;
		case 'x': // Hex
			{
				unsigned long x, stringLen, bufPos = 0;
				char hex[3] = { 0, 0, 0 };
				*bufLength  = (stringLen = strlen(string)) / 2;
				*buf        = (unsigned char *)malloc(*bufLength);

				if (!*buf)
					return;

				for (x = 0; x < stringLen; x += 2)
				{
					hex[0] = string[x];
					hex[1] = string[x+1];

					(*buf)[bufPos++] = strtoul(hex, NULL, 16) & 0xFF;
				}
			}
			break;
	}
}

unsigned long memgrep(unsigned long cmd, MEM_CTX *ctx)
{
	unsigned long ret = 0;

	if ((ret = attachProcess(ctx->pid)) <= 0)
		return ret;

	switch (cmd)
	{
		case MEMGREP_CMD_SEARCH:
			ret = memgrep_search(ctx);
			break;
		case MEMGREP_CMD_REPLACE:
			ret = memgrep_replace(ctx);
			break;
		case MEMGREP_CMD_SEARCHREPLACE:
			ret = memgrep_searchreplace(ctx);
			break;
		case MEMGREP_CMD_DUMP:
			ret = memgrep_dump(ctx);
			break;
		default:
			break;
	}

	detachProcess(ctx->pid);

	if (ctx->search.buf)
		free(ctx->search.buf);
	if (ctx->replace.buf)
		free(ctx->replace.buf);

	return ret;
}

unsigned long memgrep_search(MEM_CTX *ctx)
{
	return memgrep_searchreplace(ctx);
}

unsigned long memgrep_replace(MEM_CTX *ctx)
{
	unsigned long ret = 0, x;

	do
	{
		for (x = 0; x < ctx->numAddrs; x++)
		{
			fprintf(stdout, "Replacing memory %.8x with %lu bytes of data...\n", (unsigned int)ctx->addrs[x], ctx->replace.length);

			putBufferInAddressPid(ctx->pid, ctx->addrs[x], ctx->replace.buf, ctx->replace.length);
		}

	} while (0);

	return ret;
}

unsigned long memgrep_searchreplace(MEM_CTX *ctx)
{
	unsigned long ret = 0, x;

	do
	{
		for (x = 0; x < ctx->numAddrs; x++)
		{
			unsigned long end = (ctx->length)?ctx->addrs[x] + ctx->length:0xffffffff, inc = 4, data = 0, addr, y;
			unsigned char found = 0;

			fprintf(stdout, "Searching %.8x -> %.8x...\n", (unsigned int)ctx->addrs[x], (unsigned int)end);

			for (addr = ctx->addrs[x]; addr < end; addr += inc)
			{
				if (found)
				{
					inc   = 4;
					found = 0;
				}

				if (((data = ptrace(PTRACE_PEEKDATA, ctx->pid, addr, NULL)) == -1) && (errno == 5))
					break;

				for (y = 0; y < sizeof(unsigned long); y++)
				{
					if (((unsigned char *)&data)[y] == ctx->search.buf[0])
					{
						unsigned char *match = getBufferFromAddressPid(ctx->pid, addr + y, ctx->search.length);

						if ((match) && (memcmp(match, ctx->search.buf, ctx->search.length) == 0))
						{
							if (ctx->replace.buf)
							{
								putBufferInAddressPid(ctx->pid, addr + y, ctx->replace.buf, ctx->replace.length);
								
								fprintf(stdout, "  replaced at 0x%.8x\n", (unsigned int)(addr + y));

								inc = ctx->replace.length;
							}
							else
							{
								fprintf(stdout, "  found at 0x%.8x\n", (unsigned int)(addr + y));

								inc = ctx->search.length + y;
							}

							found = 1;
						}

						if (match)
							free(match);

						if (found)
							break;
					}

				}
			}
		}

	} while (0);

	return ret;
}

unsigned long memgrep_dump(MEM_CTX *ctx)
{
	unsigned long ret = 0, x, y, base, len;

	if (!ctx->length)
		return 0;

	for (x = 0; x < ctx->numAddrs; x++)
	{
		unsigned char *buf = getBufferFromAddressPid(ctx->pid, (base = ctx->addrs[x] - ctx->dump.padding), (len = ctx->length + (ctx->dump.padding * 2)));
		unsigned long data, count = 0, inc = sizeof(unsigned long);

		fprintf(stdout, "%lu bytes starting at %.8x (+/- %lu)...\n%.8x: ", ctx->length, (unsigned int)ctx->addrs[x], ctx->dump.padding, (unsigned int)base);

		if (!buf)
			continue;

		if (ctx->flags & MEMGREP_FLAG_DUMPCLEAN)
			inc = 1;

		for (y = 0; y < len; y += inc)
		{
			if (ctx->flags & MEMGREP_FLAG_DUMPCLEAN)
			{
				if (isdigit(buf[y]) || isalpha(buf[y]))
					fprintf(stdout, "%c", buf[y]);
				else
					fprintf(stdout, ".");

				count++;

				if (count == 16)
				{
					fprintf(stdout, "\n");
				
					if (y + 1 < len)
						fprintf(stdout, "%.8x: ", base + y + 1);

					count = 0;
				}
			}
			else
			{
				memcpy(&data, buf + y, sizeof(unsigned long));

				fprintf(stdout, "%.8x ", (unsigned int)data);

				count++;

				if (count == 4)
				{
					fprintf(stdout, "\n");
										
					if (y + 4 < len)
						fprintf(stdout, "%.8x: ", (unsigned int)(base + y + 4));

					count = 0;
				}
			}
		}

		fprintf(stdout, "\n");

		free(buf);
	}

	return ret;
}

unsigned char *getBufferFromAddressPid(int pid, unsigned long addr, unsigned long length)
{
	unsigned char *ret = (unsigned char *)malloc(length);
	unsigned long x = addr, end = addr + length, data, left = length, y, retPos = 0;

	if (!ret)
		return NULL;

	memset(ret, 0, length);

	for (; x < end; x += 4)
	{
		data = ptrace(PTRACE_PEEKDATA, pid, x, NULL);

		for (y = 0; y < sizeof(unsigned long) && left != 0; y++)
		{
			ret[retPos++] = ((unsigned char *)&data)[y];
			left--;
		}
	}

	return ret;
}

unsigned long putBufferInAddressPid(int pid, unsigned long addr, unsigned char *buf, unsigned long length)
{
	unsigned long x = addr, end = addr + length, data, pokedata, left = length, pos = 0;
	int y;

	for (; x < end; x += 4)
	{
		data = ptrace(PTRACE_PEEKDATA, pid, x, NULL);

		for (y = 0; y < sizeof(unsigned long); y++)
		{
			if (left == 0)
				((unsigned char *)&pokedata)[y] = ((unsigned char *)&data)[y];
			else
			{
				((unsigned char *)&pokedata)[y] = buf[pos++];

				left--;
			}

		}

		ptrace(PTRACE_POKEDATA, pid, x, pokedata);
	}

	return 1;
}