#include "sdb.h"
#include "logger.h"
#include "vector.h"

#include <stdlib.h>
#include <string.h>
#include <capstone/capstone.h>
#include <sys/user.h>
#include <sys/ptrace.h>
#include <sys/wait.h>

int child;
int status;

struct user_regs_struct regs;
uint64_t disassemble_addr;

long syscall_nr = 0xffff;

void sync_regs(void)
{
	if (ptrace(PTRACE_GETREGS, child, NULL, &regs) != 0) {
		ERROR("sync_regs ptrace getregs\n");
		perror("sync_regs");
		exit(1);
	}
}

uint8_t poke(uint64_t addr, uint8_t data)
{
	uint64_t ret = ptrace(PTRACE_PEEKTEXT, child, addr, NULL);
	ptrace(PTRACE_POKETEXT, child, addr, (ret & ~0xff) | data);
	return ret & 0xff;
}

void run(const char *filename)
{
	child = fork();
	if (child < 0) {
		ERROR("run fork failed\n");
		perror("run fork");
		exit(1);
	}

	if (child == 0) {
		DEBUG("forked\n");
		if (ptrace(PTRACE_TRACEME, 0, NULL, NULL) < 0) {
			perror("run traceme");
			exit(1);
		}
		execl(filename, filename, NULL);
		perror("run execl");
	}

	if (waitpid(child, &status, 0) < 0) {
		perror("run waitpid");
		exit(1);
	}
	if (WIFSTOPPED(status) == 0) {
		perror("run stopped");
		exit(1);
	}
	ptrace(PTRACE_SETOPTIONS, child, 0, 
		PTRACE_O_EXITKILL | PTRACE_O_TRACESYSGOOD);
	
	sync_regs();	
	INFO("program '%s' loaded. entry point %p.\n", filename, (void *)regs.rip);
}

void disassemble()
{
	if (disassemble_addr == 0x00) return;
	const uint64_t rip = disassemble_addr;
	long ret;
	uint8_t code[64], *ptr = (uint8_t *)&ret;

	for (int i = 0; i < 8; i++) {
		ret = ptrace(PTRACE_PEEKTEXT, child, rip + (i << 3), NULL);
		for (int j = 0; j < 8; j++)
			code[i << 3 | j] = ptr[j];
		// DEBUG("\n");
	}

	for (int i = 0; i < bps_len; i++)
		if (bps[i].addr && rip <= bps[i].addr && bps[i].addr < rip + 64)
			code[bps[i].addr - rip] = bps[i].data;
	
	csh handle;
	if (cs_open(CS_ARCH_X86, CS_MODE_64, &handle) != CS_ERR_OK) {
		perror("disassemble cs_open");
		exit(1);
	}

	cs_insn *insn;
	int count = cs_disasm(handle, code, sizeof(code) - 1, rip, 0, &insn);
	if (count > 0) {
		for (int i = 0; i < INSTRUCTION_COUNT; i++) {
			if (insn[i].size == 2 && 
					insn[i].bytes[0] == 0 && 
					insn[i].bytes[1] == 0) {
				INFO("the address is out of the range of the text section.\n");
				break;
			}
			OUTPUT("%12lx: ", insn[i].address);
			char insn_buf[25] = "", *insn_buf_ptr = insn_buf;
			for (int j = 0; j < insn[i].size; j++)
				insn_buf_ptr += sprintf(insn_buf_ptr, "%02x ", 
				                        (uint8_t)insn[i].bytes[j]);	
			OUTPUT("%-24s%-10s %s\n", insn_buf, insn[i].mnemonic, insn[i].op_str);
		}
		cs_free(insn, count);
	} else {
		perror("disassemble");
		exit(1);
	}
	cs_close(&handle);
}

void inst_load(char **filename)
{
	char *token = strtok(NULL, " \t\n");
	if (token == NULL) {
		perror("inst_load");
		disassemble_addr = 0x00;
		return;
	}
	*filename = malloc(strlen(token) + 1);
	strcpy(*filename, token);
}

void inst_si(void)
{
	if (!WIFSTOPPED(status)) {
		ERROR("program not running.\n");
		perror("inst_si not running");
		disassemble_addr = 0x00;
		return;
	}

	syscall_nr = 0xffff;

	sync_regs();
	uint64_t rip = regs.rip;
	DEBUG("rip = %lx\n", rip);

	const struct bps_node *bp = find(rip);
	if (bp)
		poke(bp->addr, bp->data);

	ptrace(PTRACE_SINGLESTEP, child, NULL, NULL);
	waitpid(child, &status, 0);

	if (bp)
		poke(bp->addr, 0xcc);
	
	sync_regs();
	rip = regs.rip, bp = find(rip);
	if (bp)
		INFO("hit a breakpoint at %p.\n", (void *)rip);

	if (!WIFSTOPPED(status)) {
		INFO("the target program terminated.\n");
		disassemble_addr = 0x00;
	} else {
		sync_regs();
		disassemble_addr = regs.rip;
	}
}

void inst_cont(void)
{
	if (!WIFSTOPPED(status)) {
		ERROR("program not running.\n");
		perror("inst_cont not running");
		exit(1);
	}

	syscall_nr = 0xffff;

	sync_regs();
	uint64_t rip = regs.rip;
	DEBUG("rip = %lx\n", rip);

	const struct bps_node *bp = find(rip);
	if (bp) {
		poke(bp->addr, bp->data);
		ptrace(PTRACE_SINGLESTEP, child, NULL, NULL);
		waitpid(child, &status, 0);
		poke(bp->addr, 0xcc);
	}

	ptrace(PTRACE_CONT, child, NULL, NULL);
	waitpid(child, &status, 0);

	if (!WIFSTOPPED(status)) {
		INFO("the target program terminated.\n");
		disassemble_addr = 0x00;
	} else {
		sync_regs();
		rip = regs.rip;
		bp = find(rip - 1);

		if (bp == NULL) {
			perror("inst_cont bp = NULL\n");
			exit(1);
		}
		INFO("hit a breakpoint at %p.\n", (void *)(rip - 1));

		poke(bp->addr, bp->data);
		regs.rip -= 1;
		ptrace(PTRACE_SETREGS, child, 0, &regs);
		poke(bp->addr, 0xcc);
		
		sync_regs();
		disassemble_addr = regs.rip;
	}
}

void inst_info(void)
{
	char *op = strtok(NULL, " \t\n");
	if (op == NULL)
		goto inst_info_invalid;

	if (strcmp(op, "reg") == 0)
		inst_info_reg();
	else if (strcmp(op, "break") == 0) 
		inst_info_break();
	else
		goto inst_info_invalid;
	
	return;

inst_info_invalid:
	ERROR("invalid command\n");
	INFO("Command: info [reg | break]\n");
}

void inst_info_reg(void)
{
	sync_regs();
#define OUTPUT_REGS(a, b, c) \
	OUTPUT("$%-7s 0x%016lx\t$%-7s 0x%016lx\t$%-7s 0x%016lx\n", \
		#a, (unsigned long)regs.a, \
		#b, (unsigned long)regs.b, \
		#c, (unsigned long)regs.c)
	OUTPUT_REGS(rax, rbx, rcx);
	OUTPUT_REGS(rdx, rsi, rdi);
	OUTPUT_REGS(rbp, rsp,  r8);
	OUTPUT_REGS( r9, r10, r11);
	OUTPUT_REGS(r12, r13, r14);
	OUTPUT_REGS(r15, rip, eflags);
#undef OUTPUT_REGS
}

void inst_info_break(void)
{
	if (bps_cnt == 0) {
		INFO("no breakpoints.\n");
		return;
	}
	
	OUTPUT("%-6s\t%-12s\n", "Num", "Address");
	for (int i = 0; i < bps_len; i++)
		if (bps[i].addr)
			OUTPUT("%-6d\t0x%lx\n", i, bps[i].addr);
}

void inst_break(void)
{
	char *token = strtok(NULL, " \t\n");
	if (token == NULL) {
		ERROR("invalid command.\n");
		INFO("Command: break [hex address]\n");
		return;
	}
	uint64_t addr;
	sscanf(token, "%lx", &addr);

	bps_push(addr, poke(addr, 0xcc));
	INFO("set a breakpoint at %p.\n", (void *)addr);
}

void inst_delete(void)
{
	char *token = strtok(NULL, " \t\n");
	if (token == NULL) {
		ERROR("invalid command.\n");
		INFO("Command: delete [id]\n");
		return;
	}

	int id;
	sscanf(token, "%d", &id);
	
	if (bps_len <= id || bps[id].addr == 0x00)
		INFO("breakpoint %d does not exist.\n", id);
	else {
		poke(bps[id].addr, bps[id].data);
		bps[id].addr = 0x00, bps_cnt--;
		INFO("delete breakpoint %d.\n", id);
	}
}

void inst_patch(void)
{
	char *token = strtok(NULL, " \t\n");
	if (token == NULL) goto inst_patch_invalid;
	uint64_t addr;
	sscanf(token, "%lx", &addr);
	
	token = strtok(NULL, " \t\n");
	if (token == NULL) goto inst_patch_invalid;
	uint64_t data;
	sscanf(token, "%lx", &data);

	token = strtok(NULL, " \t\n");
	if (token == NULL) goto inst_patch_invalid;
	int len;
	sscanf(token, "%d", &len);

	uint64_t val = ptrace(PTRACE_PEEKTEXT, child, addr, NULL);
	DEBUG("original val = %lx\n", val);
	
	uint64_t mask = ((uint64_t)1 << (len << 3)) - 1;
	DEBUG("mask = %lx\n", mask);
	val = (val & ~mask) | (data & mask);
	DEBUG("modified val = %lx\n", val);

	for (int i = 0; i < bps_len; i++)
		if (bps[i].addr && addr <= bps[i].addr && bps[i].addr < addr + len) {
			int j = bps[i].addr - addr;
			bps[i].data = (data >> j) & 0xff;
		}
	
	if (ptrace(PTRACE_POKETEXT, child, addr, val) != 0) {
		ERROR("inst_patch POKETEXT\n");
		perror("inst_patch ptrace");
		exit(1);
	}
	INFO("patch memory at address 0x%lx.\n", addr);

	return;

inst_patch_invalid:
	ERROR("invalid command.\n");
	INFO("Command: patch [hex address] [hex value] [len]\n");
}

void inst_syscall(void)
{
	if (!WIFSTOPPED(status)) {
		ERROR("program not running.\n");
		perror("inst_syscall");
		disassemble_addr = 0x00;
		return;
	}

	sync_regs();
	uint64_t rip = regs.rip;
	const struct bps_node *bp = find(rip);

	if (bp) {
		poke(bp->addr, bp->data);
		ptrace(PTRACE_SINGLESTEP, child, NULL, NULL);
		waitpid(child, &status, 0);
		poke(bp->addr, 0xcc);
	}

	ptrace(PTRACE_SYSCALL, child, NULL, NULL);
	waitpid(child, &status, 0);

	if (!WIFSTOPPED(status)) {
		INFO("the target program terminated.\n");
		disassemble_addr = 0x00;
	} else {
		sync_regs();
		rip = regs.rip;
		DEBUG("rip = %p\n", (void *)rip);

		if (WSTOPSIG(status) & 0x80) {
			if (syscall_nr == 0xffff) {
				syscall_nr = regs.orig_rax;
				INFO("enter a syscall(%ld) at %p.\n",
					syscall_nr, (void *)(rip - 2));
			} else {
				INFO("leave a syscall(%ld) = %llu at %p.\n", 
					syscall_nr, regs.rax, (void *)(rip - 2));
				syscall_nr = 0xffff;
			}
			disassemble_addr = regs.rip - 2;	
			return;
		}

		bp = find(rip - 1);
		if (bp == NULL) {
			ERROR("inst_syscall bp = NULL\n");
			perror("inst_syscall");
			exit(1);
		}
		INFO("hit a breakpoint at %p.\n", (void *)(rip - 1));

		poke(bp->addr, bp->data);
		regs.rip -= 1;
		ptrace(PTRACE_SETREGS, child, 0, &regs);
		poke(bp->addr, 0xcc);

		disassemble_addr = regs.rip;
	}
}