• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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    if key_path and not os.path.exists(key_path) and OPTIONS.search_path:
383      new_key_path = os.path.join(OPTIONS.search_path, key_path)
384      if os.path.exists(new_key_path):
385        self.key_path = new_key_path
386
387    self.salt = salt
388    self.signing_args = signing_args
389    self.image_size = None
390
391  def CalculateMinPartitionSize(self, image_size, size_calculator=None):
392    """Calculates min partition size for a given image size.
393
394    This is used when determining the partition size for a dynamic partition,
395    which should be cover the given image size (for filesystem files) as well as
396    the verity metadata size.
397
398    Args:
399      image_size: The size of the image in question.
400      size_calculator: The function to calculate max image size
401          for a given partition size.
402
403    Returns:
404      The minimum partition size required to accommodate the image size.
405    """
406    if size_calculator is None:
407      size_calculator = self.CalculateMaxImageSize
408
409    # Use image size as partition size to approximate final partition size.
410    image_ratio = size_calculator(image_size) / float(image_size)
411
412    # Prepare a binary search for the optimal partition size.
413    lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE
414
415    # Ensure lo is small enough: max_image_size should <= image_size.
416    delta = BLOCK_SIZE
417    max_image_size = size_calculator(lo)
418    while max_image_size > image_size:
419      image_ratio = max_image_size / float(lo)
420      lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta
421      delta *= 2
422      max_image_size = size_calculator(lo)
423
424    hi = lo + BLOCK_SIZE
425
426    # Ensure hi is large enough: max_image_size should >= image_size.
427    delta = BLOCK_SIZE
428    max_image_size = size_calculator(hi)
429    while max_image_size < image_size:
430      image_ratio = max_image_size / float(hi)
431      hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta
432      delta *= 2
433      max_image_size = size_calculator(hi)
434
435    partition_size = hi
436
437    # Start to binary search.
438    while lo < hi:
439      mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
440      max_image_size = size_calculator(mid)
441      if max_image_size >= image_size:  # if mid can accommodate image_size
442        if mid < partition_size:  # if a smaller partition size is found
443          partition_size = mid
444        hi = mid
445      else:
446        lo = mid + BLOCK_SIZE
447
448    logger.info(
449        "CalculateMinPartitionSize(%d): partition_size %d.", image_size,
450        partition_size)
451
452    return partition_size
453
454  def CalculateDynamicPartitionSize(self, image_size):
455    self.partition_size = self.CalculateMinPartitionSize(image_size)
456    return self.partition_size
457
458  def CalculateMaxImageSize(self, partition_size=None):
459    """Calculates max image size for a given partition size.
460
461    Args:
462      partition_size: The partition size, which defaults to self.partition_size
463          if unspecified.
464
465    Returns:
466      The maximum image size.
467
468    Raises:
469      BuildVerityImageError: On error or getting invalid image size.
470    """
471    if partition_size is None:
472      partition_size = self.partition_size
473    assert partition_size > 0, \
474        "Invalid partition size: {}".format(partition_size)
475
476    add_footer = ("add_hash_footer" if self.footer_type == self.AVB_HASH_FOOTER
477                  else "add_hashtree_footer")
478    cmd = [self.avbtool, add_footer, "--partition_size",
479           str(partition_size), "--calc_max_image_size"]
480    cmd.extend(shlex.split(self.signing_args))
481
482    proc = common.Run(cmd)
483    output, _ = proc.communicate()
484    if proc.returncode != 0:
485      raise BuildVerityImageError(
486          "Failed to calculate max image size:\n{}".format(output))
487    image_size = int(output)
488    if image_size <= 0:
489      raise BuildVerityImageError(
490          "Invalid max image size: {}".format(output))
491    self.image_size = image_size
492    return image_size
493
494  def PadSparseImage(self, out_file):
495    # No-op as the padding is taken care of by avbtool.
496    pass
497
498  def Build(self, out_file):
499    """Adds dm-verity hashtree and AVB metadata to an image.
500
501    Args:
502      out_file: Path to image to modify.
503    """
504    add_footer = ("add_hash_footer" if self.footer_type == self.AVB_HASH_FOOTER
505                  else "add_hashtree_footer")
506    cmd = [self.avbtool, add_footer,
507           "--partition_size", str(self.partition_size),
508           "--partition_name", self.partition_name,
509           "--image", out_file]
510    if self.key_path and self.algorithm:
511      cmd.extend(["--key", self.key_path, "--algorithm", self.algorithm])
512    if self.salt:
513      cmd.extend(["--salt", self.salt])
514    cmd.extend(shlex.split(self.signing_args))
515
516    proc = common.Run(cmd)
517    output, _ = proc.communicate()
518    if proc.returncode != 0:
519      raise BuildVerityImageError("Failed to add AVB footer: {}".format(output))
520
521
522class HashtreeInfoGenerationError(Exception):
523  """An Exception raised during hashtree info generation."""
524
525  def __init__(self, message):
526    Exception.__init__(self, message)
527
528
529class HashtreeInfo(object):
530  def __init__(self):
531    self.hashtree_range = None
532    self.filesystem_range = None
533    self.hash_algorithm = None
534    self.salt = None
535    self.root_hash = None
536
537
538def CreateHashtreeInfoGenerator(partition_name, block_size, info_dict):
539  generator = None
540  if (info_dict.get("verity") == "true" and
541      info_dict.get("{}_verity_block_device".format(partition_name))):
542    partition_size = info_dict["{}_size".format(partition_name)]
543    fec_supported = info_dict.get("verity_fec") == "true"
544    generator = VerifiedBootVersion1HashtreeInfoGenerator(
545        partition_size, block_size, fec_supported)
546
547  return generator
548
549
550class HashtreeInfoGenerator(object):
551  def Generate(self, image):
552    raise NotImplementedError
553
554  def DecomposeSparseImage(self, image):
555    raise NotImplementedError
556
557  def ValidateHashtree(self):
558    raise NotImplementedError
559
560
561class VerifiedBootVersion1HashtreeInfoGenerator(HashtreeInfoGenerator):
562  """A class that parses the metadata of hashtree for a given partition."""
563
564  def __init__(self, partition_size, block_size, fec_supported):
565    """Initialize VerityTreeInfo with the sparse image and input property.
566
567    Arguments:
568      partition_size: The whole size in bytes of a partition, including the
569          filesystem size, padding size, and verity size.
570      block_size: Expected size in bytes of each block for the sparse image.
571      fec_supported: True if the verity section contains fec data.
572    """
573
574    self.block_size = block_size
575    self.partition_size = partition_size
576    self.fec_supported = fec_supported
577
578    self.image = None
579    self.filesystem_size = None
580    self.hashtree_size = None
581    self.metadata_size = None
582
583    prop_dict = {
584        'partition_size': str(partition_size),
585        'verity': 'true',
586        'verity_fec': 'true' if fec_supported else None,
587        # 'verity_block_device' needs to be present to indicate a verity-enabled
588        # partition.
589        'verity_block_device': '',
590        # We don't need the following properties that are needed for signing the
591        # verity metadata.
592        'verity_key': '',
593        'verity_signer_cmd': None,
594    }
595    self.verity_image_builder = CreateVerityImageBuilder(prop_dict)
596
597    self.hashtree_info = HashtreeInfo()
598
599  def DecomposeSparseImage(self, image):
600    """Calculate the verity size based on the size of the input image.
601
602    Since we already know the structure of a verity enabled image to be:
603    [filesystem, verity_hashtree, verity_metadata, fec_data]. We can then
604    calculate the size and offset of each section.
605    """
606
607    self.image = image
608    assert self.block_size == image.blocksize
609    assert self.partition_size == image.total_blocks * self.block_size, \
610        "partition size {} doesn't match with the calculated image size." \
611        " total_blocks: {}".format(self.partition_size, image.total_blocks)
612
613    adjusted_size = self.verity_image_builder.CalculateMaxImageSize()
614    assert adjusted_size % self.block_size == 0
615
616    verity_tree_size = GetVerityTreeSize(adjusted_size)
617    assert verity_tree_size % self.block_size == 0
618
619    metadata_size = GetVerityMetadataSize(adjusted_size)
620    assert metadata_size % self.block_size == 0
621
622    self.filesystem_size = adjusted_size
623    self.hashtree_size = verity_tree_size
624    self.metadata_size = metadata_size
625
626    self.hashtree_info.filesystem_range = RangeSet(
627        data=[0, adjusted_size // self.block_size])
628    self.hashtree_info.hashtree_range = RangeSet(
629        data=[adjusted_size // self.block_size,
630              (adjusted_size + verity_tree_size) // self.block_size])
631
632  def _ParseHashtreeMetadata(self):
633    """Parses the hash_algorithm, root_hash, salt from the metadata block."""
634
635    metadata_start = self.filesystem_size + self.hashtree_size
636    metadata_range = RangeSet(
637        data=[metadata_start // self.block_size,
638              (metadata_start + self.metadata_size) // self.block_size])
639    meta_data = b''.join(self.image.ReadRangeSet(metadata_range))
640
641    # More info about the metadata structure available in:
642    # system/extras/verity/build_verity_metadata.py
643    META_HEADER_SIZE = 268
644    header_bin = meta_data[0:META_HEADER_SIZE]
645    header = struct.unpack("II256sI", header_bin)
646
647    # header: magic_number, version, signature, table_len
648    assert header[0] == 0xb001b001, header[0]
649    table_len = header[3]
650    verity_table = meta_data[META_HEADER_SIZE: META_HEADER_SIZE + table_len]
651    table_entries = verity_table.rstrip().split()
652
653    # Expected verity table format: "1 block_device block_device block_size
654    # block_size data_blocks data_blocks hash_algorithm root_hash salt"
655    assert len(table_entries) == 10, "Unexpected verity table size {}".format(
656        len(table_entries))
657    assert (int(table_entries[3]) == self.block_size and
658            int(table_entries[4]) == self.block_size)
659    assert (int(table_entries[5]) * self.block_size == self.filesystem_size and
660            int(table_entries[6]) * self.block_size == self.filesystem_size)
661
662    self.hashtree_info.hash_algorithm = table_entries[7].decode()
663    self.hashtree_info.root_hash = table_entries[8].decode()
664    self.hashtree_info.salt = table_entries[9].decode()
665
666  def ValidateHashtree(self):
667    """Checks that we can reconstruct the verity hash tree."""
668
669    # Writes the filesystem section to a temp file; and calls the executable
670    # build_verity_tree to construct the hash tree.
671    adjusted_partition = common.MakeTempFile(prefix="adjusted_partition")
672    with open(adjusted_partition, "wb") as fd:
673      self.image.WriteRangeDataToFd(self.hashtree_info.filesystem_range, fd)
674
675    generated_verity_tree = common.MakeTempFile(prefix="verity")
676    root_hash, salt = BuildVerityTree(adjusted_partition, generated_verity_tree)
677
678    # The salt should be always identical, as we use fixed value.
679    assert salt == self.hashtree_info.salt, \
680        "Calculated salt {} doesn't match the one in metadata {}".format(
681            salt, self.hashtree_info.salt)
682
683    if root_hash != self.hashtree_info.root_hash:
684      logger.warning(
685          "Calculated root hash %s doesn't match the one in metadata %s",
686          root_hash, self.hashtree_info.root_hash)
687      return False
688
689    # Reads the generated hash tree and checks if it has the exact same bytes
690    # as the one in the sparse image.
691    with open(generated_verity_tree, 'rb') as fd:
692      return fd.read() == b''.join(self.image.ReadRangeSet(
693          self.hashtree_info.hashtree_range))
694
695  def Generate(self, image):
696    """Parses and validates the hashtree info in a sparse image.
697
698    Returns:
699      hashtree_info: The information needed to reconstruct the hashtree.
700
701    Raises:
702      HashtreeInfoGenerationError: If we fail to generate the exact bytes of
703          the hashtree.
704    """
705
706    self.DecomposeSparseImage(image)
707    self._ParseHashtreeMetadata()
708
709    if not self.ValidateHashtree():
710      raise HashtreeInfoGenerationError("Failed to reconstruct the verity tree")
711
712    return self.hashtree_info
713
714
715def CreateCustomImageBuilder(info_dict, partition_name, partition_size,
716                            key_path, algorithm, signing_args):
717  builder = None
718  if info_dict.get("avb_enable") == "true":
719    builder = VerifiedBootVersion2VerityImageBuilder(
720        partition_name,
721        partition_size,
722        VerifiedBootVersion2VerityImageBuilder.AVB_HASHTREE_FOOTER,
723        info_dict.get("avb_avbtool"),
724        key_path,
725        algorithm,
726        # Salt is None because custom images have no fingerprint property to be
727        # used as the salt.
728        None,
729        signing_args)
730
731  return builder
732
733
734def GetDiskUsage(path):
735  """Returns the number of bytes that "path" occupies on host.
736
737  Args:
738    path: The directory or file to calculate size on.
739
740  Returns:
741    The number of bytes based on a 1K block_size.
742  """
743  cmd = ["du", "-b", "-k", "-s", path]
744  output = common.RunAndCheckOutput(cmd, verbose=False)
745  return int(output.split()[0]) * 1024
746
747
748def CalculateVbmetaDigest(extracted_dir, avbtool):
749  """Calculates the vbmeta digest of the images in the extracted target_file"""
750
751  images_dir = common.MakeTempDir()
752  for name in ("PREBUILT_IMAGES", "RADIO", "IMAGES"):
753    path = os.path.join(extracted_dir, name)
754    if not os.path.exists(path):
755      continue
756
757    # Create symlink for image files under PREBUILT_IMAGES, RADIO and IMAGES,
758    # and put them into one directory.
759    for filename in os.listdir(path):
760      if not filename.endswith(".img"):
761        continue
762      symlink_path = os.path.join(images_dir, filename)
763      # The files in latter directory overwrite the existing links
764      common.RunAndCheckOutput(
765        ['ln', '-sf', os.path.join(path, filename), symlink_path])
766
767  cmd = [avbtool, "calculate_vbmeta_digest", "--image",
768         os.path.join(images_dir, 'vbmeta.img')]
769  return common.RunAndCheckOutput(cmd)
770
771
772def main(argv):
773  if len(argv) != 2:
774    print(__doc__)
775    sys.exit(1)
776
777  common.InitLogging()
778
779  dict_file = argv[0]
780  out_file = argv[1]
781
782  prop_dict = {}
783  with open(dict_file, 'r') as f:
784    for line in f:
785      line = line.strip()
786      if not line or line.startswith("#"):
787        continue
788      k, v = line.split("=", 1)
789      prop_dict[k] = v
790
791  builder = CreateVerityImageBuilder(prop_dict)
792
793  if "partition_size" not in prop_dict:
794    image_size = GetDiskUsage(out_file)
795    # make sure that the image is big enough to hold vbmeta and footer
796    image_size = image_size + (MAX_VBMETA_SIZE + MAX_FOOTER_SIZE)
797    size = builder.CalculateDynamicPartitionSize(image_size)
798    prop_dict["partition_size"] = size
799
800  builder.Build(out_file)
801
802
803if __name__ == '__main__':
804  try:
805    main(sys.argv[1:])
806  finally:
807    common.Cleanup()
808