pygpt: Support 4K block images.

pygpt is a general disk image partition parser, we have to deal with
images that use non-512b blocks. Today the major format are either 512
or 4k, so the pygpt will try to read GPT in 4k location and change block
size calculation to that if needed.

This is tested by scsi_debug kernel module:

sudo modprobe scsi_debug dev_size_mb=16384 sector_size=4096
sudo lsblk -o NAME,PHY-SEC,LOG-SEC # Find the entry with NAME,4096,4096
sudo parted /dev/sdX
 mktable GPT
 mkpart
 Start=64,END=128

sudo pygpt.py show /dev/sdX  # show information properly.

BUG=chromium:832160
TEST=make test

Change-Id: I96614580a49fa39c51a8ee294d617a809793a995
Reviewed-on: https://chromium-review.googlesource.com/1011903
Commit-Ready: Hung-Te Lin <hungte@chromium.org>
Tested-by: Hung-Te Lin <hungte@chromium.org>
Reviewed-by: Pi-Hsun Shih <pihsun@chromium.org>
diff --git a/py/utils/pygpt.py b/py/utils/pygpt.py
index 052f813..608c0e8 100755
--- a/py/utils/pygpt.py
+++ b/py/utils/pygpt.py
@@ -99,7 +99,7 @@
       'Header', HEADER_FORMAT)
   PARTITION_FORMAT, PARTITION_CLASS = BuildStructFormatAndNamedTuple(
       'Partition', PARTITION_FORMAT)
-  BLOCK_SIZE = 512
+  DEFAULT_BLOCK_SIZE = 512
   HEADER_SIGNATURE = 'EFI PART'
   TYPE_GUID_UNUSED = '\x00' * 16
   TYPE_GUID_MAP = {
@@ -119,6 +119,7 @@
   def __init__(self):
     self.header = None
     self.partitions = None
+    self.block_size = self.DEFAULT_BLOCK_SIZE
 
   @staticmethod
   def GetAttributeSuccess(attrs):
@@ -168,11 +169,17 @@
   def LoadFromFile(cls, f):
     """Loads a GPT table from give disk image file object."""
     gpt = GPT()
-    f.seek(gpt.BLOCK_SIZE * 1)
-    header = gpt.ReadHeader(f)
-    if header.Signature != cls.HEADER_SIGNATURE:
+    # Try DEFAULT_BLOCK_SIZE, then 4K.
+    for block_size in [cls.DEFAULT_BLOCK_SIZE, 4096]:
+      f.seek(block_size * 1)
+      header = gpt.ReadHeader(f)
+      if header.Signature == cls.HEADER_SIGNATURE:
+        gpt.block_size = block_size
+        break
+    else:
       raise ValueError('Invalid signature in GPT header.')
-    f.seek(gpt.BLOCK_SIZE * header.PartitionEntriesStartingLBA)
+
+    f.seek(gpt.block_size * header.PartitionEntriesStartingLBA)
     partitions = [gpt.ReadPartitionEntry(f)
                   for unused_i in range(header.PartitionEntriesNumber)]
     gpt.header = header
@@ -204,8 +211,8 @@
     if header is None:
       header = self.header
     size = header.PartitionEntrySize * header.PartitionEntriesNumber
-    blocks = size / self.BLOCK_SIZE
-    if size % self.BLOCK_SIZE:
+    blocks = size / self.block_size
+    if size % self.block_size:
       blocks += 1
     return blocks
 
@@ -215,14 +222,14 @@
     Args:
       new_size: Integer for new size of disk image file.
     """
-    old_size = self.BLOCK_SIZE * (self.header.BackupLBA + 1)
-    if new_size % self.BLOCK_SIZE:
+    old_size = self.block_size * (self.header.BackupLBA + 1)
+    if new_size % self.block_size:
       raise ValueError('New file size %d is not valid for image files.' %
                        new_size)
-    new_blocks = new_size / self.BLOCK_SIZE
+    new_blocks = new_size / self.block_size
     if old_size != new_size:
       logging.warn('Image size (%d, LBA=%d) changed from %d (LBA=%d).',
-                   new_size, new_blocks, old_size, old_size / self.BLOCK_SIZE)
+                   new_size, new_blocks, old_size, old_size / self.block_size)
     else:
       logging.info('Image size (%d, LBA=%d) not changed.',
                    new_size, new_blocks)
@@ -258,7 +265,7 @@
     """Returns the free (available) space left according to LastUsableLBA."""
     max_lba = self.GetMaxUsedLBA()
     assert max_lba <= self.header.LastUsableLBA, "Partitions too large."
-    return self.BLOCK_SIZE * (self.header.LastUsableLBA - max_lba)
+    return self.block_size * (self.header.LastUsableLBA - max_lba)
 
   def ExpandPartition(self, i):
     """Expands a given partition to last usable LBA.
@@ -313,8 +320,8 @@
     def WriteData(name, blob, lba):
       """Writes a blob into given location."""
       logging.info('Writing %s in LBA %d (offset %d)',
-                   name, lba, lba * self.BLOCK_SIZE)
-      f.seek(lba * self.BLOCK_SIZE)
+                   name, lba, lba * self.block_size)
+      f.seek(lba * self.block_size)
       f.write(blob)
 
     self.UpdateChecksum()
@@ -380,7 +387,7 @@
     elif free_space:
       logging.warn('Extra space found (%d, LBA=%d), '
                    'use --expand to expand partitions.',
-                   free_space, free_space / gpt.BLOCK_SIZE)
+                   free_space, free_space / gpt.block_size)
 
     gpt.WriteToFile(args.image_file)
     print('Disk image file %s repaired.' % args.image_file.name)
@@ -507,7 +514,7 @@
 
     if do_print_gpt_blocks:
       f = args.image_file
-      f.seek(gpt.header.BackupLBA * gpt.BLOCK_SIZE)
+      f.seek(gpt.header.BackupLBA * gpt.block_size)
       backup_header = gpt.ReadHeader(f)
       print(fmt % (backup_header.PartitionEntriesStartingLBA,
                    gpt.GetPartitionTableBlocks(backup_header), '',