#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <unistd.h>
#include "util.h"
#include "vm.h"

Bytes hex2bytes(const char *hex)
{
  size_t hlen = strlen(hex);
  Bytes bytes = {
    .data = xmalloc(hlen / 2),
    .len = hlen / 2,
  };

  for (size_t i = 0; i < hlen; ++i) {
    char hdig[3] = { hex[i*2], hex[i*2+1], '\0' };
    sscanf(hdig, "%hhx", &bytes.data[i]);
  }

  return bytes;
}

void action_scan(int pid, int argc, const char *argv[])
{
  if (argc == 0) {
    ERROR("Scan: Missing argument <byte_sequence>\n");
  }

  Bytes aob = hex2bytes(argv[0]);
  MemscanResult *results_head =  memscan(pid, aob);
  MemscanResult *cur = results_head;
  while (cur) {
    void *address = cur->mapping->begin + cur->offset;
    printf("%p\n", address);
    cur = cur->next;
  }
}

void action_read(int pid, int argc, const char *argv[])
{
  if (argc == 0) {
    ERROR("Write: Missing argument <nbytes>\n");
  }

  if (argc == 1) {
    ERROR("Read: Missing argument(s) [<address> ...]");
  }

  size_t nbytes = atol(argv[0]);
  if (nbytes == 0) {
    perror("atol");
    exit(1);
  }

  void *address[argc-1];
  for (size_t i = 0; i < argc-1; ++i) {
    sscanf(argv[i+1], "%p", &address[i]);
  }

  for (size_t i = 0; i < argc-1; ++i) {
    Bytes bytes = memread(pid, address[i], nbytes);
    for (size_t j = 0; j < bytes.len; ++j) {
      printf("%02hhx", bytes.data[j]);
    }
    printf("\n");
  }
}

void action_write(int pid, int argc, const char *argv[])
{
  if (argc == 0) {
    ERROR("Write: Missing argument <bytes>\n");
  }

  if (argc == 1) {
    ERROR("Write: Missing argument [<address> ...]\n");
  }

  Bytes aob = hex2bytes(argv[0]);
  printf("%s\n", aob.data);
  void *address[argc-1];
  for (size_t i = 0; i < argc-1; ++i) {
    sscanf(argv[i+1], "%p", &address[i]);
  }

  for (size_t i = 0; i < argc-1; ++i) {
    memwrite(pid, address[i], aob);
  }
}

int main(int argc, const char *argv[])
{
  if (argc < 3) {
    ERROR("Usage: %s <tracee_pid> (scan|read|write) [args ...]\n", argv[0]);
  }

  int pid;
  if ((pid = atoi(argv[1])) == 0) {
    ERROR("Invalid pid '%s'\n", argv[1]);
  }

  ptrace(PTRACE_ATTACH, pid, NULL, NULL);
  waitpid(pid, NULL, __WALL);
  LOG("Attached to process %d\n", pid);

  if (strcmp(argv[2], "scan") == 0) {
    action_scan(pid, argc-3, &argv[3]);
  } else if (strcmp(argv[2], "read") == 0) {
    action_read(pid, argc-3, &argv[3]);
  } else if (strcmp(argv[2], "write") == 0) {
    action_write(pid, argc-3, &argv[3]);
  } else {
    ERROR("Unknown option '%s'\n", argv[1]);
  }

  ptrace(PTRACE_DETACH, pid, NULL, NULL);
  LOG("Detached from process %d\n", pid);

  return 0;
}