diff --git a/bin/cros_make_image_bootable b/bin/cros_make_image_bootable
index 6d33cf1cfd..e27862c367 100755
--- a/bin/cros_make_image_bootable
+++ b/bin/cros_make_image_bootable
@@ -35,7 +35,6 @@ assert_inside_chroot
 . /usr/lib/installer/chromeos-common.sh || exit 1
 . "${BUILD_LIBRARY_DIR}/toolchain_util.sh" || exit 1
 . "${BUILD_LIBRARY_DIR}/build_image_util.sh" || exit 1
-. "${BUILD_LIBRARY_DIR}/disk_layout_util.sh" || exit 1
 
 switch_to_strict_mode
 
diff --git a/build_library/base_image_util.sh b/build_library/base_image_util.sh
index 52331b3fab..55fc293c45 100755
--- a/build_library/base_image_util.sh
+++ b/build_library/base_image_util.sh
@@ -41,12 +41,10 @@ create_base_image() {
   local rootfs_verification_enabled=$2
 
   get_disk_layout_type
-  local image_type="${DISK_LAYOUT_TYPE}"
-
   check_valid_layout "base"
-  check_valid_layout ${image_type}
+  check_valid_layout "${DISK_LAYOUT_TYPE}"
 
-  info "Using image type ${image_type}"
+  info "Using image type ${DISK_LAYOUT_TYPE}"
 
   root_fs_dir="${BUILD_DIR}/rootfs"
   state_fs_dir="${BUILD_DIR}/state"
@@ -57,25 +55,25 @@ create_base_image() {
   cleanup_mounts &> /dev/null
 
   local root_fs_label="ROOT-A"
-  local root_fs_num=$(get_num ${image_type} ${root_fs_label})
+  local root_fs_num=$(get_num ${root_fs_label})
   local root_fs_img="${BUILD_DIR}/rootfs.image"
-  local root_fs_bytes=$(get_filesystem_size ${image_type} ${root_fs_num})
+  local root_fs_bytes=$(get_filesystem_size ${root_fs_num})
 
   local state_fs_label="STATE"
-  local state_fs_num=$(get_num ${image_type} ${state_fs_label})
+  local state_fs_num=$(get_num ${state_fs_label})
   local state_fs_img="${BUILD_DIR}/state.image"
-  local state_fs_bytes=$(get_filesystem_size ${image_type} ${state_fs_num})
+  local state_fs_bytes=$(get_filesystem_size ${state_fs_num})
   local state_fs_uuid=$(uuidgen)
 
   local esp_fs_label="EFI-SYSTEM"
-  local esp_fs_num=$(get_num ${image_type} ${esp_fs_label})
+  local esp_fs_num=$(get_num ${esp_fs_label})
   local esp_fs_img="${BUILD_DIR}/esp.image"
-  local esp_fs_bytes=$(get_filesystem_size ${image_type} ${esp_fs_num})
+  local esp_fs_bytes=$(get_filesystem_size ${esp_fs_num})
 
   local oem_fs_label="OEM"
-  local oem_fs_num=$(get_num ${image_type} ${oem_fs_label})
+  local oem_fs_num=$(get_num ${oem_fs_label})
   local oem_fs_img="${BUILD_DIR}/oem.image"
-  local oem_fs_bytes=$(get_filesystem_size ${image_type} ${oem_fs_num})
+  local oem_fs_bytes=$(get_filesystem_size ${oem_fs_num})
   local oem_fs_uuid=$(uuidgen)
 
   local fs_block_size=$(get_fs_block_size)
diff --git a/build_library/create_legacy_bootloader_templates.sh b/build_library/create_legacy_bootloader_templates.sh
index c4dd64a08b..f6b6467fa1 100755
--- a/build_library/create_legacy_bootloader_templates.sh
+++ b/build_library/create_legacy_bootloader_templates.sh
@@ -42,8 +42,8 @@ if [[ "${FLAGS_arch}" = "x86" || "${FLAGS_arch}" = "amd64"  ]]; then
   sudo mkdir -p ${FLAGS_to}
 
   # Get partition UUIDs from the json config
-  ROOTA="PARTUUID=$(get_uuid base ROOT-A)"
-  ROOTB="PARTUUID=$(get_uuid base ROOT-B)"
+  ROOTA="PARTUUID=$(get_uuid ROOT-A)"
+  ROOTB="PARTUUID=$(get_uuid ROOT-B)"
 
   # Build configuration files for pygrub/pvgrub
   GRUB_DIR="${FLAGS_to}/grub"
diff --git a/build_library/disk_layout_util.sh b/build_library/disk_layout_util.sh
index 7cd02a5f4e..915f8d3255 100644
--- a/build_library/disk_layout_util.sh
+++ b/build_library/disk_layout_util.sh
@@ -17,13 +17,8 @@ cgpt_py() {
   "${CGPT_PY}" "$@"
 }
 
-get_disk_layout_path() {
-  DISK_LAYOUT_PATH="${BUILD_LIBRARY_DIR}/legacy_disk_layout.json"
-}
-
 write_partition_table() {
-  local image_type=$1
-  local outdev=$2
+  local outdev=$1
 
   local pmbr_img
   case ${ARCH} in
@@ -39,77 +34,59 @@ write_partition_table() {
     ;;
   esac
 
-  get_disk_layout_path
-  cgpt_py write_gpt "${image_type}" "${DISK_LAYOUT_PATH}" "${outdev}"
-  cgpt_py write_mbr \
-      "${image_type}" "${DISK_LAYOUT_PATH}" "${outdev}" "${pmbr_img}"
+  cgpt_py write_gpt "${outdev}"
+  cgpt_py write_mbr "${outdev}" "${pmbr_img}"
 }
 
 get_fs_block_size() {
-  get_disk_layout_path
-
-  cgpt_py readfsblocksize "${DISK_LAYOUT_PATH}"
+  cgpt_py readfsblocksize
 }
 
 get_block_size() {
-  get_disk_layout_path
-
-  cgpt_py readblocksize "${DISK_LAYOUT_PATH}"
+  cgpt_py readblocksize
 }
 
 get_partition_size() {
-  local image_type=$1
-  local part_id=$2
-  get_disk_layout_path
+  local part_id=$1
 
-  cgpt_py readpartsize "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
+  cgpt_py readpartsize ${part_id}
 }
 
 get_filesystem_size() {
-  local image_type=$1
-  local part_id=$2
-  get_disk_layout_path
+  local part_id=$1
 
-  cgpt_py readfssize "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
+  cgpt_py readfssize ${part_id}
 }
 
 get_label() {
-  local image_type=$1
-  local part_id=$2
-  get_disk_layout_path
+  local part_id=$1
 
-  cgpt_py readlabel "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
+  cgpt_py readlabel ${part_id}
 }
 
 get_num() {
-  local image_type=$1
-  local label=$2
-  get_disk_layout_path
+  local label=$1
 
-  cgpt_py readnum "${image_type}" "${DISK_LAYOUT_PATH}" ${label}
+  cgpt_py readnum ${label}
 }
 
 get_uuid() {
-  local image_type=$1
-  local label=$2
-  get_disk_layout_path
+  local label=$1
 
-  cgpt_py readuuid "${image_type}" "${DISK_LAYOUT_PATH}" ${label}
+  cgpt_py readuuid ${label}
 }
 
 check_valid_layout() {
-  local image_type=$1
-  get_disk_layout_path
-
-  cgpt_py parseonly "${image_type}" "${DISK_LAYOUT_PATH}" > /dev/null
+  cgpt_py parseonly > /dev/null
 }
 
 get_disk_layout_type() {
-  DISK_LAYOUT_TYPE="base"
+  DISK_LAYOUT_TYPE="${1:-base}"
   if [[ -n "${FLAGS_disk_layout}" && \
       "${FLAGS_disk_layout}" != "default" ]]; then
     DISK_LAYOUT_TYPE="${FLAGS_disk_layout}"
   fi
+  export DISK_LAYOUT_TYPE
 }
 
 emit_gpt_scripts() {
@@ -180,7 +157,7 @@ build_gpt() {
   local oem_img="$5"
 
   get_disk_layout_type
-  write_partition_table "${DISK_LAYOUT_TYPE}" "${outdev}"
+  write_partition_table "${outdev}"
 
   local sudo=
   if [ ! -w "$outdev" ] ; then
diff --git a/build_library/disk_util b/build_library/disk_util
index 0299730204..6b391e754b 100755
--- a/build_library/disk_util
+++ b/build_library/disk_util
@@ -3,13 +3,13 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import argparse
 import json
 import os
 import re
 import subprocess
 import sys
 import uuid
-from optparse import OptionParser
 
 # First sector we can use.
 GPT_RESERVED_SECTORS = 34
@@ -24,11 +24,11 @@ class InvalidAdjustment(Exception):
   pass
 
 
-def LoadPartitionConfig(filename):
+def LoadPartitionConfig(options):
   """Loads a partition tables configuration file into a Python object.
 
   Args:
-    filename: Filename to load into object
+    options: Flags passed to the script
   Returns:
     Object containing disk layout configuration
   """
@@ -38,6 +38,7 @@ def LoadPartitionConfig(filename):
       '_comment', 'type', 'num', 'label', 'blocks', 'block_size', 'fs_blocks',
       'fs_block_size', 'features', 'uuid', 'alignment'))
 
+  filename = options.disk_layout_file
   if not os.path.exists(filename):
     raise ConfigNotFound('Partition config %s was not found!' % filename)
   with open(filename) as f:
@@ -99,9 +100,7 @@ def LoadPartitionConfig(filename):
   return config
 
 
-
-
-def GetPartitionTable(options, config, image_type):
+def GetPartitionTable(options, config):
   """Generates requested image_type layout from a layout configuration.
   This loads the base table and then overlays the requested layout over
   the base layout.
@@ -109,13 +108,13 @@ def GetPartitionTable(options, config, image_type):
   Args:
     options: Flags passed to the script
     config: Partition configuration file object
-    image_type: Type of image eg base/test/dev/factory_install
   Returns:
     Object representing a selected partition table
   """
 
   partitions = config['layouts']['base']
   metadata = config['metadata']
+  image_type = options.disk_layout
 
   if image_type != 'base':
     for partition_t in config['layouts'][image_type]:
@@ -198,29 +197,26 @@ def ApplyPartitionAdjustment(partitions, metadata, label, operator, operand):
     partition['bytes'] = partition['blocks'] * metadata['block_size']
 
 
-def GetPartitionTableFromConfig(options, layout_filename, image_type):
+def GetPartitionTableFromConfig(options):
   """Loads a partition table and returns a given partition table type
 
   Args:
     options: Flags passed to the script
-    layout_filename: The filename to load tables from
-    image_type: The type of partition table to return
+  Returns:
+    A list defining all known partitions.
   """
 
-  config = LoadPartitionConfig(layout_filename)
-  partitions = GetPartitionTable(options, config, image_type)
+  config = LoadPartitionConfig(options)
+  partitions = GetPartitionTable(options, config)
 
   return partitions
 
 
-def WritePartitionTable(options, image_type, layout_filename, disk_filename):
+def WritePartitionTable(options):
   """Writes the given partition table to a disk image or device.
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
-    disk_filename: Path to disk image or device file
   """
 
   def Cgpt(*args):
@@ -232,9 +228,8 @@ def WritePartitionTable(options, image_type, layout_filename, disk_filename):
       count += alignment - offset
     return count
 
-  config = LoadPartitionConfig(layout_filename)
-  metadata = config['metadata']
-  partitions = GetPartitionTable(options, config, image_type)
+  config = LoadPartitionConfig(options)
+  partitions = GetPartitionTable(options, config)
   disk_block_count = GPT_RESERVED_SECTORS
 
   for partition in partitions:
@@ -245,7 +240,7 @@ def WritePartitionTable(options, image_type, layout_filename, disk_filename):
   # Sometimes qemu-img expects disks sizes aligned to 64k
   disk_block_count = Align(disk_block_count, config['metadata']['alignment'])
 
-  Cgpt('create', '-c', '-s', disk_block_count, disk_filename)
+  Cgpt('create', '-c', '-s', disk_block_count, options.disk_image)
 
   sector = GPT_RESERVED_SECTORS
   for partition in partitions:
@@ -257,29 +252,24 @@ def WritePartitionTable(options, image_type, layout_filename, disk_filename):
                   '-t', partition['type'],
                   '-l', partition['label'],
                   '-u', partition['uuid'],
-                  disk_filename)
+                  options.disk_image)
 
     sector += partition['blocks']
 
-  Cgpt('show', disk_filename)
+  Cgpt('show', options.disk_image)
 
 
-def WriteMbrBoot(options, image_type, layout_filename,
-                 disk_filename, mbr_filename):
+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
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
-    disk_filename: Path to disk image or device file
-    mbr_filename: Path to boot code, usually gptmbr.bin from syslinux.
   """
 
-  config = LoadPartitionConfig(layout_filename)
-  partitions = GetPartitionTable(options, config, image_type)
+  config = LoadPartitionConfig(options)
+  partitions = GetPartitionTable(options, config)
 
   esp_number = None
   for partition in partitions:
@@ -289,8 +279,8 @@ def WriteMbrBoot(options, image_type, layout_filename,
   if esp_number is None:
     raise InvalidLayout('Table does not include an EFI partition.')
 
-  subprocess.check_call(['cgpt', 'boot', '-p', '-b', mbr_filename,
-      '-i', str(partition['num']), disk_filename])
+  subprocess.check_call(['cgpt', 'boot', '-p', '-b', options.mbr_boot_code,
+      '-i', str(partition['num']), options.disk_image])
 
 
 def GetPartitionByNumber(partitions, num):
@@ -329,142 +319,107 @@ def GetPartitionByLabel(partitions, label):
   raise PartitionNotFound('Partition not found')
 
 
-def GetBlockSize(options, layout_filename):
+def GetBlockSize(options):
   """Returns the partition table block size.
 
   Args:
     options: Flags passed to the script
-    layout_filename: Path to partition configuration file
-  Returns:
+  Prints:
     Block size of all partitions in the layout
   """
 
-  config = LoadPartitionConfig(layout_filename)
-  return config['metadata']['block_size']
+  config = LoadPartitionConfig(options)
+  print config['metadata']['block_size']
 
 
-def GetFilesystemBlockSize(options, layout_filename):
+def GetFilesystemBlockSize(options):
   """Returns the filesystem block size.
 
-  Args:
-    options: Flags passed to the script
-
   This is used for all partitions in the table that have filesystems.
 
   Args:
-    layout_filename: Path to partition configuration file
-  Returns:
+    options: Flags passed to the script
+  Prints:
     Block size of all filesystems in the layout
   """
 
-  config = LoadPartitionConfig(layout_filename)
-  return config['metadata']['fs_block_size']
+  config = LoadPartitionConfig(options)
+  print config['metadata']['fs_block_size']
 
 
-def GetPartitionSize(options, image_type, layout_filename, num):
+def GetPartitionSize(options):
   """Returns the partition size of a given partition for a given layout type.
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
-    num: Number of the partition you want to read from
-  Returns:
+  Prints:
     Size of selected partition in bytes
   """
 
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
-  partition = GetPartitionByNumber(partitions, num)
-
-  return partition['bytes']
+  partitions = GetPartitionTableFromConfig(options)
+  partition = GetPartitionByNumber(partitions, options.partition_num)
+  print partition['bytes']
 
 
-def GetFilesystemSize(options, image_type, layout_filename, num):
+def GetFilesystemSize(options):
   """Returns the filesystem size of a given partition for a given layout type.
 
   If no filesystem size is specified, returns the partition size.
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
-    num: Number of the partition you want to read from
-  Returns:
+  Prints:
     Size of selected partition filesystem in bytes
   """
 
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
-  partition = GetPartitionByNumber(partitions, num)
-
-  if 'fs_bytes' in partition:
-    return partition['fs_bytes']
-  else:
-    return partition['bytes']
+  partitions = GetPartitionTableFromConfig(options)
+  partition = GetPartitionByNumber(partitions, options.partition_num)
+  print partition.get('fs_bytes', partition['bytes'])
 
 
-def GetLabel(options, image_type, layout_filename, num):
+def GetLabel(options):
   """Returns the label for a given partition.
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
-    num: Number of the partition you want to read from
-  Returns:
+  Prints:
     Label of selected partition, or 'UNTITLED' if none specified
   """
 
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
-  partition = GetPartitionByNumber(partitions, num)
-
-  if 'label' in partition:
-    return partition['label']
-  else:
-    return 'UNTITLED'
+  partitions = GetPartitionTableFromConfig(options)
+  partition = GetPartitionByNumber(partitions, options.partition_num)
+  print partition.get('label', 'UNTITLED')
 
 
-def GetNum(options, image_type, layout_filename, label):
+def GetNum(options):
   """Returns the number for a given label.
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
-    label: Label of the partition you want to read from
-  Returns:
+  Prints:
     Number of selected partition, or '-1' if there is no number
   """
 
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
-  partition = GetPartitionByLabel(partitions, label)
-
-  if 'num' in partition:
-    return partition['num']
-  else:
-    return '-1'
+  partitions = GetPartitionTableFromConfig(options)
+  partition = GetPartitionByLabel(partitions, options.label)
+  print partition.get('num', '-1')
 
 
-def GetUuid(options, image_type, layout_filename, label):
+def GetUuid(options):
   """Returns the unique partition UUID for a given label.
 
-  Note: Only useful if the UUID is specified in the config file, otherwise
-  the value returned unlikely to be what is actually used in the image.
-
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/prod
-    layout_filename: Path to partition configuration file
-    label: Label of the partition you want to read from
-  Returns:
+  Prints:
     String containing the requested UUID
   """
 
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
-  partition = GetPartitionByLabel(partitions, label)
-  return partition['uuid']
+  partitions = GetPartitionTableFromConfig(options)
+  partition = GetPartitionByLabel(partitions, options.label)
+  print partition.get('uuid', '')
 
 
-def DoDebugOutput(options, image_type, layout_filename):
+def DoDebugOutput(options):
   """Prints out a human readable disk layout in on-disk order.
 
   This will round values larger than 1MB, it's exists to quickly
@@ -472,10 +427,8 @@ def DoDebugOutput(options, image_type, layout_filename):
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
   """
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
+  partitions = GetPartitionTableFromConfig(options)
 
   for partition in partitions:
     if partition['bytes'] < 1024 * 1024:
@@ -495,85 +448,73 @@ def DoDebugOutput(options, image_type, layout_filename):
       print 'blank - %s' %  size
 
 
-def DoParseOnly(options, image_type, layout_filename):
+def DoParseOnly(options):
   """Parses a layout file only, used before reading sizes to check for errors.
 
   Args:
     options: Flags passed to the script
-    image_type: Type of image eg base/test/dev/factory_install
-    layout_filename: Path to partition configuration file
   """
-  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
+  GetPartitionTableFromConfig(options)
 
 
 def main(argv):
-  action_map = {
-    'write_gpt': {
-      'usage': ['<image_type>', '<partition_config_file>', '<disk_image>'],
-      'func': WritePartitionTable,
-    },
-    'write_mbr': {
-      'usage': ['<image_type>', '<partition_config_file>', '<disk_image>',
-                '<mbr_boot_code>'],
-      'func': WriteMbrBoot,
-    },
-    'readblocksize': {
-      'usage': ['<partition_config_file>'],
-      'func': GetBlockSize,
-    },
-    'readfsblocksize': {
-      'usage': ['<partition_config_file>'],
-      'func': GetFilesystemBlockSize,
-    },
-    'readpartsize': {
-      'usage': ['<image_type>', '<partition_config_file>', '<partition_num>'],
-      'func': GetPartitionSize,
-    },
-    'readfssize': {
-      'usage': ['<image_type>', '<partition_config_file>', '<partition_num>'],
-      'func': GetFilesystemSize,
-    },
-    'readlabel': {
-      'usage': ['<image_type>', '<partition_config_file>', '<partition_num>'],
-      'func': GetLabel,
-    },
-    'readnum': {
-      'usage': ['<image_type>', '<partition_config_file>', '<label>'],
-      'func': GetNum,
-    },
-    'readuuid': {
-      'usage': ['<image_type>', '<partition_config_file>', '<label>'],
-      'func': GetUuid,
-    },
-    'debug': {
-      'usage': ['<image_type>', '<partition_config_file>'],
-      'func': DoDebugOutput,
-    },
-    'parseonly': {
-      'usage': ['<image_type>', '<partition_config_file>'],
-      'func': DoParseOnly,
-    }
-  }
+  default_layout_file = os.environ.get('DISK_LAYOUT_FILE',
+          os.path.join(os.path.dirname(__file__), 'legacy_disk_layout.json'))
+  default_layout_type = os.environ.get('DISK_LAYOUT_TYPE', 'base')
 
-  parser = OptionParser()
-  parser.add_option("--adjust_part", dest="adjust_part",
-                    help="adjust partition sizes", default="")
-  (options, args) = parser.parse_args()
+  parser = argparse.ArgumentParser(
+          formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument('--adjust_part', default='',
+          help='adjust partition sizes')
+  parser.add_argument('--disk_layout_file', default=default_layout_file,
+          help='path to disk layout json file')
+  parser.add_argument('--disk_layout', default=default_layout_type,
+          help='disk layout type from the json file')
+  actions = parser.add_subparsers(title='actions')
 
-  if len(args) < 1 or args[0] not in action_map:
-    print 'Usage: %s <action>\n' % sys.argv[0]
-    print 'Valid actions are:'
-    for action in action_map:
-      print '  %s %s' % (action, ' '.join(action_map[action]['usage']))
-    sys.exit(1)
-  else:
-    action_name = args[0]
-    action = action_map[action_name]
-    if len(action['usage']) == len(args) - 1:
-      print action['func'](options, *args[1:])
-    else:
-      sys.exit('Usage: %s %s %s' % (sys.argv[0], args[0],
-               ' '.join(action['usage'])))
+  a = actions.add_parser('write_gpt', help='write gpt to new image')
+  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)
+
+  a = actions.add_parser('readfsblocksize', help='get filesystem block size')
+  a.set_defaults(func=GetFilesystemBlockSize)
+
+  a = actions.add_parser('readpartsize', help='get partition size')
+  a.add_argument('partition_num', type=int, help='partition number')
+  a.set_defaults(func=GetPartitionSize)
+
+  a = actions.add_parser('readfssize', help='get filesystem size')
+  a.add_argument('partition_num', type=int, help='partition number')
+  a.set_defaults(func=GetFilesystemSize)
+
+  a = actions.add_parser('readlabel', help='get partition label')
+  a.add_argument('partition_num', type=int, help='partition number')
+  a.set_defaults(func=GetLabel)
+
+  a = actions.add_parser('readnum', help='get partition number')
+  a.add_argument('label', help='partition label')
+  a.set_defaults(func=GetNum)
+
+  a = actions.add_parser('readuuid', help='get partition uuid')
+  a.add_argument('label', help='partition label')
+  a.set_defaults(func=GetUuid)
+
+  a = actions.add_parser('debug', help='dump debug output')
+  a.set_defaults(func=DoDebugOutput)
+
+  a = actions.add_parser('parseonly', help='validate config')
+  a.set_defaults(func=DoParseOnly)
+
+  options = parser.parse_args(argv[1:])
+  options.func(options)
 
 
 if __name__ == '__main__':
diff --git a/build_library/vm_image_util.sh b/build_library/vm_image_util.sh
index 568e5ab402..03d1bd7289 100644
--- a/build_library/vm_image_util.sh
+++ b/build_library/vm_image_util.sh
@@ -199,8 +199,8 @@ _disk_ext() {
 # alternate filesystem image for the state partition instead of the one
 # from VM_SRC_IMG. Start new image using the given disk layout.
 unpack_source_disk() {
-    local disk_layout="${1:-$(_get_vm_opt DISK_LAYOUT)}"
-    local alternate_state_image="$2"
+    get_disk_layout_type $(_get_vm_opt DISK_LAYOUT)
+    local alternate_state_image="$1"
 
     if [[ -n "${alternate_state_image}" && ! -f "${alternate_state_image}" ]]
     then
@@ -229,13 +229,13 @@ unpack_source_disk() {
 
     if [[ $(_get_vm_opt PARTITIONED_IMG) -eq 1 ]]; then
       info "Initializing new partition table..."
-      write_partition_table "${disk_layout}" "${VM_TMP_IMG}"
+      write_partition_table "${VM_TMP_IMG}"
     fi
 }
 
 resize_state_partition() {
-    local disk_layout="${1:-$(_get_vm_opt DISK_LAYOUT)}"
-    local size_in_bytes=$(get_filesystem_size "${disk_layout}" ${NUM_STATEFUL})
+    get_disk_layout_type $(_get_vm_opt DISK_LAYOUT)
+    local size_in_bytes=$(get_filesystem_size ${NUM_STATEFUL})
     local size_in_sectors=$(( size_in_bytes / 512 ))
     local size_in_mb=$(( size_in_bytes / 1024 / 1024 ))
     local original_size=$(stat -c%s "${TEMP_STATE}")
diff --git a/image_to_vm.sh b/image_to_vm.sh
index c6727987b1..1e35812d6b 100755
--- a/image_to_vm.sh
+++ b/image_to_vm.sh
@@ -109,8 +109,8 @@ trap vm_cleanup EXIT
 
 # Unpack image, using alternate state image if defined
 # Resize to use all available space in new disk layout
-unpack_source_disk "${FLAGS_disk_layout}" "${FLAGS_state_image}"
-resize_state_partition "${FLAGS_disk_layout}"
+unpack_source_disk "${FLAGS_state_image}"
+resize_state_partition
 
 # Optionally install any OEM packages
 install_oem_package