1#!/usr/bin/env python
2#
3# Copyright (C) 2018 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""
18Signs a given image using avbtool
19
20Usage:  verity_utils properties_file output_image
21"""
22
23from __future__ import print_function
24
25import logging
26import os.path
27import shlex
28import struct
29import sys
30
31import common
32import sparse_img
33from rangelib import RangeSet
34
35logger = logging.getLogger(__name__)
36
37OPTIONS = common.OPTIONS
38BLOCK_SIZE = common.BLOCK_SIZE
39FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7"
40
41# From external/avb/avbtool.py
42MAX_VBMETA_SIZE = 64 * 1024
43MAX_FOOTER_SIZE = 4096
44
45class BuildVerityImageError(Exception):
46  """An Exception raised during verity image building."""
47
48  def __init__(self, message):
49    Exception.__init__(self, message)
50
51
52def GetVerityFECSize(image_size):
53  cmd = ["fec", "-s", str(image_size)]
54  output = common.RunAndCheckOutput(cmd, verbose=False)
55  return int(output)
56
57
58def GetVerityTreeSize(image_size):
59  cmd = ["build_verity_tree", "-s", str(image_size)]
60  output = common.RunAndCheckOutput(cmd, verbose=False)
61  return int(output)
62
63
64def GetVerityMetadataSize(image_size):
65  cmd = ["build_verity_metadata", "size", str(image_size)]
66  output = common.RunAndCheckOutput(cmd, verbose=False)
67  return int(output)
68
69
70def GetVeritySize(image_size, fec_supported):
71  verity_tree_size = GetVerityTreeSize(image_size)
72  verity_metadata_size = GetVerityMetadataSize(image_size)
73  verity_size = verity_tree_size + verity_metadata_size
74  if fec_supported:
75    fec_size = GetVerityFECSize(image_size + verity_size)
76    return verity_size + fec_size
77  return verity_size
78
79
80def GetSimgSize(image_file):
81  simg = sparse_img.SparseImage(image_file, build_map=False)
82  return simg.blocksize * simg.total_blocks
83
84
85def ZeroPadSimg(image_file, pad_size):
86  blocks = pad_size // BLOCK_SIZE
87  logger.info("Padding %d blocks (%d bytes)", blocks, pad_size)
88  simg = sparse_img.SparseImage(image_file, mode="r+b", build_map=False)
89  simg.AppendFillChunk(0, blocks)
90
91
92def BuildVerityFEC(sparse_image_path, verity_path, verity_fec_path,
93                   padding_size):
94  cmd = ["fec", "-e", "-p", str(padding_size), sparse_image_path,
95         verity_path, verity_fec_path]
96  common.RunAndCheckOutput(cmd)
97
98
99def BuildVerityTree(sparse_image_path, verity_image_path):
100  cmd = ["build_verity_tree", "-A", FIXED_SALT, sparse_image_path,
101         verity_image_path]
102  output = common.RunAndCheckOutput(cmd)
103  root, salt = output.split()
104  return root, salt
105
106
107def BuildVerityMetadata(image_size, verity_metadata_path, root_hash, salt,
108                        block_device, signer_path, key, signer_args,
109                        verity_disable):
110  cmd = ["build_verity_metadata", "build", str(image_size),
111         verity_metadata_path, root_hash, salt, block_device, signer_path, key]
112  if signer_args:
113    cmd.append("--signer_args=\"%s\"" % (' '.join(signer_args),))
114  if verity_disable:
115    cmd.append("--verity_disable")
116  common.RunAndCheckOutput(cmd)
117
118
119def Append2Simg(sparse_image_path, unsparse_image_path, error_message):
120  """Appends the unsparse image to the given sparse image.
121
122  Args:
123    sparse_image_path: the path to the (sparse) image
124    unsparse_image_path: the path to the (unsparse) image
125
126  Raises:
127    BuildVerityImageError: On error.
128  """
129  cmd = ["append2simg", sparse_image_path, unsparse_image_path]
130  try:
131    common.RunAndCheckOutput(cmd)
132  except:
133    logger.exception(error_message)
134    raise BuildVerityImageError(error_message)
135
136
137def Append(target, file_to_append, error_message):
138  """Appends file_to_append to target.
139
140  Raises:
141    BuildVerityImageError: On error.
142  """
143  try:
144    with open(target, 'ab') as out_file, \
145        open(file_to_append, 'rb') as input_file:
146      for line in input_file:
147        out_file.write(line)
148  except IOError:
149    logger.exception(error_message)
150    raise BuildVerityImageError(error_message)
151
152
153def CreateVerityImageBuilder(prop_dict):
154  """Returns a verity image builder based on the given build properties.
155
156  Args:
157    prop_dict: A dict that contains the build properties. In particular, it will
158        look for verity-related property values.
159
160  Returns:
161    A VerityImageBuilder instance for Verified Boot 1.0 or Verified Boot 2.0; or
162        None if the given build doesn't support Verified Boot.
163  """
164  partition_size = prop_dict.get("partition_size")
165  # partition_size could be None at this point, if using dynamic partitions.
166  if partition_size:
167    partition_size = int(partition_size)
168
169  # Verified Boot 1.0
170  verity_supported = prop_dict.get("verity") == "true"
171  is_verity_partition = "verity_block_device" in prop_dict
172  if verity_supported and is_verity_partition:
173    if OPTIONS.verity_signer_path is not None:
174      signer_path = OPTIONS.verity_signer_path
175    else:
176      signer_path = prop_dict["verity_signer_cmd"]
177    return Version1VerityImageBuilder(
178        partition_size,
179        prop_dict["verity_block_device"],
180        prop_dict.get("verity_fec") == "true",
181        signer_path,
182        prop_dict["verity_key"] + ".pk8",
183        OPTIONS.verity_signer_args,
184        "verity_disable" in prop_dict)
185
186  # Verified Boot 2.0
187  if (prop_dict.get("avb_hash_enable") == "true" or
188      prop_dict.get("avb_hashtree_enable") == "true"):
189    # key_path and algorithm are only available when chain partition is used.
190    key_path = prop_dict.get("avb_key_path")
191    algorithm = prop_dict.get("avb_algorithm")
192
193    # Image uses hash footer.
194    if prop_dict.get("avb_hash_enable") == "true":
195      return VerifiedBootVersion2VerityImageBuilder(
196          prop_dict["partition_name"],
197          partition_size,
198          VerifiedBootVersion2VerityImageBuilder.AVB_HASH_FOOTER,
199          prop_dict["avb_avbtool"],
200          key_path,
201          algorithm,
202          prop_dict.get("avb_salt"),
203          prop_dict["avb_add_hash_footer_args"])
204
205    # Image uses hashtree footer.
206    return VerifiedBootVersion2VerityImageBuilder(
207        prop_dict["partition_name"],
208        partition_size,
209        VerifiedBootVersion2VerityImageBuilder.AVB_HASHTREE_FOOTER,
210        prop_dict["avb_avbtool"],
211        key_path,
212        algorithm,
213        prop_dict.get("avb_salt"),
214        prop_dict["avb_add_hashtree_footer_args"])
215
216  return None
217
218
219class VerityImageBuilder(object):
220  """A builder that generates an image with verity metadata for Verified Boot.
221
222  A VerityImageBuilder instance handles the works for building an image with
223  verity metadata for supporting Android Verified Boot. This class defines the
224  common interface between Verified Boot 1.0 and Verified Boot 2.0. A matching
225  builder will be returned based on the given build properties.
226
227  More info on the verity image generation can be found at the following link.
228  https://source.android.com/security/verifiedboot/dm-verity#implementation
229  """
230
231  def CalculateMaxImageSize(self, partition_size):
232    """Calculates the filesystem image size for the given partition size."""
233    raise NotImplementedError
234
235  def CalculateDynamicPartitionSize(self, image_size):
236    """Calculates and sets the partition size for a dynamic partition."""
237    raise NotImplementedError
238
239  def PadSparseImage(self, out_file):
240    """Adds padding to the generated sparse image."""
241    raise NotImplementedError
242
243  def Build(self, out_file):
244    """Builds the verity image and writes it to the given file."""
245    raise NotImplementedError
246
247
248class Version1VerityImageBuilder(VerityImageBuilder):
249  """A VerityImageBuilder for Verified Boot 1.0."""
250
251  def __init__(self, partition_size, block_dev, fec_supported, signer_path,
252               signer_key, signer_args, verity_disable):
253    self.version = 1
254    self.partition_size = partition_size
255    self.block_device = block_dev
256    self.fec_supported = fec_supported
257    self.signer_path = signer_path
258    self.signer_key = signer_key
259    self.signer_args = signer_args
260    self.verity_disable = verity_disable
261    self.image_size = None
262    self.verity_size = None
263
264  def CalculateDynamicPartitionSize(self, image_size):
265    # This needs to be implemented. Note that returning the given image size as
266    # the partition size doesn't make sense, as it will fail later.
267    raise NotImplementedError
268
269  def CalculateMaxImageSize(self, partition_size=None):
270    """Calculates the max image size by accounting for the verity metadata.
271
272    Args:
273      partition_size: The partition size, which defaults to self.partition_size
274          if unspecified.
275
276    Returns:
277      The size of the image adjusted for verity metadata.
278    """
279    if partition_size is None:
280      partition_size = self.partition_size
281    assert partition_size > 0, \
282        "Invalid partition size: {}".format(partition_size)
283
284    hi = partition_size
285    if hi % BLOCK_SIZE != 0:
286      hi = (hi // BLOCK_SIZE) * BLOCK_SIZE
287
288    # verity tree and fec sizes depend on the partition size, which
289    # means this estimate is always going to be unnecessarily small
290    verity_size = GetVeritySize(hi, self.fec_supported)
291    lo = partition_size - verity_size
292    result = lo
293
294    # do a binary search for the optimal size
295    while lo < hi:
296      i = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
297      v = GetVeritySize(i, self.fec_supported)
298      if i + v <= partition_size:
299        if result < i:
300          result = i
301          verity_size = v
302        lo = i + BLOCK_SIZE
303      else:
304        hi = i
305
306    self.image_size = result
307    self.verity_size = verity_size
308
309    logger.info(
310        "Calculated image size for verity: partition_size %d, image_size %d, "
311        "verity_size %d", partition_size, result, verity_size)
312    return result
313
314  def Build(self, out_file):
315    """Creates an image that is verifiable using dm-verity.
316
317    Args:
318      out_file: the output image.
319
320    Returns:
321      AssertionError: On invalid partition sizes.
322      BuildVerityImageError: On other errors.
323    """
324    image_size = int(self.image_size)
325    tempdir_name = common.MakeTempDir(suffix="_verity_images")
326
327    # Get partial image paths.
328    verity_image_path = os.path.join(tempdir_name, "verity.img")
329    verity_metadata_path = os.path.join(tempdir_name, "verity_metadata.img")
330
331    # Build the verity tree and get the root hash and salt.
332    root_hash, salt = BuildVerityTree(out_file, verity_image_path)
333
334    # Build the metadata blocks.
335    BuildVerityMetadata(
336        image_size, verity_metadata_path, root_hash, salt, self.block_device,
337        self.signer_path, self.signer_key, self.signer_args,
338        self.verity_disable)
339
340    padding_size = self.partition_size - self.image_size - self.verity_size
341    assert padding_size >= 0
342
343    # Build the full verified image.
344    Append(
345        verity_image_path, verity_metadata_path,
346        "Failed to append verity metadata")
347
348    if self.fec_supported:
349      # Build FEC for the entire partition, including metadata.
350      verity_fec_path = os.path.join(tempdir_name, "verity_fec.img")
351      BuildVerityFEC(
352          out_file, verity_image_path, verity_fec_path, padding_size)
353      Append(verity_image_path, verity_fec_path, "Failed to append FEC")
354
355    Append2Simg(
356        out_file, verity_image_path, "Failed to append verity data")
357
358  def PadSparseImage(self, out_file):
359    sparse_image_size = GetSimgSize(out_file)
360    if sparse_image_size > self.image_size:
361      raise BuildVerityImageError(
362          "Error: image size of {} is larger than partition size of "
363          "{}".format(sparse_image_size, self.image_size))
364    ZeroPadSimg(out_file, self.image_size - sparse_image_size)
365
366
367class VerifiedBootVersion2VerityImageBuilder(VerityImageBuilder):
368  """A VerityImageBuilder for Verified Boot 2.0."""
369
370  AVB_HASH_FOOTER = 1
371  AVB_HASHTREE_FOOTER = 2
372
373  def __init__(self, partition_name, partition_size, footer_type, avbtool,
374               key_path, algorithm, salt, signing_args):
375    self.version = 2
376    self.partition_name = partition_name
377    self.partition_size = partition_size
378    self.footer_type = footer_type
379    self.avbtool = avbtool
380    self.algorithm = algorithm
381    self.key_path = key_path
382    self.salt = salt
383    self.signing_args = signing_args
384    self.image_size = None
385
386  def CalculateMinPartitionSize(self, image_size, size_calculator=None):
387    """Calculates min partition size for a given image size.
388
389    This is used when determining the partition size for a dynamic partition,
390    which should be cover the given image size (for filesystem files) as well as
391    the verity metadata size.
392
393    Args:
394      image_size: The size of the image in question.
395      size_calculator: The function to calculate max image size
396          for a given partition size.
397
398    Returns:
399      The minimum partition size required to accommodate the image size.
400    """
401    if size_calculator is None:
402      size_calculator = self.CalculateMaxImageSize
403
404    # Use image size as partition size to approximate final partition size.
405    image_ratio = size_calculator(image_size) / float(image_size)
406
407    # Prepare a binary search for the optimal partition size.
408    lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE
409
410    # Ensure lo is small enough: max_image_size should <= image_size.
411    delta = BLOCK_SIZE
412    max_image_size = size_calculator(lo)
413    while max_image_size > image_size:
414      image_ratio = max_image_size / float(lo)
415      lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta
416      delta *= 2
417      max_image_size = size_calculator(lo)
418
419    hi = lo + BLOCK_SIZE
420
421    # Ensure hi is large enough: max_image_size should >= image_size.
422    delta = BLOCK_SIZE
423    max_image_size = size_calculator(hi)
424    while max_image_size < image_size:
425      image_ratio = max_image_size / float(hi)
426      hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta
427      delta *= 2
428      max_image_size = size_calculator(hi)
429
430    partition_size = hi
431
432    # Start to binary search.
433    while lo < hi:
434      mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
435      max_image_size = size_calculator(mid)
436      if max_image_size >= image_size:  # if mid can accommodate image_size
437        if mid < partition_size:  # if a smaller partition size is found
438          partition_size = mid
439        hi = mid
440      else:
441        lo = mid + BLOCK_SIZE
442
443    logger.info(
444        "CalculateMinPartitionSize(%d): partition_size %d.", image_size,
445        partition_size)
446
447    return partition_size
448
449  def CalculateDynamicPartitionSize(self, image_size):
450    self.partition_size = self.CalculateMinPartitionSize(image_size)
451    return self.partition_size
452
453  def CalculateMaxImageSize(self, partition_size=None):
454    """Calculates max image size for a given partition size.
455
456    Args:
457      partition_size: The partition size, which defaults to self.partition_size
458          if unspecified.
459
460    Returns:
461      The maximum image size.
462
463    Raises:
464      BuildVerityImageError: On error or getting invalid image size.
465    """
466    if partition_size is None:
467      partition_size = self.partition_size
468    assert partition_size > 0, \
469        "Invalid partition size: {}".format(partition_size)
470
471    add_footer = ("add_hash_footer" if self.footer_type == self.AVB_HASH_FOOTER
472                  else "add_hashtree_footer")
473    cmd = [self.avbtool, add_footer, "--partition_size",
474           str(partition_size), "--calc_max_image_size"]
475    cmd.extend(shlex.split(self.signing_args))
476
477    proc = common.Run(cmd)
478    output, _ = proc.communicate()
479    if proc.returncode != 0:
480      raise BuildVerityImageError(
481          "Failed to calculate max image size:\n{}".format(output))
482    image_size = int(output)
483    if image_size <= 0:
484      raise BuildVerityImageError(
485          "Invalid max image size: {}".format(output))
486    self.image_size = image_size
487    return image_size
488
489  def PadSparseImage(self, out_file):
490    # No-op as the padding is taken care of by avbtool.
491    pass
492
493  def Build(self, out_file):
494    """Adds dm-verity hashtree and AVB metadata to an image.
495
496    Args:
497      out_file: Path to image to modify.
498    """
499    add_footer = ("add_hash_footer" if self.footer_type == self.AVB_HASH_FOOTER
500                  else "add_hashtree_footer")
501    cmd = [self.avbtool, add_footer,
502           "--partition_size", str(self.partition_size),
503           "--partition_name", self.partition_name,
504           "--image", out_file]
505    if self.key_path and self.algorithm:
506      cmd.extend(["--key", self.key_path, "--algorithm", self.algorithm])
507    if self.salt:
508      cmd.extend(["--salt", self.salt])
509    cmd.extend(shlex.split(self.signing_args))
510
511    proc = common.Run(cmd)
512    output, _ = proc.communicate()
513    if proc.returncode != 0:
514      raise BuildVerityImageError("Failed to add AVB footer: {}".format(output))
515
516
517class HashtreeInfoGenerationError(Exception):
518  """An Exception raised during hashtree info generation."""
519
520  def __init__(self, message):
521    Exception.__init__(self, message)
522
523
524class HashtreeInfo(object):
525  def __init__(self):
526    self.hashtree_range = None
527    self.filesystem_range = None
528    self.hash_algorithm = None
529    self.salt = None
530    self.root_hash = None
531
532
533def CreateHashtreeInfoGenerator(partition_name, block_size, info_dict):
534  generator = None
535  if (info_dict.get("verity") == "true" and
536      info_dict.get("{}_verity_block_device".format(partition_name))):
537    partition_size = info_dict["{}_size".format(partition_name)]
538    fec_supported = info_dict.get("verity_fec") == "true"
539    generator = VerifiedBootVersion1HashtreeInfoGenerator(
540        partition_size, block_size, fec_supported)
541
542  return generator
543
544
545class HashtreeInfoGenerator(object):
546  def Generate(self, image):
547    raise NotImplementedError
548
549  def DecomposeSparseImage(self, image):
550    raise NotImplementedError
551
552  def ValidateHashtree(self):
553    raise NotImplementedError
554
555
556class VerifiedBootVersion1HashtreeInfoGenerator(HashtreeInfoGenerator):
557  """A class that parses the metadata of hashtree for a given partition."""
558
559  def __init__(self, partition_size, block_size, fec_supported):
560    """Initialize VerityTreeInfo with the sparse image and input property.
561
562    Arguments:
563      partition_size: The whole size in bytes of a partition, including the
564          filesystem size, padding size, and verity size.
565      block_size: Expected size in bytes of each block for the sparse image.
566      fec_supported: True if the verity section contains fec data.
567    """
568
569    self.block_size = block_size
570    self.partition_size = partition_size
571    self.fec_supported = fec_supported
572
573    self.image = None
574    self.filesystem_size = None
575    self.hashtree_size = None
576    self.metadata_size = None
577
578    prop_dict = {
579        'partition_size': str(partition_size),
580        'verity': 'true',
581        'verity_fec': 'true' if fec_supported else None,
582        # 'verity_block_device' needs to be present to indicate a verity-enabled
583        # partition.
584        'verity_block_device': '',
585        # We don't need the following properties that are needed for signing the
586        # verity metadata.
587        'verity_key': '',
588        'verity_signer_cmd': None,
589    }
590    self.verity_image_builder = CreateVerityImageBuilder(prop_dict)
591
592    self.hashtree_info = HashtreeInfo()
593
594  def DecomposeSparseImage(self, image):
595    """Calculate the verity size based on the size of the input image.
596
597    Since we already know the structure of a verity enabled image to be:
598    [filesystem, verity_hashtree, verity_metadata, fec_data]. We can then
599    calculate the size and offset of each section.
600    """
601
602    self.image = image
603    assert self.block_size == image.blocksize
604    assert self.partition_size == image.total_blocks * self.block_size, \
605        "partition size {} doesn't match with the calculated image size." \
606        " total_blocks: {}".format(self.partition_size, image.total_blocks)
607
608    adjusted_size = self.verity_image_builder.CalculateMaxImageSize()
609    assert adjusted_size % self.block_size == 0
610
611    verity_tree_size = GetVerityTreeSize(adjusted_size)
612    assert verity_tree_size % self.block_size == 0
613
614    metadata_size = GetVerityMetadataSize(adjusted_size)
615    assert metadata_size % self.block_size == 0
616
617    self.filesystem_size = adjusted_size
618    self.hashtree_size = verity_tree_size
619    self.metadata_size = metadata_size
620
621    self.hashtree_info.filesystem_range = RangeSet(
622        data=[0, adjusted_size // self.block_size])
623    self.hashtree_info.hashtree_range = RangeSet(
624        data=[adjusted_size // self.block_size,
625              (adjusted_size + verity_tree_size) // self.block_size])
626
627  def _ParseHashtreeMetadata(self):
628    """Parses the hash_algorithm, root_hash, salt from the metadata block."""
629
630    metadata_start = self.filesystem_size + self.hashtree_size
631    metadata_range = RangeSet(
632        data=[metadata_start // self.block_size,
633              (metadata_start + self.metadata_size) // self.block_size])
634    meta_data = b''.join(self.image.ReadRangeSet(metadata_range))
635
636    # More info about the metadata structure available in:
637    # system/extras/verity/build_verity_metadata.py
638    META_HEADER_SIZE = 268
639    header_bin = meta_data[0:META_HEADER_SIZE]
640    header = struct.unpack("II256sI", header_bin)
641
642    # header: magic_number, version, signature, table_len
643    assert header[0] == 0xb001b001, header[0]
644    table_len = header[3]
645    verity_table = meta_data[META_HEADER_SIZE: META_HEADER_SIZE + table_len]
646    table_entries = verity_table.rstrip().split()
647
648    # Expected verity table format: "1 block_device block_device block_size
649    # block_size data_blocks data_blocks hash_algorithm root_hash salt"
650    assert len(table_entries) == 10, "Unexpected verity table size {}".format(
651        len(table_entries))
652    assert (int(table_entries[3]) == self.block_size and
653            int(table_entries[4]) == self.block_size)
654    assert (int(table_entries[5]) * self.block_size == self.filesystem_size and
655            int(table_entries[6]) * self.block_size == self.filesystem_size)
656
657    self.hashtree_info.hash_algorithm = table_entries[7].decode()
658    self.hashtree_info.root_hash = table_entries[8].decode()
659    self.hashtree_info.salt = table_entries[9].decode()
660
661  def ValidateHashtree(self):
662    """Checks that we can reconstruct the verity hash tree."""
663
664    # Writes the filesystem section to a temp file; and calls the executable
665    # build_verity_tree to construct the hash tree.
666    adjusted_partition = common.MakeTempFile(prefix="adjusted_partition")
667    with open(adjusted_partition, "wb") as fd:
668      self.image.WriteRangeDataToFd(self.hashtree_info.filesystem_range, fd)
669
670    generated_verity_tree = common.MakeTempFile(prefix="verity")
671    root_hash, salt = BuildVerityTree(adjusted_partition, generated_verity_tree)
672
673    # The salt should be always identical, as we use fixed value.
674    assert salt == self.hashtree_info.salt, \
675        "Calculated salt {} doesn't match the one in metadata {}".format(
676            salt, self.hashtree_info.salt)
677
678    if root_hash != self.hashtree_info.root_hash:
679      logger.warning(
680          "Calculated root hash %s doesn't match the one in metadata %s",
681          root_hash, self.hashtree_info.root_hash)
682      return False
683
684    # Reads the generated hash tree and checks if it has the exact same bytes
685    # as the one in the sparse image.
686    with open(generated_verity_tree, 'rb') as fd:
687      return fd.read() == b''.join(self.image.ReadRangeSet(
688          self.hashtree_info.hashtree_range))
689
690  def Generate(self, image):
691    """Parses and validates the hashtree info in a sparse image.
692
693    Returns:
694      hashtree_info: The information needed to reconstruct the hashtree.
695
696    Raises:
697      HashtreeInfoGenerationError: If we fail to generate the exact bytes of
698          the hashtree.
699    """
700
701    self.DecomposeSparseImage(image)
702    self._ParseHashtreeMetadata()
703
704    if not self.ValidateHashtree():
705      raise HashtreeInfoGenerationError("Failed to reconstruct the verity tree")
706
707    return self.hashtree_info
708
709
710def CreateCustomImageBuilder(info_dict, partition_name, partition_size,
711                            key_path, algorithm, signing_args):
712  builder = None
713  if info_dict.get("avb_enable") == "true":
714    builder = VerifiedBootVersion2VerityImageBuilder(
715        partition_name,
716        partition_size,
717        VerifiedBootVersion2VerityImageBuilder.AVB_HASHTREE_FOOTER,
718        info_dict.get("avb_avbtool"),
719        key_path,
720        algorithm,
721        # Salt is None because custom images have no fingerprint property to be
722        # used as the salt.
723        None,
724        signing_args)
725
726  return builder
727
728
729def GetDiskUsage(path):
730  """Returns the number of bytes that "path" occupies on host.
731
732  Args:
733    path: The directory or file to calculate size on.
734
735  Returns:
736    The number of bytes based on a 1K block_size.
737  """
738  cmd = ["du", "-b", "-k", "-s", path]
739  output = common.RunAndCheckOutput(cmd, verbose=False)
740  return int(output.split()[0]) * 1024
741
742
743def CalculateVbmetaDigest(extracted_dir, avbtool):
744  """Calculates the vbmeta digest of the images in the extracted target_file"""
745
746  images_dir = common.MakeTempDir()
747  for name in ("PREBUILT_IMAGES", "RADIO", "IMAGES"):
748    path = os.path.join(extracted_dir, name)
749    if not os.path.exists(path):
750      continue
751
752    # Create symlink for image files under PREBUILT_IMAGES, RADIO and IMAGES,
753    # and put them into one directory.
754    for filename in os.listdir(path):
755      if not filename.endswith(".img"):
756        continue
757      symlink_path = os.path.join(images_dir, filename)
758      # The files in latter directory overwrite the existing links
759      common.RunAndCheckOutput(
760        ['ln', '-sf', os.path.join(path, filename), symlink_path])
761
762  cmd = [avbtool, "calculate_vbmeta_digest", "--image",
763         os.path.join(images_dir, 'vbmeta.img')]
764  return common.RunAndCheckOutput(cmd)
765
766
767def main(argv):
768  if len(argv) != 2:
769    print(__doc__)
770    sys.exit(1)
771
772  common.InitLogging()
773
774  dict_file = argv[0]
775  out_file = argv[1]
776
777  prop_dict = {}
778  with open(dict_file, 'r') as f:
779    for line in f:
780      line = line.strip()
781      if not line or line.startswith("#"):
782        continue
783      k, v = line.split("=", 1)
784      prop_dict[k] = v
785
786  builder = CreateVerityImageBuilder(prop_dict)
787
788  if "partition_size" not in prop_dict:
789    image_size = GetDiskUsage(out_file)
790    # make sure that the image is big enough to hold vbmeta and footer
791    image_size = image_size + (MAX_VBMETA_SIZE + MAX_FOOTER_SIZE)
792    size = builder.CalculateDynamicPartitionSize(image_size)
793    prop_dict["partition_size"] = size
794
795  builder.Build(out_file)
796
797
798if __name__ == '__main__':
799  try:
800    main(sys.argv[1:])
801  finally:
802    common.Cleanup()
803