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

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

#include <udis86.h>

//global
unsigned long g_logaddr = 0x00000000;
unsigned long g_logsize = 0xffffffff;

unsigned long g_wrtaddr[256];
unsigned char g_wrtdata[256];
unsigned char g_reddata[256];

unsigned int g_fork_flag = 0;

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

int read_data(int pid, unsigned long addr, unsigned char *mem, int size)
{
	int i, n;
	unsigned long *out = (unsigned long *)mem;

	n = size / 4;
	
	for(i=0; i < n; i++){
		errno = 0;
		unsigned long data = ptrace(PTRACE_PEEKDATA, pid, addr, NULL);
		addr += 4;
		if(errno != 0)
			return -1;
		*out++ = data;
	}

	if((size % 4) > 0){
		errno = 0;
		unsigned long data = ptrace(PTRACE_PEEKDATA, pid, addr, NULL);
		unsigned char *ptr = (unsigned char *)&data;
		unsigned char *buff = (unsigned char *)out;
		if(errno != 0)
			return -1;
		switch(size % 4){
			case 3: buff[2] = ptr[2];
			case 2: buff[1] = ptr[1];
			case 1: buff[0] = ptr[0];
		}
	}
	
	return 0;
}

int write_data(int pid, unsigned long addr, unsigned char *mem, int size)
{
	int i, n;
	unsigned long *out = (unsigned long *)mem;

	n = size / 4;

	for(i=0; i < n; i++){
		errno = 0;
		unsigned long data = *out++;
		ptrace(PTRACE_POKEDATA, pid, addr, data);
		addr += 4;
		if(errno != 0)
			return -1;
	}

	if((size % 4) > 0){
		unsigned char buff[4];
		unsigned char *data = (unsigned char *)out;
		if(read_data(pid, addr, buff, 4) == -1)
			return -1;
		switch(size % 4){
			case 3: buff[2] = data[2];
			case 2: buff[1] = data[1];
			case 1: buff[0] = data[0];
		}
		write_data(pid, addr, buff, 4);
	}
	
	return 0;
}

int disas(int pid, unsigned long addr)
{
	ud_t ud_obj;
	unsigned char buff[32];

	if(read_data(pid, addr, buff, 32) == -1){
		printf("(Can't read)\n");
		return -1;
	}

	ud_init(&ud_obj);
	ud_set_input_buffer(&ud_obj, buff, 32);

	// 16 or 32 or 64 bit
	ud_set_mode(&ud_obj, 32);

	// UD_VEDNOR_ATT or UD_SYN_INTEL
	ud_set_syntax(&ud_obj, UD_SYN_INTEL);
	
	if(ud_disassemble(&ud_obj)){
		printf("%08lx: %14s  %s\n", addr, 
			ud_insn_hex(&ud_obj), 
			ud_insn_asm(&ud_obj));
	}
	
	return (int)ud_insn_len(&ud_obj);
}

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

void control(int pid)
{
	int flag = 0;
	int status, i, len = 0;
	struct user_regs_struct regs;
	unsigned char int80h[2];
	
	while(1){
		waitpid(pid, &status, 0);
		if(WIFEXITED(status))
			break;
		if(len != 0){
			for(i=0; g_wrtaddr[i] != 0; i++)
				write_data(pid, g_wrtaddr[i], &g_reddata[i], 1);
			len = 0;
		}
		ptrace(PTRACE_GETREGS, pid, 0, &regs);
		if(flag){
			flag = 0;
			ptrace(PTRACE_SYSCALL, pid, NULL, NULL);
			waitpid(pid, &status, 0);
			if(WIFEXITED(status))
				break;
			ptrace(PTRACE_GETREGS, pid, 0, &regs);
			if(g_fork_flag == 1)
				ptrace(PTRACE_DETACH, pid, NULL, NULL);
			if(g_fork_flag == 2)
				ptrace(PTRACE_KILL, pid, NULL, NULL);
			pid = regs.eax;
			ptrace(PTRACE_ATTACH, pid, NULL, NULL);
			continue;
		}else{
			if(g_logaddr <= regs.eip && regs.eip < g_logaddr + g_logsize)
				len = disas(pid, regs.eip);
			read_data(pid, regs.eip, int80h, 2);
			if(g_fork_flag != 0 && 
				int80h[0] == 0xcd && int80h[1] == 0x80 && 
				regs.eax == 0x78)
			{
				flag = 1;
				ptrace(PTRACE_SYSCALL, pid, NULL, NULL);
				continue;
			}
			for(i=0; g_wrtaddr[i] != 0; i++){
				if(regs.eip <= g_wrtaddr[i] &&  g_wrtaddr[i] < regs.eip + len)
					write_data(pid, g_wrtaddr[i], &g_wrtdata[i], 1);
			}
		}
		ptrace(PTRACE_SINGLESTEP, pid, 0, NULL);
	}
}

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

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

int main(int argc, char *argv[], char *argp[])
{
	int i;
	char *opt;
	
	if(argc < 2){
		fprintf(stderr, "%s <args>\n", argv[0]);
		fprintf(stderr, "%s -a <addr> <size> <args>\n", argv[0]);
		return 1;
	}

	argv++;
	opt = argv[0];
	
	memset(g_wrtaddr, 0, sizeof(g_wrtaddr));
	memset(g_wrtdata, 0, sizeof(g_wrtdata));
	
	i = 0;

	while(opt[0] == '-'){
		switch(opt[1]){
		case 'a':
			if((g_logaddr = strtoul(argv[1], NULL, 16)) == 0){
				err("STRTOUL addr");
				return -1;
			}
			if((g_logsize = strtoul(argv[2], NULL, 16)) == 0){
				err("STRTOUL size");
				return -1;
			}
			argv += 3;
			break;
		case 'w':
			if((g_wrtaddr[i] = strtoul(argv[1], NULL, 16)) == 0){
				err("STRTOUL write addr");
				return -1;
			}
			if((g_reddata[i] = (unsigned char)
				strtoul(argv[2], NULL, 16)) == 0)
			{
				err("STRTOUL read data");
				return -1;
			}
			if((g_wrtdata[i] = (unsigned char)
				strtoul(argv[3], NULL, 16)) == 0)
			{
				err("STRTOUL write data");
				return -1;
			}
			i++;
			argv += 4;
			break;
		case 'f':
			g_fork_flag = 0;
			if(strcmp("DETACH", argv[1]) == 0)
				g_fork_flag = 1;
			if(strcmp("KILL", argv[1]) == 0)
				g_fork_flag = 2;
			argv += 2;
			break;
		}
		opt = argv[0];
	}
	
	exec_prog(argv, argp);
	return 0;
}