/*
 * FuncTrace.  Copyright (c) 2005, Trent Waddington.  All rights reserved.
 * 
 * Outputs a function level trace of any unstripped target program.
 *
 */
#include <stdio.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <map>
#include <string>
#include <asm/user.h>

struct func_data {
	std::string name;
	unsigned long undercc;
};
std::map<unsigned long, struct func_data> funcs;

struct return_data {
	unsigned long proc;
	unsigned long undercc;
};
std::map<unsigned long, std::map<unsigned long, struct return_data> > returns;

void addBreakpoints(int pid, char *fname)
{
	FILE *f;
	char *buf = (char*)malloc(strlen(fname) + strlen(fname) + 1000);
	strcpy(buf, fname);
	strcat(buf, ".map");
	f = fopen(buf, "r");
	if (f == NULL) {
		fprintf(stderr, "cannot open map file, trying to create..\n");
		sprintf(buf, "nm %s > %s.map", fname, fname);
		system(buf);
		strcpy(buf, fname);
		strcat(buf, ".map");
		f = fopen(buf, "r");
		if (f == NULL) {
			fprintf(stderr, "failed to create map file.\n");
			exit(1);
		}
	}
	free(buf);
	//printf("adding breakpoints..\n");
	while (!feof(f)) {
		char line[1024], *p;
		fgets(line, 1024, f);
		p = strchr(line, '\n');
		if (p)
			*p = 0;
		p = strchr(line, ' ');
		if (p == NULL)
			continue;
		p++;
		if (*p == 'T') {
			unsigned long addr;
			if (sscanf(line, "%lx", &addr) == 1 && funcs.find(addr) == funcs.end()) {
				struct func_data d;
				d.name = p + 2;
				d.undercc = ptrace(PTRACE_PEEKTEXT, pid, (void*)addr, 0);
				if (*(unsigned char*)&d.undercc != 0xcc) {
					unsigned long undercc = d.undercc;
					*(unsigned char*)&undercc = 0xcc;
					ptrace(PTRACE_POKETEXT, pid, (void*)addr, (void*)undercc);
					funcs[addr] = d;
					//printf("added %s at addr %lx undercc=%lx\n", d.name.c_str(), addr, d.undercc);
				}
			}
		}
	}
	fclose(f);
}

int main(int argc, char **argv)
{
	if (argc < 2) {
		fprintf(stderr, "usage: functrace <program>\n");
		return 1;
	}

	int pid;
	if ((pid=fork()) == 0) {
		ptrace(PTRACE_TRACEME, 0, 0, 0);
		execl(argv[1], argv[1]);
		return 0;
	}

	int first = 1;
	int indent = 0;
	unsigned long stepfrom = 0;

	while(1) {
		int status;
		waitpid(pid, &status, 0);
		if (WIFEXITED(status))
			break;
		if (first) {
			first = 0;
			addBreakpoints(pid, argv[1]);
		} else if (WIFSTOPPED(status) && WSTOPSIG(status) == SIGTRAP) {
			struct user_regs_struct regs;
			ptrace(PTRACE_GETREGS, pid, 0, &regs);
			unsigned long bytes = ptrace(PTRACE_PEEKTEXT, pid, (void*)(regs.rip - 1), 0);
			if (*(unsigned char*)&bytes == 0xcc) {
				regs.rip--;
				ptrace(PTRACE_SETREGS, pid, 0, &regs);
				if (funcs.find(regs.rip) == funcs.end()) {
					if (returns.find(regs.rip) != returns.end()) {
						if (returns[regs.rip].find(regs.rsp-8) == returns[regs.rip].end()) {
							fprintf(stderr, "unknown return to %lx with rsp %lx\n", regs.rip, regs.rsp);
							return 1;
						}
						indent--;
						for (int i = 0; i < indent; i++)
							fprintf(stderr, "  ");
						fprintf(stderr, "<- %s\n", funcs[returns[regs.rip][regs.rsp-8].proc].name.c_str());
						ptrace(PTRACE_POKETEXT, pid, (void*)regs.rip, 
								(void*)returns[regs.rip][regs.rsp-8].undercc);
					} else {
						fprintf(stderr, "unknown breakpoint at %lx\n", regs.rip);
						return 1;
					}
				} else {
					for (int i = 0; i < indent; i++)
						fprintf(stderr, "  ");
					fprintf(stderr, "-> %s\n", funcs[regs.rip].name.c_str());
					indent++;
					ptrace(PTRACE_POKETEXT, pid, (void*)regs.rip, (void*)funcs[regs.rip].undercc);

					unsigned long retaddr = ptrace(PTRACE_PEEKTEXT, pid, (void*)regs.rsp, 0);
					struct return_data rd;
					rd.proc = regs.rip;
					rd.undercc = ptrace(PTRACE_PEEKTEXT, pid, (void*)retaddr, 0);
					returns[retaddr][regs.rsp] = rd;

					bytes = rd.undercc;
					*(unsigned char*)&bytes = 0xcc;
					ptrace(PTRACE_POKETEXT, pid, (void*)retaddr, (void*)bytes);
					//printf("adding breakpoint at return address %lx with rsp %lx\n", retaddr, regs.rsp);
					ptrace(PTRACE_SINGLESTEP, pid, 0, 0);
					stepfrom = regs.rip;
					continue;
				}
			} else if (stepfrom == 0) {
				fprintf(stderr, "unexpected single step\n");
				return 1;
			} else {
				unsigned long undercc = funcs[stepfrom].undercc;
				*(unsigned char*)&undercc = 0xcc;
				ptrace(PTRACE_POKETEXT, pid, (void*)stepfrom, (void*)undercc);
				stepfrom = 0;
			}
		} else if (WIFSTOPPED(status)) {
			fprintf(stderr, "stopped on unhandled signal %i\n", WSTOPSIG(status));
			return 1;
		}
		ptrace(PTRACE_CONT, pid, 0, 0);
	}

	return 0;
}