Adding make_disk_image subcommand.

am: 159e172053

Change-Id: If1fb30d77d1589563d51239b1473474adbbccbfe

6 files changed, 315 insertions, 3 deletions

diff --git a/README b/README
index b4e6049..785dac8 100644
--- a/README
+++ b/README
@@ -129,7 +129,7 @@ means 1 mebibyte. Strings with base-10 units (kB, MB, GB, TB, PB) and
 base-2 units (KiB, MiB, GiB, TiB, PiB) are also supported. For
 example:
 
- "size": "1 Mib"
+ "size": "1 MiB"
 
 means 1,048,576 bytes and
 
@@ -191,6 +191,30 @@ where if prev_output.bpt contained the partitions "system_a",
 "system_b" (for the default A/B suffixes) then new_output.bpt would
 contain partitions "system-A", "system-B", and "system-C".
 
+-- DISK IMAGE GENERATION
+
+Disk images may be created given an unfolded .bpt file. 'bpttool
+make_disk_image' generates the output disk image file.
+
+To generate a disk image, use the following subcommand:
+
+  $ bpttool make_disk_image \
+      --output disk-image.bin \
+      --input /path/to/bpt-file.bpt \
+      --image system_a:/path/to/system.img \
+      --image boot_a:/path/to/boot.img \
+      [...]
+
+where the 'output' argument specifies the name and location of the outputted
+disk image and the 'input' argument is the .bpt file containing valid labels and
+offsets for each partition.  The 'image' argument specifies a mapping  from
+partition name/label to the path of the corresponding image partition image.
+All partitions specified in the .bpt file must be passed in via the 'image'
+argument.
+
+Typically, each of the 'image' argument files are located in the
+ANDROID_PRODUCT_OUT directory after a build is complete.
+
 -- BUILD SYSTEM INTEGRATION NOTES
 
 To generate partition tables in the Android build system, simply add
diff --git a/bpt_unittest.py b/bpt_unittest.py
index c2aa539..cc1d925 100755
--- a/bpt_unittest.py
+++ b/bpt_unittest.py
@@ -20,6 +20,7 @@
 
 import imp
 import sys
+import tempfile
 import unittest
 
 sys.dont_write_bytecode = True
@@ -39,6 +40,17 @@ class FakeGuidGenerator(object):
     uuid = '01234567-89ab-cdef-0123-%012x' % partition_number
     return uuid
 
+class PatternPartition(object):
+  """A partition image file containing a predictable pattern.
+
+  This holds file data about a partition image file for binary pattern.
+  testing.
+  """
+  def __init__(self, char='', file=None, partition_name=None, obj=None):
+    self.char = char
+    self.file = file
+    self.partition_name = partition_name
+    self.obj = obj
 
 class RoundToMultipleTest(unittest.TestCase):
   """Unit tests for the RoundToMultiple() function."""
@@ -111,6 +123,108 @@ class ParseSizeTest(unittest.TestCase):
     self.assertEqual(bpttool.ParseSize('0.5 GiB'), 536870912)
     self.assertEqual(bpttool.ParseSize('0.1 MiB'), 104858)
 
+class MakeDiskImageTest(unittest.TestCase):
+  """Unit tests for 'bpttool make_disk_image'."""
+
+  def setUp(self):
+    """Set-up method."""
+    self.bpt = bpttool.Bpt()
+
+  def _BinaryPattern(self, bpt_file_name, partition_patterns):
+    """Checks that a binary pattern may be written to a specified partition.
+
+    This checks individual partion image writes to portions of a disk.  Known
+    patterns are written into certain partitions and are verified after each
+    pattern has been written to.
+
+    Arguments:
+      bpt_file_name: File name of bpt JSON containing partition information.
+      partition_patterns: List of tuples with each tuple having partition name
+                          as the first argument, and character pattern as the
+                          second argument.
+
+    """
+    bpt_file = open(bpt_file_name, 'r')
+    partitions_string, _ = self.bpt.make_table([bpt_file])
+    bpt_tmp = tempfile.NamedTemporaryFile()
+    bpt_tmp.write(partitions_string)
+    bpt_tmp.seek(0)
+    partitions, _ = self.bpt._read_json([bpt_tmp])
+
+    # Declare list of partition images to be written and compared on disk.
+    pattern_images = [PatternPartition(
+                      char=pp[1],
+                      file=tempfile.NamedTemporaryFile(),
+                      partition_name=pp[0])
+                      for pp in partition_patterns]
+
+    # Store partition object and write a known character pattern image.
+    for pi in pattern_images:
+      pi.obj = [p for p in partitions if str(p.label) == pi.partition_name][0]
+      pi.file.write(bytearray(pi.char * int(pi.obj.size)))
+
+    # Create the disk containing the partition filled with a known character
+    # pattern, seek to it's position and compare it to the supposed pattern.
+    with tempfile.NamedTemporaryFile() as generated_disk_image:
+      bpt_tmp.seek(0)
+      self.bpt.make_disk_image(generated_disk_image,
+                               bpt_tmp,
+                               [p.partition_name + ':' + p.file.name
+                                for p in pattern_images])
+
+      for pi in pattern_images:
+        generated_disk_image.seek(pi.obj.offset)
+        pi.file.seek(0)
+
+        self.assertEqual(generated_disk_image.read(pi.obj.size),
+                    pi.file.read())
+        pi.file.close()
+
+    bpt_file.close()
+    bpt_tmp.close()
+
+  def _LargeBinary(self, bpt_file_name):
+    """Helper function to write large partition images to disk images.
+
+    This is a simple call to make_disk_image, passing a large in an image
+    which exceeds the it's size as specfied in the bpt file.
+
+    Arguments:
+      bpt_file_name: File name of bpt JSON containing partition information.
+
+    """
+    with open(bpt_file_name, 'r') as bpt_file, \
+         tempfile.NamedTemporaryFile() as bpt_tmp, \
+         tempfile.NamedTemporaryFile() as generated_disk_image, \
+         tempfile.NamedTemporaryFile() as large_partition_image:
+        partitions_string, _ = self.bpt.make_table([bpt_file])
+        bpt_tmp.write(partitions_string)
+        bpt_tmp.seek(0)
+        partitions, _ = self.bpt._read_json([bpt_tmp])
+
+        # Create the over-sized partition image.
+        large_partition_image.write(bytearray('0' *
+          int(1.1*partitions[0].size + 1)))
+
+        bpt_tmp.seek(0)
+
+        # Expect exception here.
+        self.bpt.make_disk_image(generated_disk_image, bpt_tmp,
+          [p.label + ':' + large_partition_image.name for p in partitions])
+
+  def testBinaryPattern(self):
+    """Checks patterns written to partitions on disk images."""
+    self._BinaryPattern('test/pattern_partition_single.bpt', [('charlie', 'c')])
+    self._BinaryPattern('test/pattern_partition_multi.bpt', [('alpha', 'a'),
+                        ('beta', 'b')])
+
+  def testExceedPartitionSize(self):
+    """Checks that exceedingly large partition images are not accepted."""
+    try:
+      self._LargeBinary('test/pattern_partition_exceed_size.bpt')
+    except bpttool.BptError as e:
+      assert 'exceeds the partition size' in e.message
+
 
 class MakeTableTest(unittest.TestCase):
   """Unit tests for 'bpttool make_table'."""
diff --git a/bpttool b/bpttool
index 861fe6b..3e57695 100755
--- a/bpttool
+++ b/bpttool
@@ -677,6 +677,31 @@ class Bpt(object):
     ret = protective_mbr + primary_gpt + secondary_gpt
     return ret
 
+  def _validate_disk_partitions(self, partitions, disk_size):
+    """Check that a list of partitions have assigned offsets and fits on a
+       disk of a given size.
+
+    This function checks partition offsets and sizes to see if they may fit on
+    a disk image.
+
+    Arguments:
+      partitions: A list of Partition objects.
+      settings: Integer size of disk image.
+
+    Raises:
+      BptError: If checked condition is not satisfied.
+    """
+    for p in partitions:
+      if not p.offset or p.offset < (GPT_NUM_LBAS + 1)*DISK_SECTOR_SIZE:
+        raise BptError('Partition with label "{}" has no offset.'
+                       .format(p.label))
+      if not p.size or p.size < 0:
+        raise BptError('Partition with label "{}" has no size.'
+                        .format(p.label))
+      if (p.offset + p.size) > (disk_size - GPT_NUM_LBAS*DISK_SECTOR_SIZE):
+        raise BptError('Partition with label "{}" exceeds the disk '
+                       'image size.'.format(p.label))
+
   def make_table(self,
                  inputs,
                  ab_suffixes=None,
@@ -798,7 +823,7 @@ class Bpt(object):
                      'totaling {} bytes.\n'.format(
                          settings.disk_size, offset))
 
-    # If we have an grow partition, it'll starts at the next
+    # If we have a grow partition, it'll starts at the next
     # available alignment offset and we can calculate its size as
     # follows.
     if grow_part:
@@ -827,6 +852,78 @@ class Bpt(object):
 
     return json_str, gpt_bin
 
+  def make_disk_image(self, output, bpt, images, allow_empty_partitions=False):
+    """Implementation of the 'make_disk_image' command.
+
+    This function takes in a list of partitions images and a bpt file
+    for the purpose of creating a raw disk image with a protective MBR,
+    primary and secondary GPT, and content for each partition as specified.
+
+    Arguments:
+      output: Output file where disk image is to be written to.
+      bpt: BPT JSON file to parse.
+      images: List of partition image paths to be combined (as specified by
+              bpt).  Each element is of the form.
+              'PARTITION_NAME:/PATH/TO/PARTITION_IMAGE'
+      allow_empty_partitions: If True, partitions defined in |bpt| need not to
+                              be present in |images|. Otherwise an exception is
+                              thrown if a partition is referenced in |bpt| but
+                              not in |images|.
+
+    Raises:
+      BptParsingError: If an image file has an error.
+      BptError: If another application-specific error occurs.
+    """
+    # Generate partition list and settings.
+    partitions, settings = self._read_json([bpt], ab_collapse=False)
+
+    # Validated partition sizes and offsets.
+    self._validate_disk_partitions(partitions, settings.disk_size)
+
+    # Sort according to 'offset' attribute.
+    partitions = sorted(partitions, cmp=lambda x, y: cmp(x.offset, y.offset))
+
+    # Create necessary tables.
+    protective_mbr = self._generate_protective_mbr(settings)
+    primary_gpt = self._generate_gpt(partitions, settings)
+    secondary_gpt = self._generate_gpt(partitions, settings, primary=False)
+
+    # Start at 0 offset for mbr and primary gpt.
+    output.seek(0)
+    output.write(protective_mbr)
+    output.write(primary_gpt)
+
+    # Create mapping of partition name to partition image file.
+    image_file_names = {}
+    try:
+      for name_path in images:
+        name, path = name_path.split(":")
+        image_file_names[name] = path
+    except ValueError as e:
+      raise BptParsingError(name_path, 'Bad image argument {}.'.format(
+                            images[i]))
+
+    # Read image and insert in correct offset.
+    for p in partitions:
+      if p.label not in image_file_names:
+        if allow_empty_partitions:
+          continue
+        else:
+          raise BptParsingError(bpt.name, 'No content specified for partition'
+                                ' with label {}'.format(p.label))
+
+      with open(image_file_names[p.label], 'rb') as partition_image:
+        output.seek(p.offset)
+        partition_blob = partition_image.read()
+        if len(partition_blob) > p.size:
+          raise BptError('Partition image content with label "{}" exceeds the '
+                         'partition size.'.format(p.label))
+        output.write(partition_blob)
+
+    # Put secondary GPT and end of disk.
+    output.seek(settings.disk_size - len(secondary_gpt))
+    output.write(secondary_gpt)
+
   def query_partition(self, input_file, part_label, query_type, ab_collapse):
     """Implementation of the 'query_partition' command.
 
@@ -849,7 +946,7 @@ class Bpt(object):
       BptError: If another application-specific error occurs
     """
 
-    partitions, _ = self._read_json([input_file], ab_collapse)
+    partitions, _ = self._read_json([input_file], 'ab_collapse')
 
     part = None
     for p in partitions:
@@ -915,6 +1012,26 @@ class BptTool(object):
     sub_parser.set_defaults(func=self.make_table)
 
     sub_parser = subparsers.add_parser(
+        'make_disk_image',
+        help='Creates disk image for loaded with partitions.')
+    sub_parser.add_argument('--output',
+                            help='Path to image output.',
+                            type=argparse.FileType('w'),
+                            required=True)
+    sub_parser.add_argument('--input',
+                            help='Path to bpt file input.',
+                            type=argparse.FileType('r'),
+                            required=True)
+    sub_parser.add_argument('--image',
+                            help='Partition name and path to image file.',
+                            metavar='PARTITION_NAME:PATH',
+                            action='append')
+    sub_parser.add_argument('--allow_empty_partitions',
+                            help='Allow skipping partitions in bpt file.',
+                            action='store_false')
+    sub_parser.set_defaults(func=self.make_disk_image)
+
+    sub_parser = subparsers.add_parser(
         'query_partition',
         help='Looks up informtion about a partition.')
     sub_parser.add_argument('--input',
@@ -983,6 +1100,26 @@ class BptTool(object):
     if args.output_gpt:
       args.output_gpt.write(gpt_bin)
 
+  def make_disk_image(self, args):
+    """Implements the 'make_disk_image' sub-command."""
+    if not args.input:
+      sys.stderr.write('Option --input is required.\n')
+      sys.exit(1)
+    if not args.output:
+      sys.stderr.write('Option --ouptut is required.\n')
+      sys.exit(1)
+
+    try:
+      self.bpt.make_disk_image(args.output,
+                               args.input,
+                               args.image,
+                               args.allow_empty_partitions)
+    except BptParsingError as e:
+      sys.stderr.write('{}: Error parsing: {}\n'.format(e.filename, e.message))
+      sys.exit(1)
+    except 'BptError' as e:
+      sys.stderr.write('{}\n'.format(e.message))
+      sys.exit(1)
 
 if __name__ == '__main__':
   tool = BptTool()
diff --git a/test/pattern_partition_exceed_size.bpt b/test/pattern_partition_exceed_size.bpt
new file mode 100644
index 0000000..c1c54c0
--- /dev/null
+++ b/test/pattern_partition_exceed_size.bpt
@@ -0,0 +1,11 @@
+{
+    "settings": {
+        "disk_size": "50 MiB"
+    },
+    "partitions": [
+        {
+            "label": "delta",
+            "size": "20 MiB"
+        }
+    ]
+}
diff --git a/test/pattern_partition_multi.bpt b/test/pattern_partition_multi.bpt
new file mode 100644
index 0000000..56cd061
--- /dev/null
+++ b/test/pattern_partition_multi.bpt
@@ -0,0 +1,15 @@
+{
+    "settings": {
+        "disk_size": "80 MiB"
+    },
+    "partitions": [
+        {
+            "label": "alpha",
+            "size": "10 MiB"
+        },
+        {
+            "label": "beta",
+            "size": "50 MiB"
+        }
+    ]
+}
diff --git a/test/pattern_partition_single.bpt b/test/pattern_partition_single.bpt
new file mode 100644
index 0000000..036de77
--- /dev/null
+++ b/test/pattern_partition_single.bpt
@@ -0,0 +1,11 @@
+{
+    "settings": {
+        "disk_size": "40 GiB"
+    },
+    "partitions": [
+        {
+            "label": "charlie",
+            "size": "10 MiB"
+        }
+    ]
+}