#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/reg.h>
#include <sys/uio.h>
#include <strings.h>
#include <string.h>
#include <unistd.h>
#include <elf.h>
#include "sandbox.h"

/* DEBUG HEADER */
#include <stdio.h>

static inline void get_args(struct s_sandboxenv *env)
{
#ifdef __x86_64__
	REGISTER_TYPE result[] = { env->registers.rdi, env->registers.rsi,
			env->registers.rdx, env->registers.r10,
			env->registers.r8, env->registers.r9 };
#else
	REGISTER_TYPE result[] = { env->registers.ebx, env->registers.ecx,
			env->registers.edx, env->registers.esi,
			env->registers.edi, env->registers.ebp };
#endif
	memcpy(env->syscall_args, result, sizeof(REGISTER_TYPE) * 6);
}

static inline void init_syscalls(struct s_sandboxenv *env)
{
	bzero(env->functions, sizeof(*(env->functions)) * NR_syscalls);
	env->functions[__NR_write] = ovr_write;
	env->functions[__NR_open] = ovr_open;
	env->functions[__NR_close] = ovr_close;
}

static inline int waitForSyscall(const int pid)
{
	int status;

	ptrace(PTRACE_SYSCALL, pid, NULL, 0);
	if (waitpid(pid, &status, 0) == -1)
		return -1;
	return 0;
}

int manageSyscall(struct s_sandboxenv *env)
{
	t_syscall_fnc ovr_fnc;
	int syscall_nr;

#ifdef __x86_64__
	syscall_nr = env->registers.orig_rax;
#else
	syscall_nr = env->registers.orig_eax;
#endif

	if (env->registers.orig_rax >= NR_syscalls ||
			!(ovr_fnc = env->functions[(int) env->registers.orig_rax]))
		return 0;
	get_args(env);
	(ovr_fnc)(env);
	return 1;
}

void doTrace(int pid, const t_param *params)
{
	int status;
	struct s_sandboxenv sandbox_env;
	struct iovec iov = { &(sandbox_env.registers), sizeof(sandbox_env.registers) };

	ptrace(PTRACE_ATTACH, pid, 0, 0);
	kill(pid, SIGTRAP);
	waitpid(pid, &status, 0);
	sandbox_env.params = params;
	sandbox_env.child_pid = pid;
	init_syscalls(&sandbox_env);

	while (1)
	{
		if (waitForSyscall(pid))
			break;
		ptrace(PTRACE_GETREGSET, pid, NT_PRSTATUS, &iov);
		if (manageSyscall(&sandbox_env))
			ptrace(PTRACE_SETREGS, pid, 0, &(sandbox_env.registers));
		if (waitForSyscall(pid))
			break;
		iov.iov_len = sizeof(sandbox_env.registers);
		ptrace(PTRACE_GETREGSET, pid, NT_PRSTATUS, &iov);
		manageSyscall(&sandbox_env);
	}
}

int launch_program(const t_param *params)
{
	int child = fork();

	if (child == -1)
		return -1;
	else if (child == 0)
		doExec(getpid(), params);
	else
		doTrace(child, params);
	return 0;
}