#include <stdio.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/reg.h>
#include <sys/syscall.h>   /* For SYS_write etc */

/* Too lazy to figure out the proper way to detect x86 and x86_64 */
#if __WORDSIZE == 64
#	define REG_WIDTH 8
#	define REG_IO_PORT RDX
#	define REG_IO_VALUE RAX
#	define REG_IO_VALUE_INS RDI
#	define REG_IO_VALUE_OUTS RSI
#	define REG_IP RIP
#else
#	define REG_WIDTH 4
#	define REG_IO_PORT EDX
#	define REG_IO_VALUE EAX
#	define REG_IO_VALUE_INS EDI
#	define REG_IO_VALUE_OUTS ESI
#	define REG_IP EIP
#endif

int main(int argc, char **argv)
{
	pid_t child;
	unsigned long rip, ins, tmp, size, i, port, value, pending;
	unsigned char *code = (unsigned char *)&ins;
	char *dir;
	int status;

	if(argc < 2) {
		printf("usage: %s PROGRAM [ARGUMENTS]\n", argv[0]);
		return -1;
	}

	child = fork();
	if(child == 0) {
		ptrace(PTRACE_TRACEME, 0, NULL, NULL);
		if(execvp(argv[1], &argv[1]) == -1) {
			perror("Could not run the specified command");
		}
	} else {
		pending = 0;
		while(1) {
			wait(&status);
			if(WIFEXITED(status))
				break;

			/* If we saw a good-looking instruction at the end of the loop last time, we now read the value we possibly got */
			if(pending) {
				pending = 0;

				/* Only read from register if they haven't been loaded from eg. immediates yet */
				if(port == -1) {
					port = ptrace(PTRACE_PEEKUSER, child, REG_WIDTH * REG_IO_PORT, NULL) & 0xffff;
				}
				if(value == -1) {
					value = ptrace(PTRACE_PEEKUSER, child, REG_WIDTH * REG_IO_VALUE, NULL);
				}

				/* Mask the value according to the instruction size */
				switch(size) {
					case 8:
						value = value & 0xff;
						break;
					case 16:
						value = value & 0xffff;
						break;
					case 32:
						value = value & 0xffffffff;
						break;
					default:
						printf("ERROR: Invalid size!\n");
						return -1;
				}

				printf("port: %ld; %s: %lu\n", port, dir, value);
			}

			/* Get next instruction */
			rip = ptrace(PTRACE_PEEKUSER, child, REG_WIDTH * REG_IP, NULL);
			ins = ptrace(PTRACE_PEEKTEXT, child, rip, NULL);

			/* Decode the instruction: check for operand size -prefix */
			i = 0;
			size = 32;
			if(code[0] == 0x66) {
				i = 1;
				size = 16;
			}

			/* Decode the instruction: check if IN/OUT, check size, get immediate/read value from memory */
			/* XXX: Check the operand size -prefix isn't used on an 8bit instruction */
			port = value = -1;
			switch(code[i]) {
				case 0xe4:
					size = 8;
				case 0xe5:
					port = code[i+1];
					//printf("IN.i%ld A*, %02lu\n", size, port);
					dir = "in";
					break;
				case 0xec:
					size = 8;
				case 0xed:
					//printf("IN.r%ld A*, DX\n", size);
					dir = "in";
					break;
				case 0xe6:
					size = 8;
				case 0xe7:
					port = code[i+1];
					//printf("OUT.i%ld %02lu, A*\n", size, port);
					dir = "out";
					break;
				case 0xee:
					size = 8;
				case 0xef:
					//printf("OUT.r%ld DX, A*\n", size);
					dir = "out";
					break;
				case 0x6c:
					size = 8;
				case 0x6d: /* XXX: I SMELL A BUG (This should be done after the instruction is executed...) */
					tmp = ptrace(PTRACE_PEEKUSER, child, REG_WIDTH * REG_IO_VALUE_INS, NULL);
					value = ptrace(PTRACE_PEEKTEXT, child, tmp, NULL);
					//printf("INS.m%ld ES:RDI, DX\n", size);
					printf("BUG: INS will return invalid data at the moment\n");
					dir = "in";
					break;
				case 0x6e:
					size = 8;
				case 0x6f:
					tmp = ptrace(PTRACE_PEEKUSER, child, REG_WIDTH * REG_IO_VALUE_OUTS, NULL);
					value = ptrace(PTRACE_PEEKTEXT, child, tmp, NULL);
					//printf("OUTS.m%ld DX, DS:RSI\n", size);
					dir = "out";
					break;
				default:
					size = 0;
					break;
			}

			/* The instruction is IN or OUT, print port and value at next step (because for in, we can't see the value before that) */
			if(size) {
				pending = 1;
			}

			ptrace(PTRACE_SINGLESTEP, child, NULL, NULL);
		}
	}

	return 0;
}