From cc341f856b78d7f3733be59a5ec2042851cffac2 Mon Sep 17 00:00:00 2001
From: Michael Marineau <michael.marineau@coreos.com>
Date: Thu, 19 Dec 2013 16:53:14 -0800
Subject: [PATCH] fix(disk_util): Merge write_mbr into write_gpt

No need for this to be a different command.
---
 build_library/disk_layout_util.sh |  3 +--
 build_library/disk_util           | 40 ++++++++++---------------------
 2 files changed, 14 insertions(+), 29 deletions(-)

diff --git a/build_library/disk_layout_util.sh b/build_library/disk_layout_util.sh
index 915f8d3255..59b869bf39 100644
--- a/build_library/disk_layout_util.sh
+++ b/build_library/disk_layout_util.sh
@@ -34,8 +34,7 @@ write_partition_table() {
     ;;
   esac
 
-  cgpt_py write_gpt "${outdev}"
-  cgpt_py write_mbr "${outdev}" "${pmbr_img}"
+  cgpt_py write_gpt --mbr_boot_code="${pmbr_img}" "${outdev}"
 }
 
 get_fs_block_size() {
diff --git a/build_library/disk_util b/build_library/disk_util
index 6b391e754b..3751a9c70f 100755
--- a/build_library/disk_util
+++ b/build_library/disk_util
@@ -243,6 +243,7 @@ def WritePartitionTable(options):
   Cgpt('create', '-c', '-s', disk_block_count, options.disk_image)
 
   sector = GPT_RESERVED_SECTORS
+  esp_number = None
   for partition in partitions:
     sector = Align(sector, partition['alignment'])
     if partition['type'] != 'blank':
@@ -254,33 +255,21 @@ def WritePartitionTable(options):
                   '-u', partition['uuid'],
                   options.disk_image)
 
-    sector += partition['blocks']
-
-  Cgpt('show', options.disk_image)
-
-
-def WriteMbrBoot(options):
-  """Writes the protective MBR with the given boot code.
-
-  The EFI System Partition will be marked as the 'boot' partition.
-
-  Args:
-    options: Flags passed to the script
-  """
-
-  config = LoadPartitionConfig(options)
-  partitions = GetPartitionTable(options, config)
-
-  esp_number = None
-  for partition in partitions:
     if partition['type'] == 'efi':
       esp_number = partition['num']
-      break
+
+    sector += partition['blocks']
+
   if esp_number is None:
     raise InvalidLayout('Table does not include an EFI partition.')
 
-  subprocess.check_call(['cgpt', 'boot', '-p', '-b', options.mbr_boot_code,
-      '-i', str(partition['num']), options.disk_image])
+  if options.mbr_boot_code:
+    Cgpt('boot', '-p',
+                 '-b', options.mbr_boot_code,
+                 '-i', esp_number,
+                 options.disk_image)
+
+  Cgpt('show', options.disk_image)
 
 
 def GetPartitionByNumber(partitions, num):
@@ -473,14 +462,11 @@ def main(argv):
   actions = parser.add_subparsers(title='actions')
 
   a = actions.add_parser('write_gpt', help='write gpt to new image')
+  a.add_argument('--mbr_boot_code',
+          help='path to mbr boot block, such as syslinux/gptmbr.bin')
   a.add_argument('disk_image', help='path to disk image file')
   a.set_defaults(func=WritePartitionTable)
 
-  a = actions.add_parser('write_mbr', help='write mbr to existing image')
-  a.add_argument('disk_image', help='path to disk image file')
-  a.add_argument('mbr_boot_code', help='path to mbr boot block')
-  a.set_defaults(func=WriteMbrBoot)
-
   a = actions.add_parser('readblocksize', help='get device block size')
   a.set_defaults(func=GetBlockSize)