/* Copyright Yann Weber <asmsh@yannweb.net>
   This file is part of asmsh.

   asmsh is free software: you can redistribute it and/or modify it under the
   terms of the GNU General Public License as published by the Free Software
   Foundation, either version 3 of the License, or any later version.
   
   asmsh is distributed in the hope that it will be useful, but WITHOUT ANY
   WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
   FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
   details.
   
   You should have received a copy of the GNU General Public License along
   with asmsh. If not, see <https://www.gnu.org/licenses/>.
*/
#include "asm_env.h"

static int _asmsh_env_spawn(asmsh_env_t *asmenv);
static void _asmsh_env_child(const char *childpath);
/** Return a path (that should be freed) of a temporary executable
 *  child that can be exec on */
static char *asmsh_env_tmpexec();

/* binary buffer of the child elf */
extern unsigned char _binary_child_start;
extern unsigned char _binary_child_end;

asmsh_env_t* asmsh_env(const char *childpath)
{
	asmsh_env_t *res;
	int err;

	if((res = malloc(sizeof(*res))) == NULL)
	{
		err = errno;
		asmsh_log_perror("Unable to allocate env");
		errno = err;
		return NULL;
	}
	child_mmap_init(&(res->mmap));

	asmsh_brk_init(&res->brks);

	res->childpath = NULL;
	if(childpath && (res->childpath = strdup(childpath)) == NULL)
	{
		err=errno;
		goto err_pathdup;
	}

	if(_asmsh_env_spawn(res) < 0)
	{
		err = errno;
		goto err;
	}

	if(asmsh_env_update(res) < 0)
	{
		err=errno;
		goto err;
	}
	
	res->txt_map_addr = (void*)res->regs.rax;
	res->txt_map_sz = res->regs.r14;
	res->stack_addr = (void*)res->regs.r15;
	res->stack_sz = (size_t)res->stack_addr - res->regs.rsp;

	res->code_write_ptr = res->txt_map_addr;


	return res;
err:
	if(res->childpath)
	{
		free(res->childpath);
	}
err_pathdup:
	free(res);
	errno = err;
	return NULL;
}


void asmsh_env_free(asmsh_env_t *asmenv)
{
	if(asmenv->mmap.maps)
	{
		free(asmenv->mmap.maps);
	}
	kill(asmenv->pid, SIGKILL);
	asmsh_brk_free(&asmenv->brks);
	free(asmenv->childpath);
	free(asmenv);
}

int asmsh_env_write_mem(asmsh_env_t *env, void *addr, const unsigned char *buf, size_t buf_sz)
{
	int err;
	u_int64_t data;
	unsigned char written;
	char bleft  = (u_int64_t)addr % 8;

	written = 0;
	if(bleft)
	{
		// First write to correct further write allignement
		void *wr_addr = (void*)(((u_int64_t)addr / 8) * 8);
		char towrite = 8 - bleft;
		towrite = (towrite > buf_sz)?buf_sz:towrite;
		
		errno = 0;
		data = ptrace(PTRACE_PEEKTEXT, env->pid, wr_addr, NULL);
		if(errno)
		{
			err = errno;
			asmsh_log_perror("Unable to peektext in order to allign write");
			errno = err;
			return -1;
		}
		memcpy(((unsigned char*)&data)+bleft, &(buf[written]), towrite);
		if(ptrace(PTRACE_POKETEXT, env->pid, wr_addr, data) < 0)
		{
			err = errno;
			asmsh_log_perror("Unable to poketext in order to allign write");
			errno = err;
			return -1;
		}
		written += towrite;
		env->code_write_ptr += towrite;
	}
	if(written >= buf_sz)
	{
		return 0;
	}
	// env->code_write_ptr is now word alligned and "written" bytes are
	// allready written
	while(written < buf_sz)
	{
		char towrite = buf_sz - written;
		if(towrite >= 8)
		{
			data = *(u_int64_t*)(&buf[written]);
		}
		else
		{
			data = 0;
			memcpy(&data, &buf[written], towrite);
		}

		if(ptrace(PTRACE_POKETEXT, env->pid, env->code_write_ptr, data) < 0)
		{
			err = errno;
			asmsh_log_perror("Unable to poketext");
			errno = err;
			return -1;
		}
		written += towrite;
		env->code_write_ptr += towrite;
	}
	return 0;
}


int asmsh_env_write_code(asmsh_env_t *env, asmsh_bytecode_t *bcode)
{
	return asmsh_env_write_mem(env, env->code_write_ptr,
			bcode->bytes, bcode->size);
}


int asmsh_env_step(asmsh_env_t *env, int *status)
{
	int err;

	if(status) { *status = 0; }

	if(ptrace(PTRACE_SINGLESTEP, env->pid, NULL, 0) < 0)
	{
		err = errno;
		asmsh_log_perror("Unable to ptrace singlestep");
		goto err;
	}
	if(waitpid(env->pid, &env->status, 0) < 0)
	{
		err = errno;
		asmsh_log_perror("Unable to wait for child process to stop on step");
		goto err;
	}
	if(status) { *status = env->status; }

	if(!WIFSTOPPED(env->status) || env->status >> 8 != 5)
	{
		goto err_wstatus;
	}
	
	return 0;

/// TODO replace by an utility function that logs ?
err_wstatus:
	if(WIFEXITED(env->status))
	{
		dprintf(2, "Child exited with status %d\n", WEXITSTATUS(env->status));
	}
	else if(WIFSIGNALED(env->status))
	{
		if(WCOREDUMP(env->status))
		{
			dprintf(2, "Child segfault\n");
		}
		else
		{
			dprintf(2, "Child killed by sig %d\n", WTERMSIG(env->status));
		}
	}
	else if(WIFSTOPPED(env->status))
	{
		dprintf(2, "Child stopped by %s(sig#%d)\n",
				strsignal(WSTOPSIG(env->status)),
				WSTOPSIG(env->status));
	}
	else
	{
		dprintf(2, "Unexpected child status 0x%04X\n", env->status);
	}
	errno = ECHILD;
	return 1;
err:
	errno = err;
	return -1;
}


int asmsh_env_run(asmsh_env_t *env, int *status)
{
	while(1)
	{
		int ret;
		if((ret = asmsh_env_update_regs(env)) < 0)
		{
			return ret;
		}
		const unsigned long rip = env->regs.rip;
		if(asmsh_brk_isset(&(env->brks), rip))
		{
			asmsh_log_info("Breakpoint %016lx reached", rip);
			*status = 0;
			return 0;
		}
		if((ret = asmsh_env_step(env, status)) != 0)
		{
			return ret;
		}
	}
}


int asmsh_env_update(asmsh_env_t *asmenv)
{
	if(asmsh_env_update_regs(asmenv) < 0)
	{
		return -1;
	}
	return child_mmap_get(asmenv->pid, &(asmenv->mmap));
}


int asmsh_env_update_maps(asmsh_env_t *asmenv)
{
	if(child_mmap_get(asmenv->pid, &(asmenv->mmap)) < 0)
	{
		return -1;
	}

	return 0;
}


int asmsh_env_update_regs(asmsh_env_t *asmenv)
{
	bzero(&(asmenv->regs), sizeof(asmenv->regs));
	if(ptrace(PTRACE_GETREGS, asmenv->pid, NULL, &(asmenv->regs)) < 0)
	{
		int err = errno;
		asmsh_log_perror("ptrace getregs error");
		errno = err;
		return -1;
	}
	return 0;
}


static int _asmsh_env_spawn(asmsh_env_t *env)
{
	int err;
	int wstatus;

	const char *childpath = env->childpath?env->childpath:asmsh_env_tmpexec();

	if(!childpath)
	{
		return -1; // Error in asmsh_env_tmpexec()
	}

	if((env->pid = fork()) == -1)
	{
		err = errno;
		asmsh_log_perror("Unable to fork!");
		goto err_fork;
	}
	else if(env->pid == 0)
	{
		_asmsh_env_child(childpath);
	}

	if(waitpid(env->pid, &wstatus, WUNTRACED) < 0)
	{
		err=errno;
		asmsh_log_perror("Unable to wait for child process");
		goto err;
	}
	if(!WIFSTOPPED(wstatus))
	{
		dprintf(2, "child didn't stop as expected");
		goto err_wstatus;
	}

	if(ptrace(PTRACE_ATTACH, env->pid, 0, 0) == -1)
	{
		err=errno;
		asmsh_log_perror("Unable to attach to child process");
		goto err;
	}

	if(waitpid(env->pid, &wstatus, 0) < 0)
	{
		err=errno;
		asmsh_log_perror("Unable to wait for child process");
		goto err;
	}
	if(!WIFSTOPPED(wstatus))
	{
		dprintf(2, "child didn't stop as expected");
		goto err_wstatus;
	}

	// Attached to child
	// tell child to stop on exec
	if(ptrace(PTRACE_SETOPTIONS, env->pid, NULL, PTRACE_O_TRACEEXEC) < 0)
	{
		err = errno;
		asmsh_log_perror("ptrace setoptions failed");
		goto err;
	}

	for(int i=0; i<2; i++) // cont kill & ptrace
	{
		if(ptrace(PTRACE_CONT, env->pid, NULL, 0) < 0)
		{
			err = errno;
			asmsh_log_perror("ptrace CONT failed after attach");
			goto err;
		}

		if(waitpid(env->pid, &wstatus, 0) < 0)
		{
			err = errno;
			asmsh_log_perror("Unable to wait for child process to stop after exec");
			goto err;
		}
		if(wstatus >> 8 != (SIGTRAP | (PTRACE_EVENT_EXEC<<8)) && \
				!WIFSTOPPED(wstatus))
		{
			goto err_wstatus;
		}
	}

	// next child events are : execve ret, mmap in, mmap out
	for(int i=0; i<3; i++)
	{
		if(ptrace(PTRACE_SYSCALL, env->pid, NULL, 0) < 0)
		{
			err = errno;
			asmsh_log_error("Unable to ptrace syscall on %dth time : %s",
					i+1, strerror(errno));
			goto err;
		}
		if(waitpid(env->pid, &wstatus, 0) < 0)
		{
			err = errno;
			asmsh_log_perror("Unable to wait for child process to stop on syscall");
			goto err;
		}
		if(wstatus != 1407 && wstatus != 263551)
		{
			dprintf(2, "unexpected status after ptrace syscall %d\n", wstatus);
			goto err_wstatus;
		}
	}
	// mmap done by child process
	
	// WARNING totally depends on child.s source
	// right now there is only 4 instructions (the fourth is the jmp to the txt_map)
	// before reaching the start of the code mmap
	for(int i=0; i<4; i++)
	{
		/* // DEBUG to monitor placement in child exec
		asmsh_env_update_regs(env);
		dprintf(2, "%d) rax: %08X rip : %08X mmap_addr = %08X\n",
				i, env->regs.rax, env->regs.rip, env->txt_map_ptr);
		*/
		if(asmsh_env_step(env, NULL))
		{
			err=errno;
			goto err;
		}
	}

	if(!env->childpath) { unlink(childpath); } // rm tmp child exec

	return 0;

/// TODO replace by an utility function that logs ?
err_wstatus:
	if(WIFEXITED(wstatus))
	{
		dprintf(2, "Child exited with status %d\n", WEXITSTATUS(wstatus));
	}
	else if(WIFSIGNALED(wstatus))
	{
		if(WCOREDUMP(wstatus))
		{
			dprintf(2, "Child segfault\n");
		}
		else
		{
			dprintf(2, "Child killed by sig %d\n", WTERMSIG(wstatus));
		}
	}
	else if(WIFSTOPPED(wstatus))
	{
		dprintf(2, "Child stopped by sig %d\n", WSTOPSIG(wstatus));
	}
	else
	{
		dprintf(2, "Unexpected child status 0x%04X\n", wstatus);
	}
	err = ECHILD;
err:
	kill(env->pid, SIGKILL);
err_fork:

	if(!env->childpath) { unlink(childpath); } // rm tmp child exec

	errno = err;
	return -1;
}

static void _asmsh_env_child(const char *childpath)
{

	char *argv[] = {NULL, NULL};
	char *envp[] = {NULL};
	if(!(argv[0] = strdupa(childpath)))
	{
		int err = errno;
		perror("Unable to strdupa asmsh childpath :/");
		//asmsh_log_perror("Unable to strdupa asmsh childpath :/");
		///! @todo dump logs before exit !
		exit(err?err:-1);
	}

	kill(getpid(), SIGSTOP);
	execve(childpath, argv, envp);
	int err = errno;
	perror("Unable to execve");
	//asmsh_log_perror("Child is unable to execve");
	///! @todo dump logs before exit !
	exit(err?err:-1);
}

static char *asmsh_env_tmpexec()
{
	char *ret = strdup("asmsh_child_XXXXXXXXX");
	int err;
	if(!ret)
	{
		err = errno;
		asmsh_log_perror("getting a temporary file name");
		errno = err;
		return NULL;
	}
	int tmpfd = mkstemp(ret);
	if(tmpfd < 0)
	{
		err = errno;
		asmsh_log_perror("Unable to mk temporary file");
		errno = err;
		return NULL;
	}
	const int sz = &_binary_child_end - &_binary_child_start;
	int rsz = write(tmpfd, &_binary_child_start, sz);
	if(rsz<sz)
	{
		// TODO : differenciate incomplete write & error !
		err = errno;
		asmsh_log_perror("Unable to write the child executable");
		free(ret);
		errno = err;
		return NULL;
	}
	fchmod(tmpfd, 0555);
	close(tmpfd);
	return ret;
}