releasetools: Move verity related functions into verity_utils.py.

This prepares for upcoming changes that refactor verity-related
functions into classes.

This CL makes minimal changes to the moved functions, by creating a new
BuildVerityImageError class in verity_utils.py replacing the former
BuildImageError that's specific to build_image.py. As part of the
change, it also moves the tests for AVBCalcMinPartitionSize into
test_verity_utils.py.

Test: python -m unittest test_verity_utils
Test: `m dist` with aosp_marlin-userdebug (Verified Boot 1.0)
Test: `m dist` with aosp_taimen-userdebug (Verified Boot 2.0)
Change-Id: I19b52714d8980705ea1f9484ac03eb0af9483240
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 42f05a7..43c91da 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -29,18 +29,14 @@
 import os
 import os.path
 import re
-import shlex
 import shutil
 import sys
 
 import common
-import sparse_img
-
+import verity_utils
 
 OPTIONS = common.OPTIONS
-
-FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7"
-BLOCK_SIZE = 4096
+BLOCK_SIZE = common.BLOCK_SIZE
 BYTES_IN_MB = 1024 * 1024
 
 
@@ -51,34 +47,6 @@
     Exception.__init__(self, message)
 
 
-def GetVerityFECSize(partition_size):
-  cmd = ["fec", "-s", str(partition_size)]
-  output = common.RunAndCheckOutput(cmd, verbose=False)
-  return int(output)
-
-
-def GetVerityTreeSize(partition_size):
-  cmd = ["build_verity_tree", "-s", str(partition_size)]
-  output = common.RunAndCheckOutput(cmd, verbose=False)
-  return int(output)
-
-
-def GetVerityMetadataSize(partition_size):
-  cmd = ["build_verity_metadata.py", "size", str(partition_size)]
-  output = common.RunAndCheckOutput(cmd, verbose=False)
-  return int(output)
-
-
-def GetVeritySize(partition_size, fec_supported):
-  verity_tree_size = GetVerityTreeSize(partition_size)
-  verity_metadata_size = GetVerityMetadataSize(partition_size)
-  verity_size = verity_tree_size + verity_metadata_size
-  if fec_supported:
-    fec_size = GetVerityFECSize(partition_size + verity_size)
-    return verity_size + fec_size
-  return verity_size
-
-
 def GetDiskUsage(path):
   """Returns the number of bytes that "path" occupies on host.
 
@@ -102,258 +70,6 @@
   return int(output.split()[0]) * 512
 
 
-def GetSimgSize(image_file):
-  simg = sparse_img.SparseImage(image_file, build_map=False)
-  return simg.blocksize * simg.total_blocks
-
-
-def ZeroPadSimg(image_file, pad_size):
-  blocks = pad_size // BLOCK_SIZE
-  print("Padding %d blocks (%d bytes)" % (blocks, pad_size))
-  simg = sparse_img.SparseImage(image_file, mode="r+b", build_map=False)
-  simg.AppendFillChunk(0, blocks)
-
-
-def AVBCalcMaxImageSize(avbtool, footer_type, partition_size, additional_args):
-  """Calculates max image size for a given partition size.
-
-  Args:
-    avbtool: String with path to avbtool.
-    footer_type: 'hash' or 'hashtree' for generating footer.
-    partition_size: The size of the partition in question.
-    additional_args: Additional arguments to pass to "avbtool add_hash_footer"
-        or "avbtool add_hashtree_footer".
-
-  Returns:
-    The maximum image size.
-
-  Raises:
-    BuildImageError: On invalid image size.
-  """
-  cmd = [avbtool, "add_%s_footer" % footer_type,
-         "--partition_size", str(partition_size), "--calc_max_image_size"]
-  cmd.extend(shlex.split(additional_args))
-
-  output = common.RunAndCheckOutput(cmd)
-  image_size = int(output)
-  if image_size <= 0:
-    raise BuildImageError(
-        "Invalid max image size: {}".format(output))
-  return image_size
-
-
-def AVBCalcMinPartitionSize(image_size, size_calculator):
-  """Calculates min partition size for a given image size.
-
-  Args:
-    image_size: The size of the image in question.
-    size_calculator: The function to calculate max image size
-        for a given partition size.
-
-  Returns:
-    The minimum partition size required to accommodate the image size.
-  """
-  # Use image size as partition size to approximate final partition size.
-  image_ratio = size_calculator(image_size) / float(image_size)
-
-  # Prepare a binary search for the optimal partition size.
-  lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE
-
-  # Ensure lo is small enough: max_image_size should <= image_size.
-  delta = BLOCK_SIZE
-  max_image_size = size_calculator(lo)
-  while max_image_size > image_size:
-    image_ratio = max_image_size / float(lo)
-    lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta
-    delta *= 2
-    max_image_size = size_calculator(lo)
-
-  hi = lo + BLOCK_SIZE
-
-  # Ensure hi is large enough: max_image_size should >= image_size.
-  delta = BLOCK_SIZE
-  max_image_size = size_calculator(hi)
-  while max_image_size < image_size:
-    image_ratio = max_image_size / float(hi)
-    hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta
-    delta *= 2
-    max_image_size = size_calculator(hi)
-
-  partition_size = hi
-
-  # Start to binary search.
-  while lo < hi:
-    mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
-    max_image_size = size_calculator(mid)
-    if max_image_size >= image_size:  # if mid can accommodate image_size
-      if mid < partition_size:  # if a smaller partition size is found
-        partition_size = mid
-      hi = mid
-    else:
-      lo = mid + BLOCK_SIZE
-
-  if OPTIONS.verbose:
-    print("AVBCalcMinPartitionSize({}): partition_size: {}.".format(
-        image_size, partition_size))
-
-  return partition_size
-
-
-def AVBAddFooter(image_path, avbtool, footer_type, partition_size,
-                 partition_name, key_path, algorithm, salt,
-                 additional_args):
-  """Adds dm-verity hashtree and AVB metadata to an image.
-
-  Args:
-    image_path: Path to image to modify.
-    avbtool: String with path to avbtool.
-    footer_type: 'hash' or 'hashtree' for generating footer.
-    partition_size: The size of the partition in question.
-    partition_name: The name of the partition - will be embedded in metadata.
-    key_path: Path to key to use or None.
-    algorithm: Name of algorithm to use or None.
-    salt: The salt to use (a hexadecimal string) or None.
-    additional_args: Additional arguments to pass to "avbtool add_hash_footer"
-        or "avbtool add_hashtree_footer".
-  """
-  cmd = [avbtool, "add_%s_footer" % footer_type,
-         "--partition_size", partition_size,
-         "--partition_name", partition_name,
-         "--image", image_path]
-
-  if key_path and algorithm:
-    cmd.extend(["--key", key_path, "--algorithm", algorithm])
-  if salt:
-    cmd.extend(["--salt", salt])
-
-  cmd.extend(shlex.split(additional_args))
-
-  common.RunAndCheckOutput(cmd)
-
-
-def AdjustPartitionSizeForVerity(partition_size, fec_supported):
-  """Modifies the provided partition size to account for the verity metadata.
-
-  This information is used to size the created image appropriately.
-
-  Args:
-    partition_size: the size of the partition to be verified.
-
-  Returns:
-    A tuple of the size of the partition adjusted for verity metadata, and
-    the size of verity metadata.
-  """
-  key = "%d %d" % (partition_size, fec_supported)
-  if key in AdjustPartitionSizeForVerity.results:
-    return AdjustPartitionSizeForVerity.results[key]
-
-  hi = partition_size
-  if hi % BLOCK_SIZE != 0:
-    hi = (hi // BLOCK_SIZE) * BLOCK_SIZE
-
-  # verity tree and fec sizes depend on the partition size, which
-  # means this estimate is always going to be unnecessarily small
-  verity_size = GetVeritySize(hi, fec_supported)
-  lo = partition_size - verity_size
-  result = lo
-
-  # do a binary search for the optimal size
-  while lo < hi:
-    i = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
-    v = GetVeritySize(i, fec_supported)
-    if i + v <= partition_size:
-      if result < i:
-        result = i
-        verity_size = v
-      lo = i + BLOCK_SIZE
-    else:
-      hi = i
-
-  if OPTIONS.verbose:
-    print("Adjusted partition size for verity, partition_size: {},"
-          " verity_size: {}".format(result, verity_size))
-  AdjustPartitionSizeForVerity.results[key] = (result, verity_size)
-  return (result, verity_size)
-
-
-AdjustPartitionSizeForVerity.results = {}
-
-
-def BuildVerityFEC(sparse_image_path, verity_path, verity_fec_path,
-                   padding_size):
-  cmd = ["fec", "-e", "-p", str(padding_size), sparse_image_path,
-         verity_path, verity_fec_path]
-  common.RunAndCheckOutput(cmd)
-
-
-def BuildVerityTree(sparse_image_path, verity_image_path):
-  cmd = ["build_verity_tree", "-A", FIXED_SALT, sparse_image_path,
-         verity_image_path]
-  output = common.RunAndCheckOutput(cmd)
-  root, salt = output.split()
-  return root, salt
-
-
-def BuildVerityMetadata(image_size, verity_metadata_path, root_hash, salt,
-                        block_device, signer_path, key, signer_args,
-                        verity_disable):
-  cmd = ["build_verity_metadata.py", "build", str(image_size),
-         verity_metadata_path, root_hash, salt, block_device, signer_path, key]
-  if signer_args:
-    cmd.append("--signer_args=\"%s\"" % (' '.join(signer_args),))
-  if verity_disable:
-    cmd.append("--verity_disable")
-  common.RunAndCheckOutput(cmd)
-
-
-def Append2Simg(sparse_image_path, unsparse_image_path, error_message):
-  """Appends the unsparse image to the given sparse image.
-
-  Args:
-    sparse_image_path: the path to the (sparse) image
-    unsparse_image_path: the path to the (unsparse) image
-
-  Raises:
-    BuildImageError: On error.
-  """
-  cmd = ["append2simg", sparse_image_path, unsparse_image_path]
-  try:
-    common.RunAndCheckOutput(cmd)
-  except:
-    raise BuildImageError(error_message)
-
-
-def Append(target, file_to_append, error_message):
-  """Appends file_to_append to target.
-
-  Raises:
-    BuildImageError: On error.
-  """
-  try:
-    with open(target, "a") as out_file, open(file_to_append, "r") as input_file:
-      for line in input_file:
-        out_file.write(line)
-  except IOError:
-    raise BuildImageError(error_message)
-
-
-def BuildVerifiedImage(data_image_path, verity_image_path,
-                       verity_metadata_path, verity_fec_path,
-                       padding_size, fec_supported):
-  Append(
-      verity_image_path, verity_metadata_path,
-      "Could not append verity metadata!")
-
-  if fec_supported:
-    # Build FEC for the entire partition, including metadata.
-    BuildVerityFEC(
-        data_image_path, verity_image_path, verity_fec_path, padding_size)
-    Append(verity_image_path, verity_fec_path, "Could not append FEC!")
-
-  Append2Simg(
-      data_image_path, verity_image_path, "Could not append verity data!")
-
-
 def UnsparseImage(sparse_image_path, replace=True):
   img_dir = os.path.dirname(sparse_image_path)
   unsparse_image_path = "unsparse_" + os.path.basename(sparse_image_path)
@@ -372,56 +88,6 @@
   return unsparse_image_path
 
 
-def MakeVerityEnabledImage(out_file, fec_supported, prop_dict):
-  """Creates an image that is verifiable using dm-verity.
-
-  Args:
-    out_file: the location to write the verifiable image at
-    prop_dict: a dictionary of properties required for image creation and
-               verification
-
-  Raises:
-    AssertionError: On invalid partition sizes.
-    BuildImageError: On other errors.
-  """
-  # get properties
-  image_size = int(prop_dict["image_size"])
-  block_dev = prop_dict["verity_block_device"]
-  signer_key = prop_dict["verity_key"] + ".pk8"
-  if OPTIONS.verity_signer_path is not None:
-    signer_path = OPTIONS.verity_signer_path
-  else:
-    signer_path = prop_dict["verity_signer_cmd"]
-  signer_args = OPTIONS.verity_signer_args
-
-  tempdir_name = common.MakeTempDir(suffix="_verity_images")
-
-  # Get partial image paths.
-  verity_image_path = os.path.join(tempdir_name, "verity.img")
-  verity_metadata_path = os.path.join(tempdir_name, "verity_metadata.img")
-  verity_fec_path = os.path.join(tempdir_name, "verity_fec.img")
-
-  # Build the verity tree and get the root hash and salt.
-  root_hash, salt = BuildVerityTree(out_file, verity_image_path)
-
-  # Build the metadata blocks.
-  verity_disable = "verity_disable" in prop_dict
-  BuildVerityMetadata(
-      image_size, verity_metadata_path, root_hash, salt, block_dev, signer_path,
-      signer_key, signer_args, verity_disable)
-
-  # Build the full verified image.
-  partition_size = int(prop_dict["partition_size"])
-  verity_size = int(prop_dict["verity_size"])
-
-  padding_size = partition_size - image_size - verity_size
-  assert padding_size >= 0
-
-  BuildVerifiedImage(
-      out_file, verity_image_path, verity_metadata_path, verity_fec_path,
-      padding_size, fec_supported)
-
-
 def ConvertBlockMapToBaseFs(block_map_file):
   base_fs_file = common.MakeTempFile(prefix="script_gen_", suffix=".base_fs")
   convert_command = ["blk_alloc_to_base_fs", block_map_file, base_fs_file]
@@ -570,9 +236,9 @@
     # Adjust partition_size to add more space for AVB footer, to prevent
     # it from consuming partition_reserved_size.
     if avb_footer_type:
-      size = AVBCalcMinPartitionSize(
+      size = verity_utils.AVBCalcMinPartitionSize(
           size,
-          lambda x: AVBCalcMaxImageSize(
+          lambda x: verity_utils.AVBCalcMaxImageSize(
               avbtool, avb_footer_type, x, avb_signing_args))
     prop_dict["partition_size"] = str(size)
     if OPTIONS.verbose:
@@ -583,7 +249,7 @@
   # Adjust the image size to make room for the hashes if this is to be verified.
   if verity_supported and is_verity_partition:
     partition_size = int(prop_dict.get("partition_size"))
-    image_size, verity_size = AdjustPartitionSizeForVerity(
+    image_size, verity_size = verity_utils.AdjustPartitionSizeForVerity(
         partition_size, verity_fec_supported)
     prop_dict["image_size"] = str(image_size)
     prop_dict["verity_size"] = str(verity_size)
@@ -592,7 +258,7 @@
   if avb_footer_type:
     partition_size = prop_dict["partition_size"]
     # avb_add_hash_footer_args or avb_add_hashtree_footer_args.
-    max_image_size = AVBCalcMaxImageSize(
+    max_image_size = verity_utils.AVBCalcMaxImageSize(
         avbtool, avb_footer_type, partition_size, avb_signing_args)
     prop_dict["image_size"] = str(max_image_size)
 
@@ -709,17 +375,18 @@
   if not fs_spans_partition:
     mount_point = prop_dict.get("mount_point")
     image_size = int(prop_dict["image_size"])
-    sparse_image_size = GetSimgSize(out_file)
+    sparse_image_size = verity_utils.GetSimgSize(out_file)
     if sparse_image_size > image_size:
       raise BuildImageError(
           "Error: {} image size of {} is larger than partition size of "
           "{}".format(mount_point, sparse_image_size, image_size))
     if verity_supported and is_verity_partition:
-      ZeroPadSimg(out_file, image_size - sparse_image_size)
+      verity_utils.ZeroPadSimg(out_file, image_size - sparse_image_size)
 
   # Create the verified image if this is to be verified.
   if verity_supported and is_verity_partition:
-    MakeVerityEnabledImage(out_file, verity_fec_supported, prop_dict)
+    verity_utils.MakeVerityEnabledImage(
+        out_file, verity_fec_supported, prop_dict)
 
   # Add AVB HASH or HASHTREE footer (metadata).
   if avb_footer_type:
@@ -729,7 +396,7 @@
     key_path = prop_dict.get("avb_key_path")
     algorithm = prop_dict.get("avb_algorithm")
     salt = prop_dict.get("avb_salt")
-    AVBAddFooter(
+    verity_utils.AVBAddFooter(
         out_file, avbtool, avb_footer_type, partition_size, partition_name,
         key_path, algorithm, salt, avb_signing_args)