From c79bbdb0448501987c0c16c2877c780143156d1e Mon Sep 17 00:00:00 2001
From: Orfeas <38209077+0xfea5@users.noreply.github.com>
Date: Sun, 21 Apr 2024 16:27:55 +0300
Subject: Changes in trainer interface, bug fixes & helper scripts

---
 src/main.c | 125 +++++++++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 88 insertions(+), 37 deletions(-)

(limited to 'src/main.c')

diff --git a/src/main.c b/src/main.c
index 667ca4c..623d767 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1,5 +1,4 @@
 #include <string.h>
-#include <stdint.h>
 #include <stdlib.h>
 #include <stdio.h>
 #include <sys/ptrace.h>
@@ -8,24 +7,94 @@
 #include "util.h"
 #include "vm.h"
 
-void hex2bytes(char *hex)
+Bytes hex2bytes(const char *hex)
 {
-  size_t len = strlen(hex);
-  char bytes[len + 1];
+  size_t hlen = strlen(hex);
+  Bytes bytes = {
+    .data = xmalloc(hlen / 2),
+    .len = hlen / 2,
+  };
 
-  for (size_t i = 0; i < len; ++i) {
+  for (size_t i = 0; i < hlen; ++i) {
     char hdig[3] = { hex[i*2], hex[i*2+1], '\0' };
-    sscanf(hdig, "%hhx", &bytes[i]);
+    sscanf(hdig, "%hhx", &bytes.data[i]);
   }
 
-  memcpy(hex, bytes, len);
-  hex[len] = '\0';
+  return bytes;
 }
 
-int main(int argc, char *argv[])
+void action_scan(int pid, int argc, const char *argv[])
 {
-  if (argc < 2) {
-    ERROR("Usage: %s <tracee_pid>\n", argv[0]);
+  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;
@@ -37,32 +106,14 @@ int main(int argc, char *argv[])
   waitpid(pid, NULL, __WALL);
   LOG("Attached to process %d\n", pid);
 
-  char *byte_seq = "DEADBEEF";
-  size_t byte_seq_len = strlen(byte_seq);
-  MemscanResult *head = memscan(pid, (uint8_t*)byte_seq, byte_seq_len);
-  MemscanResult *cur = head;
-
-  printf("\n\n\nMemory scan results:\n");
-  printf("%-16s|%-16s|%-10s|%-s\n", "Address", "Base", "Offset", "Name");
-  puts("--------------------------------------------------");
-  while (cur) {
-    printf("%-16p|%-16p|%#-10lx|%-s\n",
-	   cur->mapping->begin + cur->offset,
-	   cur->mapping->begin,
-	   cur->offset,
-	   cur->mapping->name);
-    cur = cur->next;
-  }
-
-  printf("\n\n");
-
-  char *buf = "CAFEBABE";
-  size_t len = strlen(buf);
-  cur = head;
-  while (cur) {
-    void *address = cur->mapping->begin + cur->offset;
-    memwrite(pid, address, (uint8_t*)buf, len);
-    cur = cur->next;
+  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);
-- 
cgit v1.2.3