tplink-safeloader: add support to split & extract firmwares

Split the oem firmware upgrade images into seperate files.
Useful when analysing oem firmware files.

Signed-off-by: Alexander Couzens <lynxis@fe80.eu>

diff --git a/src/tplink-safeloader.c b/src/tplink-safeloader.c
index c96edb4609c4e36e987bb1eb29e1240568038dc7..0a11287f875009537eef27a9cbf4e4998421ebdf 100644 (file)
--- a/src/tplink-safeloader.c
+++ b/src/tplink-safeloader.c
@@ -46,6 +46,7 @@
 
 #include <sys/types.h>
 #include <sys/stat.h>
+#include <limits.h>
 
 #include "md5.h"
 
@@ -1379,6 +1380,9 @@ static void usage(const char *argv0) {
                "Usage: %s [OPTIONS...]\n"
                "\n"
                "Options:\n"
+               "  -h              show this help\n"
+               "\n"
+               "Create a new image:\n"
                "  -B <board>      create image for the board specified with <board>\n"
                "  -k <file>       read kernel image from the file <file>\n"
                "  -r <file>       read rootfs image from the file <file>\n"
@@ -1386,7 +1390,9 @@ static void usage(const char *argv0) {
                "  -V <rev>        sets the revision number to <rev>\n"
                "  -j              add jffs2 end-of-filesystem markers\n"
                "  -S              create sysupgrade instead of factory image\n"
-               "  -h              show this help\n",
+               "Extract an old image:\n"
+               "  -x <file>       extract all oem firmware partition\n"
+               "  -d <dir>        destination to extract the firmware partition\n",
                argv0
        );
 };
@@ -1403,8 +1409,232 @@ static const struct device_info *find_board(const char *id)
        return NULL;
 }
 
+static int add_flash_partition(
+               struct flash_partition_entry *part_list,
+               size_t max_entries,
+               const char *name,
+               unsigned long base,
+               unsigned long size)
+{
+       int ptr = 0;
+       /* check if the list has a free entry */
+       for (int ptr=0; ptr<max_entries; ptr++, part_list++) {
+               if (part_list->name == NULL &&
+                               part_list->base == 0 &&
+                               part_list->size == 0)
+                       break;
+       }
+
+       if (ptr == max_entries) {
+               error(1, 0, "No free flash part entry available.");
+       }
+
+       part_list->name = calloc(1, strlen(name) + 1);
+       memcpy((char *)part_list->name, name, strlen(name));
+       part_list->base = base;
+       part_list->size = size;
+
+       return 0;
+}
+
+/** read the partition table into struct flash_partition_entry */
+static int read_partition_table(
+               FILE *file, long offset,
+               struct flash_partition_entry *entries, size_t max_entries,
+               int type)
+{
+       char buf[2048];
+       char *ptr, *end;
+       const char *parthdr = NULL;
+       const char *fwuphdr = "fwup-ptn";
+       const char *flashhdr = "partition";
+
+       /* TODO: search for the partition table */
+
+       switch(type) {
+               case 0:
+                       parthdr = fwuphdr;
+                       break;
+               case 1:
+                       parthdr = flashhdr;
+                       break;
+               default:
+                       error(1, 0, "Invalid partition table");
+       }
+
+       if (fseek(file, offset, SEEK_SET) < 0)
+               error(1, errno, "Can not seek in the firmware");
+
+       if (fread(buf, 1, 2048, file) < 0)
+               error(1, errno, "Can not read fwup-ptn from the firmware");
+
+       buf[2047] = '\0';
+
+       /* look for the partition header */
+       if (memcmp(buf, parthdr, strlen(parthdr)) != 0) {
+               fprintf(stderr, "DEBUG: can not find fwuphdr\n");
+               return 1;
+       }
+
+       ptr = buf;
+       end = buf + sizeof(buf);
+       while ((ptr + strlen(parthdr)) < end &&
+                       memcmp(ptr, parthdr, strlen(parthdr)) == 0) {
+               char *end_part;
+               char *end_element;
+
+               char name[32] = { 0 };
+               int name_len = 0;
+               unsigned long base = 0;
+               unsigned long size = 0;
+
+               end_part = memchr(ptr, '\n', (end - ptr));
+               if (end_part == NULL) {
+                       /* in theory this should never happen, because a partition always ends with 0x09, 0x0D, 0x0A */
+                       break;
+               }
+
+               for (int i=0; i<=4; i++) {
+                       if (end_part <= ptr)
+                               break;
+
+                       end_element = memchr(ptr, 0x20, (end_part - ptr));
+                       if (end_element == NULL) {
+                               error(1, errno, "Ignoring the rest of the partition entries.");
+                               break;
+                       }
+
+                       switch (i) {
+                               /* partition header */
+                               case 0:
+                                       ptr = end_element + 1;
+                                       continue;
+                               /* name */
+                               case 1:
+                                       name_len = (end_element - ptr) > 31 ? 31 : (end_element - ptr);
+                                       strncpy(name, ptr, name_len);
+                                       name[name_len] = '\0';
+                                       ptr = end_element + 1;
+                                       continue;
+
+                               /* string "base" */
+                               case 2:
+                                       ptr = end_element + 1;
+                                       continue;
+
+                               /* actual base */
+                               case 3:
+                                       base = strtoul(ptr, NULL, 16);
+                                       ptr = end_element + 1;
+                                       continue;
+
+                               /* string "size" */
+                               case 4:
+                                       ptr = end_element + 1;
+                                       /* actual size. The last element doesn't have a sepeartor */
+                                       size = strtoul(ptr, NULL, 16);
+                                       /* the part ends with 0x09, 0x0d, 0x0a */
+                                       ptr = end_part + 1;
+                                       add_flash_partition(entries, max_entries, name, base, size);
+                                       continue;
+                       }
+               }
+       }
+
+       return 0;
+}
+
+static void write_partition(
+               FILE *input_file,
+               size_t firmware_offset,
+               struct flash_partition_entry *entry,
+               FILE *output_file)
+{
+       char buf[4096];
+       size_t offset;
+
+       fseek(input_file, entry->base + firmware_offset, SEEK_SET);
+
+       for (offset = 0; sizeof(buf) + offset <= entry->size; offset += sizeof(buf)) {
+               if (fread(buf, sizeof(buf), 1, input_file) < 0)
+                       error(1, errno, "Can not read partition from input_file");
+
+               if (fwrite(buf, sizeof(buf), 1, output_file) < 0)
+                       error(1, errno, "Can not write partition to output_file");
+       }
+       /* write last chunk smaller than buffer */
+       if (offset < entry->size) {
+               offset = entry->size - offset;
+               if (fread(buf, offset, 1, input_file) < 0)
+                       error(1, errno, "Can not read partition from input_file");
+               if (fwrite(buf, offset, 1, output_file) < 0)
+                       error(1, errno, "Can not write partition to output_file");
+       }
+}
+
+static int extract_firmware_partition(FILE *input_file, size_t firmware_offset, struct flash_partition_entry *entry, const char *output_directory)
+{
+       FILE *output_file;
+       char output[PATH_MAX];
+
+       snprintf(output, PATH_MAX, "%s/%s", output_directory, entry->name);
+       output_file = fopen(output, "wb+");
+       if (output_file == NULL) {
+               error(1, errno, "Can not open output file %s", output);
+       }
+
+       write_partition(input_file, firmware_offset, entry, output_file);
+
+       fclose(output_file);
+
+       return 0;
+}
+
+/** extract all partitions from the firmware file */
+static int extract_firmware(const char *input, const char *output_directory)
+{
+       struct flash_partition_entry entries[16] = { 0 };
+       size_t max_entries = 16;
+       size_t firmware_offset = 0x1014;
+       FILE *input_file;
+
+       struct stat statbuf;
+
+       /* check input file */
+       if (stat(input, &statbuf)) {
+               error(1, errno, "Can not read input firmware %s", input);
+       }
+
+       /* check if output directory exists */
+       if (stat(output_directory, &statbuf)) {
+               error(1, errno, "Failed to stat output directory %s", output_directory);
+       }
+
+       if ((statbuf.st_mode & S_IFMT) != S_IFDIR) {
+               error(1, errno, "Given output directory is not a directory %s", output_directory);
+       }
+
+       input_file = fopen(input, "rb");
+
+       if (read_partition_table(input_file, firmware_offset, entries, 16, 0) != 0) {
+               error(1, 0, "Error can not read the partition table (fwup-ptn)");
+       }
+
+       for (int i=0; i<max_entries; i++) {
+               if (entries[i].name == NULL &&
+                               entries[i].base == 0 &&
+                               entries[i].size == 0)
+                       continue;
+
+               extract_firmware_partition(input_file, firmware_offset, &entries[i], output_directory);
+       }
+
+       return 0;
+}
+
 int main(int argc, char *argv[]) {
        const char *board = NULL, *kernel_image = NULL, *rootfs_image = NULL, *output = NULL;
+       const char *extract_image = NULL, *output_directory = NULL;
        bool add_jffs2_eof = false, sysupgrade = false;
        unsigned rev = 0;
        const struct device_info *info;
@@ -1413,7 +1643,7 @@ int main(int argc, char *argv[]) {
        while (true) {
                int c;
 
-               c = getopt(argc, argv, "B:k:r:o:V:jSh");
+               c = getopt(argc, argv, "B:k:r:o:V:jSh:x:d:");
                if (c == -1)
                        break;
 
@@ -1450,27 +1680,43 @@ int main(int argc, char *argv[]) {
                        usage(argv[0]);
                        return 0;
 
+               case 'd':
+                       output_directory = optarg;
+                       break;
+
+               case 'x':
+                       extract_image = optarg;
+                       break;
+
                default:
                        usage(argv[0]);
                        return 1;
                }
        }
 
-       if (!board)
-               error(1, 0, "no board has been specified");
-       if (!kernel_image)
-               error(1, 0, "no kernel image has been specified");
-       if (!rootfs_image)
-               error(1, 0, "no rootfs image has been specified");
-       if (!output)
-               error(1, 0, "no output filename has been specified");
-
-       info = find_board(board);
-
-       if (info == NULL)
-               error(1, 0, "unsupported board %s", board);
-
-       build_image(output, kernel_image, rootfs_image, rev, add_jffs2_eof, sysupgrade, info);
+       if (extract_image || output_directory) {
+               if (!extract_image)
+                       error(1, 0, "No factory/oem image given via -x <file>. Output directory is only valid with -x");
+               if (!output_directory)
+                       error(1, 0, "Can not extract an image without output directory. Use -d <dir>");
+               extract_firmware(extract_image, output_directory);
+       } else {
+               if (!board)
+                       error(1, 0, "no board has been specified");
+               if (!kernel_image)
+                       error(1, 0, "no kernel image has been specified");
+               if (!rootfs_image)
+                       error(1, 0, "no rootfs image has been specified");
+               if (!output)
+                       error(1, 0, "no output filename has been specified");
+
+               info = find_board(board);
+
+               if (info == NULL)
+                       error(1, 0, "unsupported board %s", board);
+
+               build_image(output, kernel_image, rootfs_image, rev, add_jffs2_eof, sysupgrade, info);
+       }
 
        return 0;
 }