/* 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, const char *childpath);
static void _asmsh_env_child(const char *childpath);

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

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

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

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

	if(asmsh_env_update(res) < 0)
	{
		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:
	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);
	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;
	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;
			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;
			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;
			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, ret;

	if(status) { *status = 0; }

	if(ptrace(PTRACE_SINGLESTEP, env->pid, NULL, 0) < 0)
	{
		err = errno;
		perror("Unable to ptrace singlestep");
		goto err;
	}
	if(waitpid(env->pid, &env->status, 0) < 0)
	{
		err = errno;
		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_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;
		perror("ptrace getregs error");
		errno = err;
		return -1;
	}
	return 0;
}


static int _asmsh_env_spawn(asmsh_env_t *env, const char *childpath)
{
	int err;
	int wstatus;

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

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

	if(waitpid(env->pid, &wstatus, 0) < 0)
	{
		perror("Unable to wait for child process");
		goto err;
	}
	if(!WIFSTOPPED(wstatus))
	{
		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;
		perror("ptrace setoptions failed");
		goto err;
	}
	if(ptrace(PTRACE_CONT, env->pid, NULL, 0) < 0)
	{
		err = errno;
		perror("ptrace CONT failed after attach");
		goto err;
	}

	if(waitpid(env->pid, &wstatus, 0) < 0)
	{
		err = errno;
		perror("Unable to wait for child process to stop after exec");
		goto err;
	}
	if(wstatus >> 8 != (SIGTRAP | (PTRACE_EVENT_EXEC<<8)))
	{
		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;
			perror("Unable to ptrace syscall");
			dprintf(2, "ptrace syscall failed on %d time\n", i+1);
			goto err;
		}
		if(waitpid(env->pid, &wstatus, 0) < 0)
		{
			err = errno;
			perror("Unable to wait for child process to stop on syscall");
			goto err;
		}
		if(wstatus != 1407)
		{
			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))
		{
			goto err;
		}
	}

	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:
	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 :/");
		exit(err?err:-1);
	}
	execve(childpath, argv, envp);
	int err = errno;
	perror("Unable to execve");
	exit(err?err:-1);
}