/* This code is in the Public Domain. Any use for any purpose is explicitly allowed. */

#include <stdio.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/ptrace.h>
#include <assert.h>
#include <mach/mach.h>
#include <errno.h>
#include <stdlib.h>

void check(int cond, char* msg) {
  if (!cond) {
    printf("%s\n", msg);
    exit(-1);
  }
}

int main() {
  pid_t child = fork();

  if (child == 0) {
    /* PT_TRACE_ME will stop the process after the execl is executed and allows the parent
       to take control. */
    check(!ptrace(PT_TRACE_ME, 0, 0, 0), "PT_TRACE_ME failed.");
    execl("./hello", "hello", NULL);
  } else {
    /* Get the task for this pid. Seems to require superuser privileges or some gid hack.
       Go Apple! */
    mach_port_t task;
    check(task_for_pid(mach_task_self(), child, &task) == KERN_SUCCESS, "task_for_pid failed.");

    /* Get the list of threads in that process (we expect one thread exactly.) */
    thread_act_port_array_t threadList;
    mach_msg_type_number_t threadCount;
    check(task_threads(task, &threadList, &threadCount) == KERN_SUCCESS, "task_threads failed.");
    check(threadCount == 1, "task has more than one thread.");

    /* If we find the syscall we are looking for we set this flag, continue and then check
       the result of the mmap syscall, which is the address of the new mapping. */
    int after_syscall = 0;

    while (1) {
      /* Wait for updates from the child process we are tracing */
      int status;
      pid_t w = wait(&status);

      /* Read the register state of the child via Mach (since we are single-stepping the child is guaranteed to be suspended at the moment. */
      x86_thread_state64_t state;
      mach_msg_type_number_t stateCount = x86_THREAD_STATE64_COUNT;
      check(thread_get_state(threadList[0],
			     x86_THREAD_STATE64,
			     (thread_state_t)&state,
			     &stateCount) == KERN_SUCCESS, "thread_get_state failed.");

      /* If the previous step found an interesting system call, and we just stepped over it, analyze it now. */
      if (after_syscall) {
	printf("I think the file was mapped to %lx.\n", state.__rax);
	/* remote process is now stopped at the instruction after the syscall, and we can take control of it */
	exit(0); 
      }

      /* 0xc5 is mmap on Darwin, and we also check for the length of the mmap we are looking for (0x88000). Check for other criteria
	 (i.e. rsi > some_large_size_of_image) at your leasure. */
      if (((state.__rax & 0xffff) == 0xc5) && (state.__rsi == 0x88000)) {
	/* To confirm this is a syscall we have to read the machince code at that address */
	char data[16];
	vm_size_t data_count;
	check(vm_read_overwrite(task, (vm_address_t)state.__rip, 16, (vm_address_t)data, &data_count) == KERN_SUCCESS, "vm_read_overwrite failed.");
	check(data_count == 16, "vm_read_overwrite didn't return 16 bytes as expected.");

	/* System calls on Darwin use SYSCALL (0x0F 0x05) */
	if ((data[0] == 0x0f) && (data[1] == 0x05)) {
	  /* Step over the syscall and analyze the result next time around. */
	  after_syscall = 1;
	}
      }
      check(ptrace(PT_STEP, child, (char*)1, 0) == 0, "PT_STEP failed.");
    }
  }
}