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
34from hashlib import sha256
35
36logger = logging.getLogger(__name__)
37
38OPTIONS = common.OPTIONS
39BLOCK_SIZE = common.BLOCK_SIZE
40FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7"
41
42# From external/avb/avbtool.py
43MAX_VBMETA_SIZE = 64 * 1024
44MAX_FOOTER_SIZE = 4096
45
46
47class BuildVerityImageError(Exception):
48  """An Exception raised during verity image building."""
49
50  def __init__(self, message):
51    Exception.__init__(self, message)
52
53
54def CreateVerityImageBuilder(prop_dict):
55  """Returns a verity image builder based on the given build properties.
56
57  Args:
58    prop_dict: A dict that contains the build properties. In particular, it will
59        look for verity-related property values.
60
61  Returns:
62    A VerityImageBuilder instance for Verified Boot 1.0 or Verified Boot 2.0; or
63        None if the given build doesn't support Verified Boot.
64  """
65  partition_size = prop_dict.get("partition_size")
66  # partition_size could be None at this point, if using dynamic partitions.
67  if partition_size:
68    partition_size = int(partition_size)
69  # Set up the salt (based on fingerprint) that will be used when adding AVB
70  # hash / hashtree footers.
71  salt = prop_dict.get("avb_salt")
72  if salt is None:
73    salt = sha256(prop_dict.get("fingerprint", "").encode()).hexdigest()
74
75  # Verified Boot 2.0
76  if (prop_dict.get("avb_hash_enable") == "true" or
77      prop_dict.get("avb_hashtree_enable") == "true"):
78    # key_path and algorithm are only available when chain partition is used.
79    key_path = prop_dict.get("avb_key_path")
80    algorithm = prop_dict.get("avb_algorithm")
81
82    # Image uses hash footer.
83    if prop_dict.get("avb_hash_enable") == "true":
84      return VerifiedBootVersion2VerityImageBuilder(
85          prop_dict["partition_name"],
86          partition_size,
87          VerifiedBootVersion2VerityImageBuilder.AVB_HASH_FOOTER,
88          prop_dict["avb_avbtool"],
89          key_path,
90          algorithm,
91          salt,
92          prop_dict["avb_add_hash_footer_args"])
93
94    # Image uses hashtree footer.
95    return VerifiedBootVersion2VerityImageBuilder(
96        prop_dict["partition_name"],
97        partition_size,
98        VerifiedBootVersion2VerityImageBuilder.AVB_HASHTREE_FOOTER,
99        prop_dict["avb_avbtool"],
100        key_path,
101        algorithm,
102        salt,
103        prop_dict["avb_add_hashtree_footer_args"])
104
105  return None
106
107
108class VerityImageBuilder(object):
109  """A builder that generates an image with verity metadata for Verified Boot.
110
111  A VerityImageBuilder instance handles the works for building an image with
112  verity metadata for supporting Android Verified Boot. This class defines the
113  common interface between Verified Boot 1.0 and Verified Boot 2.0. A matching
114  builder will be returned based on the given build properties.
115
116  More info on the verity image generation can be found at the following link.
117  https://source.android.com/security/verifiedboot/dm-verity#implementation
118  """
119
120  def CalculateMaxImageSize(self, partition_size):
121    """Calculates the filesystem image size for the given partition size."""
122    raise NotImplementedError
123
124  def CalculateDynamicPartitionSize(self, image_size):
125    """Calculates and sets the partition size for a dynamic partition."""
126    raise NotImplementedError
127
128  def PadSparseImage(self, out_file):
129    """Adds padding to the generated sparse image."""
130    raise NotImplementedError
131
132  def Build(self, out_file):
133    """Builds the verity image and writes it to the given file."""
134    raise NotImplementedError
135
136
137class VerifiedBootVersion2VerityImageBuilder(VerityImageBuilder):
138  """A VerityImageBuilder for Verified Boot 2.0."""
139
140  AVB_HASH_FOOTER = 1
141  AVB_HASHTREE_FOOTER = 2
142
143  def __init__(self, partition_name, partition_size, footer_type, avbtool,
144               key_path, algorithm, salt, signing_args):
145    self.version = 2
146    self.partition_name = partition_name
147    self.partition_size = partition_size
148    self.footer_type = footer_type
149    self.avbtool = avbtool
150    self.algorithm = algorithm
151    self.key_path = common.ResolveAVBSigningPathArgs(key_path)
152
153    self.salt = salt
154    self.signing_args = signing_args
155    self.image_size = None
156
157  def CalculateMinPartitionSize(self, image_size, size_calculator=None):
158    """Calculates min partition size for a given image size.
159
160    This is used when determining the partition size for a dynamic partition,
161    which should be cover the given image size (for filesystem files) as well as
162    the verity metadata size.
163
164    Args:
165      image_size: The size of the image in question.
166      size_calculator: The function to calculate max image size
167          for a given partition size.
168
169    Returns:
170      The minimum partition size required to accommodate the image size.
171    """
172    if size_calculator is None:
173      size_calculator = self.CalculateMaxImageSize
174
175    # Use image size as partition size to approximate final partition size.
176    image_ratio = size_calculator(image_size) / float(image_size)
177
178    # Prepare a binary search for the optimal partition size.
179    lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE
180
181    # Ensure lo is small enough: max_image_size should <= image_size.
182    delta = BLOCK_SIZE
183    max_image_size = size_calculator(lo)
184    while max_image_size > image_size:
185      image_ratio = max_image_size / float(lo)
186      lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta
187      delta *= 2
188      max_image_size = size_calculator(lo)
189
190    hi = lo + BLOCK_SIZE
191
192    # Ensure hi is large enough: max_image_size should >= image_size.
193    delta = BLOCK_SIZE
194    max_image_size = size_calculator(hi)
195    while max_image_size < image_size:
196      image_ratio = max_image_size / float(hi)
197      hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta
198      delta *= 2
199      max_image_size = size_calculator(hi)
200
201    partition_size = hi
202
203    # Start to binary search.
204    while lo < hi:
205      mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
206      max_image_size = size_calculator(mid)
207      if max_image_size >= image_size:  # if mid can accommodate image_size
208        if mid < partition_size:  # if a smaller partition size is found
209          partition_size = mid
210        hi = mid
211      else:
212        lo = mid + BLOCK_SIZE
213
214    logger.info(
215        "CalculateMinPartitionSize(%d): partition_size %d.", image_size,
216        partition_size)
217
218    return partition_size
219
220  def CalculateDynamicPartitionSize(self, image_size):
221    self.partition_size = self.CalculateMinPartitionSize(image_size)
222    return self.partition_size
223
224  def CalculateMaxImageSize(self, partition_size=None):
225    """Calculates max image size for a given partition size.
226
227    Args:
228      partition_size: The partition size, which defaults to self.partition_size
229          if unspecified.
230
231    Returns:
232      The maximum image size.
233
234    Raises:
235      BuildVerityImageError: On error or getting invalid image size.
236    """
237    if partition_size is None:
238      partition_size = self.partition_size
239    assert partition_size > 0, \
240        "Invalid partition size: {}".format(partition_size)
241
242    add_footer = ("add_hash_footer" if self.footer_type == self.AVB_HASH_FOOTER
243                  else "add_hashtree_footer")
244    cmd = [self.avbtool, add_footer, "--partition_size",
245           str(partition_size), "--calc_max_image_size"]
246    cmd.extend(shlex.split(self.signing_args))
247
248    proc = common.Run(cmd)
249    output, _ = proc.communicate()
250    if proc.returncode != 0:
251      raise BuildVerityImageError(
252          "Failed to calculate max image size:\n{}".format(output))
253    image_size = int(output)
254    if image_size <= 0:
255      raise BuildVerityImageError(
256          "Invalid max image size: {}".format(output))
257    self.image_size = image_size
258    return image_size
259
260  def PadSparseImage(self, out_file):
261    # No-op as the padding is taken care of by avbtool.
262    pass
263
264  def Build(self, out_file):
265    """Adds dm-verity hashtree and AVB metadata to an image.
266
267    Args:
268      out_file: Path to image to modify.
269    """
270    add_footer = ("add_hash_footer" if self.footer_type == self.AVB_HASH_FOOTER
271                  else "add_hashtree_footer")
272    cmd = [self.avbtool, add_footer,
273           "--partition_size", str(self.partition_size),
274           "--partition_name", self.partition_name,
275           "--image", out_file]
276    if self.key_path and self.algorithm:
277      cmd.extend(["--key", self.key_path, "--algorithm", self.algorithm])
278    if self.salt:
279      cmd.extend(["--salt", self.salt])
280    cmd.extend(shlex.split(self.signing_args))
281
282    proc = common.Run(cmd)
283    output, _ = proc.communicate()
284    if proc.returncode != 0:
285      raise BuildVerityImageError("Failed to add AVB footer: {}".format(output))
286
287
288def CreateCustomImageBuilder(info_dict, partition_name, partition_size,
289                             key_path, algorithm, signing_args):
290  builder = None
291  if info_dict.get("avb_enable") == "true":
292    builder = VerifiedBootVersion2VerityImageBuilder(
293        partition_name,
294        partition_size,
295        VerifiedBootVersion2VerityImageBuilder.AVB_HASHTREE_FOOTER,
296        info_dict.get("avb_avbtool"),
297        key_path,
298        algorithm,
299        # Salt is None because custom images have no fingerprint property to be
300        # used as the salt.
301        None,
302        signing_args)
303
304  return builder
305
306
307def GetDiskUsage(path):
308  """Returns the number of bytes that "path" occupies on host.
309
310  Args:
311    path: The directory or file to calculate size on.
312
313  Returns:
314    The number of bytes based on a 1K block_size.
315  """
316  cmd = ["du", "-b", "-k", "-s", path]
317  output = common.RunAndCheckOutput(cmd, verbose=False)
318  return int(output.split()[0]) * 1024
319
320
321def CalculateVbmetaDigest(extracted_dir, avbtool):
322  """Calculates the vbmeta digest of the images in the extracted target_file"""
323
324  images_dir = common.MakeTempDir()
325  for name in ("PREBUILT_IMAGES", "RADIO", "IMAGES"):
326    path = os.path.join(extracted_dir, name)
327    if not os.path.exists(path):
328      continue
329
330    # Create symlink for image files under PREBUILT_IMAGES, RADIO and IMAGES,
331    # and put them into one directory.
332    for filename in os.listdir(path):
333      if not filename.endswith(".img"):
334        continue
335      symlink_path = os.path.join(images_dir, filename)
336      # The files in latter directory overwrite the existing links
337      common.RunAndCheckOutput(
338        ['ln', '-sf', os.path.join(path, filename), symlink_path])
339
340  cmd = [avbtool, "calculate_vbmeta_digest", "--image",
341         os.path.join(images_dir, 'vbmeta.img')]
342  return common.RunAndCheckOutput(cmd)
343
344
345def main(argv):
346  if len(argv) != 2:
347    print(__doc__)
348    sys.exit(1)
349
350  common.InitLogging()
351
352  dict_file = argv[0]
353  out_file = argv[1]
354
355  prop_dict = {}
356  with open(dict_file, 'r') as f:
357    for line in f:
358      line = line.strip()
359      if not line or line.startswith("#"):
360        continue
361      k, v = line.split("=", 1)
362      prop_dict[k] = v
363
364  builder = CreateVerityImageBuilder(prop_dict)
365
366  if "partition_size" not in prop_dict:
367    image_size = GetDiskUsage(out_file)
368    # make sure that the image is big enough to hold vbmeta and footer
369    image_size = image_size + (MAX_VBMETA_SIZE + MAX_FOOTER_SIZE)
370    size = builder.CalculateDynamicPartitionSize(image_size)
371    prop_dict["partition_size"] = size
372
373  builder.Build(out_file)
374
375
376if __name__ == '__main__':
377  try:
378    main(sys.argv[1:])
379  finally:
380    common.Cleanup()
381