imxtools: completely rework patching logic to prepare unpatching

The old code was working but a mess to maintain. The new code is cleaner
and always simpler handling of all the different options. Extraction of the
OF is no longer a standalone function but just one particular output type.
This commit prepares the ground for firmware "unpatching" (aka OF extraction
from patched OF). The patching code itself did not change so this commit
should still produce the exact same images as before.

Change-Id: I3840793d4b78b8435e38c08f558840925085ead1

3 files changed, 211 insertions, 178 deletions

diff --git a/rbutil/mkimxboot/main.c b/rbutil/mkimxboot/main.c
index 611715f7c1..90f2f20a5f 100644
--- a/rbutil/mkimxboot/main.c
+++ b/rbutil/mkimxboot/main.c
@@ -72,10 +72,9 @@ static void usage(void)
     printf("  -i <file>   Set input file\n");
     printf("  -b <file>   Set boot file\n");
     printf("  -d/--debug  Enable debug output\n");
-    printf("  -t <type>   Set type (dualboot, singleboot, recovery, charge)\n");
+    printf("  -t <type>   Set type (dualboot, singleboot, recovery, origfw, charge)\n");
     printf("  -v <v>      Set variant\n");
     printf("  -x          Dump device informations\n");
-    printf("  -w          Extract the original firmware\n");
     printf("  -p <ver>    Force product and component version\n");
     printf("  -5 <type>   Compute <type> MD5 sum of the input file\n");
     printf("  -m <model>  Specify model (useful for soft MD5 sum)\n");
@@ -108,14 +107,14 @@ static void usage(void)
     exit(1);
 }
 
-static int print_md5(const char *file, enum imx_model_t model, const char *type)
+static int print_md5(const char *file, const char *type)
 {
     uint8_t md5sum[16];
     enum imx_error_t err;
     if(strcmp(type, "full") == 0)
         err = compute_md5sum(file, md5sum);
     else if(strcmp(type, "soft") == 0)
-        err = compute_soft_md5sum(file, model, md5sum);
+        err = compute_soft_md5sum(file, md5sum);
     else
     {
         printf("Invalid md5sum type '%s'\n", type);
@@ -142,7 +141,6 @@ int main(int argc, char *argv[])
     enum imx_output_type_t type = IMX_DUALBOOT;
     enum imx_model_t model = MODEL_UNKNOWN;
     bool debug = false;
-    bool extract_of = false;
     const char *md5type = NULL;
     const char *force_version = NULL;
 
@@ -153,7 +151,7 @@ int main(int argc, char *argv[])
     {
         static struct option long_options[] =
         {
-            {"help", no_argument, 0, '?'},
+            {"help", no_argument, 0, 'h'},
             {"in-file", no_argument, 0, 'i'},
             {"out-file", required_argument, 0, 'o'},
             {"boot-file", required_argument, 0, 'b'},
@@ -166,7 +164,7 @@ int main(int argc, char *argv[])
             {0, 0, 0, 0}
         };
 
-        int c = getopt_long(argc, argv, "?di:o:b:t:v:xwp:m:5:", long_options, NULL);
+        int c = getopt_long(argc, argv, "hdi:o:b:t:v:xp:m:5:", long_options, NULL);
         if(c == -1)
             break;
         switch(c)
@@ -174,7 +172,7 @@ int main(int argc, char *argv[])
             case 'd':
                 debug = true;
                 break;
-            case '?':
+            case 'h':
                 usage();
                 break;
             case 'o':
@@ -195,6 +193,8 @@ int main(int argc, char *argv[])
                     type = IMX_RECOVERY;
                 else if(strcmp(optarg, "charge") == 0)
                     type = IMX_CHARGE;
+                else if(strcmp(optarg, "origfw") == 0)
+                    type = IMX_ORIG_FW;
                 else
                 {
                     printf("Invalid boot type '%s'\n", optarg);
@@ -223,9 +223,6 @@ int main(int argc, char *argv[])
                 for(int i = 0; i < sizeof(imx_variants) / sizeof(imx_variants[0]); i++)
                     printf("  %s -> variant=%d\n", imx_variants[i].name, imx_variants[i].variant);
                 break;
-            case 'w':
-                extract_of = true;
-                break;
             case 'p':
                 force_version = optarg;
                 break;
@@ -266,24 +263,19 @@ int main(int argc, char *argv[])
         printf("You must specify an output file\n");
         return 1;
     }
-    if(!bootfile && !extract_of)
+
+    if(!bootfile && type != IMX_ORIG_FW)
     {
         printf("You must specify an boot file\n");
         return 1;
     }
+
     if(optind != argc)
     {
         printf("Extra arguments on command line\n");
         return 1;
     }
 
-    if(extract_of)
-    {
-        enum imx_error_t err = extract_firmware(infile, variant, outfile);
-        printf("Result: %d\n", err);
-        return 0;
-    }
-
     struct imx_option_t opt;
     memset(&opt, 0, sizeof(opt));
     opt.debug = debug;
diff --git a/rbutil/mkimxboot/mkimxboot.c b/rbutil/mkimxboot/mkimxboot.c
index 909dad8641..1d1ec54e10 100644
--- a/rbutil/mkimxboot/mkimxboot.c
+++ b/rbutil/mkimxboot/mkimxboot.c
@@ -38,6 +38,8 @@ struct rb_fw_t
     int entry_idx;
 };
 
+/* A firmware upgrade can contains several variants like recovery image, or
+ * images for different models */
 struct imx_fw_variant_desc_t
 {
     /* Offset within file */
@@ -46,6 +48,7 @@ struct imx_fw_variant_desc_t
     size_t size;
 };
 
+/* Map a MD5 sum of the whole file to a model and describe the variants in it */
 struct imx_md5sum_t
 {
     /* Device model */
@@ -58,6 +61,7 @@ struct imx_md5sum_t
     struct imx_fw_variant_desc_t fw_variants[VARIANT_COUNT];
 };
 
+/* Describe how to produce a bootloader image for a specific model */
 struct imx_model_desc_t
 {
     /* Descriptive name of this model */
@@ -72,16 +76,15 @@ struct imx_model_desc_t
     /* Model number used to initialise the checksum in the Rockbox header in
        ".sansa" files - these are the same as MODEL_NUMBER in config-target.h */
     const int rb_model_num;
-    /* Number of keys needed to decrypt/encrypt */
-    int nr_keys;
-    /* Array of keys */
-    struct crypto_key_t *keys;
+    /* Array of NULL-terminated keys */
+    struct crypto_key_t **keys;
     /* Dualboot load address */
     uint32_t dualboot_addr;
     /* Bootloader load address */
     uint32_t bootloader_addr;
 };
 
+/* Friendly names for variants */
 static const char *imx_fw_variant[] =
 {
     [VARIANT_DEFAULT] = "default",
@@ -92,6 +95,7 @@ static const char *imx_fw_variant[] =
     [VARIANT_ZENSTYLE_RECOVERY] = "ZEN Style 100/300 Recovery",
 };
 
+/* List of known MD5 sums for firmware upgrades */
 static const struct imx_md5sum_t imx_sums[] =
 {
     /** Fuze+ */
@@ -196,22 +200,24 @@ static struct crypto_key_t zero_key =
     .u.key = {0}
 };
 
+static struct crypto_key_t *list_zero_key[] = { &zero_key, NULL };
+static struct crypto_key_t *list_all_keys[] = { &zero_key, NULL };
+
 static const struct imx_model_desc_t imx_models[] =
 {
-    [MODEL_FUZEPLUS] = {"Fuze+",  dualboot_fuzeplus, sizeof(dualboot_fuzeplus), "fuz+", 72,
-                          1, &zero_key, 0, 0x40000000 },
-    [MODEL_ZENXFI2] = {"Zen X-Fi2", dualboot_zenxfi2, sizeof(dualboot_zenxfi2), "zxf2", 82,
-                       1, &zero_key, 0, 0x40000000 },
-    [MODEL_ZENXFI3] = {"Zen X-Fi3", dualboot_zenxfi3, sizeof(dualboot_zenxfi3), "zxf3", 83,
-                       1, &zero_key, 0, 0x40000000 },
-    [MODEL_ZENXFISTYLE] = {"Zen X-Fi Style", dualboot_zenxfistyle, sizeof(dualboot_zenxfistyle), "zxfs", 94,
-                       1, &zero_key, 0, 0x40000000 },
-    [MODEL_ZENSTYLE] = {"Zen Style 100/300", NULL, 0, "", -1,
-                       1, &zero_key, 0, 0x40000000 },
-    [MODEL_NWZE370] = {"NWZ-E370", dualboot_nwze370, sizeof(dualboot_nwze370), "e370", 88,
-                       1, &zero_key, 0, 0x40000000 },
-    [MODEL_NWZE360] = {"NWZ-E360", dualboot_nwze360, sizeof(dualboot_nwze360), "e360", 89,
-                       1, &zero_key, 0, 0x40000000 },
+    [MODEL_FUZEPLUS] = {"Fuze+",  dualboot_fuzeplus, sizeof(dualboot_fuzeplus),
+        "fuz+", 72, list_zero_key, 0, 0x40000000 },
+    [MODEL_ZENXFI2] = {"Zen X-Fi2", dualboot_zenxfi2, sizeof(dualboot_zenxfi2),
+        "zxf2", 82, list_zero_key, 0, 0x40000000 },
+    [MODEL_ZENXFI3] = {"Zen X-Fi3", dualboot_zenxfi3, sizeof(dualboot_zenxfi3),
+        "zxf3", 83, list_zero_key, 0, 0x40000000 },
+    [MODEL_ZENXFISTYLE] = {"Zen X-Fi Style", dualboot_zenxfistyle, sizeof(dualboot_zenxfistyle),
+        "zxfs", 94, list_zero_key, 0, 0x40000000 },
+    [MODEL_ZENSTYLE] = {"Zen Style 100/300", NULL, 0, "", -1, list_zero_key, 0, 0x40000000 },
+    [MODEL_NWZE370] = {"NWZ-E370", dualboot_nwze370, sizeof(dualboot_nwze370),
+        "e370", 88, list_zero_key, 0, 0x40000000 },
+    [MODEL_NWZE360] = {"NWZ-E360", dualboot_nwze360, sizeof(dualboot_nwze360),
+        "e360", 89, list_zero_key, 0, 0x40000000 },
 };
 
 #define NR_IMX_SUMS     (sizeof(imx_sums) / sizeof(imx_sums[0]))
@@ -222,6 +228,12 @@ static const struct imx_model_desc_t imx_models[] =
 #define MAGIC_NORMAL    0xcafebabe
 #define MAGIC_CHARGE    0x67726863 /* 'chrg' */
 
+static void add_key_list(struct crypto_key_t **list)
+{
+    while(*list != NULL)
+        add_keys(*list++, 1);
+}
+
 static int rb_fw_get_sb_inst_count(struct rb_fw_t *fw)
 {
     return fw->nr_insts;
@@ -240,8 +252,8 @@ static void rb_fw_fill_sb(struct rb_fw_t *fw, struct sb_inst_t *inst,
     inst[fw->entry_idx].argument = entry_arg;
 }
 
-static enum imx_error_t patch_std_zero_host_play(int jump_before, int model,
-    enum imx_output_type_t type, struct sb_file_t *sb_file, struct rb_fw_t boot_fw)
+static enum imx_error_t patch_std_zero_host_play(int jump_before,
+    struct imx_option_t opt, struct sb_file_t *sb_file, struct rb_fw_t boot_fw)
 {
     /* We assume the file has three boot sections: ____, host, play and one
      * resource section rsrc.
@@ -277,7 +289,7 @@ static enum imx_error_t patch_std_zero_host_play(int jump_before, int model,
         return IMX_DONT_KNOW_HOW_TO_PATCH;
     }
 
-    if(type == IMX_DUALBOOT)
+    if(opt.output == IMX_DUALBOOT)
     {
         /* create a new instruction array with a hole for two instructions */
         struct sb_inst_t *new_insts = xmalloc(sizeof(struct sb_inst_t) * (sec->nr_insts + 2));
@@ -288,15 +300,16 @@ static enum imx_error_t patch_std_zero_host_play(int jump_before, int model,
         struct sb_inst_t *load = &new_insts[jump_idx];
         memset(load, 0, sizeof(struct sb_inst_t));
         load->inst = SB_INST_LOAD;
-        load->size = imx_models[model].dualboot_size;
-        load->addr = imx_models[model].dualboot_addr;
+        load->size = imx_models[opt.model].dualboot_size;
+        load->addr = imx_models[opt.model].dualboot_addr;
         /* duplicate memory because it will be free'd */
-        load->data = memdup(imx_models[model].dualboot, imx_models[model].dualboot_size);
+        load->data = memdup(imx_models[opt.model].dualboot,
+            imx_models[opt.model].dualboot_size);
         /* second instruction is a call */
         struct sb_inst_t *call = &new_insts[jump_idx + 1];
         memset(call, 0, sizeof(struct sb_inst_t));
         call->inst = SB_INST_CALL;
-        call->addr = imx_models[model].dualboot_addr;
+        call->addr = imx_models[opt.model].dualboot_addr;
         call->argument = MAGIC_ROCK;
         /* free old instruction array */
         free(sec->insts);
@@ -320,9 +333,9 @@ static enum imx_error_t patch_std_zero_host_play(int jump_before, int model,
 
         return IMX_SUCCESS;
     }
-    else if(type == IMX_SINGLEBOOT || type == IMX_RECOVERY)
+    else if(opt.output == IMX_SINGLEBOOT || opt.output == IMX_RECOVERY)
     {
-        bool recovery = type == IMX_RECOVERY;
+        bool recovery = (opt.output == IMX_RECOVERY);
         /* remove everything after the call and add instructions for firmware */
         struct sb_inst_t *new_insts = xmalloc(sizeof(struct sb_inst_t) * (jump_idx + nr_boot_inst));
         memcpy(new_insts, sec->insts, sizeof(struct sb_inst_t) * jump_idx);
@@ -344,7 +357,7 @@ static enum imx_error_t patch_std_zero_host_play(int jump_before, int model,
 
         return IMX_SUCCESS;
     }
-    else if(type == IMX_CHARGE)
+    else if(opt.output == IMX_CHARGE)
     {
         /* throw away everything except the dualboot stub with a special argument */
         struct sb_inst_t *new_insts = xmalloc(sizeof(struct sb_inst_t) * 2);
@@ -352,15 +365,16 @@ static enum imx_error_t patch_std_zero_host_play(int jump_before, int model,
         struct sb_inst_t *load = &new_insts[0];
         memset(load, 0, sizeof(struct sb_inst_t));
         load->inst = SB_INST_LOAD;
-        load->size = imx_models[model].dualboot_size;
-        load->addr = imx_models[model].dualboot_addr;
+        load->size = imx_models[opt.model].dualboot_size;
+        load->addr = imx_models[opt.model].dualboot_addr;
         /* duplicate memory because it will be free'd */
-        load->data = memdup(imx_models[model].dualboot, imx_models[model].dualboot_size);
+        load->data = memdup(imx_models[opt.model].dualboot,
+            imx_models[opt.model].dualboot_size);
         /* second instruction is a call */
         struct sb_inst_t *call = &new_insts[1];
         memset(call, 0, sizeof(struct sb_inst_t));
         call->inst = SB_INST_CALL;
-        call->addr = imx_models[model].dualboot_addr;
+        call->addr = imx_models[opt.model].dualboot_addr;
         call->argument = MAGIC_CHARGE;
         /* free old instruction array */
         free(sec->insts);
@@ -428,42 +442,42 @@ static enum imx_error_t parse_version(const char *s, struct sb_version_t *ver)
     return IMX_SUCCESS;
 }
 
-static enum imx_error_t patch_firmware(enum imx_model_t model,
-    enum imx_firmware_variant_t variant, enum imx_output_type_t type,
-    struct sb_file_t *sb_file, struct rb_fw_t boot_fw,
-    const char *force_version)
+static enum imx_error_t patch_firmware(struct imx_option_t opt,
+    struct sb_file_t *sb_file, struct rb_fw_t boot_fw)
 {
-    if(force_version)
+    if(opt.force_version)
     {
-        enum imx_error_t err = parse_version(force_version, &sb_file->product_ver);
-        if(err != IMX_SUCCESS) return err;
-        err = parse_version(force_version, &sb_file->component_ver);
-        if(err != IMX_SUCCESS) return err;
+        enum imx_error_t err = parse_version(opt.force_version, &sb_file->product_ver);
+        if(err != IMX_SUCCESS)
+            return err;
+        err = parse_version(opt.force_version, &sb_file->component_ver);
+        if(err != IMX_SUCCESS)
+            return err;
     }
-    switch(model)
+    switch(opt.model)
     {
         case MODEL_FUZEPLUS:
             /* The Fuze+ uses the standard ____, host, play sections, patch after third
              * call in ____ section */
-            return patch_std_zero_host_play(3, model, type, sb_file, boot_fw);
+            return patch_std_zero_host_play(3, opt, sb_file, boot_fw);
         case MODEL_ZENXFI3:
             /* The ZEN X-Fi3 uses the standard ____, hSst, pSay sections, patch after third
              * call in ____ section. Although sections names use the S variant, they are standard. */
-            return patch_std_zero_host_play(3, model, type, sb_file, boot_fw);
+            return patch_std_zero_host_play(3, opt, sb_file, boot_fw);
         case MODEL_NWZE360:
         case MODEL_NWZE370:
             /* The NWZ-E360/E370 uses the standard ____, host, play sections, patch after first
              * call in ____ section. */
-            return patch_std_zero_host_play(1, model, type, sb_file, boot_fw);
+            return patch_std_zero_host_play(1, opt, sb_file, boot_fw);
         case MODEL_ZENXFI2:
             /* The ZEN X-Fi2 has two types of firmware: recovery and normal.
              * Normal uses the standard ___, host, play sections and recovery only ____ */
-            switch(variant)
+            switch(opt.fw_variant)
             {
                 case VARIANT_ZENXFI2_RECOVERY:
                 case VARIANT_ZENXFI2_NAND:
                 case VARIANT_ZENXFI2_SD:
-                    return patch_std_zero_host_play(1, model, type, sb_file, boot_fw);
+                    return patch_std_zero_host_play(1, opt, sb_file, boot_fw);
                 default:
                     return IMX_DONT_KNOW_HOW_TO_PATCH;
             }
@@ -471,12 +485,19 @@ static enum imx_error_t patch_firmware(enum imx_model_t model,
         case MODEL_ZENXFISTYLE:
             /* The ZEN X-Fi Style uses the standard ____, host, play sections, patch after first
              * call in ____ section. */
-            return patch_std_zero_host_play(1, model, type, sb_file, boot_fw);
+            return patch_std_zero_host_play(1, opt, sb_file, boot_fw);
         default:
             return IMX_DONT_KNOW_HOW_TO_PATCH;
     }
 }
 
+static enum imx_error_t unpatch_firmware(struct imx_option_t opt,
+    struct sb_file_t *sb_file)
+{
+    printf("[ERR] Unimplemented\n");
+    return IMX_ERROR;
+}
+
 static uint32_t get_uint32be(unsigned char *p)
 {
     return (p[0] << 24) | (p[1] << 16) | (p[2] << 8) | p[3];
@@ -680,16 +701,10 @@ enum imx_error_t compute_md5sum(const char *file, uint8_t file_md5sum[16])
 }
 
 /* compute soft MD5 of a file */
-enum imx_error_t compute_soft_md5sum(const char *file, enum imx_model_t model,
-    uint8_t soft_md5sum[16])
+enum imx_error_t compute_soft_md5sum(const char *file, uint8_t soft_md5sum[16])
 {
-    if(model == MODEL_UNKNOWN)
-    {
-        printf("[ERR] Cannot compute soft MD5 without knowing the model\n");
-        return IMX_ERROR;
-    }
     clear_keys();
-    add_keys(imx_models[model].keys, imx_models[model].nr_keys);
+    add_key_list(list_all_keys);
     /* read file */
     enum sb_error_t err;
     struct sb_file_t *sb = sb_read_file(file, false, NULL, generic_std_printf, &err);
@@ -705,29 +720,6 @@ enum imx_error_t compute_soft_md5sum(const char *file, enum imx_model_t model,
     return err;
 }
 
-static enum imx_error_t load_sb_file(const char *file, int md5_idx,
-    struct imx_option_t opt, struct sb_file_t **sb_file)
-{
-    if(imx_sums[md5_idx].fw_variants[opt.fw_variant].size == 0)
-    {
-        printf("[ERR] Input file does not contain variant '%s'\n", imx_fw_variant[opt.fw_variant]);
-        return IMX_VARIANT_MISMATCH;
-    }
-    enum imx_model_t model = imx_sums[md5_idx].model;
-    enum sb_error_t err;
-    g_debug = opt.debug;
-    clear_keys();
-    add_keys(imx_models[model].keys, imx_models[model].nr_keys);
-    *sb_file = sb_read_file_ex(file, imx_sums[md5_idx].fw_variants[opt.fw_variant].offset,
-                              imx_sums[md5_idx].fw_variants[opt.fw_variant].size, false, NULL, generic_std_printf, &err);
-    if(*sb_file == NULL)
-    {
-        clear_keys();
-        return IMX_FIRST_SB_ERROR + err;
-    }
-    return IMX_SUCCESS;
-}
-
 /* Load a rockbox firwmare from a buffer. Data is copied. Assume firmware is
  * using our scramble format. */
 static enum imx_error_t rb_fw_load_buf_scramble(struct rb_fw_t *fw, uint8_t *buf,
@@ -873,101 +865,145 @@ static void rb_fw_free(struct rb_fw_t *fw)
     memset(fw, 0, sizeof(struct rb_fw_t));
 }
 
-enum imx_error_t mkimxboot(const char *infile, const char *bootfile,
-    const char *outfile, struct imx_option_t opt)
+static bool contains_rockbox_bootloader(struct sb_file_t *sb_file)
 {
-    /* sanity check */
-    if(opt.fw_variant > VARIANT_COUNT)
-        return IMX_ERROR;
-    /* Dump tables */
-    dump_imx_dev_info("[INFO] ");
-    /* compute MD5 sum of the file */
-    uint8_t file_md5sum[16];
-    enum imx_error_t ret = compute_md5sum(infile, file_md5sum);
-    if(ret != IMX_SUCCESS)
-        return ret;
-    printf("[INFO] MD5 sum of the file: ");
-    for(int i = 0; i < 16; i++)
-        printf("%02x", file_md5sum[i]);
-    printf("\n");
-    /* find model */
-    int md5_idx;
-    ret = find_model_by_md5sum(file_md5sum, &md5_idx);
-    if(ret != IMX_SUCCESS)
-        return ret;
-    enum imx_model_t model = imx_sums[md5_idx].model;
-    printf("[INFO] File is for model %d (%s, version %s)\n", model,
-        imx_models[model].model_name, imx_sums[md5_idx].version);
+    for(int i = 0; i < sb_file->nr_sections; i++)
+        if(sb_file->sections[i].identifier == MAGIC_ROCK)
+            return true;
+    return false;
+}
+
+/* modify sb_file to produce requested boot image */
+static enum imx_error_t make_boot(struct sb_file_t *sb_file, const char *bootfile,
+    struct imx_option_t opt)
+{
+    /* things went smoothly, we have a SB image but it may not be suitable as an
+     * input image: if it contains a rockbox bootloader, we need to remove it */
+    if(contains_rockbox_bootloader(sb_file))
+    {
+        printf("[INFO] SB file contains a Rockbox bootloader, trying to remove it...\n");
+        enum imx_error_t ret = unpatch_firmware(opt, sb_file);
+        if(ret != IMX_SUCCESS)
+            return ret;
+    }
     /* load rockbox file */
     struct rb_fw_t boot_fw;
-    ret = rb_fw_load(&boot_fw, bootfile, model);
-    if(ret != IMX_SUCCESS)
-        return ret;
-    /* load OF file */
-    struct sb_file_t *sb_file;
-    ret = load_sb_file(infile, md5_idx, opt, &sb_file);
+    enum imx_error_t ret = rb_fw_load(&boot_fw, bootfile, opt.model);
     if(ret != IMX_SUCCESS)
-    {
-        rb_fw_free(&boot_fw);
         return ret;
-    }
     /* produce file */
-    ret = patch_firmware(model, opt.fw_variant, opt.output,
-        sb_file, boot_fw, opt.force_version);
-    if(ret == IMX_SUCCESS)
-        ret = sb_write_file(sb_file, outfile, NULL, generic_std_printf);
-
-    clear_keys();
+    ret = patch_firmware(opt, sb_file, boot_fw);
     rb_fw_free(&boot_fw);
-    sb_free(sb_file);
-    return ret;
+    return IMX_SUCCESS;
 }
 
-enum imx_error_t extract_firmware(const char *infile,
-    enum imx_firmware_variant_t fw_variant, const char *outfile)
+enum imx_error_t mkimxboot(const char *infile, const char *bootfile,
+    const char *outfile, struct imx_option_t opt)
 {
     /* sanity check */
-    if(fw_variant > VARIANT_COUNT)
+    if(opt.fw_variant >= VARIANT_COUNT || opt.model >= MODEL_COUNT)
         return IMX_ERROR;
     /* dump tables */
     dump_imx_dev_info("[INFO] ");
-    /* compute MD5 sum of the file */
+    /* load file */
     void *buf;
-    size_t sz;
-    uint8_t file_md5sum[16];
-    int ret = read_file(infile, &buf, &sz);
+    size_t offset = 0, size = 0;
+    enum imx_error_t ret = read_file(infile, &buf, &size);
     if(ret != IMX_SUCCESS)
         return ret;
-    ret = compute_md5sum_buf(buf, sz, file_md5sum);
-    if(ret != IMX_SUCCESS)
-    {
-        free(buf);
-        return ret;
-    }
+    /* compute MD5 sum of the file */
+    uint8_t file_md5sum[16];
+    compute_md5sum_buf(buf, size, file_md5sum);
     printf("[INFO] MD5 sum of the file: ");
-    print_hex(NULL, misc_std_printf, file_md5sum, 16, true);
+    for(int i = 0; i < 16; i++)
+        printf("%02x", file_md5sum[i]);
+    printf("\n");
     /* find model */
     int md5_idx;
     ret = find_model_by_md5sum(file_md5sum, &md5_idx);
-    if(ret != IMX_SUCCESS)
+    /* is this a known firmware upgrade ? */
+    if(ret == IMX_SUCCESS)
     {
+        enum imx_model_t model = imx_sums[md5_idx].model;
+        printf("[INFO] File is for model %d (%s, version %s)\n", model,
+            imx_models[model].model_name, imx_sums[md5_idx].version);
+        /* check the model is the expected one */
+        if(opt.model == MODEL_UNKNOWN)
+            opt.model = model;
+        else if(opt.model != model)
+        {
+            printf("[ERR] Model mismatch, was expecting model %d (%s)\n",
+                opt.model, imx_models[opt.model].model_name);
+            free(buf);
+            return IMX_MODEL_MISMATCH;
+        }
+        /* use database values */
+        offset = imx_sums[md5_idx].fw_variants[opt.fw_variant].offset;
+        size = imx_sums[md5_idx].fw_variants[opt.fw_variant].size;
+        if(size == 0)
+        {
+            printf("[ERR] Input file does not contain variant '%s'\n", imx_fw_variant[opt.fw_variant]);
+            free(buf);
+            return IMX_VARIANT_MISMATCH;
+        }
+        /* special case: if we need to produce the OF, just bypass read/write of
+         * the SB file and output this chunk of the file. This is faster and it
+         * also avoids modifying the OF by reconstructing it */
+        if(opt.output == IMX_ORIG_FW)
+        {
+            printf("[INFO] Extracting original firmware...\n");
+            ret = write_file(outfile, buf + offset, size);
+            free(buf);
+            return ret;
+        }
+    }
+    else
+    {
+        printf("[INFO] File doesn't have a known MD5 sum, assuming it's a SB image...\n");
+        /* image didn't match, so we expect the file to be a raw SB image, either
+        * produced by mkimxboot when uninstalling bootloader or after installing RB,
+        * so load all known keys and go on */
+
+        /* To be more user friendly, give a nice error message if we detect
+         * the file is not a SB file */
+        if(guess_sb_version(infile) == SB_VERSION_UNK)
+        {
+            printf("[ERR] Your firmware doesn't look like a SB file\n");
+            printf("[ERR] This is probably a firmware upgrade\n");
+            printf("[ERR] Unfortunately, this tool doesn't know about it yet\n");
+            printf("[ERR] Please report to the developers to add it\n");
+            free(buf);
+            return IMX_ERROR;
+        }
+    }
+    /* to proceed further, we need to know the model */
+    if(opt.model == MODEL_UNKNOWN)
+    {
+        printf("[ERR] Cannot do processing of soft image without knowing the model\n");
         free(buf);
-        return ret;
+        return IMX_MODEL_MISMATCH;
     }
-    enum imx_model_t model = imx_sums[md5_idx].model;
-    printf("[INFO] File is for model %d (%s, version %s)\n", model,
-        imx_models[model].model_name, imx_sums[md5_idx].version);
-    /* extract firmware */
-    if(imx_sums[md5_idx].fw_variants[fw_variant].size == 0)
+    /* load image */
+    g_debug = opt.debug;
+    clear_keys();
+    add_key_list(imx_models[opt.model].keys);
+    enum sb_error_t err;
+    struct sb_file_t *sb_file = sb_read_memory(buf + offset, size, false, NULL, generic_std_printf, &err);
+    if(sb_file == NULL)
     {
-        printf("[ERR] Input file does not contain variant '%s'\n", imx_fw_variant[fw_variant]);
+        printf("[ERR] Cannot open firmware as SB file: %d\n", err);
         free(buf);
-        return IMX_VARIANT_MISMATCH;
+        return IMX_FIRST_SB_ERROR + err;
     }
-
-    ret = write_file(outfile,
-        buf + imx_sums[md5_idx].fw_variants[fw_variant].offset,
-        imx_sums[md5_idx].fw_variants[fw_variant].size);
+    /* modify image */
+    ret = make_boot(sb_file, bootfile, opt);
+    if(ret == IMX_SUCCESS)
+    {
+        /* write image */
+        ret = sb_write_file(sb_file, outfile, NULL, generic_std_printf);
+    }
+    /* cleanup */
+    sb_free(sb_file);
     free(buf);
     return ret;
 }
diff --git a/rbutil/mkimxboot/mkimxboot.h b/rbutil/mkimxboot/mkimxboot.h
index 3157bd7daa..59f28a98a0 100644
--- a/rbutil/mkimxboot/mkimxboot.h
+++ b/rbutil/mkimxboot/mkimxboot.h
@@ -29,6 +29,7 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+
 enum imx_error_t
 {
     IMX_SUCCESS = 0,
@@ -43,30 +44,31 @@ enum imx_error_t
     IMX_VARIANT_MISMATCH = -9,
     IMX_WRITE_ERROR = -10,
     IMX_FIRST_SB_ERROR = -11,
+    IMX_MODEL_MISMATCH = -12,
 };
 
 enum imx_output_type_t
 {
     IMX_DUALBOOT = 0,
-    IMX_RECOVERY = 1,
-    IMX_SINGLEBOOT = 2,
-    IMX_CHARGE = 3,
+    IMX_RECOVERY,
+    IMX_SINGLEBOOT,
+    IMX_CHARGE,
+    IMX_ORIG_FW,
 };
 
 /* Supported models */
 enum imx_model_t
 {
-    MODEL_UNKNOWN = -1,
-    MODEL_FUZEPLUS = 0,
+    MODEL_UNKNOWN = 0,
+    MODEL_FUZEPLUS,
     MODEL_ZENXFI2,
     MODEL_ZENXFI3,
     MODEL_ZENXFISTYLE,
     MODEL_ZENSTYLE, /* Style 100 and Style 300 */
     MODEL_NWZE370,
     MODEL_NWZE360,
-    /* new models go here */
-
-    NUM_MODELS
+    /* Last */
+    MODEL_COUNT
 };
 
 /* Supported firmware variants */
@@ -88,19 +90,22 @@ enum imx_firmware_variant_t
 struct imx_option_t
 {
     bool debug;
+    enum imx_model_t model;
     enum imx_output_type_t output;
     enum imx_firmware_variant_t fw_variant;
     const char *force_version; // set to NULL to ignore
 };
 
+/* Print internal information to stdout about device database */
 void dump_imx_dev_info(const char *prefix);
+/* Build a SB image from an input firmware and a bootloader, input firmware
+ * can either be a firmware update or another SB file produced by this tool */
 enum imx_error_t mkimxboot(const char *infile, const char *bootfile,
     const char *outfile, struct imx_option_t opt);
-enum imx_error_t extract_firmware(const char *infile,
-    enum imx_firmware_variant_t fw_variant, const char *outfile);
+/* Compute MD5 sum of an entire file */
 enum imx_error_t compute_md5sum(const char *file, uint8_t file_md5sum[16]);
-enum imx_error_t compute_soft_md5sum(const char *file, enum imx_model_t model,
-    uint8_t soft_md5sum[16]);
+/* Compute "soft" MD5 sum of a SB file */
+enum imx_error_t compute_soft_md5sum(const char *file, uint8_t soft_md5sum[16]);
 
 #ifdef __cplusplus
 }