Nearly complete rewrite of cgpt tool.

cgpt/cmd_add.c

+// Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "cgpt.h"
+
+#include <getopt.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <uuid/uuid.h>
+
+#include "cgptlib_internal.h"
+
+static void Usage(void)
+{
+  printf("\nUsage: %s add [OPTIONS] DRIVE\n\n"
+         "Add, edit, or remove a partition entry.\n\n"
+         "Options:\n"
+         "  -i NUM       Specify partition (default is next available)\n"
+         "  -b NUM       Beginning sector\n"
+         "  -s NUM       Size in sectors\n"
+         "  -t GUID      Partition Type GUID\n"
+         "  -u GUID      Partition Unique ID\n"
+         "  -l LABEL     Label\n"
+         "  -S NUM       set Successful flag (0|1)\n"
+         "  -T NUM       set Tries flag (0-15)\n"
+         "  -P NUM       set Priority flag (0-15)\n"
+         "  -A NUM       set raw 64-bit attribute value\n"
+         "\n"
+         "Use the -i option to modify an existing partition.\n"
+         "The -b, -s, and -t options must be given for new partitions.\n"
+         "\n", progname);
+  PrintTypes();
+}
+
+int cmd_add(int argc, char *argv[]) {
+  struct drive drive;
+  int partition = 0;
+  uint64_t begin = 0;
+  uint64_t size = 0;
+  Guid type_guid;
+  Guid unique_guid;
+  char *label = 0;
+  int successful = 0;
+  int tries = 0;
+  int priority = 0;
+  uint64_t raw_value = 0;
+  int set_begin = 0;
+  int set_size = 0;
+  int set_type = 0;
+  int set_unique = 0;
+  int set_successful = 0;
+  int set_tries = 0;
+  int set_priority = 0;
+  int set_raw = 0;
+
+  int gpt_retval;
+  GptEntry *entry;
+  int index;
+  
+  int c;
+  int errorcnt = 0;
+  char *e = 0;
+
+  opterr = 0;                     // quiet, you
+  while ((c=getopt(argc, argv, ":hi:b:s:t:u:l:S:T:P:A:")) != -1)
+  {
+    switch (c)
+    {
+    case 'i':
+      partition = (uint32_t)strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
+    case 'b':
+      set_begin = 1;
+      begin = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
+    case 's':
+      set_size = 1;
+      size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
+    case 't':
+      set_type = 1;
+      if (CGPT_OK != SupportedType(optarg, &type_guid) &&
+          CGPT_OK != StrToGuid(optarg, &type_guid)) {
+        Error("invalid argument to -%c: %s\n", c, optarg);
+        errorcnt++;
+      }
+      break;
+    case 'u':
+      set_unique = 1;
+      if (CGPT_OK != StrToGuid(optarg, &unique_guid)) {
+        Error("invalid argument to -%c: %s\n", c, optarg);
+        errorcnt++;
+      }
+      break;
+    case 'l':
+      label = optarg;
+      break;
+    case 'S':
+      set_successful = 1;
+      successful = (uint32_t)strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      if (successful < 0 || successful > 1) {
+        Error("value for -%c must be between 0 and 1", c);
+        errorcnt++;
+      }
+      break;
+    case 'T':
+      set_tries = 1;
+      tries = (uint32_t)strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        fprintf(stderr, "%s: invalid argument to -%c: \"%s\"\n",
+                progname, c, optarg);
+        errorcnt++;
+      }
+      if (tries < 0 || tries > 15) {
+        Error("value for -%c must be between 0 and 15", c);
+        errorcnt++;
+      }
+      break;
+    case 'P':
+      set_priority = 1;
+      priority = (uint32_t)strtoul(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      if (priority < 0 || priority > 15) {
+        Error("value for -%c must be between 0 and 15", c);
+        errorcnt++;
+      }
+      break;
+    case 'A':
+      set_raw = 1;
+      raw_value = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
+
+    case 'h':
+      Usage();
+      return CGPT_OK;
+    case '?':
+      Error("unrecognized option: -%c\n", optopt);
+      errorcnt++;
+      break;
+    case ':':
+      Error("missing argument to -%c\n", optopt);
+      errorcnt++;
+      break;
+    default:
+      errorcnt++;
+      break;
+    }
+  }
+  if (errorcnt)
+  {
+    Usage();
+    return CGPT_FAILED;
+  }
+
+  if (optind >= argc) {
+    Error("missing drive argument\n");
+    return CGPT_FAILED;
+  }
+
+  if (CGPT_OK != DriveOpen(argv[optind], &drive))
+    return CGPT_FAILED;
+
+  if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
+    Error("GptSanityCheck() returned %d: %s\n",
+          gpt_retval, GptError(gpt_retval));
+    return CGPT_FAILED;
+  }
+
+  int max_part = GetNumberOfEntries(&drive.gpt);
+  if (partition) {
+    if (partition > max_part) {
+      Error("invalid partition number: %d\n", partition);
+      goto bad;
+    }
+    index = partition - 1;
+    entry = GetEntry(&drive.gpt, PRIMARY, index);
+  } else {
+    // find next empty partition
+    for (index = 0; index < max_part; index++) {
+      entry = GetEntry(&drive.gpt, PRIMARY, index);
+      if (IsZero(&entry->type)) {
+        partition = index + 1;
+        break;
+      }
+    }
+    if (index >= max_part) {
+      Error("no unused partitions available\n");
+      goto bad;
+    }
+  }
+
+  // New partitions must specify type, begin, and size.
+  if (IsZero(&entry->type)) {
+    if (!set_begin || !set_size || !set_type) {
+      Error("-t, -b, and -s options are required for new partitions\n");
+      goto bad;
+    }
+    if (IsZero(&type_guid)) {
+      Error("New partitions must have a type other than \"unused\"\n");
+      goto bad;
+    }
+    if (!set_unique)
+      uuid_generate((uint8_t *)&entry->unique);
+  }
+
+  if (set_begin)
+    entry->starting_lba = begin;
+  if (set_size)
+    entry->ending_lba = begin + size - 1;
+  if (set_type)
+    memcpy(&entry->type, &type_guid, sizeof(Guid));
+  if (set_unique)
+    memcpy(&entry->unique, &unique_guid, sizeof(Guid));
+  if (label) {
+    uint16_t buf[128];
+    UTF8ToUTF16((uint8_t *)label, buf);
+    memcpy(entry->name, buf, sizeof(entry->name));
+  }
+  if (set_raw) {
+    entry->attributes = raw_value;
+  } else {
+    if (set_successful)
+      SetSuccessful(&drive.gpt, PRIMARY, index, successful);
+    if (set_tries)
+      SetTries(&drive.gpt, PRIMARY, index, tries);
+    if (set_priority)
+      SetPriority(&drive.gpt, PRIMARY, index, priority);
+  }
+
+  RepairEntries(&drive.gpt, MASK_PRIMARY);
+  RepairHeader(&drive.gpt, MASK_PRIMARY);
+
+  drive.gpt.modified |= (GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1 |
+                         GPT_MODIFIED_HEADER2 | GPT_MODIFIED_ENTRIES2);
+  UpdateCrc(&drive.gpt);
+
+
+  // Write it all out
+  return DriveClose(&drive, 1);
+
+bad:
+  (void) DriveClose(&drive, 0);
+  return CGPT_FAILED;
+}