Generate the compute_hash_tree command

Generate the transfer command "compute_hash_tree" for incremental
updates of the non-A/B devices that enable verified boot 1.0

Other changes include:
i.  factor out verity_utils to use both in build_image and blockimgdiff
ii. add a new flag 'hashtree_info_generator' in sparse_image to generate
    the hashtree information.

Bug: 25170618
Test: generate a package for aosp_angler; and run simulator
Change-Id: I4d4d7a3e41dc3da810d2cbf8988e85d187c9ab0e
diff --git a/tools/releasetools/verity_utils.py b/tools/releasetools/verity_utils.py
new file mode 100644
index 0000000..0e605b1
--- /dev/null
+++ b/tools/releasetools/verity_utils.py
@@ -0,0 +1,204 @@
+#!/usr/bin/env python
+#
+# Copyright (C) 2018 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import struct
+
+import common
+from build_image import (AdjustPartitionSizeForVerity, GetVerityTreeSize,
+                         GetVerityMetadataSize, BuildVerityTree)
+from rangelib import RangeSet
+
+
+class HashtreeInfoGenerationError(Exception):
+  """An Exception raised during hashtree info generation."""
+
+  def __init__(self, message):
+    Exception.__init__(self, message)
+
+
+class HashtreeInfo(object):
+  def __init__(self):
+    self.hashtree_range = None
+    self.filesystem_range = None
+    self.hash_algorithm = None
+    self.salt = None
+    self.root_hash = None
+
+
+def CreateHashtreeInfoGenerator(partition_name, block_size, info_dict):
+  generator = None
+  if (info_dict.get("verity") == "true" and
+      info_dict.get("{}_verity_block_device".format(partition_name))):
+    partition_size = info_dict["{}_size".format(partition_name)]
+    fec_supported = info_dict.get("verity_fec") == "true"
+    generator = VerifiedBootVersion1HashtreeInfoGenerator(
+        partition_size, block_size, fec_supported)
+
+  return generator
+
+
+class HashtreeInfoGenerator(object):
+  def Generate(self, image):
+    raise NotImplementedError
+
+  def DecomposeSparseImage(self, image):
+    raise NotImplementedError
+
+  def ValidateHashtree(self):
+    raise NotImplementedError
+
+
+class VerifiedBootVersion2HashtreeInfoGenerator(HashtreeInfoGenerator):
+  pass
+
+
+class VerifiedBootVersion1HashtreeInfoGenerator(HashtreeInfoGenerator):
+  """A class that parses the metadata of hashtree for a given partition."""
+
+  def __init__(self, partition_size, block_size, fec_supported):
+    """Initialize VerityTreeInfo with the sparse image and input property.
+
+    Arguments:
+      partition_size: The whole size in bytes of a partition, including the
+        filesystem size, padding size, and verity size.
+      block_size: Expected size in bytes of each block for the sparse image.
+      fec_supported: True if the verity section contains fec data.
+    """
+
+    self.block_size = block_size
+    self.partition_size = partition_size
+    self.fec_supported = fec_supported
+
+    self.image = None
+    self.filesystem_size = None
+    self.hashtree_size = None
+    self.metadata_size = None
+
+    self.hashtree_info = HashtreeInfo()
+
+  def DecomposeSparseImage(self, image):
+    """Calculate the verity size based on the size of the input image.
+
+    Since we already know the structure of a verity enabled image to be:
+    [filesystem, verity_hashtree, verity_metadata, fec_data]. We can then
+    calculate the size and offset of each section.
+    """
+
+    self.image = image
+    assert self.block_size == image.blocksize
+    assert self.partition_size == image.total_blocks * self.block_size, \
+        "partition size {} doesn't match with the calculated image size." \
+        " total_blocks: {}".format(self.partition_size, image.total_blocks)
+
+    adjusted_size, _ = AdjustPartitionSizeForVerity(
+        self.partition_size, self.fec_supported)
+    assert adjusted_size % self.block_size == 0
+
+    verity_tree_size = GetVerityTreeSize(adjusted_size)
+    assert verity_tree_size % self.block_size == 0
+
+    metadata_size = GetVerityMetadataSize(adjusted_size)
+    assert metadata_size % self.block_size == 0
+
+    self.filesystem_size = adjusted_size
+    self.hashtree_size = verity_tree_size
+    self.metadata_size = metadata_size
+
+    self.hashtree_info.filesystem_range = RangeSet(
+        data=[0, adjusted_size / self.block_size])
+    self.hashtree_info.hashtree_range = RangeSet(
+        data=[adjusted_size / self.block_size,
+              (adjusted_size + verity_tree_size) / self.block_size])
+
+  def _ParseHashtreeMetadata(self):
+    """Parses the hash_algorithm, root_hash, salt from the metadata block."""
+
+    metadata_start = self.filesystem_size + self.hashtree_size
+    metadata_range = RangeSet(
+        data=[metadata_start / self.block_size,
+              (metadata_start + self.metadata_size) / self.block_size])
+    meta_data = ''.join(self.image.ReadRangeSet(metadata_range))
+
+    # More info about the metadata structure available in:
+    # system/extras/verity/build_verity_metadata.py
+    META_HEADER_SIZE = 268
+    header_bin = meta_data[0:META_HEADER_SIZE]
+    header = struct.unpack("II256sI", header_bin)
+
+    # header: magic_number, version, signature, table_len
+    assert header[0] == 0xb001b001, header[0]
+    table_len = header[3]
+    verity_table = meta_data[META_HEADER_SIZE: META_HEADER_SIZE + table_len]
+    table_entries = verity_table.rstrip().split()
+
+    # Expected verity table format: "1 block_device block_device block_size
+    # block_size data_blocks data_blocks hash_algorithm root_hash salt"
+    assert len(table_entries) == 10, "Unexpected verity table size {}".format(
+        len(table_entries))
+    assert (int(table_entries[3]) == self.block_size and
+            int(table_entries[4]) == self.block_size)
+    assert (int(table_entries[5]) * self.block_size == self.filesystem_size and
+            int(table_entries[6]) * self.block_size == self.filesystem_size)
+
+    self.hashtree_info.hash_algorithm = table_entries[7]
+    self.hashtree_info.root_hash = table_entries[8]
+    self.hashtree_info.salt = table_entries[9]
+
+  def ValidateHashtree(self):
+    """Checks that we can reconstruct the verity hash tree."""
+
+    # Writes the file system section to a temp file; and calls the executable
+    # build_verity_tree to construct the hash tree.
+    adjusted_partition = common.MakeTempFile(prefix="adjusted_partition")
+    with open(adjusted_partition, "wb") as fd:
+      self.image.WriteRangeDataToFd(self.hashtree_info.filesystem_range, fd)
+
+    generated_verity_tree = common.MakeTempFile(prefix="verity")
+    prop_dict = {}
+    BuildVerityTree(adjusted_partition, generated_verity_tree, prop_dict)
+
+    assert prop_dict["verity_salt"] == self.hashtree_info.salt
+    if prop_dict["verity_root_hash"] != self.hashtree_info.root_hash:
+      print("Calculated verty root hash {} doesn't match the one in metadata"
+            " {}".format(prop_dict["verity_root_hash"],
+                         self.hashtree_info.root_hash))
+      return False
+
+    # Reads the generated hash tree and checks if it has the exact same bytes
+    # as the one in the sparse image.
+    with open(generated_verity_tree, "rb") as fd:
+      return fd.read() == ''.join(self.image.ReadRangeSet(
+          self.hashtree_info.hashtree_range))
+
+  def Generate(self, image):
+    """Parses and validates the hashtree info in a sparse image.
+
+    Returns:
+      hashtree_info: The information needed to reconstruct the hashtree.
+    Raises:
+      HashtreeInfoGenerationError: If we fail to generate the exact bytes of
+          the hashtree.
+    """
+
+    self.DecomposeSparseImage(image)
+    self._ParseHashtreeMetadata()
+
+    if not self.ValidateHashtree():
+      raise HashtreeInfoGenerationError("Failed to reconstruct the verity tree")
+
+    return self.hashtree_info