From 5ee05ff3919e212745e3e5f93c7cac6935719b41 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Wed, 18 Oct 2023 12:52:18 +0200
Subject: [PATCH] feat(core/bootloader_emu): add support for setting device
 variant

---
 core/embed/bootloader/emulator.c | 101 ++++++++++++++++++++++++++-----
 1 file changed, 85 insertions(+), 16 deletions(-)

diff --git a/core/embed/bootloader/emulator.c b/core/embed/bootloader/emulator.c
index 10bbcbdf7..5a35afb82 100644
--- a/core/embed/bootloader/emulator.c
+++ b/core/embed/bootloader/emulator.c
@@ -39,6 +39,41 @@ bool sector_is_empty(uint16_t sector) {
   return true;
 }
 
+void usage(void) {
+  printf("Usage: ./build/bootloader/bootloader_emu [options]\n");
+  printf("Options:\n");
+  printf("  -s  stay in bootloader\n");
+  printf("  -e MESSAGE [TITLE [FOOTER]]  display error screen and stop\n");
+  printf("  -c COLOR_VARIANT  set color variant\n");
+  printf("  -b BITCOIN_ONLY  set bitcoin only flag\n");
+#ifdef USE_OPTIGA
+  printf("  -l  lock bootloader\n");
+#endif
+  printf("  -h  show this help\n");
+}
+
+__attribute__((noreturn)) void display_error_and_die(const char *message,
+                                                     const char *title,
+                                                     const char *footer) {
+  if (footer == NULL) {
+    footer = "PLEASE VISIT\nTREZOR.IO/RSOD";
+  }
+  if (title == NULL) {
+    title = "INTERNAL ERROR";
+  }
+  display_init();
+  display_backlight(180);
+  screen_fatal_error_rust(title, message, footer);
+  display_refresh();
+#if USE_TOUCH
+  printf("Click screen to exit.\n");
+#elif USE_BUTTON
+  printf("Press both buttons to exit.\n");
+#endif
+  ui_click();
+  exit(0);
+}
+
 __attribute__((noreturn)) int main(int argc, char **argv) {
   flash_init();
   FIRMWARE_START = (uint8_t *)flash_area_get_address(&FIRMWARE_AREA, 0, 0);
@@ -49,26 +84,60 @@ __attribute__((noreturn)) int main(int argc, char **argv) {
     (void)ret;
   }
 
-  if (argc == 2) {
-    if (argv[1][0] == 's') {
-      // Run the firmware
-      g_boot_command = BOOT_COMMAND_STOP_AND_WAIT;
-    }
+  int opt;
+  bool display_error = false;
+  uint8_t set_variant = 0xff;
+  uint8_t color_variant = 0;
+  uint8_t bitcoin_only = 0;
+  while ((opt = getopt(argc, argv, "hslec:b:")) != -1) {
+    switch (opt) {
+      case 's':
+        g_boot_command = BOOT_COMMAND_STOP_AND_WAIT;
+        break;
+      case 'e':
+        display_error = true;
+        break;
+      case 'c':
+        set_variant = 1;
+        color_variant = atoi(optarg);
+        break;
+      case 'b':
+        set_variant = 1;
+        bitcoin_only = atoi(optarg);
+        break;
 #ifdef USE_OPTIGA
-    else if (argv[1][0] == 'l') {
-      // write bootloader-lock secret
-      secret_write_header();
-    }
+      case 'l':
+        // write bootloader-lock secret
+        secret_write_header();
+        break;
 #endif
-  } else if (argc == 4) {
-    display_init();
-    display_backlight(180);
-    screen_fatal_error_rust(argv[1], argv[2], argv[3]);
-    display_refresh();
-    ui_click();
-    exit(0);
+      default:
+        usage();
+        exit(1);
+    }
   }
 
+  if (display_error) {
+    const char *message, *title = NULL, *footer = NULL;
+    if (optind < argc) {
+      message = argv[optind++];
+      if (optind < argc) {
+        title = argv[optind++];
+        if (optind < argc) {
+          footer = argv[optind++];
+        }
+      }
+    } else {
+      message = "No message specified";
+    }
+    display_error_and_die(message, title, footer);
+  }
+
+  // write variant to OTP
+  const uint8_t otp_data[] = {set_variant, color_variant, bitcoin_only};
+  (void)!flash_otp_write(FLASH_OTP_BLOCK_DEVICE_VARIANT, 0, otp_data,
+                         sizeof(otp_data));
+
   bootloader_main();
   hal_delay(3000);
   jump_to(NULL);