cgpt: add edit command

This change adds a command to cgpt to
change the GUID of the drive.

BRANCH=none
BUG=None
TEST=Compiled and ran utility to verify that GUID changes.
Also verified that the new and existing tests completed
successfully.

Change-Id: Ia8a815447509626312e2b06c6f293901290c73c3
Signed-off-by: Matt Delco <delco@google.com>
Reviewed-on: https://chromium-review.googlesource.com/1171834
Reviewed-by: Julius Werner <jwerner@chromium.org>
diff --git a/Makefile b/Makefile
index 136d48e..2c4dd7f 100644
--- a/Makefile
+++ b/Makefile
@@ -550,6 +550,7 @@
 	cgpt/cgpt_boot.c \
 	cgpt/cgpt_common.c \
 	cgpt/cgpt_create.c \
+	cgpt/cgpt_edit.c \
 	cgpt/cgpt_find.c \
 	cgpt/cgpt_legacy.c \
 	cgpt/cgpt_nor.c \
@@ -559,6 +560,7 @@
 	cgpt/cmd_add.c \
 	cgpt/cmd_boot.c \
 	cgpt/cmd_create.c \
+	cgpt/cmd_edit.c \
 	cgpt/cmd_find.c \
 	cgpt/cmd_legacy.c \
 	cgpt/cmd_prioritize.c \
diff --git a/cgpt/cgpt.c b/cgpt/cgpt.c
index 3809200..0977b65 100644
--- a/cgpt/cgpt.c
+++ b/cgpt/cgpt.c
@@ -34,6 +34,7 @@
   {"repair", cmd_repair, "Repair damaged GPT headers and tables"},
   {"boot", cmd_boot, "Edit the PMBR sector for legacy BIOSes"},
   {"find", cmd_find, "Locate a partition by its GUID"},
+  {"edit", cmd_edit, "Edit a drive entry"},
   {"prioritize", cmd_prioritize,
    "Reorder the priority of all kernel partitions"},
   {"legacy", cmd_legacy, "Switch between GPT and Legacy GPT"},
diff --git a/cgpt/cgpt.h b/cgpt/cgpt.h
index 86f4b85..c53f0f3 100644
--- a/cgpt/cgpt.h
+++ b/cgpt/cgpt.h
@@ -201,6 +201,7 @@
 int cmd_add(int argc, char *argv[]);
 int cmd_boot(int argc, char *argv[]);
 int cmd_find(int argc, char *argv[]);
+int cmd_edit(int argc, char *argv[]);
 int cmd_prioritize(int argc, char *argv[]);
 int cmd_legacy(int argc, char *argv[]);
 
diff --git a/cgpt/cgpt_edit.c b/cgpt/cgpt_edit.c
new file mode 100644
index 0000000..098187f
--- /dev/null
+++ b/cgpt/cgpt_edit.c
@@ -0,0 +1,47 @@
+// Copyright 2018 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 "cgpt_params.h"
+#include "cgptlib_internal.h"
+#include "vboot_host.h"
+
+int CgptEdit(CgptEditParams *params) {
+  struct drive drive;
+  GptHeader *h;
+  int gpt_retval;
+
+  if (params == NULL)
+    return CGPT_FAILED;
+
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR, 0))
+    return CGPT_FAILED;
+
+  if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
+    Error("GptSanityCheck() returned %d: %s\n",
+          gpt_retval, GptError(gpt_retval));
+    goto bad;
+  }
+
+  if (CGPT_OK != CheckValid(&drive)) {
+    Error("Please run 'cgpt repair' before changing settings.\n");
+    goto bad;
+  }
+
+  h = (GptHeader *)drive.gpt.primary_header;
+  memcpy(&h->disk_uuid, &params->unique_guid, sizeof(h->disk_uuid));
+  // Copy to secondary
+  RepairHeader(&drive.gpt, MASK_PRIMARY);
+  drive.gpt.modified |= (GPT_MODIFIED_HEADER1 | GPT_MODIFIED_HEADER2);
+
+  UpdateCrc(&drive.gpt);
+
+  // Write it all out.
+  return DriveClose(&drive, 1);
+
+bad:
+
+  DriveClose(&drive, 0);
+  return CGPT_FAILED;
+}
diff --git a/cgpt/cmd_edit.c b/cgpt/cmd_edit.c
new file mode 100644
index 0000000..4f4290b
--- /dev/null
+++ b/cgpt/cmd_edit.c
@@ -0,0 +1,78 @@
+// Copyright 2018 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 <getopt.h>
+
+#include "cgpt.h"
+#include "vboot_host.h"
+
+extern const char* progname;
+
+static void Usage(void)
+{
+  printf("\nUsage: %s edit [OPTIONS] DRIVE\n\n"
+         "Edit a drive's parameters.\n\n"
+         "Options:\n"
+         "  -u GUID      Drive Unique ID\n"
+         "\n", progname);
+}
+
+int cmd_edit(int argc, char *argv[]) {
+
+  CgptEditParams params;
+  memset(&params, 0, sizeof(params));
+
+  int c;
+  int errorcnt = 0;
+
+  opterr = 0;                     // quiet, you
+  while ((c=getopt(argc, argv, ":hu:")) != -1)
+  {
+    switch (c)
+    {
+    case 'u':
+      params.set_unique = 1;
+      if (CGPT_OK != StrToGuid(optarg, &params.unique_guid)) {
+        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;
+  }
+
+  params.drive_name = argv[optind];
+
+  if (!params.set_unique)
+  {
+    Error("no parameters were edited\n");
+    return CGPT_FAILED;
+  }
+
+  return CgptEdit(&params);
+}
diff --git a/host/include/cgpt_params.h b/host/include/cgpt_params.h
index 7c4a4b2..5bd1f76 100644
--- a/host/include/cgpt_params.h
+++ b/host/include/cgpt_params.h
@@ -48,6 +48,12 @@
 	int set_raw;
 } CgptAddParams;
 
+typedef struct CgptEditParams {
+	char *drive_name;
+	Guid unique_guid;
+	int set_unique;
+} CgptEditParams;
+
 typedef struct CgptShowParams {
 	char *drive_name;
 	uint64_t drive_size;
diff --git a/host/include/vboot_host.h b/host/include/vboot_host.h
index 9ebe404..cb0d0f3 100644
--- a/host/include/vboot_host.h
+++ b/host/include/vboot_host.h
@@ -19,6 +19,7 @@
 /* partition table manipulation */
 int CgptCreate(CgptCreateParams *params);
 int CgptAdd(CgptAddParams *params);
+int CgptEdit(CgptEditParams *params);
 int CgptSetAttributes(CgptAddParams *params);
 int CgptGetPartitionDetails(CgptAddParams *params);
 int CgptBoot(CgptBootParams *params);
diff --git a/tests/run_cgpt_tests.sh b/tests/run_cgpt_tests.sh
index 88799ea..81e5d3b 100755
--- a/tests/run_cgpt_tests.sh
+++ b/tests/run_cgpt_tests.sh
@@ -74,6 +74,8 @@
 RANDOM_GUID='2364a860-bf63-42fb-a83d-9ad3e057fcf5'
 RANDOM_NUM=6
 
+RANDOM_DRIVE_GUID='12345678-0000-1111-2222-123456789ABC'
+
 $CGPT create $MTD ${DEV}
 
 run_basic_tests() {
@@ -144,9 +146,32 @@
   $CGPT add $MTD -i 1 -t data ${DEV} || error
   X=$($CGPT show $MTD -t -i 1 ${DEV} | tr 'A-Z' 'a-z')
   [ "$X" = "$DATA_GUID" ] || error
+
+  ORIG_ID=$($CGPT show $MTD -v ${DEV} | \
+    grep -i "disk uuid" | head -1 | awk ' { print $3 } ' )
+  [ ! "$ORIG_ID" = "$RANDOM_DRIVE_GUID" ] || error
+  $CGPT edit $MTD -u ${RANDOM_DRIVE_GUID} ${DEV} || error
+  X=$($CGPT show $MTD -v ${DEV} | grep -i "disk uuid" | \
+    head -1 | awk ' { print $3 } ' )
+  [ "$X" = "${RANDOM_DRIVE_GUID}" ] || error
+  $CGPT edit $MTD -u ${ORIG_ID} ${DEV} || error
+  X=$($CGPT show $MTD -v ${DEV} | grep -i "disk uuid" | \
+    head -1 | awk ' { print $3 } ' )
+  [ "$X" = "${ORIG_ID}" ] || error
 }
 run_basic_tests
 
+ORIG_ID=$($CGPT show $MTD -v ${DEV} | \
+  grep -i "disk uuid" | awk ' { print $3 } ' )
+[ ! "$ORIG_ID" = "$RANDOM_DRIVE_GUID" ] || error
+$CGPT edit $MTD -u ${RANDOM_DRIVE_GUID} ${DEV} || error
+X=$($CGPT show $MTD -v ${DEV} | grep -i "disk uuid" | \
+  head -1 | awk ' { print $3 } ' )
+[ "$X" = "${RANDOM_DRIVE_GUID}" ] || error
+$CGPT edit $MTD -u ${ORIG_ID} ${DEV} || error
+X=$($CGPT show $MTD -v ${DEV} | grep -i "disk uuid" | \
+  awk ' { print $3 } ' )
+[ "$X" = "${ORIG_ID}" ] || error
 
 echo "Set the boot partition.."
 $CGPT boot $MTD -i ${KERN_NUM} ${DEV} >/dev/null