Merge "releasetools: Move verity related functions into verity_utils.py."

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

    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.

  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)

  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]

    # 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:

  # 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)

  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)

  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:

    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)

OPTIONS = Options()

# The block size that's used across the releasetools scripts.

BLOCK_SIZE = 4096

# Values for "certificate" in apkcerts that mean special things.

SPECIAL_CERT_STRINGS = ("PRESIGNED", "EXTERNAL")

#

import filecmp

import math

import os.path

import random

import common

from build_image import (

    AVBCalcMinPartitionSize, BLOCK_SIZE, BuildImageError, CheckHeadroom,

    SetUpInDirAndFsConfig)

    BuildImageError, CheckHeadroom, SetUpInDirAndFsConfig)

from test_utils import ReleaseToolsTestCase

  EXT4FS_OUTPUT = (

      "Created filesystem with 2777/129024 inodes and 515099/516099 blocks")

  def setUp(self):

    # To test AVBCalcMinPartitionSize(), by using 200MB to 2GB image size.

    #   -  51200 = 200MB * 1024 * 1024 / 4096

    #   - 524288 = 2GB * 1024 * 1024 * 1024 / 4096

    self._image_sizes = [BLOCK_SIZE * random.randint(51200, 524288) + offset

                         for offset in range(BLOCK_SIZE)]

  def test_CheckHeadroom_SizeUnderLimit(self):

    # Required headroom: 1000 blocks.

    prop_dict = {

    self.assertIn('fs-config-system\n', fs_config_data)

    self.assertIn('fs-config-root\n', fs_config_data)

    self.assertEqual('/', prop_dict['mount_point'])

  def test_AVBCalcMinPartitionSize_LinearFooterSize(self):

    """Tests with footer size which is linear to partition size."""

    for image_size in self._image_sizes:

      for ratio in 0.95, 0.56, 0.22:

        expected_size = common.RoundUpTo4K(int(math.ceil(image_size / ratio)))

        self.assertEqual(

            expected_size,

            AVBCalcMinPartitionSize(image_size, lambda x: int(x * ratio)))

  def test_AVBCalcMinPartitionSize_SlowerGrowthFooterSize(self):

    """Tests with footer size which grows slower than partition size."""

    def _SizeCalculator(partition_size):

      """Footer size is the power of 0.95 of partition size."""

      # Minus footer size to return max image size.

      return partition_size - int(math.pow(partition_size, 0.95))

    for image_size in self._image_sizes:

      min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator)

      # Checks min_partition_size can accommodate image_size.

      self.assertGreaterEqual(

          _SizeCalculator(min_partition_size),

          image_size)

      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.

      self.assertLess(

          _SizeCalculator(min_partition_size - BLOCK_SIZE),

          image_size)

  def test_AVBCalcMinPartitionSize_FasterGrowthFooterSize(self):

    """Tests with footer size which grows faster than partition size."""

    def _SizeCalculator(partition_size):

      """Max image size is the power of 0.95 of partition size."""

      # Max image size grows less than partition size, which means

      # footer size grows faster than partition size.

      return int(math.pow(partition_size, 0.95))

    for image_size in self._image_sizes:

      min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator)

      # Checks min_partition_size can accommodate image_size.

      self.assertGreaterEqual(

          _SizeCalculator(min_partition_size),

          image_size)

      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.

      self.assertLess(

          _SizeCalculator(min_partition_size - BLOCK_SIZE),

          image_size)

"""Unittests for verity_utils.py."""

import math

import os.path

import random

import build_image

import common

import sparse_img

from rangelib import RangeSet

from test_utils import get_testdata_dir, ReleaseToolsTestCase

from verity_utils import (

    CreateHashtreeInfoGenerator, HashtreeInfo,

    AdjustPartitionSizeForVerity, AVBCalcMinPartitionSize, BLOCK_SIZE,

    CreateHashtreeInfoGenerator, HashtreeInfo, MakeVerityEnabledImage,

    VerifiedBootVersion1HashtreeInfoGenerator)

  def _generate_image(self):

    partition_size = 1024 * 1024

    adjusted_size, verity_size = build_image.AdjustPartitionSizeForVerity(

    adjusted_size, verity_size = AdjustPartitionSizeForVerity(

        partition_size, True)

    raw_image = ""

        'verity_signer_cmd': 'verity_signer',

        'verity_size': str(verity_size),

    }

    build_image.MakeVerityEnabledImage(output_file, True, prop_dict)

    MakeVerityEnabledImage(output_file, True, prop_dict)

    return output_file

    self.assertEqual(self.hash_algorithm, info.hash_algorithm)

    self.assertEqual(self.fixed_salt, info.salt)

    self.assertEqual(self.expected_root_hash, info.root_hash)

class VerityUtilsTest(ReleaseToolsTestCase):

  def setUp(self):

    # To test AVBCalcMinPartitionSize(), by using 200MB to 2GB image size.

    #   -  51200 = 200MB * 1024 * 1024 / 4096

    #   - 524288 = 2GB * 1024 * 1024 * 1024 / 4096

    self._image_sizes = [BLOCK_SIZE * random.randint(51200, 524288) + offset

                         for offset in range(BLOCK_SIZE)]

  def test_AVBCalcMinPartitionSize_LinearFooterSize(self):

    """Tests with footer size which is linear to partition size."""

    for image_size in self._image_sizes:

      for ratio in 0.95, 0.56, 0.22:

        expected_size = common.RoundUpTo4K(int(math.ceil(image_size / ratio)))

        self.assertEqual(

            expected_size,

            AVBCalcMinPartitionSize(

                image_size, lambda x, ratio=ratio: int(x * ratio)))

  def test_AVBCalcMinPartitionSize_SlowerGrowthFooterSize(self):

    """Tests with footer size which grows slower than partition size."""

    def _SizeCalculator(partition_size):

      """Footer size is the power of 0.95 of partition size."""

      # Minus footer size to return max image size.

      return partition_size - int(math.pow(partition_size, 0.95))

    for image_size in self._image_sizes:

      min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator)

      # Checks min_partition_size can accommodate image_size.

      self.assertGreaterEqual(

          _SizeCalculator(min_partition_size),

          image_size)

      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.

      self.assertLess(

          _SizeCalculator(min_partition_size - BLOCK_SIZE),

          image_size)

  def test_AVBCalcMinPartitionSize_FasterGrowthFooterSize(self):

    """Tests with footer size which grows faster than partition size."""

    def _SizeCalculator(partition_size):

      """Max image size is the power of 0.95 of partition size."""

      # Max image size grows less than partition size, which means

      # footer size grows faster than partition size.

      return int(math.pow(partition_size, 0.95))

    for image_size in self._image_sizes:

      min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator)

      # Checks min_partition_size can accommodate image_size.

      self.assertGreaterEqual(

          _SizeCalculator(min_partition_size),

          image_size)

      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.

      self.assertLess(

          _SizeCalculator(min_partition_size - BLOCK_SIZE),

          image_size)