// execlog1.c
// $ gcc -Wall execlog1.c -o execlog -ludis86
//
// http://udis86.sourceforge.net/
// http://udis86.sourceforge.net/manual/manual.html

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/user.h>
#include <sys/syscall.h>

#include <udis86.h>


#define TEXT_ADDRESS  0x400e70
#define TEXT_SIZE     0x2bb8



void err(char *str)
{
	fprintf(stderr, "ERROR: %s\n", str);
}

void target(char *argv[], char *argp[])
{
	if(ptrace(PTRACE_TRACEME, 0, NULL, NULL) != -1)
		execve(argv[0], argv, NULL);
	else
		err("PTRACE_TRACEME");
	exit(0);
}

int read_data(int pid, unsigned long addr, unsigned long *mem)
{
	errno = 0;
	unsigned long data = ptrace(PTRACE_PEEKDATA, pid, addr, NULL);
	if(errno != 0)
		return -1;
	*mem = (unsigned long)data;
	return 0;
}

int wait_fork(int pid)
{
	int status, flag = 0;
	struct user_regs_struct regs;

	waitpid(pid, &status, 0);

	while(1){
		if(ptrace(PTRACE_SYSCALL, pid, NULL, NULL) == 0){
			waitpid(pid, &status, 0);
			if(WIFEXITED(status))
				break;
			ptrace(PTRACE_GETREGS, pid, 0, &regs);
			if(regs.orig_rax != (SYS_fork - 1))
				continue;
			if(flag == 0){
				flag = 1;
				continue;
			}
			// ptrace(PTRACE_DETACH, pid, NULL, NULL);
			ptrace(PTRACE_KILL, pid, NULL, NULL);
			pid = regs.rax;
			ptrace(PTRACE_ATTACH, pid, NULL, NULL);
			break;
		}
        }

	return pid;
}

char * print_value(int pid, unsigned long addr)
{
	static char buff[256];
	unsigned char data[8];
	if(read_data(pid, addr, (unsigned long *)data) == -1){
		sprintf(buff, " ");
	}else{
		sprintf(buff, "( %02x %02x %02x %02x %02x %02x %02x %02x )", 
			data[0], data[1], data[2], data[3],
			data[4], data[5], data[6], data[7]);
	}
	return buff;
}

void output_asm(unsigned long rip, char *asm_str, int indent)
{
	int i;
	for(i=0; i < indent; i++)
		printf(" ");

	printf("%lx: %s\n", rip, asm_str);
}

int g_regs_len = -1;
struct user_regs_struct g_regs[256];

void push(struct user_regs_struct *pregs)
{
	if(g_regs_len == 255)
		return;
	g_regs_len++;
	memcpy(&g_regs[g_regs_len], pregs, 
		sizeof(struct user_regs_struct));
}

struct user_regs_struct * pop(void)
{
	if(g_regs_len <= -1)
		return NULL;
	g_regs_len--;
	return &g_regs[g_regs_len + 1];
}

void output_call(struct user_regs_struct *pregs, int pid)
{
	printf("-- call --------------------------------------------\n");
	printf("rdi = %lx\t%s\n", pregs->rdi, print_value(pid, pregs->rdi));
	printf("rsi = %lx\t%s\n", pregs->rsi, print_value(pid, pregs->rsi));
	printf("rdx = %lx\t%s\n", pregs->rdx, print_value(pid, pregs->rdx));
	printf("rcx = %lx\t%s\n", pregs->rcx, print_value(pid, pregs->rcx));
	printf("r8  = %lx\t%s\n", pregs->r8,  print_value(pid, pregs->r8));
	printf("----------------------------------------------------\n");
	push(pregs);
}

void output_retn(struct user_regs_struct *pregs, int pid)
{
	struct user_regs_struct *stack = pop();
	if(stack == NULL){
		stack = pregs;
		printf("-- retn (regs) -------------------------------------\n");
	}else{
		printf("-- retn (args) -------------------------------------\n");
	}
	printf("rdi = %lx\t%s\n", stack->rdi, print_value(pid, stack->rdi));
	printf("rsi = %lx\t%s\n", stack->rsi, print_value(pid, stack->rsi));
	printf("rdx = %lx\t%s\n", stack->rdx, print_value(pid, stack->rdx));
	printf("rcx = %lx\t%s\n", stack->rcx, print_value(pid, stack->rcx));
	printf("r8  = %lx\t%s\n", stack->r8,  print_value(pid, stack->r8));
	printf("rax = %lx\t%s\n", pregs->rax, print_value(pid, pregs->rax));
	printf("----------------------------------------------------\n");
}

void analyze(int pid, struct user_regs_struct *pregs)
{
	int i;
	unsigned char code[64];
	char *asm_str;
	ud_t ud_obj;
	static int flag = 0;
	static int indent = 0;

	for(i=0; i < 8; i++){
		unsigned long data;
		read_data(pid, pregs->rip + (i * 4), &data);
		memcpy(code + (i * 4), &data, 4);
	}
	
	ud_init(&ud_obj);
	ud_set_input_buffer(&ud_obj, code, 32);
	ud_set_mode(&ud_obj, 64);
	ud_set_syntax(&ud_obj, UD_SYN_INTEL);

	while(ud_disassemble(&ud_obj)){
		asm_str = ud_insn_asm(&ud_obj);
		if(TEXT_ADDRESS <= pregs->rip && 
			pregs->rip < (TEXT_ADDRESS + TEXT_SIZE))
		{
			if(flag == 2){
				indent -= ((indent >= 2) ? 2 : 0);
				output_retn(pregs, pid);
			}
			output_asm(pregs->rip, asm_str, indent);
			if(strncmp(asm_str, "call", 4) == 0){
				indent += 2;
				output_call(pregs, pid);
			}
			if(strncmp(asm_str, "ret", 3) == 0 && indent >= 2){
				indent -= ((indent >= 2) ? 2 : 0);
				output_retn(pregs, pid);
			}
			flag = 1;
		}else{
			flag = 2;
		}
		break;
	}
}

int control(int pid)
{
	int status, pid_back;
	struct user_regs_struct regs;

	pid_back = pid;
	if((pid = wait_fork(pid)) == pid_back)
		return -1;
	
	while(1){
		waitpid(pid, &status, 0);
		if(WIFEXITED(status))
			break;
		ptrace(PTRACE_GETREGS, pid, 0, &regs);
		analyze(pid, &regs);
		ptrace(PTRACE_SINGLESTEP, pid, 0, NULL);
	}
	return 0;
}

int exec_prog(char *argv[], char *argp[])
{
	int pid;

	switch(pid = fork())
	{
	case 0:
		target(argv, argp);
		break;
	case -1:
		err("ERROR");
		break;
	default:
		control(pid);
		break;
	}
	return 0;
}

int main(int argc, char *argv[], char *argp[])
{
	if(argc < 2){
		fprintf(stderr, "%s <args>\n", argv[0]);
		return 1;
	}
	argv++;
	exec_prog(argv, argp);
	return 0;
}