flashmap: Bring over changes from factory repo

fmap.py is identical to the one in factory/setup/fmap.py except:
 - no test functionality in main function
 - module docstring and license

Added some unit tests for the new functionality

BUG=chromium:726356
TEST=ran unittests

Change-Id: Ib67bd4260cf4e1cc08543b822521a4adcf0cbce6
Reviewed-on: https://chromium-review.googlesource.com/516382
Commit-Ready: Drew Davenport <ddavenport@chromium.org>
Tested-by: Drew Davenport <ddavenport@chromium.org>
Reviewed-by: Drew Davenport <ddavenport@chromium.org>

fmap.py

diff --git a/fmap.py b/fmap.py
index 21519a4..830a64e 100644
--- a/fmap.py
+++ b/fmap.py
@@ -52,6 +52,7 @@
 """
 
 
+import logging
 import struct
 import sys
 
@@ -59,8 +60,10 @@
 # constants imported from lib/fmap.h
 FMAP_SIGNATURE = '__FMAP__'
 FMAP_VER_MAJOR = 1
-FMAP_VER_MINOR = 0
+FMAP_VER_MINOR_MIN = 0
+FMAP_VER_MINOR_MAX = 1
 FMAP_STRLEN = 32
+FMAP_SEARCH_STRIDE = 4
 
 FMAP_FLAGS = {
     'FMAP_AREA_STATIC': 1 << 0,
@@ -102,7 +105,8 @@
   if header['signature'] != FMAP_SIGNATURE:
     raise struct.error('Invalid signature')
   if (header['ver_major'] != FMAP_VER_MAJOR or
-      header['ver_minor'] != FMAP_VER_MINOR):
+      header['ver_minor'] < FMAP_VER_MINOR_MIN or
+      header['ver_minor'] > FMAP_VER_MINOR_MAX):
     raise struct.error('Incompatible version')
 
   # convert null-terminated names
@@ -128,7 +132,58 @@
   return tuple([name for name in FMAP_FLAGS if area_flags & FMAP_FLAGS[name]])
 
 
-def fmap_decode(blob, offset=None):
+def _fmap_check_name(fmap, name):
+  """Checks if the FMAP structure has correct name.
+
+  Args:
+    fmap: A decoded FMAP structure.
+    name: A string to specify expected FMAP name.
+
+  Raises:
+    struct.error if the name does not match.
+  """
+  if fmap['name'] != name:
+    raise struct.error('Incorrect FMAP (found: "%s", expected: "%s")' %
+                       (fmap['name'], name))
+
+
+def _fmap_search_header(blob, fmap_name=None):
+  """Searches FMAP headers in given blob.
+
+  Uses same logic from vboot_reference/host/lib/fmap.c.
+
+  Args:
+    blob: A string containing FMAP data.
+    fmap_name: A string to specify target FMAP name.
+
+  Returns:
+    A tuple of (fmap, size, offset).
+  """
+  lim = len(blob) - struct.calcsize(FMAP_HEADER_FORMAT)
+  align = FMAP_SEARCH_STRIDE
+
+  # Search large alignments before small ones to find "right" FMAP.
+  while align <= lim:
+    align *= 2
+
+  while align >= FMAP_SEARCH_STRIDE:
+    for offset in xrange(align, lim + 1, align * 2):
+      if not blob.startswith(FMAP_SIGNATURE, offset):
+        continue
+      try:
+        (fmap, size) = _fmap_decode_header(blob, offset)
+        if fmap_name is not None:
+          _fmap_check_name(fmap, fmap_name)
+        return (fmap, size, offset)
+      except struct.error as e:
+        # Search for next FMAP candidate.
+        logging.debug('Continue searching FMAP due to exception %r', e)
+        pass
+    align /= 2
+  raise struct.error('No valid FMAP signatures.')
+
+
+def fmap_decode(blob, offset=None, fmap_name=None):
   """ Decodes a blob to FMAP dictionary object.
 
   Arguments:
@@ -137,10 +192,13 @@
             the blob.
   """
   fmap = {}
+
   if offset is None:
-    # try search magic in fmap
-    offset = blob.find(FMAP_SIGNATURE)
-  (fmap, size) = _fmap_decode_header(blob, offset)
+    (fmap, size, offset) = _fmap_search_header(blob, fmap_name)
+  else:
+    (fmap, size) = _fmap_decode_header(blob, offset)
+    if fmap_name is not None:
+      _fmap_check_name(fmap, fmap_name)
   fmap['areas'] = []
   offset = offset + size
   for _ in range(fmap['nareas']):

fmap_unittest.py

diff --git a/fmap_unittest.py b/fmap_unittest.py
index 7006520..4473da9 100644
--- a/fmap_unittest.py
+++ b/fmap_unittest.py
@@ -5,6 +5,7 @@
 # found in the LICENSE file.
 """Unit test for fmap module."""
 
+import struct
 import unittest
 
 import fmap
@@ -57,6 +58,30 @@
     decoded = fmap.fmap_decode(self.example_blob)
     self.assertEquals(_EXAMPLE_BIN_FMAP, decoded)
 
+  def testDecodeWithOffset(self):
+    decoded = fmap.fmap_decode(self.example_blob, 512)
+    self.assertEquals(_EXAMPLE_BIN_FMAP, decoded)
+
+  def testDecodeWithName(self):
+    decoded = fmap.fmap_decode(self.example_blob, fmap_name='example')
+    self.assertEquals(_EXAMPLE_BIN_FMAP, decoded)
+    decoded = fmap.fmap_decode(self.example_blob, 512, 'example')
+    self.assertEquals(_EXAMPLE_BIN_FMAP, decoded)
+
+  def testDecodeWithWrongName(self):
+    with self.assertRaises(struct.error):
+      decoded = fmap.fmap_decode(self.example_blob, fmap_name='banana')
+    with self.assertRaises(struct.error):
+      decoded = fmap.fmap_decode(self.example_blob, 512, 'banana')
+
+  def testDecodeWithOffset(self):
+    decoded = fmap.fmap_decode(self.example_blob, 512)
+    self.assertEquals(_EXAMPLE_BIN_FMAP, decoded)
+
+  def testDecodeWithWrongOffset(self):
+    with self.assertRaises(struct.error):
+      fmap.fmap_decode(self.example_blob, 42)
+
   def testEncode(self):
     encoded = fmap.fmap_encode(_EXAMPLE_BIN_FMAP)
     # example.bin contains other binary data besides the fmap