From c94c776a61a0ab9af867d77f8d8dd22b6ec405cd Mon Sep 17 00:00:00 2001
From: Gregor Haas <gregorhaas1997@gmail.com>
Date: Sun, 12 Jan 2020 17:14:21 -0500
Subject: [PATCH] Also place payload lengths in the header... not known at
 compile time

---
 c8_libpayload/scripts/headerize.py | 10 +++++++++-
 c8_libpayload/scripts/librarize.py |  2 +-
 c8_remote/CMakeLists.txt           |  4 ++--
 c8_remote/src/payload.c            |  9 ++++++++-
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/c8_libpayload/scripts/headerize.py b/c8_libpayload/scripts/headerize.py
index a0cef6e..fa9ec78 100644
--- a/c8_libpayload/scripts/headerize.py
+++ b/c8_libpayload/scripts/headerize.py
@@ -21,13 +21,21 @@ if __name__ == '__main__':
                     '#define CHECKM8_TOOL_LIBPAYLOAD_H\n',
                     '\n']
 
+    name_lines = []
+    size_lines = []
+
     for n in lib_names:
         with open(n, 'r') as f:
             line = f.readline() # looks like "const unsigned char PAYLOAD_NAME[PAYLOAD_SIZE] = "
             name = line.split(' ')[3].split('[')[0]
             size = line.split(' ')[3].split('[')[1][:-1]
 
-            header_lines.append('extern const unsigned char %s[%s];\n' % (name, size))
+            name_lines.append('extern const unsigned char %s[%s_SZ];\n' % (name, name.upper()))
+            size_lines.append('#define %s_SZ %s\n' % (name.upper(), size))
+
+    header_lines.extend(size_lines)
+    header_lines.append('\n')
+    header_lines.extend(name_lines)
 
     header_lines.append('\n')
     header_lines.append('#endif //CHECKM8_TOOL_LIBPAYLOAD_H\n')
diff --git a/c8_libpayload/scripts/librarize.py b/c8_libpayload/scripts/librarize.py
index a94eeae..78d62b4 100644
--- a/c8_libpayload/scripts/librarize.py
+++ b/c8_libpayload/scripts/librarize.py
@@ -24,7 +24,7 @@ if __name__ == '__main__':
     for n in bin_names:
         payload_name = os.path.basename(n).split('.')[0]
         with open(n, 'rb') as fbin:
-                fbytes = fbin.read()
+            fbytes = fbin.read()
 
         source_lines[payload_name].append('const unsigned char %s[%i] =\n' % (payload_name, len(fbytes)))
         source_lines[payload_name].append('\t{')
diff --git a/c8_remote/CMakeLists.txt b/c8_remote/CMakeLists.txt
index 6763206..b94b9b9 100644
--- a/c8_remote/CMakeLists.txt
+++ b/c8_remote/CMakeLists.txt
@@ -1,10 +1,10 @@
+cmake_minimum_required(VERSION 3.10)
 project(checkm8_remote C)
 
 set(CMAKE_C_STANDARD 99)
 set(CMAKE_C_FLAGS "-g -Wall")
 
-
 include_directories(include)
-add_executable(checkm8_remote main.c src/usb_helpers.c src/exploit.c src/payload.c src/command.c)
 
+add_executable(checkm8_remote main.c src/usb_helpers.c src/exploit.c src/payload.c src/command.c)
 target_link_libraries(checkm8_remote usb-1.0 pthread udev payload)
\ No newline at end of file
diff --git a/c8_remote/src/payload.c b/c8_remote/src/payload.c
index 24e6e97..4a57b55 100644
--- a/c8_remote/src/payload.c
+++ b/c8_remote/src/payload.c
@@ -24,31 +24,38 @@ struct payload *get_payload(PAYLOAD_T p)
 {
     struct payload *res;
     const unsigned char *pl;
+    int len;
 
     switch(p)
     {
         case PAYLOAD_AES:
             pl = payload_aes;
+            len = PAYLOAD_AES_SZ;
             break;
 
         case PAYLOAD_AES_BUSY:
             pl = payload_aes_busy;
+            len = PAYLOAD_AES_BUSY_SZ;
             break;
 
         case PAYLOAD_AES_SW:
             pl = payload_aes_sw;
+            len = PAYLOAD_AES_SW_SZ;
             break;
 
         case PAYLOAD_SYNC:
             pl = payload_sync;
+            len = PAYLOAD_SYNC_SZ;
             break;
 
         case PAYLOAD_SYSREG:
             pl = payload_sysreg;
+            len = PAYLOAD_SYSREG_SZ;
             break;
 
         case PAYLOAD_TASK_SLEEP_TEST:
             pl = payload_task_sleep_test;
+            len = PAYLOAD_TASK_SLEEP_TEST_SZ;
             break;
 
         default:
@@ -60,7 +67,7 @@ struct payload *get_payload(PAYLOAD_T p)
     if(res == NULL) return NULL;
 
     res->type = p;
-    res->len = sizeof(pl);
+    res->len = len;
     res->data = pl;
     res->install_base = -1;
     res->next = NULL;