Add password support to citadel_updater

This adds the following actions:

  --enable_ro   Mark new RO image as good
  --enable_rw   Mark new RW image as good

  --change_pw   Change update password

  --erase=CODE  Erase all user secrets and reboot.
                This skips all other actions.

Change-Id: I5499d5f5262bc49282d5670dcf69bb969932ecce
Signed-off-by: Bill Richardson <wfrichar@google.com>
diff --git a/citadel/updater/updater.cpp b/citadel/updater/updater.cpp
index 75c4aaf..c2ee303 100644
--- a/citadel/updater/updater.cpp
+++ b/citadel/updater/updater.cpp
@@ -55,6 +55,10 @@
   int ro;
   int rw;
   int reboot;
+  int enable_ro;
+  int enable_rw;
+  int change_pw;
+  uint32_t erase_code;
   /* generic connection options */
   const char *device;
 } options;
@@ -64,6 +68,10 @@
   OPT_RO,
   OPT_RW,
   OPT_REBOOT,
+  OPT_ENABLE_RO,
+  OPT_ENABLE_RW,
+  OPT_CHANGE_PW,
+  OPT_ERASE,
 };
 
 const char *short_opts = ":hv";
@@ -73,6 +81,10 @@
   {"ro",          0, NULL, OPT_RO},
   {"rw",          0, NULL, OPT_RW},
   {"reboot",      0, NULL, OPT_REBOOT},
+  {"enable_ro",   0, NULL, OPT_ENABLE_RO},
+  {"enable_rw",   0, NULL, OPT_ENABLE_RW},
+  {"change_pw",   0, NULL, OPT_CHANGE_PW},
+  {"erase",       1, NULL, OPT_ERASE},
   {"device",      1, NULL, OPT_DEVICE},
   {"help",        0, NULL, 'h'},
   {NULL, 0, NULL, 0},
@@ -107,6 +119,14 @@
     "      --rw          Update RW firmware from the image file\n"
     "      --ro          Update RO firmware from the image file\n"
     "      --reboot      Tell Citadel to reboot\n"
+    "\n"
+    "      --enable_ro   Mark new RO image as good\n"
+    "      --enable_rw   Mark new RW image as good\n"
+    "\n"
+    "      --change_pw   Change update password\n"
+    "\n\n"
+    "      --erase=CODE  Erase all user secrets and reboot.\n"
+    "                    This skips all other actions.\n"
     "\n",
     progname);
 }
@@ -123,11 +143,11 @@
 {
   va_list ap;
 
-        va_start(ap, format);
-        fprintf(stderr, "ERROR: ");
-        vfprintf(stderr, format, ap);
-        fprintf(stderr, "\n");
-        va_end(ap);
+  va_start(ap, format);
+  fprintf(stderr, "ERROR: ");
+  vfprintf(stderr, format, ap);
+  fprintf(stderr, "\n");
+  va_end(ap);
 
   errorcnt++;
 }
@@ -210,24 +230,30 @@
   return buf;
 }
 
-uint32_t compute_digest(struct nugget_app_flash_block *blk)
+uint32_t compute_digest(void *ptr, size_t len)
 {
-  uint8_t *start_here = ((uint8_t *)blk) +
-    offsetof(struct nugget_app_flash_block, offset);
-  size_t size_to_hash = sizeof(*blk) -
-    offsetof(struct nugget_app_flash_block, offset);
   SHA_CTX ctx;
   uint8_t digest[SHA_DIGEST_LENGTH];
   uint32_t retval;
 
   SHA1_Init(&ctx);
-  SHA1_Update(&ctx, start_here, size_to_hash);
+  SHA1_Update(&ctx, ptr, len);
   SHA1_Final(digest, &ctx);
 
   memcpy(&retval, digest, sizeof(retval));
   return retval;
 }
 
+uint32_t compute_fb_digest(struct nugget_app_flash_block *blk)
+{
+  uint8_t *start_here = ((uint8_t *)blk) +
+    offsetof(struct nugget_app_flash_block, offset);
+  size_t size_to_hash = sizeof(*blk) -
+    offsetof(struct nugget_app_flash_block, offset);
+
+  return compute_digest(start_here, size_to_hash);
+}
+
 uint32_t try_update(AppClient &app, const std::vector<uint8_t> &image,
         uint32_t offset, uint32_t imagesize)
 {
@@ -245,7 +271,7 @@
 
     fb->offset = offset;
     memcpy(fb->payload, image.data() + offset, CHIP_FLASH_BANK_SIZE);
-    fb->block_digest = compute_digest(fb);
+    fb->block_digest = compute_fb_digest(fb);
 
     printf("writing 0x%05x / 0x%05x",
            CHIP_FLASH_BASE + offset, CHIP_FLASH_BASE + stop);
@@ -254,15 +280,19 @@
       if (rv == NUGGET_ERROR_RETRY)
         printf(" retrying");
     } while (rv == NUGGET_ERROR_RETRY && retries--);
-    printf(" %s\n", rv ? "fail" : "ok");
-    if (rv)
+    if (rv) {
+      if (rv == NUGGET_ERROR_LOCKED)
+        printf(" locked\n");
+      else
+        printf(" fail %d\n", rv);
       break;
+    }
+    printf(" ok\n");
   }
 
   return rv;
 }
 
-
 uint32_t do_update(AppClient &app, const std::vector<uint8_t> &image,
        uint32_t offset_A, uint32_t offset_B)
 {
@@ -314,6 +344,78 @@
   return retval;
 }
 
+static uint32_t do_change_pw(AppClient &app,
+                             const char *old_pw, const char *new_pw)
+{
+  std::vector<uint8_t> data(sizeof(struct nugget_app_change_update_password));
+  struct nugget_app_change_update_password *s =
+    (struct nugget_app_change_update_password*)data.data();
+
+
+  memset(s, 0xff, sizeof(*s));
+  if (old_pw && *old_pw) {
+    int len = strlen(old_pw);
+    memcpy(&s->old_password.password, old_pw, len);
+    s->old_password.digest =
+      compute_digest(&s->old_password.password,
+                     sizeof(s->old_password.password));
+  }
+
+  if (new_pw && *new_pw) {
+    int len = strlen(new_pw);
+    memcpy(&s->new_password.password, new_pw, len);
+    s->new_password.digest =
+      compute_digest(&s->new_password.password,
+                     sizeof(s->new_password.password));
+  }
+
+  uint32_t rv = app.Call(NUGGET_PARAM_CHANGE_UPDATE_PASSWORD, data, nullptr);
+
+  if (is_app_success(rv))
+    printf("Password changed\n");
+
+  return rv;
+}
+
+static uint32_t do_enable(AppClient &app, const char *pw)
+{
+
+  std::vector<uint8_t> data(sizeof(struct nugget_app_enable_update));
+  struct nugget_app_enable_update *s =
+    (struct nugget_app_enable_update*)data.data();
+
+  memset(&s->password, 0xff, sizeof(s->password));
+  if (pw && *pw) {
+    int len = strlen(pw);
+    memcpy(&s->password.password, pw, len);
+    s->password.digest =
+      compute_digest(&s->password.password,
+                     sizeof(s->password.password));
+  }
+  s->which_headers = options.enable_ro ? NUGGET_ENABLE_HEADER_RO : 0;
+  s->which_headers |= options.enable_rw ? NUGGET_ENABLE_HEADER_RW : 0;
+
+  uint32_t rv = app.Call(NUGGET_PARAM_ENABLE_UPDATE, data, nullptr);
+
+  if (is_app_success(rv))
+    printf("Update enabled\n");
+
+  return rv;
+}
+
+static uint32_t do_erase(AppClient &app)
+{
+  std::vector<uint8_t> data(sizeof(uint32_t));
+  memcpy(data.data(), &options.erase_code, data.size());
+
+  uint32_t rv = app.Call(NUGGET_PARAM_NUKE_FROM_ORBIT, data, nullptr);
+
+  if (is_app_success(rv))
+    printf("Citadel erase and reboot requested\n");
+
+  return rv;
+}
+
 std::unique_ptr<NuggetClientInterface> select_client()
 {
 #ifdef ANDROID
@@ -324,7 +426,8 @@
 #endif
 }
 
-int update_to_image(const std::vector<uint8_t> &image)
+int execute_commands(const std::vector<uint8_t> &image,
+                     const char *old_passwd, const char *passwd)
 {
   auto client = select_client();
   client->Open();
@@ -336,6 +439,11 @@
 
   /* Try all requested actions in reasonable order, bail out on error */
 
+  if (options.erase_code) {                     /* zero doesn't count */
+    /* whether we succeed or not, only do this */
+    return do_erase(app);
+  }
+
   if (options.version &&
       do_version(app) != APP_SUCCESS) {
     return 2;
@@ -353,10 +461,19 @@
     return 4;
   }
 
+  if (options.change_pw &&
+      do_change_pw(app, old_passwd, passwd) != APP_SUCCESS)
+    return 5;
+
+  if ((options.enable_ro || options.enable_rw) &&
+      do_enable(app, passwd) != APP_SUCCESS)
+    return 6;
+
   if (options.reboot &&
       do_reboot(app) != APP_SUCCESS) {
-    return 5;
+    return 7;
   }
+
   return 0;
 }
 
@@ -367,8 +484,11 @@
   int i;
   int idx = 0;
   char *this_prog;
+  char *old_passwd = 0;
+  char *passwd = 0;
   std::vector<uint8_t> image;
   int got_action = 0;
+  char *e = 0;
 
   this_prog= strrchr(argv[0], '/');
         if (this_prog) {
@@ -398,6 +518,26 @@
       options.reboot = 1;
       got_action = 1;
       break;
+    case OPT_ENABLE_RO:
+      options.enable_ro = 1;
+      got_action = 1;
+      break;
+    case OPT_ENABLE_RW:
+      options.enable_rw = 1;
+      got_action = 1;
+      break;
+    case OPT_CHANGE_PW:
+      options.change_pw = 1;
+      got_action = 1;
+      break;
+    case OPT_ERASE:
+      options.erase_code = (uint32_t)strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e)) {
+        Error("Invalid argument: \"%s\"\n", optarg);
+        errorcnt++;
+      }
+      got_action = 1;
+      break;
 
       /* generic options below */
     case OPT_DEVICE:
@@ -438,16 +578,41 @@
     if (optind < argc) {
       /* Sets errorcnt on failure */
       image = read_image_from_file(argv[optind]);
+      if (errorcnt)
+        goto out;
     } else {
       Error("An image file is required with --ro and --rw");
+      goto out;
     }
   }
 
-  if (errorcnt)
-    goto out;
+  if (options.change_pw) {
+    /* one arg provided, maybe the old password isn't set */
+    if (optind < argc) {
+      passwd = argv[optind++];
+    } else {
+      Error("Need a new password at least."
+            " Use '' to clear it.");
+      goto out;
+    }
+    /* two args provided, use both old & new passwords */
+    if (optind < argc) {
+      old_passwd = passwd;
+      passwd = argv[optind++];
+    }
+  }
 
-  /* Okay, let's do something */
-  (void) update_to_image(image);
+  if ((options.enable_ro || options.enable_rw) && !passwd) {
+    if (optind < argc)
+      passwd = argv[optind++];
+    else {
+      Error("Need a password to enable images. Use '' if none.");
+      goto out;
+    }
+  }
+
+  /* Okay, let's do it! */
+  (void) execute_commands(image, old_passwd, passwd);
   /* This is the last action, so fall through either way */
 
 out: