Make vbutil_ec operate on the entire image at once.

Instead of taking images apart, signing bits, and reassembling them, this
just operates on the entire image at once. The image can be built without
regard to the signing process, and then the signing can be done in one step
afterwards.

BUG=chrome-os-partner:7459
TEST=none

No test at the moment, since we don't have the rest of the vboot stuff
working yet.

Change-Id: Icbde9cbb89d0ef85c0f6b8ac0637e0a51a894199
Reviewed-on: https://gerrit.chromium.org/gerrit/22116
Tested-by: Bill Richardson <wfrichar@chromium.org>
Reviewed-by: Randall Spangler <rspangler@chromium.org>
Commit-Ready: Bill Richardson <wfrichar@chromium.org>

1 files changed, 411 insertions, 210 deletions

diff --git a/utility/vbutil_ec.c b/utility/vbutil_ec.c
index 86747c0a..c1b280f6 100644
--- a/utility/vbutil_ec.c
+++ b/utility/vbutil_ec.c
@@ -5,22 +5,28 @@
  * Verified boot utility for EC firmware
  */
 
+#include <errno.h>
+#include <fcntl.h>
 #include <getopt.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
 #include <unistd.h>
 
 #include "cryptolib.h"
+#include "fmap.h"
 #include "host_common.h"
 #include "vboot_common.h"
 
 
 /* Command line options */
 enum {
-  OPT_MODE_VBLOCK = 1000,
+  OPT_MODE_SIGN = 1000,
   OPT_MODE_VERIFY,
   OPT_KEYBLOCK,
   OPT_SIGNPUBKEY,
@@ -33,13 +39,12 @@ enum {
 };
 
 static struct option long_opts[] = {
-  {"vblock", 1, 0,                    OPT_MODE_VBLOCK             },
+  {"sign", 1, 0,                      OPT_MODE_SIGN               },
   {"verify", 1, 0,                    OPT_MODE_VERIFY             },
   {"keyblock", 1, 0,                  OPT_KEYBLOCK                },
   {"signpubkey", 1, 0,                OPT_SIGNPUBKEY              },
   {"signprivate", 1, 0,               OPT_SIGNPRIVATE             },
   {"version", 1, 0,                   OPT_VERSION                 },
-  {"fv", 1, 0,                        OPT_FV                      },
   {"flags", 1, 0,                     OPT_FLAGS                   },
   {"name", 1, 0,                      OPT_NAME                    },
   {NULL, 0, 0, 0}
@@ -51,255 +56,435 @@ static int PrintHelp(void) {
 
   puts("vbutil_ec - Verified boot signing utility for EC firmware\n"
        "\n"
-       "Usage:  vbutil_ec <--vblock|--verify> <file> [OPTIONS]\n"
+       "This will sign, re-sign, or test a complete EC firmware image.\n"
+       "The EC image is initially completely unsigned. To make it bootable\n"
+       "the pubic root key must be installed in the RO section, and each RW\n"
+       "section must be signed with the appropriate private keys.\n"
+       "\n"
+       "To sign an image:   vbutil_ec --sign <file> [OPTIONS]\n"
+       "\n"
+       "For signing, these options are required:\n"
        "\n"
-       "For '--vblock <file>', required OPTIONS are:\n"
        "  --keyblock <file>           Key block in .keyblock format\n"
        "  --signprivate <file>        Signing private key in .vbprivk format\n"
        "  --version <number>          Firmware version\n"
-       "  --fv <file>                 Firmware volume to sign\n"
-       "optional OPTIONS are:\n"
+       "\n"
+       "If the RO public key has not been installed, you will also need\n"
+       "\n"
+       "  --signpubkey <file>         Signing public key in .vbpubk format\n"
+       "\n"
+       "Optional args are:\n"
+       "\n"
        "  --flags <number>            Preamble flags (defaults to 0)\n"
        "  --name <string>             Human-readable description\n"
        "\n"
-       "For '--verify <file>', required OPTIONS are:\n"
-       "  --fv <file>                 Firmware volume to verify\n"
-       "optional OPTIONS are:\n"
-       "  --signpubkey <file>         Signing public key in .vbpubk format\n"
+       "\n"
+       "To verify an image:   vbutil_ec --verify <file>\n"
        "\n");
   return 1;
 }
 
 
-/* Create an EC firmware .vblock */
-static int Vblock(const char* outfile, const char* keyblock_file,
-                  const char* signprivate, uint64_t version,
-                  const char* fv_file, uint32_t preamble_flags,
-                  const char* name) {
-  VbPrivateKey* signing_key;
-  VbSignature* body_digest;
-  VbECPreambleHeader* preamble;
-  VbKeyBlockHeader* key_block;
-  uint64_t key_block_size;
-  uint8_t* fv_data;
-  uint64_t fv_size;
-  FILE* f;
-  uint64_t i;
+static int FindInFmap(FmapHeader *fh, const char *name,
+                      uint8_t *base, uint64_t base_size,
+                      uint8_t **data, uint64_t *size) {
+  const FmapAreaHeader *ah;
+  int i;
+
+  ah = (FmapAreaHeader *)(fh + 1);
+  for (i = 0; i < fh->fmap_nareas; i++)
+    if (!strncmp(ah[i].area_name, name, FMAP_NAMELEN)) {
+      if (ah[i].area_size + ah[i].area_offset > base_size) {
+        printf("FMAP region %s extends off image file\n", name);
+        return 0;
+      }
+      if (data)
+        *data = base + ah[i].area_offset;
+      if (size)
+        *size = ah[i].area_size;
+      return 1;
+    }
+
+  return 0;
+}
+
+static int GoodKey(VbPublicKey *key, uint64_t region_size)
+{
+  uint64_t key_size;
+
+  if (0 != VerifyPublicKeyInside(key, region_size, key))
+    return 0;
 
-  if (!outfile)
-    VbExError("Must specify output filename\n");
+  if (key->algorithm >= kNumAlgorithms)
+    return 0;
 
-  if (!keyblock_file || !signprivate)
-    VbExError("Must specify all keys\n");
+  /* Currently, TPM only supports 16-bit version */
+  if (key->key_version > 0xFFFF)
+    return 0;
 
-  if (!fv_file)
-    VbExError("Must specify firmware volume\n");
+  if (!RSAProcessedKeySize(key->algorithm, &key_size) ||
+      key_size != key->key_size)
+    return 0;
 
-  /* Read the key block and keys */
-  key_block = (VbKeyBlockHeader*)ReadFile(keyblock_file, &key_block_size);
-  if (!key_block)
-    VbExError("Error reading key block.\n");
+  return 1;
+}
 
-  signing_key = PrivateKeyRead(signprivate);
-  if (!signing_key)
-    VbExError("Error reading signing key.\n");
 
-  /* Read and sign the firmware volume */
-  fv_data = ReadFile(fv_file, &fv_size);
-  if (!fv_data)
-    return 1;
-  if (!fv_size)
-    VbExError("Empty firmware volume file\n");
+/* We build the image file with a non-FF byte at the end of each RW firmware,
+ * just so we can do this. */
+static uint64_t FindImageEnd(uint8_t *data, uint64_t size)
+{
+  for (size-- ; size && data[size] == 0xff; size--)
+    ;
+  return size;
+}
+
+static void SignImage(const char *filename,
+                      VbKeyBlockHeader *key_block, uint64_t key_block_size,
+                      VbPrivateKey *privkey, uint64_t version,
+                      VbPublicKey *pubkey, uint32_t preamble_flags,
+                      const char *name) {
+  struct stat sb;
+  int fd;
+  void *image;
+  uint64_t image_size;
+  FmapHeader* fmap;
+  VbECPreambleHeader *preamble;
+  uint8_t *fv_data = 0;
+  uint8_t *vblock_data = 0;
+  uint64_t fv_size, vblock_size;
+  VbSignature* body_digest;
 
   if (name && strlen(name)+1 > sizeof(preamble->name))
     VbExError("Name string is too long\n");
 
-  body_digest = CalculateHash(fv_data, fv_size, signing_key);
+  if (0 != stat(filename, &sb))
+    VbExError("Can't stat %s: %s\n", filename, strerror(errno));
+
+  fd = open(filename, O_RDWR);
+  if (fd < 0)
+    VbExError("Can't open %s: %s\n", filename, strerror(errno));
+
+  image = mmap(0, sb.st_size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
+  if (image == (void *)-1)
+    VbExError("Can't mmap %s: %s\n", filename, strerror(errno));
+  close(fd);                            /* done with this now */
+
+  fmap = (FmapHeader *)FmapFind(image, sb.st_size);
+  if (!fmap)
+    VbExError("File %s doesn't have an FMAP - can't continue.\n");
+
+  if (fmap->fmap_size > sb.st_size)
+    VbExError("FMAP is bigger than file size (%ld vs %ld)\n",
+              fmap->fmap_size, sb.st_size);
+
+  image_size = sb.st_size;
+
+  /* Install pubkey if provided */
+  if (pubkey) {
+    if (!FindInFmap(fmap, "ROOT_KEY", image, image_size,
+                    &vblock_data, &vblock_size))
+      VbExError("Can't find ROOT_KEY in %s\n", filename);
+
+    if (pubkey->key_offset + pubkey->key_size > vblock_size)
+      VbExError("ROOT_KEY is too small for pubkey (%d bytes, needs %d)\n",
+                vblock_size, pubkey->key_offset + pubkey->key_size);
+
+    memcpy(vblock_data, pubkey, pubkey->key_offset + pubkey->key_size);
+  }
+
+
+  /* Sign FW A */
+  if (!FindInFmap(fmap, "FW_MAIN_A", image, image_size, &fv_data, &fv_size))
+    VbExError("Can't find FW_MAIN_A in %s\n", filename);
+
+  if (!FindInFmap(fmap, "VBLOCK_A", image, image_size,
+                  &vblock_data, &vblock_size))
+    VbExError("Can't find VBLOCK_A in %s\n", filename);
+
+  fv_size = FindImageEnd(fv_data, fv_size);
+
+  body_digest = CalculateHash(fv_data, fv_size, privkey);
   if (!body_digest)
     VbExError("Error calculating body digest\n");
-  free(fv_data);
 
-  /* Create preamble */
-  preamble = CreateECPreamble(version, body_digest, signing_key,
+  preamble = CreateECPreamble(version, body_digest, privkey,
                               preamble_flags, name);
   if (!preamble)
     VbExError("Error creating preamble.\n");
 
-  /* Write the output file */
-  f = fopen(outfile, "wb");
-  if (!f)
-    VbExError("Can't open output file %s\n", outfile);
-
-  i = ((1 != fwrite(key_block, key_block_size, 1, f)) ||
-       (1 != fwrite(preamble, preamble->preamble_size, 1, f)));
-  fclose(f);
-  if (i) {
-    unlink(outfile);
-    VbExError("Can't write output file %s\n", outfile);
+  if (key_block_size + preamble->preamble_size > vblock_size)
+    VbExError("VBLOCK_A is too small for digest (%d bytes, needs %d)\n",
+              vblock_size, key_block_size + preamble->preamble_size);
+
+  memcpy(vblock_data, key_block, key_block_size);
+  memcpy(vblock_data + key_block_size, preamble, preamble->preamble_size);
+
+  free(body_digest);
+  free(preamble);
+
+
+  /* Sign FW B - skip if there isn't one */
+  if (!FindInFmap(fmap, "FW_MAIN_B", image, image_size, &fv_data, &fv_size) ||
+      !FindInFmap(fmap, "VBLOCK_B", image, image_size,
+                  &vblock_data, &vblock_size)) {
+    printf("Image does not contain FW B - ignoring that part\n");
+  } else {
+    fv_size = FindImageEnd(fv_data, fv_size);
+
+    body_digest = CalculateHash(fv_data, fv_size, privkey);
+    if (!body_digest)
+      VbExError("Error calculating body digest\n");
+
+    preamble = CreateECPreamble(version, body_digest, privkey,
+                                preamble_flags, name);
+    if (!preamble)
+      VbExError("Error creating preamble.\n");
+
+    if (key_block_size + preamble->preamble_size > vblock_size)
+      VbExError("VBLOCK_B is too small for digest (%d bytes, needs %d)\n",
+                vblock_size, key_block_size + preamble->preamble_size);
+
+    memcpy(vblock_data, key_block, key_block_size);
+    memcpy(vblock_data + key_block_size, preamble, preamble->preamble_size);
+
+    free(body_digest);
+    free(preamble);
   }
 
-  /* Success */
-  return 0;
+  /* Unmap to write changes to disk. */
+  if (0 != munmap(image, sb.st_size))
+    VbExError("Can't munmap %s: %s\n", filename, strerror(errno));
+
+  printf("Image signing completed\n");
+
 }
 
-static int Verify(const char* infile,
-                  const char* signpubkey,
-                  const char* fv_file) {
-  VbKeyBlockHeader* key_block;
-  VbECPreambleHeader* preamble;
-  VbPublicKey* data_key;
-  VbPublicKey* sign_key = 0;
-  RSAPublicKey* rsa;
-  uint8_t* blob;
-  uint64_t blob_size;
-  uint8_t* fv_data;
+static int Verify(const char *filename) {
+  struct stat sb;
+  int fd;
+  void *image;
+  uint64_t image_size;
+  FmapHeader* fmap;
+  VbECPreambleHeader *preamble;
+  VbPublicKey *pubkey;
+  uint64_t pubkey_size;
+  VbKeyBlockHeader *key_block;
+  uint64_t key_block_size;
+  uint8_t *fv_data = 0;
   uint64_t fv_size;
-  uint64_t now = 0;
+  VbPublicKey *data_key;
+  RSAPublicKey* rsa;
+  int errorcnt = 0;
+  char buf[80];
+  int i;
 
-  if (!infile || !fv_file) {
-    VbExError("Must specify filename and fv\n");
-    return 1;
-  }
+  if (0 != stat(filename, &sb))
+    VbExError("Can't stat %s: %s\n", filename, strerror(errno));
+
+  fd = open(filename, O_RDONLY);
+  if (fd < 0)
+    VbExError("Can't open %s: %s\n", filename, strerror(errno));
+
+  image = mmap(0, sb.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
+  if (image == (void *)-1)
+    VbExError("Can't mmap %s: %s\n", filename, strerror(errno));
+  close(fd);                            /* done with this now */
+
+  fmap = (FmapHeader *)FmapFind(image, sb.st_size);
+  if (!fmap)
+    VbExError("File %s doesn't have an FMAP - can't continue.\n");
+
+  if (fmap->fmap_size > sb.st_size)
+    VbExError("FMAP is bigger than file size (%ld vs %ld)\n",
+              fmap->fmap_size, sb.st_size);
+
+  image_size = sb.st_size;
 
-  /* Read public signing key */
-  if (signpubkey) {
-    sign_key = PublicKeyRead(signpubkey);
-    if (!sign_key)
-      VbExError("Error reading signpubkey.\n");
+  /* Read pubkey */
+  if (!FindInFmap(fmap, "ROOT_KEY", image, image_size,
+                  (uint8_t **)&pubkey, &pubkey_size)) {
+    printf("Can't find ROOT_KEY in %s\n", filename);
+    errorcnt++;
+  } else if (!GoodKey(pubkey, pubkey_size)) {
+    printf("ROOT_KEY is invalid\n");
+    errorcnt++;
   } else {
-    printf("WARNING:  No public key given - signature is not checked\n");
+    printf("ROOT_KEY\n");
+    printf("  Algorithm:         %" PRIu64 " %s\n", pubkey->algorithm,
+           (pubkey->algorithm < kNumAlgorithms ?
+            algo_strings[pubkey->algorithm] : "(invalid)"));
+    printf("  Key Version:       %" PRIu64 "\n", pubkey->key_version);
+    printf("  Key sha1sum:       ");
+    PrintPubKeySha1Sum(pubkey);
+    printf("\n");
   }
 
-  /* Read blob */
-  blob = ReadFile(infile, &blob_size);
-  if (!blob)
-    VbExError("Error reading input file\n");
-
-  /* Read firmware volume */
-  fv_data = ReadFile(fv_file, &fv_size);
-  if (!fv_data)
-    VbExError("Error reading firmware volume\n");
-
-  /* Verify key block */
-  key_block = (VbKeyBlockHeader*)blob;
-  if (0 != KeyBlockVerify(key_block, blob_size, sign_key, !signpubkey))
-    VbExError("Error verifying key block.\n");
-
-  if (sign_key)
-    free(sign_key);
-  now += key_block->key_block_size;
-
-  printf("Key block:\n");
-  data_key = &key_block->data_key;
-  printf("  Size:                %" PRIu64 "\n", key_block->key_block_size);
-  printf("  Flags:               %" PRIu64 " (ignored)\n",
-         key_block->key_block_flags);
-  printf("  Data key algorithm:  %" PRIu64 " %s\n", data_key->algorithm,
-         (data_key->algorithm < kNumAlgorithms ?
-          algo_strings[data_key->algorithm] : "(invalid)"));
-  printf("  Data key version:    %" PRIu64 "\n", data_key->key_version);
-  printf("  Data key sha1sum:    ");
-  PrintPubKeySha1Sum(data_key);
-  printf("\n");
-
-  rsa = PublicKeyToRSA(&key_block->data_key);
-  if (!rsa)
-    VbExError("Error parsing data key.\n");
-
-  /* Verify preamble */
-  preamble = (VbECPreambleHeader*)(blob + now);
-  if (0 != VerifyECPreamble(preamble, blob_size - now, rsa))
-    VbExError("Error verifying preamble.\n");
-
-  now += preamble->preamble_size;
-
-  printf("Preamble:\n");
-  printf("  Size:                  %" PRIu64 "\n", preamble->preamble_size);
-  printf("  Header version:        %" PRIu32 ".%" PRIu32"\n",
-         preamble->header_version_major, preamble->header_version_minor);
-  printf("  Firmware version:      %" PRIu64 "\n", preamble->firmware_version);
-  printf("  Firmware body size:    %" PRIu64 "\n",
-         preamble->body_digest.data_size);
-  printf("  Preamble flags:        %" PRIu32 "\n", preamble->flags);
-  printf("  Preamble name:         %s\n", preamble->name);
-
-  /* TODO: verify body size same as signature size */
-
-  /* Verify body */
-  if (preamble->flags & VB_FIRMWARE_PREAMBLE_USE_RO_NORMAL) {
-    printf("Preamble requests USE_RO_NORMAL; skipping body verification.\n");
-  } else {
-    if (0 != EqualData(fv_data, fv_size, &preamble->body_digest, rsa))
-      VbExError("Error verifying firmware body.\n");
-    printf("Body verification succeeded.\n");
+  for (i = 'A'; i <= 'B'; i++) {
+
+    fv_data = 0;
+    key_block = 0;
+    preamble = 0;
+
+    printf("FW %c\n", i);
+    sprintf(buf, "FW_MAIN_%c", i);
+    if (!FindInFmap(fmap, buf, image, image_size, &fv_data, &fv_size)) {
+      printf("Can't find %s in %s\n", buf, filename);
+      errorcnt++;
+      continue;
+    }
+
+    sprintf(buf, "VBLOCK_%c", i);
+    if (!FindInFmap(fmap, buf, image, image_size,
+                    (uint8_t **)&key_block, &key_block_size)) {
+      printf("Can't find %s in %s\n", buf, filename);
+      errorcnt++;
+      continue;
+    }
+
+    if (0 != KeyBlockVerify(key_block, key_block_size, pubkey, !pubkey)) {
+      printf("Error verifying key block for %s.\n", buf);
+      errorcnt++;
+      continue;
+    }
+    printf("  Key block:\n");
+    data_key = &key_block->data_key;
+    printf("    Size:                %" PRIu64 "\n",
+           key_block->key_block_size);
+    printf("    Flags:               %" PRIu64 " (ignored)\n",
+           key_block->key_block_flags);
+    printf("    Data key algorithm:  %" PRIu64 " %s\n", data_key->algorithm,
+           (data_key->algorithm < kNumAlgorithms ?
+            algo_strings[data_key->algorithm] : "(invalid)"));
+    printf("    Data key version:    %" PRIu64 "\n", data_key->key_version);
+    printf("    Data key sha1sum:    ");
+    PrintPubKeySha1Sum(data_key);
+    printf("\n");
+
+    preamble = (VbECPreambleHeader*)
+      ((uint8_t *)key_block + key_block->key_block_size);
+
+    rsa = PublicKeyToRSA(&key_block->data_key);
+    if (!rsa) {
+      printf("Error parsing data key.\n");
+      errorcnt++;
+    }
+    /* Verify preamble */
+    if (0 != VerifyECPreamble(preamble,
+                              key_block_size - key_block->key_block_size,
+                              rsa)) {
+      printf("Error verifying preamble.\n");
+      errorcnt++;
+      free(rsa);
+      continue;
+    }
+    printf("  Preamble:\n");
+    printf("    Size:                  %" PRIu64 "\n",
+           preamble->preamble_size);
+    printf("    Header version:        %" PRIu32 ".%" PRIu32"\n",
+           preamble->header_version_major,
+           preamble->header_version_minor);
+    printf("    Firmware version:      %" PRIu64 "\n",
+           preamble->firmware_version);
+    printf("    Firmware body size:    %" PRIu64 "\n",
+           preamble->body_digest.data_size);
+    printf("    Preamble flags:        %" PRIu32 "\n", preamble->flags);
+    printf("    Preamble name:         %s\n", preamble->name);
+
+    /* TODO: verify body size same as signature size */
+
+    /* Verify body */
+    if (preamble->flags & VB_FIRMWARE_PREAMBLE_USE_RO_NORMAL) {
+      printf("Preamble requests USE_RO_NORMAL; skipping verification.\n");
+    } else {
+      if (0 != EqualData(fv_data, fv_size,
+                         &preamble->body_digest, rsa)) {
+        printf("Error verifying firmware body.\n");
+        errorcnt++;
+      }
+    }
+    free(rsa);
   }
 
-  return 0;
-}
+  /* Done */
+  if (0 != munmap(image, sb.st_size))
+    VbExError("Can't munmap %s: %s\n", filename, strerror(errno));
 
+  printf("Done\n");
+  return errorcnt;
+}
 
 int main(int argc, char* argv[]) {
 
   char* filename = NULL;
-  char* key_block_file = NULL;
-  char* signpubkey = NULL;
-  char* signprivate = NULL;
   uint64_t version = 0;
   int got_version = 0;
-  char* fv_file = NULL;
   uint32_t preamble_flags = 0;
   char *name = NULL;
   int mode = 0;
-  int parse_error = 0;
+  VbKeyBlockHeader* key_block = 0;
+  VbPrivateKey* privkey = 0;
+  VbPublicKey* pubkey = 0;
+  uint64_t key_block_size;
+  int errorcnt = 0;
   char* e;
   int i;
 
   while ((i = getopt_long(argc, argv, "", long_opts, NULL)) != -1) {
     switch (i) {
-      case '?':
-        /* Unhandled option */
-        printf("Unknown option\n");
-        parse_error = 1;
-        break;
-
-      case OPT_MODE_VBLOCK:
-      case OPT_MODE_VERIFY:
-        mode = i;
-        filename = optarg;
-        break;
-
-      case OPT_KEYBLOCK:
-        key_block_file = optarg;
-        break;
-
-      case OPT_SIGNPUBKEY:
-        signpubkey = optarg;
-        break;
-
-      case OPT_SIGNPRIVATE:
-        signprivate = optarg;
-        break;
-
-      case OPT_FV:
-        fv_file = optarg;
-        break;
-
-      case OPT_VERSION:
-        version = strtoul(optarg, &e, 0);
-        if (!*optarg || (e && *e)) {
-          printf("Invalid --version\n");
-          parse_error = 1;
-        }
-        got_version = 1;
-        break;
-
-      case OPT_FLAGS:
-        preamble_flags = strtoul(optarg, &e, 0);
-        if (!*optarg || (e && *e)) {
-          printf("Invalid --flags\n");
-          parse_error = 1;
-        }
-        break;
+    case '?':
+      /* Unhandled option */
+      printf("Unknown option\n");
+      errorcnt++;
+      break;
+
+    case OPT_MODE_SIGN:
+    case OPT_MODE_VERIFY:
+      mode = i;
+      filename = optarg;
+      break;
+
+    case OPT_KEYBLOCK:
+      /* Read the key block and keys */
+      key_block = (VbKeyBlockHeader*)ReadFile(optarg, &key_block_size);
+      if (!key_block) {
+        printf("Error reading key block from %s\n", optarg);
+        errorcnt++;
+      }
+      break;
+
+    case OPT_SIGNPUBKEY:
+      pubkey = PublicKeyRead(optarg);
+      if (!pubkey) {
+        printf("Error reading public key from %s\n", optarg);
+        errorcnt++;
+      }
+      break;
+
+    case OPT_SIGNPRIVATE:
+      privkey = PrivateKeyRead(optarg);
+      if (!privkey) {
+        printf("Error reading private key from %s\n", optarg);
+        errorcnt++;
+      }
+      break;
+
+    case OPT_VERSION:
+      version = strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e)) {
+        printf("Invalid --version argument: \"%s\"\n", optarg);
+        errorcnt++;
+      }
+      got_version = 1;
+      break;
+
+    case OPT_FLAGS:
+      preamble_flags = strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e)) {
+        printf("Invalid --flags argument: \"%s\"\n", optarg);
+        errorcnt++;
+      }
+      break;
 
     case OPT_NAME:
       name = optarg;
@@ -307,21 +492,37 @@ int main(int argc, char* argv[]) {
     }
   }
 
-  if (parse_error)
-    return PrintHelp();
-
   switch(mode) {
-    case OPT_MODE_VBLOCK:
-      if (!got_version) {
-        printf("Must specify a version\n");
-        return PrintHelp();
-      }
-      return Vblock(filename, key_block_file, signprivate, version,
-                    fv_file, preamble_flags, name);
-    case OPT_MODE_VERIFY:
-      return Verify(filename, signpubkey, fv_file);
-    default:
-      printf("Must specify a mode.\n");
+
+  case OPT_MODE_SIGN:
+    /* Check required args */
+    if (!key_block) {
+      printf("The ----keyblock arg is required when signing\n");
+      errorcnt++;
+    }
+    if (!privkey) {
+      printf("The --signprivate arg is required when signing\n");
+      errorcnt++;
+    }
+    if (!got_version) {
+      printf("The --version arg is required when signing\n");
+      errorcnt++;
+    }
+
+    if (errorcnt)
       return PrintHelp();
+
+    /* Sign or die */
+    SignImage(filename, key_block, key_block_size,
+              privkey, version, pubkey, preamble_flags, name);
+
+    /* fall through and verify what we've just done */
+
+  case OPT_MODE_VERIFY:
+    return Verify(filename);
+
+  default:
+    printf("\nMust specify a mode, either --sign or --verify.\n\n");
+    return PrintHelp();
   }
 }