#include <stdio.h>
#include <sys/types.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <stdlib.h>
#include <err.h>
#include <machine/reg.h>

#define	CODE_SIZE 128
#define	STACK_SIZE 128

extern void injected(void);
extern int injected_len;

int
ptrace_wait(pid_t pid)
{
	int status;

	while (waitpid(pid, &status, 0) != -1) {
		if (WIFEXITED(status) || WIFEXITED(status)) {
			printf("exited\n");
			return (1);
		}
		if (WIFSTOPPED(status))
			break;
	}
	return (0);
}

void
ptrace_io(pid_t pid, int op, void *off, void *addr, size_t len)
{
	struct ptrace_io_desc io_desc;

	io_desc.piod_op = op;
	io_desc.piod_offs = off;
	io_desc.piod_addr = addr;
	io_desc.piod_len = len;
	if (ptrace(PT_IO, pid, (caddr_t)&io_desc, 0) < 0)
		err(1, "ptrace(PT_IO)");
}

int
main(int argc, char **argv)
{
	struct reg oregs;
	char ocode[CODE_SIZE], ostack[STACK_SIZE];
	int pid;

	if (argc < 2)
		errx(1, "Usage: %s <pid>\n", argv[0]);

	pid = atoi(argv[1]);

	printf("injecting in %d\n", pid);

	if (ptrace(PT_ATTACH, pid, NULL, 0) < 0)
		err(1, "ptrace(PT_ATTACH)");

	if (ptrace_wait(pid))
		return (1);

	if (ptrace(PT_GETREGS, pid, (caddr_t)&oregs, 0) < 0)
		err(1, "ptrace(PT_GETREGS)");

	ptrace_io(pid, PIOD_READ_I, (void *)oregs.r_eip, &ocode, CODE_SIZE);
	ptrace_io(pid, PIOD_READ_D, (void *)oregs.r_esp, &ostack, STACK_SIZE);

	ptrace_io(pid, PIOD_WRITE_I, (void *)oregs.r_eip, injected,
	    injected_len);

	if (ptrace(PT_CONTINUE, pid, (caddr_t)1, 0) < 0)
		err(1, "ptrace(PT_DETACH)");

	if (ptrace_wait(pid))
		return (1);

	ptrace_io(pid, PIOD_WRITE_I, (void *)oregs.r_eip, &ocode, CODE_SIZE);
	ptrace_io(pid, PIOD_WRITE_D, (void *)oregs.r_esp, &ostack, STACK_SIZE);

	if (ptrace(PT_SETREGS, pid, (caddr_t)&oregs, 0) < 0)
		err(1, "ptrace(PT_GETREGS)");

	if (ptrace(PT_DETACH, pid, (caddr_t)1, 0) < 0)
		err(1, "ptrace(PT_DETACH)");
	return (0);
}