#include <sys/ptrace.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <stdint.h>
#include <sys/mman.h>
#include "vm.h"
#include "util.h"

static void strtrim (char *str)
{
  char *begin = str;
  while (isspace(*begin)) {
    begin++;
  }
  size_t len = strlen(begin);
  char *end = begin + len;
  while (end > begin && isspace(*end)) {
    end++;
  }

  memmove(str, begin, end - begin + 1);
}

VMMapping* parse_vmmap (int pid)
{
  char fmaps_path[128];
  sprintf(fmaps_path, "/proc/%d/maps", pid);
  FILE *fmaps;

  if ((fmaps = fopen(fmaps_path, "r")) == NULL) {
    perror("fopen");
    exit(1);
  }

  uint64_t begin, end;
  char perms[16];
  char filename[1024];
  VMMapping *head = NULL;
  VMMapping *cur = NULL;

  while(
    fscanf(fmaps, "%lx-%lx %15s %*x %*x:%*x %*u%1023[^\n]",
	   &begin, &end, perms, filename) != EOF) {
    strtrim(filename);

    VMMapping *new_mapping = xmalloc(sizeof(VMMapping));
    *new_mapping = (VMMapping) {
      .begin = (void*)begin,
      .end = (void*)end,
      .r = perms[0] == 'r',
      .w = perms[1] == 'w',
      .x = perms[2] == 'x',
      .s = perms[3] == 's',
      .p = perms[3] == 'p',
      .name = strdup(filename),
      .next = NULL,
    };

    if (head) {
      cur->next = new_mapping;
      cur = new_mapping;
    } else {
      head = cur = new_mapping;
    }
  }

  return head;
}

static off_t memfind(const uint8_t *hay, size_t hay_size, const uint8_t *needle, size_t needle_size)
{
  for (off_t i = 0 ; i < hay_size - needle_size; ++i) {
    if (memcmp(&hay[i], needle, needle_size) == 0) {
      return i;
    }
  }
  return hay_size;
}

MemscanResult* memscan(int pid, uint8_t *byte_seq, uint64_t byte_seq_len)
{
  char fmem_path[1024] = {0};
  sprintf(fmem_path, "/proc/%d/mem", pid);
  FILE *fmem = fopen(fmem_path, "rb");
  VMMapping *vmmaps_head = parse_vmmap(pid);
  VMMapping *cur_vmmap = vmmaps_head;
  MemscanResult *cur = NULL, *head = NULL;

  while (cur_vmmap) {
    if (!cur_vmmap->r) {
      cur_vmmap = cur_vmmap->next;
      continue;
    }

    size_t region_size = cur_vmmap->end - cur_vmmap->begin;
    uint8_t region_data[region_size];
    fseek(fmem, (off_t)cur_vmmap->begin, SEEK_SET);
    fread(region_data, 1, region_size, fmem);

    off_t offset = 0;
    while (offset += memfind(region_data + offset,
			     region_size - offset,
			     byte_seq,
			     byte_seq_len),
	   offset < region_size) {
      LOG("Matched pattern at [%p]\n", cur_vmmap->begin + (off_t)offset);
      MemscanResult *new_result = xmalloc(sizeof(MemscanResult));
      *new_result = (MemscanResult) {
	.mapping = cur_vmmap,
	.offset = offset,
	.next = NULL,
      };

      if (head) {
	cur->next = new_result;
	cur = new_result;
      }
      else {
	head = cur = new_result;
      }

      offset += byte_seq_len;
    }

    cur_vmmap = cur_vmmap->next;
  }
  fclose(fmem);

  return head;
}

void memwrite(int pid, void *address, uint8_t *data, size_t data_len)
{
  char fmem_path[1024] = {0};
  sprintf(fmem_path, "/proc/%d/mem", pid);
  FILE *fmem = fopen(fmem_path, "rb+");

  fseek(fmem, (off_t)address, SEEK_SET);
  fwrite(data, 1, data_len, fmem);
  fclose(fmem);

  LOG("Data written successfully at address %p\n", address);
}