#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 pid) {
  int status = 0;

  // preliminary stop
  // works for both SIGSTOP of the main tracee
  // as well as their children when stopped by ptrace trace fork
  if (waitpid(pid, &status, 0) < 0)
    die("waitpid failed: %m");
  if (!WIFSTOPPED(status) || WSTOPSIG(status) != SIGSTOP) {
    kill(pid, SIGKILL);
    die("tracer: unexpected wait status: %x", status);
  }

  // capture these "fork" syscalls
  // and then capture execve syscall
  // fork vfork clone clone3
  unsigned long ptrace_options = PTRACE_O_TRACECLONE | PTRACE_O_TRACEFORK |
                                 PTRACE_O_TRACEVFORK | PTRACE_O_TRACESYSGOOD;

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

  ptrace(PTRACE_CONT, pid, 0, WSTOPSIG(status));

  int sysenter = 0;
  while (1) {
    ptrace(PTRACE_SYSCALL, pid, 0, 0);
    waitpid(pid, &status, __WALL);

    if (WIFEXITED(status) || WIFSIGNALED(status)) {
      printf("[%d] child ded\n", pid);
      break;
    }

    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))) {
      printf("[%d] clone\n", pid);
      // ptrace(PTRACE_CONT, pid, 0, 0);
      continue;
    }
    if (status >> 8 == (SIGTRAP | (PTRACE_EVENT_FORK << 8))) {
      printf("[%d] fork\n", pid);
      // ptrace(PTRACE_CONT, pid, 0, 0);
      continue;
    }
    if (status >> 8 == (SIGTRAP | (PTRACE_EVENT_VFORK << 8))) {
      printf("[%d] vfork\n", pid);
      int spawn_pid;
      if (ptrace(PTRACE_GETEVENTMSG, pid, NULL, &spawn_pid) == -1) {
        die("can't get lower child pid\n");
      }
      printf("=> %d\n", spawn_pid);
      tracer(spawn_pid);

      // ptrace(PTRACE_CONT, 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)) {
      if (sysenter) {
        sysenter = 1;
        continue;
      }

      struct user_regs_struct regs;
      if (ptrace(PTRACE_GETREGS, pid, NULL, &regs) == -1) {
        die("Cannot get registers");
        return;
      }

      int syscall = regs.orig_rax;
      const int EXECVE_SYSCALL = 59;
      if (syscall == EXECVE_SYSCALL) {
        sysenter = 1;
        printf("[%d] execve %d\n", pid, syscall);

        // RDI: const char *filename
        // RSI: char *const argv[]
        // RDX: char *const envp[]

        char *program;
        read_string_ptrace(pid, (void *)regs.rdi, &program);

        char **argv;
        read_string_array_ptrace(pid, (void *)regs.rsi, &argv);

        char **envp;
        read_string_array_ptrace(pid, (void *)regs.rdx, &envp);

        hook_process_invoke(program, argv, envp);

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

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;
}