#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ptrace.h>
#include <sys/reg.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <unistd.h>

#define die(...)                                                               \
  do {                                                                         \
    fprintf(stderr, __VA_ARGS__);                                              \
    fprintf(stderr, "\n");                                                     \
    exit(1);                                                                   \
  } while (0)

void read_string_ptrace(pid_t pid, void *addr, char **strout) {
  size_t strlength = 0;
  for (size_t i = 0;; i++) {
    strlength += sizeof(long);
    long d = ptrace(PTRACE_PEEKDATA, pid, addr + i * sizeof(long), NULL);
    if ((d & 0xff) == 0x00 || ((d >> 8) & 0xff) == 0x00 ||
        ((d >> 16) & 0xff) == 0x00 || ((d >> 24) & 0xff) == 0x00 ||
        ((d >> 32) & 0xff) == 0x00 || ((d >> 40) & 0xff) == 0x00 ||
        ((d >> 48) & 0xff) == 0x00 || ((d >> 56) & 0xff) == 0x00) {
      break;
    }
  }

  // add extra bytes to malloc to prevent str exactly strlength
  // for NULL byte set
  char *str = malloc(strlength + sizeof(long));
  if (str == NULL) {
    *strout = NULL;
    return;
  }
  *strout = str;

  // endianess is a problem here btw
  for (size_t i = 0; i < (strlength / sizeof(long)); i++) {
    long d = ptrace(PTRACE_PEEKDATA, pid, addr + i * sizeof(long), NULL);
    str[i * sizeof(long) + 0] = d & 0xff;
    str[i * sizeof(long) + 1] = (d >> 8) & 0xff;
    str[i * sizeof(long) + 2] = (d >> 16) & 0xff;
    str[i * sizeof(long) + 3] = (d >> 24) & 0xff;
    str[i * sizeof(long) + 4] = (d >> 32) & 0xff;
    str[i * sizeof(long) + 5] = (d >> 40) & 0xff;
    str[i * sizeof(long) + 6] = (d >> 48) & 0xff;
    str[i * sizeof(long) + 7] = (d >> 56) & 0xff;
  }
  // if the string is exactly strlength size
  str[strlength + 1] = '\0';
}

void read_string_array_ptrace(pid_t pid, void *addr, char ***strarray_out) {
  size_t num_element = 0;
  for (size_t i = 0;; i++) {
    long buffer = ptrace(PTRACE_PEEKDATA, pid, addr + i * sizeof(void *), NULL);
    if (buffer == 0) {
      break;
    }
    num_element += 1;
  }

  char **strarray = malloc(sizeof(char *) * (num_element + 1));
  *strarray_out = strarray;

  for (size_t i = 0; i < num_element; i++) {
    long buffer_addr =
        ptrace(PTRACE_PEEKDATA, pid, addr + i * sizeof(void *), NULL);
    if (buffer_addr == 0) {
      break;
    }
    read_string_ptrace(pid, (void *)buffer_addr, &strarray[i]);
  }
  strarray[num_element] = NULL;
}

void hook_process_invoke(char *program, char **argv, char **envp) {
  printf("==> %s\n", program);
  for (size_t i = 0;; i++) {
    if (argv[i] == NULL)
      break;
    printf("    argv[%zu]=%s\n", i, argv[i]);
  }
  // for (size_t i = 0;; i++) {
  //   if (envp[i] == NULL)
  //     break;
  //   printf("==> envp[%zu]=%s\n", i, envp[i]);
  // }

  for (size_t i = 0;; i++) {
    if (argv[i] == NULL)
      break;
    free(argv[i]);
  }
  for (size_t i = 0;; i++) {
    if (envp[i] == NULL)
      break;
    free(envp[i]);
  }
}

static void tracee(int argc, char **argv) {
  if (ptrace(PTRACE_TRACEME, 0, NULL, NULL) < 0)
    die("child: ptrace(traceme) failed: %m");

  /* Make sure tracer starts tracing us. */
  if (raise(SIGSTOP))
    die("child: raise(SIGSTOP) failed: %m");

  printf("[%d] spawned\n", getpid());
  /* Start the process. */
  execvp(argv[1], &argv[1]);

  /* Should never be reached. */
  die("tracee start failed: %m");
}

static void tracer(pid_t ppid) {
  int status = 0;

  if (waitpid(ppid, &status, 0) < 0)
    die("waitpid failed: %m");

  unsigned long ptrace_options = PTRACE_O_TRACECLONE | PTRACE_O_TRACEFORK |
                                 PTRACE_O_TRACEVFORK | PTRACE_O_TRACESYSGOOD;

  if (ptrace(PTRACE_SETOPTIONS, ppid, NULL, ptrace_options) < 0) {
    perror("ptrace(PTRACE_SETOPTIONS)");
    ptrace(PTRACE_DETACH, ppid, NULL, NULL);
    die("bruh bruh cannot trace pid %d", ppid);
  }

  ptrace(PTRACE_SYSCALL, ppid, 0, 0);

  while (1) {
    pid_t pid = waitpid(-1, &status, __WALL);

    if (WIFEXITED(status) || WIFSIGNALED(status)) {
      // printf("[%d] child ded\n", pid);
      if (pid == ppid) {
        // the main process dies
        break;
      }
      continue;
    }

    if (!WIFSTOPPED(status)) {
      printf("[%d] what's this?\n", pid);
      ptrace(PTRACE_CONT, pid, 0, 0);
      continue;
    }

    // special events handled first
    if (status >> 8 == (SIGTRAP | (PTRACE_EVENT_CLONE << 8)) ||
        status >> 8 == (SIGTRAP | (PTRACE_EVENT_VFORK << 8)) ||
        status >> 8 == (SIGTRAP | (PTRACE_EVENT_FORK << 8))) {
      int spawn_pid;
      if (ptrace(PTRACE_GETEVENTMSG, pid, NULL, &spawn_pid) == -1) {
        die("can't get lower child pid\n");
      }
      // printf("[%d] spawn => %d\n", pid, spawn_pid);

      ptrace(PTRACE_SYSCALL, spawn_pid, 0, 0);
      ptrace(PTRACE_SYSCALL, pid, 0, 0);
      continue;
    }

    // generic syscall events handled later
    int event_msg = 0;
    ptrace(PTRACE_GETEVENTMSG, pid, 0, &event_msg);

    if (WSTOPSIG(status) == (SIGTRAP | 0x80)) {
      struct user_regs_struct regs;
      if (ptrace(PTRACE_GETREGS, pid, NULL, &regs) == -1) {
        // TODO: this reaches, even though the stop reason is syscall
        // printf("error=%d pid=%d\n", errno, pid);
        printf("Cannot get registers\n");
        continue;
      }

      if (regs.rax != -ENOSYS) {
        // syscall exit
        ptrace(PTRACE_SYSCALL, pid, 0, 0);
        continue;
      }

      const int EXECVE_SYSCALL = 59;
      const int EXECVEAT_SYSCALL = 322;

      int syscall = regs.orig_rax;
      if (syscall == EXECVE_SYSCALL || syscall == EXECVEAT_SYSCALL) {
        printf("[%d] execve %d\n", pid, syscall);

        char *program;
        char **argv;
        char **envp;

        if (syscall == EXECVE_SYSCALL) {
          read_string_ptrace(pid, (void *)regs.rdi, &program);
          read_string_array_ptrace(pid, (void *)regs.rsi, &argv);
          read_string_array_ptrace(pid, (void *)regs.rdx, &envp);
        }
        if (syscall == EXECVEAT_SYSCALL) {
          read_string_ptrace(pid, (void *)regs.rsi, &program);
          read_string_array_ptrace(pid, (void *)regs.rdx, &argv);
          read_string_array_ptrace(pid, (void *)regs.r10, &envp);
        }

        hook_process_invoke(program, argv, envp);

        free(program);
        free(argv);
        free(envp);
      }

      // TODO: capture other syscalls also
      // renameat -> track file renames
      // printf("syscall %d\n", syscall);
    }

    // ask to wait for syscall of the currently stopped process
    ptrace(PTRACE_SYSCALL, pid, 0, 0);
  }
}

int main(int argc, char **argv) {
  pid_t pid = fork();
  if (pid < 0)
    die("couldn't fork: %m");

  if (pid == 0)
    tracee(argc, argv);
  else
    tracer(pid);

  return 0;
}