xref: /OK3568_Linux_fs/tools/linux/Linux_SecurityAVB/scripts/avbtool (revision 4882a59341e53eb6f0b4789bf948001014eff981)
1*4882a593Smuzhiyun#!/usr/bin/env python
2*4882a593Smuzhiyun
3*4882a593Smuzhiyun# Copyright 2016, The Android Open Source Project
4*4882a593Smuzhiyun#
5*4882a593Smuzhiyun# Permission is hereby granted, free of charge, to any person
6*4882a593Smuzhiyun# obtaining a copy of this software and associated documentation
7*4882a593Smuzhiyun# files (the "Software"), to deal in the Software without
8*4882a593Smuzhiyun# restriction, including without limitation the rights to use, copy,
9*4882a593Smuzhiyun# modify, merge, publish, distribute, sublicense, and/or sell copies
10*4882a593Smuzhiyun# of the Software, and to permit persons to whom the Software is
11*4882a593Smuzhiyun# furnished to do so, subject to the following conditions:
12*4882a593Smuzhiyun#
13*4882a593Smuzhiyun# The above copyright notice and this permission notice shall be
14*4882a593Smuzhiyun# included in all copies or substantial portions of the Software.
15*4882a593Smuzhiyun#
16*4882a593Smuzhiyun# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
17*4882a593Smuzhiyun# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
18*4882a593Smuzhiyun# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
19*4882a593Smuzhiyun# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
20*4882a593Smuzhiyun# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
21*4882a593Smuzhiyun# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
22*4882a593Smuzhiyun# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23*4882a593Smuzhiyun# SOFTWARE.
24*4882a593Smuzhiyun#
25*4882a593Smuzhiyun"""Command-line tool for working with Android Verified Boot images."""
26*4882a593Smuzhiyun
27*4882a593Smuzhiyunimport argparse
28*4882a593Smuzhiyunimport binascii
29*4882a593Smuzhiyunimport bisect
30*4882a593Smuzhiyunimport hashlib
31*4882a593Smuzhiyunimport math
32*4882a593Smuzhiyunimport os
33*4882a593Smuzhiyunimport struct
34*4882a593Smuzhiyunimport subprocess
35*4882a593Smuzhiyunimport sys
36*4882a593Smuzhiyunimport tempfile
37*4882a593Smuzhiyunimport time
38*4882a593Smuzhiyun
39*4882a593Smuzhiyun# Keep in sync with libavb/avb_version.h.
40*4882a593SmuzhiyunAVB_VERSION_MAJOR = 1
41*4882a593SmuzhiyunAVB_VERSION_MINOR = 1
42*4882a593SmuzhiyunAVB_VERSION_SUB = 0
43*4882a593Smuzhiyun
44*4882a593Smuzhiyun# Keep in sync with libavb/avb_footer.h.
45*4882a593SmuzhiyunAVB_FOOTER_VERSION_MAJOR = 1
46*4882a593SmuzhiyunAVB_FOOTER_VERSION_MINOR = 0
47*4882a593Smuzhiyun
48*4882a593SmuzhiyunAVB_VBMETA_IMAGE_FLAGS_HASHTREE_DISABLED = 1
49*4882a593Smuzhiyun
50*4882a593Smuzhiyun
51*4882a593Smuzhiyunclass AvbError(Exception):
52*4882a593Smuzhiyun  """Application-specific errors.
53*4882a593Smuzhiyun
54*4882a593Smuzhiyun  These errors represent issues for which a stack-trace should not be
55*4882a593Smuzhiyun  presented.
56*4882a593Smuzhiyun
57*4882a593Smuzhiyun  Attributes:
58*4882a593Smuzhiyun    message: Error message.
59*4882a593Smuzhiyun  """
60*4882a593Smuzhiyun
61*4882a593Smuzhiyun  def __init__(self, message):
62*4882a593Smuzhiyun    Exception.__init__(self, message)
63*4882a593Smuzhiyun
64*4882a593Smuzhiyun
65*4882a593Smuzhiyunclass Algorithm(object):
66*4882a593Smuzhiyun  """Contains details about an algorithm.
67*4882a593Smuzhiyun
68*4882a593Smuzhiyun  See the avb_vbmeta_header.h file for more details about
69*4882a593Smuzhiyun  algorithms.
70*4882a593Smuzhiyun
71*4882a593Smuzhiyun  The constant |ALGORITHMS| is a dictionary from human-readable
72*4882a593Smuzhiyun  names (e.g 'SHA256_RSA2048') to instances of this class.
73*4882a593Smuzhiyun
74*4882a593Smuzhiyun  Attributes:
75*4882a593Smuzhiyun    algorithm_type: Integer code corresponding to |AvbAlgorithmType|.
76*4882a593Smuzhiyun    hash_name: Empty or a name from |hashlib.algorithms|.
77*4882a593Smuzhiyun    hash_num_bytes: Number of bytes used to store the hash.
78*4882a593Smuzhiyun    signature_num_bytes: Number of bytes used to store the signature.
79*4882a593Smuzhiyun    public_key_num_bytes: Number of bytes used to store the public key.
80*4882a593Smuzhiyun    padding: Padding used for signature, if any.
81*4882a593Smuzhiyun  """
82*4882a593Smuzhiyun
83*4882a593Smuzhiyun  def __init__(self, algorithm_type, hash_name, hash_num_bytes,
84*4882a593Smuzhiyun               signature_num_bytes, public_key_num_bytes, padding):
85*4882a593Smuzhiyun    self.algorithm_type = algorithm_type
86*4882a593Smuzhiyun    self.hash_name = hash_name
87*4882a593Smuzhiyun    self.hash_num_bytes = hash_num_bytes
88*4882a593Smuzhiyun    self.signature_num_bytes = signature_num_bytes
89*4882a593Smuzhiyun    self.public_key_num_bytes = public_key_num_bytes
90*4882a593Smuzhiyun    self.padding = padding
91*4882a593Smuzhiyun
92*4882a593Smuzhiyun
93*4882a593Smuzhiyun# This must be kept in sync with the avb_crypto.h file.
94*4882a593Smuzhiyun#
95*4882a593Smuzhiyun# The PKC1-v1.5 padding is a blob of binary DER of ASN.1 and is
96*4882a593Smuzhiyun# obtained from section 5.2.2 of RFC 4880.
97*4882a593SmuzhiyunALGORITHMS = {
98*4882a593Smuzhiyun    'NONE': Algorithm(
99*4882a593Smuzhiyun        algorithm_type=0,        # AVB_ALGORITHM_TYPE_NONE
100*4882a593Smuzhiyun        hash_name='',
101*4882a593Smuzhiyun        hash_num_bytes=0,
102*4882a593Smuzhiyun        signature_num_bytes=0,
103*4882a593Smuzhiyun        public_key_num_bytes=0,
104*4882a593Smuzhiyun        padding=[]),
105*4882a593Smuzhiyun    'SHA256_RSA2048': Algorithm(
106*4882a593Smuzhiyun        algorithm_type=1,        # AVB_ALGORITHM_TYPE_SHA256_RSA2048
107*4882a593Smuzhiyun        hash_name='sha256',
108*4882a593Smuzhiyun        hash_num_bytes=32,
109*4882a593Smuzhiyun        signature_num_bytes=256,
110*4882a593Smuzhiyun        public_key_num_bytes=8 + 2*2048/8,
111*4882a593Smuzhiyun        padding=[
112*4882a593Smuzhiyun            # PKCS1-v1_5 padding
113*4882a593Smuzhiyun            0x00, 0x01] + [0xff]*202 + [0x00] + [
114*4882a593Smuzhiyun                # ASN.1 header
115*4882a593Smuzhiyun                0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
116*4882a593Smuzhiyun                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05,
117*4882a593Smuzhiyun                0x00, 0x04, 0x20,
118*4882a593Smuzhiyun            ]),
119*4882a593Smuzhiyun    'SHA256_RSA4096': Algorithm(
120*4882a593Smuzhiyun        algorithm_type=2,        # AVB_ALGORITHM_TYPE_SHA256_RSA4096
121*4882a593Smuzhiyun        hash_name='sha256',
122*4882a593Smuzhiyun        hash_num_bytes=32,
123*4882a593Smuzhiyun        signature_num_bytes=512,
124*4882a593Smuzhiyun        public_key_num_bytes=8 + 2*4096/8,
125*4882a593Smuzhiyun        padding=[
126*4882a593Smuzhiyun            # PKCS1-v1_5 padding
127*4882a593Smuzhiyun            0x00, 0x01] + [0xff]*458 + [0x00] + [
128*4882a593Smuzhiyun                # ASN.1 header
129*4882a593Smuzhiyun                0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
130*4882a593Smuzhiyun                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05,
131*4882a593Smuzhiyun                0x00, 0x04, 0x20,
132*4882a593Smuzhiyun            ]),
133*4882a593Smuzhiyun    'SHA256_RSA8192': Algorithm(
134*4882a593Smuzhiyun        algorithm_type=3,        # AVB_ALGORITHM_TYPE_SHA256_RSA8192
135*4882a593Smuzhiyun        hash_name='sha256',
136*4882a593Smuzhiyun        hash_num_bytes=32,
137*4882a593Smuzhiyun        signature_num_bytes=1024,
138*4882a593Smuzhiyun        public_key_num_bytes=8 + 2*8192/8,
139*4882a593Smuzhiyun        padding=[
140*4882a593Smuzhiyun            # PKCS1-v1_5 padding
141*4882a593Smuzhiyun            0x00, 0x01] + [0xff]*970 + [0x00] + [
142*4882a593Smuzhiyun                # ASN.1 header
143*4882a593Smuzhiyun                0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
144*4882a593Smuzhiyun                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05,
145*4882a593Smuzhiyun                0x00, 0x04, 0x20,
146*4882a593Smuzhiyun            ]),
147*4882a593Smuzhiyun    'SHA512_RSA2048': Algorithm(
148*4882a593Smuzhiyun        algorithm_type=4,        # AVB_ALGORITHM_TYPE_SHA512_RSA2048
149*4882a593Smuzhiyun        hash_name='sha512',
150*4882a593Smuzhiyun        hash_num_bytes=64,
151*4882a593Smuzhiyun        signature_num_bytes=256,
152*4882a593Smuzhiyun        public_key_num_bytes=8 + 2*2048/8,
153*4882a593Smuzhiyun        padding=[
154*4882a593Smuzhiyun            # PKCS1-v1_5 padding
155*4882a593Smuzhiyun            0x00, 0x01] + [0xff]*170 + [0x00] + [
156*4882a593Smuzhiyun                # ASN.1 header
157*4882a593Smuzhiyun                0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
158*4882a593Smuzhiyun                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05,
159*4882a593Smuzhiyun                0x00, 0x04, 0x40
160*4882a593Smuzhiyun            ]),
161*4882a593Smuzhiyun    'SHA512_RSA4096': Algorithm(
162*4882a593Smuzhiyun        algorithm_type=5,        # AVB_ALGORITHM_TYPE_SHA512_RSA4096
163*4882a593Smuzhiyun        hash_name='sha512',
164*4882a593Smuzhiyun        hash_num_bytes=64,
165*4882a593Smuzhiyun        signature_num_bytes=512,
166*4882a593Smuzhiyun        public_key_num_bytes=8 + 2*4096/8,
167*4882a593Smuzhiyun        padding=[
168*4882a593Smuzhiyun            # PKCS1-v1_5 padding
169*4882a593Smuzhiyun            0x00, 0x01] + [0xff]*426 + [0x00] + [
170*4882a593Smuzhiyun                # ASN.1 header
171*4882a593Smuzhiyun                0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
172*4882a593Smuzhiyun                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05,
173*4882a593Smuzhiyun                0x00, 0x04, 0x40
174*4882a593Smuzhiyun            ]),
175*4882a593Smuzhiyun    'SHA512_RSA8192': Algorithm(
176*4882a593Smuzhiyun        algorithm_type=6,        # AVB_ALGORITHM_TYPE_SHA512_RSA8192
177*4882a593Smuzhiyun        hash_name='sha512',
178*4882a593Smuzhiyun        hash_num_bytes=64,
179*4882a593Smuzhiyun        signature_num_bytes=1024,
180*4882a593Smuzhiyun        public_key_num_bytes=8 + 2*8192/8,
181*4882a593Smuzhiyun        padding=[
182*4882a593Smuzhiyun            # PKCS1-v1_5 padding
183*4882a593Smuzhiyun            0x00, 0x01] + [0xff]*938 + [0x00] + [
184*4882a593Smuzhiyun                # ASN.1 header
185*4882a593Smuzhiyun                0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
186*4882a593Smuzhiyun                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05,
187*4882a593Smuzhiyun                0x00, 0x04, 0x40
188*4882a593Smuzhiyun            ]),
189*4882a593Smuzhiyun}
190*4882a593Smuzhiyun
191*4882a593Smuzhiyun
192*4882a593Smuzhiyundef get_release_string():
193*4882a593Smuzhiyun  """Calculates the release string to use in the VBMeta struct."""
194*4882a593Smuzhiyun  # Keep in sync with libavb/avb_version.c:avb_version_string().
195*4882a593Smuzhiyun  return 'avbtool {}.{}.{}'.format(AVB_VERSION_MAJOR,
196*4882a593Smuzhiyun                                   AVB_VERSION_MINOR,
197*4882a593Smuzhiyun                                   AVB_VERSION_SUB)
198*4882a593Smuzhiyun
199*4882a593Smuzhiyun
200*4882a593Smuzhiyundef round_to_multiple(number, size):
201*4882a593Smuzhiyun  """Rounds a number up to nearest multiple of another number.
202*4882a593Smuzhiyun
203*4882a593Smuzhiyun  Args:
204*4882a593Smuzhiyun    number: The number to round up.
205*4882a593Smuzhiyun    size: The multiple to round up to.
206*4882a593Smuzhiyun
207*4882a593Smuzhiyun  Returns:
208*4882a593Smuzhiyun    If |number| is a multiple of |size|, returns |number|, otherwise
209*4882a593Smuzhiyun    returns |number| + |size|.
210*4882a593Smuzhiyun  """
211*4882a593Smuzhiyun  remainder = number % size
212*4882a593Smuzhiyun  if remainder == 0:
213*4882a593Smuzhiyun    return number
214*4882a593Smuzhiyun  return number + size - remainder
215*4882a593Smuzhiyun
216*4882a593Smuzhiyun
217*4882a593Smuzhiyundef round_to_pow2(number):
218*4882a593Smuzhiyun  """Rounds a number up to the next power of 2.
219*4882a593Smuzhiyun
220*4882a593Smuzhiyun  Args:
221*4882a593Smuzhiyun    number: The number to round up.
222*4882a593Smuzhiyun
223*4882a593Smuzhiyun  Returns:
224*4882a593Smuzhiyun    If |number| is already a power of 2 then |number| is
225*4882a593Smuzhiyun    returned. Otherwise the smallest power of 2 greater than |number|
226*4882a593Smuzhiyun    is returned.
227*4882a593Smuzhiyun  """
228*4882a593Smuzhiyun  return 2**((number - 1).bit_length())
229*4882a593Smuzhiyun
230*4882a593Smuzhiyun
231*4882a593Smuzhiyundef encode_long(num_bits, value):
232*4882a593Smuzhiyun  """Encodes a long to a bytearray() using a given amount of bits.
233*4882a593Smuzhiyun
234*4882a593Smuzhiyun  This number is written big-endian, e.g. with the most significant
235*4882a593Smuzhiyun  bit first.
236*4882a593Smuzhiyun
237*4882a593Smuzhiyun  This is the reverse of decode_long().
238*4882a593Smuzhiyun
239*4882a593Smuzhiyun  Arguments:
240*4882a593Smuzhiyun    num_bits: The number of bits to write, e.g. 2048.
241*4882a593Smuzhiyun    value: The value to write.
242*4882a593Smuzhiyun
243*4882a593Smuzhiyun  Returns:
244*4882a593Smuzhiyun    A bytearray() with the encoded long.
245*4882a593Smuzhiyun  """
246*4882a593Smuzhiyun  ret = bytearray()
247*4882a593Smuzhiyun  for bit_pos in range(num_bits, 0, -8):
248*4882a593Smuzhiyun    octet = (value >> (bit_pos - 8)) & 0xff
249*4882a593Smuzhiyun    ret.extend(struct.pack('!B', octet))
250*4882a593Smuzhiyun  return ret
251*4882a593Smuzhiyun
252*4882a593Smuzhiyun
253*4882a593Smuzhiyundef decode_long(blob):
254*4882a593Smuzhiyun  """Decodes a long from a bytearray() using a given amount of bits.
255*4882a593Smuzhiyun
256*4882a593Smuzhiyun  This number is expected to be in big-endian, e.g. with the most
257*4882a593Smuzhiyun  significant bit first.
258*4882a593Smuzhiyun
259*4882a593Smuzhiyun  This is the reverse of encode_long().
260*4882a593Smuzhiyun
261*4882a593Smuzhiyun  Arguments:
262*4882a593Smuzhiyun    value: A bytearray() with the encoded long.
263*4882a593Smuzhiyun
264*4882a593Smuzhiyun  Returns:
265*4882a593Smuzhiyun    The decoded value.
266*4882a593Smuzhiyun  """
267*4882a593Smuzhiyun  ret = 0
268*4882a593Smuzhiyun  for b in bytearray(blob):
269*4882a593Smuzhiyun    ret *= 256
270*4882a593Smuzhiyun    ret += b
271*4882a593Smuzhiyun  return ret
272*4882a593Smuzhiyun
273*4882a593Smuzhiyun
274*4882a593Smuzhiyundef egcd(a, b):
275*4882a593Smuzhiyun  """Calculate greatest common divisor of two numbers.
276*4882a593Smuzhiyun
277*4882a593Smuzhiyun  This implementation uses a recursive version of the extended
278*4882a593Smuzhiyun  Euclidian algorithm.
279*4882a593Smuzhiyun
280*4882a593Smuzhiyun  Arguments:
281*4882a593Smuzhiyun    a: First number.
282*4882a593Smuzhiyun    b: Second number.
283*4882a593Smuzhiyun
284*4882a593Smuzhiyun  Returns:
285*4882a593Smuzhiyun    A tuple (gcd, x, y) that where |gcd| is the greatest common
286*4882a593Smuzhiyun    divisor of |a| and |b| and |a|*|x| + |b|*|y| = |gcd|.
287*4882a593Smuzhiyun  """
288*4882a593Smuzhiyun  if a == 0:
289*4882a593Smuzhiyun    return (b, 0, 1)
290*4882a593Smuzhiyun  else:
291*4882a593Smuzhiyun    g, y, x = egcd(b % a, a)
292*4882a593Smuzhiyun    return (g, x - (b // a) * y, y)
293*4882a593Smuzhiyun
294*4882a593Smuzhiyun
295*4882a593Smuzhiyundef modinv(a, m):
296*4882a593Smuzhiyun  """Calculate modular multiplicative inverse of |a| modulo |m|.
297*4882a593Smuzhiyun
298*4882a593Smuzhiyun  This calculates the number |x| such that |a| * |x| == 1 (modulo
299*4882a593Smuzhiyun  |m|). This number only exists if |a| and |m| are co-prime - |None|
300*4882a593Smuzhiyun  is returned if this isn't true.
301*4882a593Smuzhiyun
302*4882a593Smuzhiyun  Arguments:
303*4882a593Smuzhiyun    a: The number to calculate a modular inverse of.
304*4882a593Smuzhiyun    m: The modulo to use.
305*4882a593Smuzhiyun
306*4882a593Smuzhiyun  Returns:
307*4882a593Smuzhiyun    The modular multiplicative inverse of |a| and |m| or |None| if
308*4882a593Smuzhiyun    these numbers are not co-prime.
309*4882a593Smuzhiyun  """
310*4882a593Smuzhiyun  gcd, x, _ = egcd(a, m)
311*4882a593Smuzhiyun  if gcd != 1:
312*4882a593Smuzhiyun    return None  # modular inverse does not exist
313*4882a593Smuzhiyun  else:
314*4882a593Smuzhiyun    return x % m
315*4882a593Smuzhiyun
316*4882a593Smuzhiyun
317*4882a593Smuzhiyundef parse_number(string):
318*4882a593Smuzhiyun  """Parse a string as a number.
319*4882a593Smuzhiyun
320*4882a593Smuzhiyun  This is just a short-hand for int(string, 0) suitable for use in the
321*4882a593Smuzhiyun  |type| parameter of |ArgumentParser|'s add_argument() function. An
322*4882a593Smuzhiyun  improvement to just using type=int is that this function supports
323*4882a593Smuzhiyun  numbers in other bases, e.g. "0x1234".
324*4882a593Smuzhiyun
325*4882a593Smuzhiyun  Arguments:
326*4882a593Smuzhiyun    string: The string to parse.
327*4882a593Smuzhiyun
328*4882a593Smuzhiyun  Returns:
329*4882a593Smuzhiyun    The parsed integer.
330*4882a593Smuzhiyun
331*4882a593Smuzhiyun  Raises:
332*4882a593Smuzhiyun    ValueError: If the number could not be parsed.
333*4882a593Smuzhiyun  """
334*4882a593Smuzhiyun  return int(string, 0)
335*4882a593Smuzhiyun
336*4882a593Smuzhiyun
337*4882a593Smuzhiyunclass RSAPublicKey(object):
338*4882a593Smuzhiyun  """Data structure used for a RSA public key.
339*4882a593Smuzhiyun
340*4882a593Smuzhiyun  Attributes:
341*4882a593Smuzhiyun    exponent: The key exponent.
342*4882a593Smuzhiyun    modulus: The key modulus.
343*4882a593Smuzhiyun    num_bits: The key size.
344*4882a593Smuzhiyun  """
345*4882a593Smuzhiyun
346*4882a593Smuzhiyun  MODULUS_PREFIX = 'modulus='
347*4882a593Smuzhiyun
348*4882a593Smuzhiyun  def __init__(self, key_path):
349*4882a593Smuzhiyun    """Loads and parses an RSA key from either a private or public key file.
350*4882a593Smuzhiyun
351*4882a593Smuzhiyun    Arguments:
352*4882a593Smuzhiyun      key_path: The path to a key file.
353*4882a593Smuzhiyun    """
354*4882a593Smuzhiyun    # We used to have something as simple as this:
355*4882a593Smuzhiyun    #
356*4882a593Smuzhiyun    #  key = Crypto.PublicKey.RSA.importKey(open(key_path).read())
357*4882a593Smuzhiyun    #  self.exponent = key.e
358*4882a593Smuzhiyun    #  self.modulus = key.n
359*4882a593Smuzhiyun    #  self.num_bits = key.size() + 1
360*4882a593Smuzhiyun    #
361*4882a593Smuzhiyun    # but unfortunately PyCrypto is not available in the builder. So
362*4882a593Smuzhiyun    # instead just parse openssl(1) output to get this
363*4882a593Smuzhiyun    # information. It's ugly but...
364*4882a593Smuzhiyun    args = ['openssl', 'rsa', '-in', key_path, '-modulus', '-noout']
365*4882a593Smuzhiyun    p = subprocess.Popen(args,
366*4882a593Smuzhiyun                         stdin=subprocess.PIPE,
367*4882a593Smuzhiyun                         stdout=subprocess.PIPE,
368*4882a593Smuzhiyun                         stderr=subprocess.PIPE)
369*4882a593Smuzhiyun    (pout, perr) = p.communicate()
370*4882a593Smuzhiyun    if p.wait() != 0:
371*4882a593Smuzhiyun      # Could be just a public key is passed, try that.
372*4882a593Smuzhiyun      args.append('-pubin')
373*4882a593Smuzhiyun      p = subprocess.Popen(args,
374*4882a593Smuzhiyun                           stdin=subprocess.PIPE,
375*4882a593Smuzhiyun                           stdout=subprocess.PIPE,
376*4882a593Smuzhiyun                           stderr=subprocess.PIPE)
377*4882a593Smuzhiyun      (pout, perr) = p.communicate()
378*4882a593Smuzhiyun      if p.wait() != 0:
379*4882a593Smuzhiyun        raise AvbError('Error getting public key: {}'.format(perr))
380*4882a593Smuzhiyun
381*4882a593Smuzhiyun    if not pout.lower().startswith(self.MODULUS_PREFIX):
382*4882a593Smuzhiyun      raise AvbError('Unexpected modulus output')
383*4882a593Smuzhiyun
384*4882a593Smuzhiyun    modulus_hexstr = pout[len(self.MODULUS_PREFIX):]
385*4882a593Smuzhiyun
386*4882a593Smuzhiyun    # The exponent is assumed to always be 65537 and the number of
387*4882a593Smuzhiyun    # bits can be derived from the modulus by rounding up to the
388*4882a593Smuzhiyun    # nearest power of 2.
389*4882a593Smuzhiyun    self.modulus = int(modulus_hexstr, 16)
390*4882a593Smuzhiyun    self.num_bits = round_to_pow2(int(math.ceil(math.log(self.modulus, 2))))
391*4882a593Smuzhiyun    self.exponent = 65537
392*4882a593Smuzhiyun
393*4882a593Smuzhiyun
394*4882a593Smuzhiyundef encode_rsa_key(key_path):
395*4882a593Smuzhiyun  """Encodes a public RSA key in |AvbRSAPublicKeyHeader| format.
396*4882a593Smuzhiyun
397*4882a593Smuzhiyun  This creates a |AvbRSAPublicKeyHeader| as well as the two large
398*4882a593Smuzhiyun  numbers (|key_num_bits| bits long) following it.
399*4882a593Smuzhiyun
400*4882a593Smuzhiyun  Arguments:
401*4882a593Smuzhiyun    key_path: The path to a key file.
402*4882a593Smuzhiyun
403*4882a593Smuzhiyun  Returns:
404*4882a593Smuzhiyun    A bytearray() with the |AvbRSAPublicKeyHeader|.
405*4882a593Smuzhiyun  """
406*4882a593Smuzhiyun  key = RSAPublicKey(key_path)
407*4882a593Smuzhiyun  if key.exponent != 65537:
408*4882a593Smuzhiyun    raise AvbError('Only RSA keys with exponent 65537 are supported.')
409*4882a593Smuzhiyun  ret = bytearray()
410*4882a593Smuzhiyun  # Calculate n0inv = -1/n[0] (mod 2^32)
411*4882a593Smuzhiyun  b = 2L**32
412*4882a593Smuzhiyun  n0inv = b - modinv(key.modulus, b)
413*4882a593Smuzhiyun  # Calculate rr = r^2 (mod N), where r = 2^(# of key bits)
414*4882a593Smuzhiyun  r = 2L**key.modulus.bit_length()
415*4882a593Smuzhiyun  rrmodn = r * r % key.modulus
416*4882a593Smuzhiyun  ret.extend(struct.pack('!II', key.num_bits, n0inv))
417*4882a593Smuzhiyun  ret.extend(encode_long(key.num_bits, key.modulus))
418*4882a593Smuzhiyun  ret.extend(encode_long(key.num_bits, rrmodn))
419*4882a593Smuzhiyun  return ret
420*4882a593Smuzhiyun
421*4882a593Smuzhiyun
422*4882a593Smuzhiyundef lookup_algorithm_by_type(alg_type):
423*4882a593Smuzhiyun  """Looks up algorithm by type.
424*4882a593Smuzhiyun
425*4882a593Smuzhiyun  Arguments:
426*4882a593Smuzhiyun    alg_type: The integer representing the type.
427*4882a593Smuzhiyun
428*4882a593Smuzhiyun  Returns:
429*4882a593Smuzhiyun    A tuple with the algorithm name and an |Algorithm| instance.
430*4882a593Smuzhiyun
431*4882a593Smuzhiyun  Raises:
432*4882a593Smuzhiyun    Exception: If the algorithm cannot be found
433*4882a593Smuzhiyun  """
434*4882a593Smuzhiyun  for alg_name in ALGORITHMS:
435*4882a593Smuzhiyun    alg_data = ALGORITHMS[alg_name]
436*4882a593Smuzhiyun    if alg_data.algorithm_type == alg_type:
437*4882a593Smuzhiyun      return (alg_name, alg_data)
438*4882a593Smuzhiyun  raise AvbError('Unknown algorithm type {}'.format(alg_type))
439*4882a593Smuzhiyun
440*4882a593Smuzhiyun
441*4882a593Smuzhiyundef raw_sign(signing_helper, signing_helper_with_files,
442*4882a593Smuzhiyun             algorithm_name, signature_num_bytes, key_path,
443*4882a593Smuzhiyun             raw_data_to_sign):
444*4882a593Smuzhiyun  """Computes a raw RSA signature using |signing_helper| or openssl.
445*4882a593Smuzhiyun
446*4882a593Smuzhiyun  Arguments:
447*4882a593Smuzhiyun    signing_helper: Program which signs a hash and returns the signature.
448*4882a593Smuzhiyun    signing_helper_with_files: Same as signing_helper but uses files instead.
449*4882a593Smuzhiyun    algorithm_name: The algorithm name as per the ALGORITHMS dict.
450*4882a593Smuzhiyun    signature_num_bytes: Number of bytes used to store the signature.
451*4882a593Smuzhiyun    key_path: Path to the private key file. Must be PEM format.
452*4882a593Smuzhiyun    raw_data_to_sign: Data to sign (bytearray or str expected).
453*4882a593Smuzhiyun
454*4882a593Smuzhiyun  Returns:
455*4882a593Smuzhiyun    A bytearray containing the signature.
456*4882a593Smuzhiyun
457*4882a593Smuzhiyun  Raises:
458*4882a593Smuzhiyun    Exception: If an error occurs.
459*4882a593Smuzhiyun  """
460*4882a593Smuzhiyun  p = None
461*4882a593Smuzhiyun  if signing_helper_with_files is not None:
462*4882a593Smuzhiyun    signing_file = tempfile.NamedTemporaryFile()
463*4882a593Smuzhiyun    signing_file.write(str(raw_data_to_sign))
464*4882a593Smuzhiyun    signing_file.flush()
465*4882a593Smuzhiyun    p = subprocess.Popen(
466*4882a593Smuzhiyun      [signing_helper_with_files, algorithm_name, key_path, signing_file.name])
467*4882a593Smuzhiyun    retcode = p.wait()
468*4882a593Smuzhiyun    if retcode != 0:
469*4882a593Smuzhiyun      raise AvbError('Error signing')
470*4882a593Smuzhiyun    signing_file.seek(0)
471*4882a593Smuzhiyun    signature = bytearray(signing_file.read())
472*4882a593Smuzhiyun  else:
473*4882a593Smuzhiyun    if signing_helper is not None:
474*4882a593Smuzhiyun      p = subprocess.Popen(
475*4882a593Smuzhiyun          [signing_helper, algorithm_name, key_path],
476*4882a593Smuzhiyun          stdin=subprocess.PIPE,
477*4882a593Smuzhiyun          stdout=subprocess.PIPE,
478*4882a593Smuzhiyun          stderr=subprocess.PIPE)
479*4882a593Smuzhiyun    else:
480*4882a593Smuzhiyun      p = subprocess.Popen(
481*4882a593Smuzhiyun          ['openssl', 'rsautl', '-sign', '-inkey', key_path, '-raw'],
482*4882a593Smuzhiyun          stdin=subprocess.PIPE,
483*4882a593Smuzhiyun          stdout=subprocess.PIPE,
484*4882a593Smuzhiyun          stderr=subprocess.PIPE)
485*4882a593Smuzhiyun    (pout, perr) = p.communicate(str(raw_data_to_sign))
486*4882a593Smuzhiyun    retcode = p.wait()
487*4882a593Smuzhiyun    if retcode != 0:
488*4882a593Smuzhiyun      raise AvbError('Error signing: {}'.format(perr))
489*4882a593Smuzhiyun    signature = bytearray(pout)
490*4882a593Smuzhiyun  if len(signature) != signature_num_bytes:
491*4882a593Smuzhiyun    raise AvbError('Error signing: Invalid length of signature')
492*4882a593Smuzhiyun  return signature
493*4882a593Smuzhiyun
494*4882a593Smuzhiyun
495*4882a593Smuzhiyundef verify_vbmeta_signature(vbmeta_header, vbmeta_blob):
496*4882a593Smuzhiyun  """Checks that the signature in a vbmeta blob was made by
497*4882a593Smuzhiyun     the embedded public key.
498*4882a593Smuzhiyun
499*4882a593Smuzhiyun  Arguments:
500*4882a593Smuzhiyun    vbmeta_header: A AvbVBMetaHeader.
501*4882a593Smuzhiyun    vbmeta_blob: The whole vbmeta blob, including the header.
502*4882a593Smuzhiyun
503*4882a593Smuzhiyun  Returns:
504*4882a593Smuzhiyun    True if the signature is valid and corresponds to the embedded
505*4882a593Smuzhiyun    public key. Also returns True if the vbmeta blob is not signed.
506*4882a593Smuzhiyun  """
507*4882a593Smuzhiyun  (_, alg) = lookup_algorithm_by_type(vbmeta_header.algorithm_type)
508*4882a593Smuzhiyun  if alg.hash_name == '':
509*4882a593Smuzhiyun    return True
510*4882a593Smuzhiyun  header_blob = vbmeta_blob[0:256]
511*4882a593Smuzhiyun  auth_offset = 256
512*4882a593Smuzhiyun  aux_offset = auth_offset + vbmeta_header.authentication_data_block_size
513*4882a593Smuzhiyun  aux_size = vbmeta_header.auxiliary_data_block_size
514*4882a593Smuzhiyun  aux_blob = vbmeta_blob[aux_offset:aux_offset + aux_size]
515*4882a593Smuzhiyun  pubkey_offset = aux_offset + vbmeta_header.public_key_offset
516*4882a593Smuzhiyun  pubkey_size = vbmeta_header.public_key_size
517*4882a593Smuzhiyun  pubkey_blob = vbmeta_blob[pubkey_offset:pubkey_offset + pubkey_size]
518*4882a593Smuzhiyun
519*4882a593Smuzhiyun  digest_offset = auth_offset + vbmeta_header.hash_offset
520*4882a593Smuzhiyun  digest_size = vbmeta_header.hash_size
521*4882a593Smuzhiyun  digest_blob = vbmeta_blob[digest_offset:digest_offset + digest_size]
522*4882a593Smuzhiyun
523*4882a593Smuzhiyun  sig_offset = auth_offset + vbmeta_header.signature_offset
524*4882a593Smuzhiyun  sig_size = vbmeta_header.signature_size
525*4882a593Smuzhiyun  sig_blob = vbmeta_blob[sig_offset:sig_offset + sig_size]
526*4882a593Smuzhiyun
527*4882a593Smuzhiyun  # Now that we've got the stored digest, public key, and signature
528*4882a593Smuzhiyun  # all we need to do is to verify. This is the exactly the same
529*4882a593Smuzhiyun  # steps as performed in the avb_vbmeta_image_verify() function in
530*4882a593Smuzhiyun  # libavb/avb_vbmeta_image.c.
531*4882a593Smuzhiyun
532*4882a593Smuzhiyun  ha = hashlib.new(alg.hash_name)
533*4882a593Smuzhiyun  ha.update(header_blob)
534*4882a593Smuzhiyun  ha.update(aux_blob)
535*4882a593Smuzhiyun  computed_digest = ha.digest()
536*4882a593Smuzhiyun
537*4882a593Smuzhiyun  if computed_digest != digest_blob:
538*4882a593Smuzhiyun    return False
539*4882a593Smuzhiyun
540*4882a593Smuzhiyun  padding_and_digest = bytearray(alg.padding)
541*4882a593Smuzhiyun  padding_and_digest.extend(computed_digest)
542*4882a593Smuzhiyun
543*4882a593Smuzhiyun  (num_bits,) = struct.unpack('!I', pubkey_blob[0:4])
544*4882a593Smuzhiyun  modulus_blob = pubkey_blob[8:8 + num_bits/8]
545*4882a593Smuzhiyun  modulus = decode_long(modulus_blob)
546*4882a593Smuzhiyun  exponent = 65537
547*4882a593Smuzhiyun
548*4882a593Smuzhiyun  # For now, just use Crypto.PublicKey.RSA to verify the signature. This
549*4882a593Smuzhiyun  # is OK since 'avbtool verify_image' is not expected to run on the
550*4882a593Smuzhiyun  # Android builders (see bug #36809096).
551*4882a593Smuzhiyun  import Crypto.PublicKey.RSA
552*4882a593Smuzhiyun  key = Crypto.PublicKey.RSA.construct((modulus, long(exponent)))
553*4882a593Smuzhiyun  if not key.verify(decode_long(padding_and_digest),
554*4882a593Smuzhiyun                    (decode_long(sig_blob), None)):
555*4882a593Smuzhiyun    return False
556*4882a593Smuzhiyun  return True
557*4882a593Smuzhiyun
558*4882a593Smuzhiyun
559*4882a593Smuzhiyunclass ImageChunk(object):
560*4882a593Smuzhiyun  """Data structure used for representing chunks in Android sparse files.
561*4882a593Smuzhiyun
562*4882a593Smuzhiyun  Attributes:
563*4882a593Smuzhiyun    chunk_type: One of TYPE_RAW, TYPE_FILL, or TYPE_DONT_CARE.
564*4882a593Smuzhiyun    chunk_offset: Offset in the sparse file where this chunk begins.
565*4882a593Smuzhiyun    output_offset: Offset in de-sparsified file where output begins.
566*4882a593Smuzhiyun    output_size: Number of bytes in output.
567*4882a593Smuzhiyun    input_offset: Offset in sparse file for data if TYPE_RAW otherwise None.
568*4882a593Smuzhiyun    fill_data: Blob with data to fill if TYPE_FILL otherwise None.
569*4882a593Smuzhiyun  """
570*4882a593Smuzhiyun
571*4882a593Smuzhiyun  FORMAT = '<2H2I'
572*4882a593Smuzhiyun  TYPE_RAW = 0xcac1
573*4882a593Smuzhiyun  TYPE_FILL = 0xcac2
574*4882a593Smuzhiyun  TYPE_DONT_CARE = 0xcac3
575*4882a593Smuzhiyun  TYPE_CRC32 = 0xcac4
576*4882a593Smuzhiyun
577*4882a593Smuzhiyun  def __init__(self, chunk_type, chunk_offset, output_offset, output_size,
578*4882a593Smuzhiyun               input_offset, fill_data):
579*4882a593Smuzhiyun    """Initializes an ImageChunk object.
580*4882a593Smuzhiyun
581*4882a593Smuzhiyun    Arguments:
582*4882a593Smuzhiyun      chunk_type: One of TYPE_RAW, TYPE_FILL, or TYPE_DONT_CARE.
583*4882a593Smuzhiyun      chunk_offset: Offset in the sparse file where this chunk begins.
584*4882a593Smuzhiyun      output_offset: Offset in de-sparsified file.
585*4882a593Smuzhiyun      output_size: Number of bytes in output.
586*4882a593Smuzhiyun      input_offset: Offset in sparse file if TYPE_RAW otherwise None.
587*4882a593Smuzhiyun      fill_data: Blob with data to fill if TYPE_FILL otherwise None.
588*4882a593Smuzhiyun
589*4882a593Smuzhiyun    Raises:
590*4882a593Smuzhiyun      ValueError: If data is not well-formed.
591*4882a593Smuzhiyun    """
592*4882a593Smuzhiyun    self.chunk_type = chunk_type
593*4882a593Smuzhiyun    self.chunk_offset = chunk_offset
594*4882a593Smuzhiyun    self.output_offset = output_offset
595*4882a593Smuzhiyun    self.output_size = output_size
596*4882a593Smuzhiyun    self.input_offset = input_offset
597*4882a593Smuzhiyun    self.fill_data = fill_data
598*4882a593Smuzhiyun    # Check invariants.
599*4882a593Smuzhiyun    if self.chunk_type == self.TYPE_RAW:
600*4882a593Smuzhiyun      if self.fill_data is not None:
601*4882a593Smuzhiyun        raise ValueError('RAW chunk cannot have fill_data set.')
602*4882a593Smuzhiyun      if not self.input_offset:
603*4882a593Smuzhiyun        raise ValueError('RAW chunk must have input_offset set.')
604*4882a593Smuzhiyun    elif self.chunk_type == self.TYPE_FILL:
605*4882a593Smuzhiyun      if self.fill_data is None:
606*4882a593Smuzhiyun        raise ValueError('FILL chunk must have fill_data set.')
607*4882a593Smuzhiyun      if self.input_offset:
608*4882a593Smuzhiyun        raise ValueError('FILL chunk cannot have input_offset set.')
609*4882a593Smuzhiyun    elif self.chunk_type == self.TYPE_DONT_CARE:
610*4882a593Smuzhiyun      if self.fill_data is not None:
611*4882a593Smuzhiyun        raise ValueError('DONT_CARE chunk cannot have fill_data set.')
612*4882a593Smuzhiyun      if self.input_offset:
613*4882a593Smuzhiyun        raise ValueError('DONT_CARE chunk cannot have input_offset set.')
614*4882a593Smuzhiyun    else:
615*4882a593Smuzhiyun      raise ValueError('Invalid chunk type')
616*4882a593Smuzhiyun
617*4882a593Smuzhiyun
618*4882a593Smuzhiyunclass ImageHandler(object):
619*4882a593Smuzhiyun  """Abstraction for image I/O with support for Android sparse images.
620*4882a593Smuzhiyun
621*4882a593Smuzhiyun  This class provides an interface for working with image files that
622*4882a593Smuzhiyun  may be using the Android Sparse Image format. When an instance is
623*4882a593Smuzhiyun  constructed, we test whether it's an Android sparse file. If so,
624*4882a593Smuzhiyun  operations will be on the sparse file by interpreting the sparse
625*4882a593Smuzhiyun  format, otherwise they will be directly on the file. Either way the
626*4882a593Smuzhiyun  operations do the same.
627*4882a593Smuzhiyun
628*4882a593Smuzhiyun  For reading, this interface mimics a file object - it has seek(),
629*4882a593Smuzhiyun  tell(), and read() methods. For writing, only truncation
630*4882a593Smuzhiyun  (truncate()) and appending is supported (append_raw() and
631*4882a593Smuzhiyun  append_dont_care()). Additionally, data can only be written in units
632*4882a593Smuzhiyun  of the block size.
633*4882a593Smuzhiyun
634*4882a593Smuzhiyun  Attributes:
635*4882a593Smuzhiyun    is_sparse: Whether the file being operated on is sparse.
636*4882a593Smuzhiyun    block_size: The block size, typically 4096.
637*4882a593Smuzhiyun    image_size: The size of the unsparsified file.
638*4882a593Smuzhiyun  """
639*4882a593Smuzhiyun  # See system/core/libsparse/sparse_format.h for details.
640*4882a593Smuzhiyun  MAGIC = 0xed26ff3a
641*4882a593Smuzhiyun  HEADER_FORMAT = '<I4H4I'
642*4882a593Smuzhiyun
643*4882a593Smuzhiyun  # These are formats and offset of just the |total_chunks| and
644*4882a593Smuzhiyun  # |total_blocks| fields.
645*4882a593Smuzhiyun  NUM_CHUNKS_AND_BLOCKS_FORMAT = '<II'
646*4882a593Smuzhiyun  NUM_CHUNKS_AND_BLOCKS_OFFSET = 16
647*4882a593Smuzhiyun
648*4882a593Smuzhiyun  def __init__(self, image_filename):
649*4882a593Smuzhiyun    """Initializes an image handler.
650*4882a593Smuzhiyun
651*4882a593Smuzhiyun    Arguments:
652*4882a593Smuzhiyun      image_filename: The name of the file to operate on.
653*4882a593Smuzhiyun
654*4882a593Smuzhiyun    Raises:
655*4882a593Smuzhiyun      ValueError: If data in the file is invalid.
656*4882a593Smuzhiyun    """
657*4882a593Smuzhiyun    self._image_filename = image_filename
658*4882a593Smuzhiyun    self._read_header()
659*4882a593Smuzhiyun
660*4882a593Smuzhiyun  def _read_header(self):
661*4882a593Smuzhiyun    """Initializes internal data structures used for reading file.
662*4882a593Smuzhiyun
663*4882a593Smuzhiyun    This may be called multiple times and is typically called after
664*4882a593Smuzhiyun    modifying the file (e.g. appending, truncation).
665*4882a593Smuzhiyun
666*4882a593Smuzhiyun    Raises:
667*4882a593Smuzhiyun      ValueError: If data in the file is invalid.
668*4882a593Smuzhiyun    """
669*4882a593Smuzhiyun    self.is_sparse = False
670*4882a593Smuzhiyun    self.block_size = 4096
671*4882a593Smuzhiyun    self._file_pos = 0
672*4882a593Smuzhiyun    self._image = open(self._image_filename, 'r+b')
673*4882a593Smuzhiyun    self._image.seek(0, os.SEEK_END)
674*4882a593Smuzhiyun    self.image_size = self._image.tell()
675*4882a593Smuzhiyun
676*4882a593Smuzhiyun    self._image.seek(0, os.SEEK_SET)
677*4882a593Smuzhiyun    header_bin = self._image.read(struct.calcsize(self.HEADER_FORMAT))
678*4882a593Smuzhiyun    (magic, major_version, minor_version, file_hdr_sz, chunk_hdr_sz,
679*4882a593Smuzhiyun     block_size, self._num_total_blocks, self._num_total_chunks,
680*4882a593Smuzhiyun     _) = struct.unpack(self.HEADER_FORMAT, header_bin)
681*4882a593Smuzhiyun    if magic != self.MAGIC:
682*4882a593Smuzhiyun      # Not a sparse image, our job here is done.
683*4882a593Smuzhiyun      return
684*4882a593Smuzhiyun    if not (major_version == 1 and minor_version == 0):
685*4882a593Smuzhiyun      raise ValueError('Encountered sparse image format version {}.{} but '
686*4882a593Smuzhiyun                       'only 1.0 is supported'.format(major_version,
687*4882a593Smuzhiyun                                                      minor_version))
688*4882a593Smuzhiyun    if file_hdr_sz != struct.calcsize(self.HEADER_FORMAT):
689*4882a593Smuzhiyun      raise ValueError('Unexpected file_hdr_sz value {}.'.
690*4882a593Smuzhiyun                       format(file_hdr_sz))
691*4882a593Smuzhiyun    if chunk_hdr_sz != struct.calcsize(ImageChunk.FORMAT):
692*4882a593Smuzhiyun      raise ValueError('Unexpected chunk_hdr_sz value {}.'.
693*4882a593Smuzhiyun                       format(chunk_hdr_sz))
694*4882a593Smuzhiyun
695*4882a593Smuzhiyun    self.block_size = block_size
696*4882a593Smuzhiyun
697*4882a593Smuzhiyun    # Build an list of chunks by parsing the file.
698*4882a593Smuzhiyun    self._chunks = []
699*4882a593Smuzhiyun
700*4882a593Smuzhiyun    # Find the smallest offset where only "Don't care" chunks
701*4882a593Smuzhiyun    # follow. This will be the size of the content in the sparse
702*4882a593Smuzhiyun    # image.
703*4882a593Smuzhiyun    offset = 0
704*4882a593Smuzhiyun    output_offset = 0
705*4882a593Smuzhiyun    for _ in xrange(1, self._num_total_chunks + 1):
706*4882a593Smuzhiyun      chunk_offset = self._image.tell()
707*4882a593Smuzhiyun
708*4882a593Smuzhiyun      header_bin = self._image.read(struct.calcsize(ImageChunk.FORMAT))
709*4882a593Smuzhiyun      (chunk_type, _, chunk_sz, total_sz) = struct.unpack(ImageChunk.FORMAT,
710*4882a593Smuzhiyun                                                          header_bin)
711*4882a593Smuzhiyun      data_sz = total_sz - struct.calcsize(ImageChunk.FORMAT)
712*4882a593Smuzhiyun
713*4882a593Smuzhiyun      if chunk_type == ImageChunk.TYPE_RAW:
714*4882a593Smuzhiyun        if data_sz != (chunk_sz * self.block_size):
715*4882a593Smuzhiyun          raise ValueError('Raw chunk input size ({}) does not match output '
716*4882a593Smuzhiyun                           'size ({})'.
717*4882a593Smuzhiyun                           format(data_sz, chunk_sz*self.block_size))
718*4882a593Smuzhiyun        self._chunks.append(ImageChunk(ImageChunk.TYPE_RAW,
719*4882a593Smuzhiyun                                       chunk_offset,
720*4882a593Smuzhiyun                                       output_offset,
721*4882a593Smuzhiyun                                       chunk_sz*self.block_size,
722*4882a593Smuzhiyun                                       self._image.tell(),
723*4882a593Smuzhiyun                                       None))
724*4882a593Smuzhiyun        self._image.read(data_sz)
725*4882a593Smuzhiyun
726*4882a593Smuzhiyun      elif chunk_type == ImageChunk.TYPE_FILL:
727*4882a593Smuzhiyun        if data_sz != 4:
728*4882a593Smuzhiyun          raise ValueError('Fill chunk should have 4 bytes of fill, but this '
729*4882a593Smuzhiyun                           'has {}'.format(data_sz))
730*4882a593Smuzhiyun        fill_data = self._image.read(4)
731*4882a593Smuzhiyun        self._chunks.append(ImageChunk(ImageChunk.TYPE_FILL,
732*4882a593Smuzhiyun                                       chunk_offset,
733*4882a593Smuzhiyun                                       output_offset,
734*4882a593Smuzhiyun                                       chunk_sz*self.block_size,
735*4882a593Smuzhiyun                                       None,
736*4882a593Smuzhiyun                                       fill_data))
737*4882a593Smuzhiyun      elif chunk_type == ImageChunk.TYPE_DONT_CARE:
738*4882a593Smuzhiyun        if data_sz != 0:
739*4882a593Smuzhiyun          raise ValueError('Don\'t care chunk input size is non-zero ({})'.
740*4882a593Smuzhiyun                           format(data_sz))
741*4882a593Smuzhiyun        self._chunks.append(ImageChunk(ImageChunk.TYPE_DONT_CARE,
742*4882a593Smuzhiyun                                       chunk_offset,
743*4882a593Smuzhiyun                                       output_offset,
744*4882a593Smuzhiyun                                       chunk_sz*self.block_size,
745*4882a593Smuzhiyun                                       None,
746*4882a593Smuzhiyun                                       None))
747*4882a593Smuzhiyun      elif chunk_type == ImageChunk.TYPE_CRC32:
748*4882a593Smuzhiyun        if data_sz != 4:
749*4882a593Smuzhiyun          raise ValueError('CRC32 chunk should have 4 bytes of CRC, but '
750*4882a593Smuzhiyun                           'this has {}'.format(data_sz))
751*4882a593Smuzhiyun        self._image.read(4)
752*4882a593Smuzhiyun      else:
753*4882a593Smuzhiyun        raise ValueError('Unknown chunk type {}'.format(chunk_type))
754*4882a593Smuzhiyun
755*4882a593Smuzhiyun      offset += chunk_sz
756*4882a593Smuzhiyun      output_offset += chunk_sz*self.block_size
757*4882a593Smuzhiyun
758*4882a593Smuzhiyun    # Record where sparse data end.
759*4882a593Smuzhiyun    self._sparse_end = self._image.tell()
760*4882a593Smuzhiyun
761*4882a593Smuzhiyun    # Now that we've traversed all chunks, sanity check.
762*4882a593Smuzhiyun    if self._num_total_blocks != offset:
763*4882a593Smuzhiyun      raise ValueError('The header said we should have {} output blocks, '
764*4882a593Smuzhiyun                       'but we saw {}'.format(self._num_total_blocks, offset))
765*4882a593Smuzhiyun    junk_len = len(self._image.read())
766*4882a593Smuzhiyun    if junk_len > 0:
767*4882a593Smuzhiyun      raise ValueError('There were {} bytes of extra data at the end of the '
768*4882a593Smuzhiyun                       'file.'.format(junk_len))
769*4882a593Smuzhiyun
770*4882a593Smuzhiyun    # Assign |image_size|.
771*4882a593Smuzhiyun    self.image_size = output_offset
772*4882a593Smuzhiyun
773*4882a593Smuzhiyun    # This is used when bisecting in read() to find the initial slice.
774*4882a593Smuzhiyun    self._chunk_output_offsets = [i.output_offset for i in self._chunks]
775*4882a593Smuzhiyun
776*4882a593Smuzhiyun    self.is_sparse = True
777*4882a593Smuzhiyun
778*4882a593Smuzhiyun  def _update_chunks_and_blocks(self):
779*4882a593Smuzhiyun    """Helper function to update the image header.
780*4882a593Smuzhiyun
781*4882a593Smuzhiyun    The the |total_chunks| and |total_blocks| fields in the header
782*4882a593Smuzhiyun    will be set to value of the |_num_total_blocks| and
783*4882a593Smuzhiyun    |_num_total_chunks| attributes.
784*4882a593Smuzhiyun
785*4882a593Smuzhiyun    """
786*4882a593Smuzhiyun    self._image.seek(self.NUM_CHUNKS_AND_BLOCKS_OFFSET, os.SEEK_SET)
787*4882a593Smuzhiyun    self._image.write(struct.pack(self.NUM_CHUNKS_AND_BLOCKS_FORMAT,
788*4882a593Smuzhiyun                                  self._num_total_blocks,
789*4882a593Smuzhiyun                                  self._num_total_chunks))
790*4882a593Smuzhiyun
791*4882a593Smuzhiyun  def append_dont_care(self, num_bytes):
792*4882a593Smuzhiyun    """Appends a DONT_CARE chunk to the sparse file.
793*4882a593Smuzhiyun
794*4882a593Smuzhiyun    The given number of bytes must be a multiple of the block size.
795*4882a593Smuzhiyun
796*4882a593Smuzhiyun    Arguments:
797*4882a593Smuzhiyun      num_bytes: Size in number of bytes of the DONT_CARE chunk.
798*4882a593Smuzhiyun    """
799*4882a593Smuzhiyun    assert num_bytes % self.block_size == 0
800*4882a593Smuzhiyun
801*4882a593Smuzhiyun    if not self.is_sparse:
802*4882a593Smuzhiyun      self._image.seek(0, os.SEEK_END)
803*4882a593Smuzhiyun      # This is more efficient that writing NUL bytes since it'll add
804*4882a593Smuzhiyun      # a hole on file systems that support sparse files (native
805*4882a593Smuzhiyun      # sparse, not Android sparse).
806*4882a593Smuzhiyun      self._image.truncate(self._image.tell() + num_bytes)
807*4882a593Smuzhiyun      self._read_header()
808*4882a593Smuzhiyun      return
809*4882a593Smuzhiyun
810*4882a593Smuzhiyun    self._num_total_chunks += 1
811*4882a593Smuzhiyun    self._num_total_blocks += num_bytes / self.block_size
812*4882a593Smuzhiyun    self._update_chunks_and_blocks()
813*4882a593Smuzhiyun
814*4882a593Smuzhiyun    self._image.seek(self._sparse_end, os.SEEK_SET)
815*4882a593Smuzhiyun    self._image.write(struct.pack(ImageChunk.FORMAT,
816*4882a593Smuzhiyun                                  ImageChunk.TYPE_DONT_CARE,
817*4882a593Smuzhiyun                                  0,  # Reserved
818*4882a593Smuzhiyun                                  num_bytes / self.block_size,
819*4882a593Smuzhiyun                                  struct.calcsize(ImageChunk.FORMAT)))
820*4882a593Smuzhiyun    self._read_header()
821*4882a593Smuzhiyun
822*4882a593Smuzhiyun  def append_raw(self, data):
823*4882a593Smuzhiyun    """Appends a RAW chunk to the sparse file.
824*4882a593Smuzhiyun
825*4882a593Smuzhiyun    The length of the given data must be a multiple of the block size.
826*4882a593Smuzhiyun
827*4882a593Smuzhiyun    Arguments:
828*4882a593Smuzhiyun      data: Data to append.
829*4882a593Smuzhiyun    """
830*4882a593Smuzhiyun    assert len(data) % self.block_size == 0
831*4882a593Smuzhiyun
832*4882a593Smuzhiyun    if not self.is_sparse:
833*4882a593Smuzhiyun      self._image.seek(0, os.SEEK_END)
834*4882a593Smuzhiyun      self._image.write(data)
835*4882a593Smuzhiyun      self._read_header()
836*4882a593Smuzhiyun      return
837*4882a593Smuzhiyun
838*4882a593Smuzhiyun    self._num_total_chunks += 1
839*4882a593Smuzhiyun    self._num_total_blocks += len(data) / self.block_size
840*4882a593Smuzhiyun    self._update_chunks_and_blocks()
841*4882a593Smuzhiyun
842*4882a593Smuzhiyun    self._image.seek(self._sparse_end, os.SEEK_SET)
843*4882a593Smuzhiyun    self._image.write(struct.pack(ImageChunk.FORMAT,
844*4882a593Smuzhiyun                                  ImageChunk.TYPE_RAW,
845*4882a593Smuzhiyun                                  0,  # Reserved
846*4882a593Smuzhiyun                                  len(data) / self.block_size,
847*4882a593Smuzhiyun                                  len(data) +
848*4882a593Smuzhiyun                                  struct.calcsize(ImageChunk.FORMAT)))
849*4882a593Smuzhiyun    self._image.write(data)
850*4882a593Smuzhiyun    self._read_header()
851*4882a593Smuzhiyun
852*4882a593Smuzhiyun  def append_fill(self, fill_data, size):
853*4882a593Smuzhiyun    """Appends a fill chunk to the sparse file.
854*4882a593Smuzhiyun
855*4882a593Smuzhiyun    The total length of the fill data must be a multiple of the block size.
856*4882a593Smuzhiyun
857*4882a593Smuzhiyun    Arguments:
858*4882a593Smuzhiyun      fill_data: Fill data to append - must be four bytes.
859*4882a593Smuzhiyun      size: Number of chunk - must be a multiple of four and the block size.
860*4882a593Smuzhiyun    """
861*4882a593Smuzhiyun    assert len(fill_data) == 4
862*4882a593Smuzhiyun    assert size % 4 == 0
863*4882a593Smuzhiyun    assert size % self.block_size == 0
864*4882a593Smuzhiyun
865*4882a593Smuzhiyun    if not self.is_sparse:
866*4882a593Smuzhiyun      self._image.seek(0, os.SEEK_END)
867*4882a593Smuzhiyun      self._image.write(fill_data * (size/4))
868*4882a593Smuzhiyun      self._read_header()
869*4882a593Smuzhiyun      return
870*4882a593Smuzhiyun
871*4882a593Smuzhiyun    self._num_total_chunks += 1
872*4882a593Smuzhiyun    self._num_total_blocks += size / self.block_size
873*4882a593Smuzhiyun    self._update_chunks_and_blocks()
874*4882a593Smuzhiyun
875*4882a593Smuzhiyun    self._image.seek(self._sparse_end, os.SEEK_SET)
876*4882a593Smuzhiyun    self._image.write(struct.pack(ImageChunk.FORMAT,
877*4882a593Smuzhiyun                                  ImageChunk.TYPE_FILL,
878*4882a593Smuzhiyun                                  0,  # Reserved
879*4882a593Smuzhiyun                                  size / self.block_size,
880*4882a593Smuzhiyun                                  4 + struct.calcsize(ImageChunk.FORMAT)))
881*4882a593Smuzhiyun    self._image.write(fill_data)
882*4882a593Smuzhiyun    self._read_header()
883*4882a593Smuzhiyun
884*4882a593Smuzhiyun  def seek(self, offset):
885*4882a593Smuzhiyun    """Sets the cursor position for reading from unsparsified file.
886*4882a593Smuzhiyun
887*4882a593Smuzhiyun    Arguments:
888*4882a593Smuzhiyun      offset: Offset to seek to from the beginning of the file.
889*4882a593Smuzhiyun    """
890*4882a593Smuzhiyun    if offset < 0:
891*4882a593Smuzhiyun      raise RuntimeError("Seeking with negative offset: %d" % offset)
892*4882a593Smuzhiyun    self._file_pos = offset
893*4882a593Smuzhiyun
894*4882a593Smuzhiyun  def read(self, size):
895*4882a593Smuzhiyun    """Reads data from the unsparsified file.
896*4882a593Smuzhiyun
897*4882a593Smuzhiyun    This method may return fewer than |size| bytes of data if the end
898*4882a593Smuzhiyun    of the file was encountered.
899*4882a593Smuzhiyun
900*4882a593Smuzhiyun    The file cursor for reading is advanced by the number of bytes
901*4882a593Smuzhiyun    read.
902*4882a593Smuzhiyun
903*4882a593Smuzhiyun    Arguments:
904*4882a593Smuzhiyun      size: Number of bytes to read.
905*4882a593Smuzhiyun
906*4882a593Smuzhiyun    Returns:
907*4882a593Smuzhiyun      The data.
908*4882a593Smuzhiyun
909*4882a593Smuzhiyun    """
910*4882a593Smuzhiyun    if not self.is_sparse:
911*4882a593Smuzhiyun      self._image.seek(self._file_pos)
912*4882a593Smuzhiyun      data = self._image.read(size)
913*4882a593Smuzhiyun      self._file_pos += len(data)
914*4882a593Smuzhiyun      return data
915*4882a593Smuzhiyun
916*4882a593Smuzhiyun    # Iterate over all chunks.
917*4882a593Smuzhiyun    chunk_idx = bisect.bisect_right(self._chunk_output_offsets,
918*4882a593Smuzhiyun                                    self._file_pos) - 1
919*4882a593Smuzhiyun    data = bytearray()
920*4882a593Smuzhiyun    to_go = size
921*4882a593Smuzhiyun    while to_go > 0:
922*4882a593Smuzhiyun      chunk = self._chunks[chunk_idx]
923*4882a593Smuzhiyun      chunk_pos_offset = self._file_pos - chunk.output_offset
924*4882a593Smuzhiyun      chunk_pos_to_go = min(chunk.output_size - chunk_pos_offset, to_go)
925*4882a593Smuzhiyun
926*4882a593Smuzhiyun      if chunk.chunk_type == ImageChunk.TYPE_RAW:
927*4882a593Smuzhiyun        self._image.seek(chunk.input_offset + chunk_pos_offset)
928*4882a593Smuzhiyun        data.extend(self._image.read(chunk_pos_to_go))
929*4882a593Smuzhiyun      elif chunk.chunk_type == ImageChunk.TYPE_FILL:
930*4882a593Smuzhiyun        all_data = chunk.fill_data*(chunk_pos_to_go/len(chunk.fill_data) + 2)
931*4882a593Smuzhiyun        offset_mod = chunk_pos_offset % len(chunk.fill_data)
932*4882a593Smuzhiyun        data.extend(all_data[offset_mod:(offset_mod + chunk_pos_to_go)])
933*4882a593Smuzhiyun      else:
934*4882a593Smuzhiyun        assert chunk.chunk_type == ImageChunk.TYPE_DONT_CARE
935*4882a593Smuzhiyun        data.extend('\0' * chunk_pos_to_go)
936*4882a593Smuzhiyun
937*4882a593Smuzhiyun      to_go -= chunk_pos_to_go
938*4882a593Smuzhiyun      self._file_pos += chunk_pos_to_go
939*4882a593Smuzhiyun      chunk_idx += 1
940*4882a593Smuzhiyun      # Generate partial read in case of EOF.
941*4882a593Smuzhiyun      if chunk_idx >= len(self._chunks):
942*4882a593Smuzhiyun        break
943*4882a593Smuzhiyun
944*4882a593Smuzhiyun    return data
945*4882a593Smuzhiyun
946*4882a593Smuzhiyun  def tell(self):
947*4882a593Smuzhiyun    """Returns the file cursor position for reading from unsparsified file.
948*4882a593Smuzhiyun
949*4882a593Smuzhiyun    Returns:
950*4882a593Smuzhiyun      The file cursor position for reading.
951*4882a593Smuzhiyun    """
952*4882a593Smuzhiyun    return self._file_pos
953*4882a593Smuzhiyun
954*4882a593Smuzhiyun  def truncate(self, size):
955*4882a593Smuzhiyun    """Truncates the unsparsified file.
956*4882a593Smuzhiyun
957*4882a593Smuzhiyun    Arguments:
958*4882a593Smuzhiyun      size: Desired size of unsparsified file.
959*4882a593Smuzhiyun
960*4882a593Smuzhiyun    Raises:
961*4882a593Smuzhiyun      ValueError: If desired size isn't a multiple of the block size.
962*4882a593Smuzhiyun    """
963*4882a593Smuzhiyun    if not self.is_sparse:
964*4882a593Smuzhiyun      self._image.truncate(size)
965*4882a593Smuzhiyun      self._read_header()
966*4882a593Smuzhiyun      return
967*4882a593Smuzhiyun
968*4882a593Smuzhiyun    if size % self.block_size != 0:
969*4882a593Smuzhiyun      raise ValueError('Cannot truncate to a size which is not a multiple '
970*4882a593Smuzhiyun                       'of the block size')
971*4882a593Smuzhiyun
972*4882a593Smuzhiyun    if size == self.image_size:
973*4882a593Smuzhiyun      # Trivial where there's nothing to do.
974*4882a593Smuzhiyun      return
975*4882a593Smuzhiyun    elif size < self.image_size:
976*4882a593Smuzhiyun      chunk_idx = bisect.bisect_right(self._chunk_output_offsets, size) - 1
977*4882a593Smuzhiyun      chunk = self._chunks[chunk_idx]
978*4882a593Smuzhiyun      if chunk.output_offset != size:
979*4882a593Smuzhiyun        # Truncation in the middle of a trunk - need to keep the chunk
980*4882a593Smuzhiyun        # and modify it.
981*4882a593Smuzhiyun        chunk_idx_for_update = chunk_idx + 1
982*4882a593Smuzhiyun        num_to_keep = size - chunk.output_offset
983*4882a593Smuzhiyun        assert num_to_keep % self.block_size == 0
984*4882a593Smuzhiyun        if chunk.chunk_type == ImageChunk.TYPE_RAW:
985*4882a593Smuzhiyun          truncate_at = (chunk.chunk_offset +
986*4882a593Smuzhiyun                         struct.calcsize(ImageChunk.FORMAT) + num_to_keep)
987*4882a593Smuzhiyun          data_sz = num_to_keep
988*4882a593Smuzhiyun        elif chunk.chunk_type == ImageChunk.TYPE_FILL:
989*4882a593Smuzhiyun          truncate_at = (chunk.chunk_offset +
990*4882a593Smuzhiyun                         struct.calcsize(ImageChunk.FORMAT) + 4)
991*4882a593Smuzhiyun          data_sz = 4
992*4882a593Smuzhiyun        else:
993*4882a593Smuzhiyun          assert chunk.chunk_type == ImageChunk.TYPE_DONT_CARE
994*4882a593Smuzhiyun          truncate_at = chunk.chunk_offset + struct.calcsize(ImageChunk.FORMAT)
995*4882a593Smuzhiyun          data_sz = 0
996*4882a593Smuzhiyun        chunk_sz = num_to_keep/self.block_size
997*4882a593Smuzhiyun        total_sz = data_sz + struct.calcsize(ImageChunk.FORMAT)
998*4882a593Smuzhiyun        self._image.seek(chunk.chunk_offset)
999*4882a593Smuzhiyun        self._image.write(struct.pack(ImageChunk.FORMAT,
1000*4882a593Smuzhiyun                                      chunk.chunk_type,
1001*4882a593Smuzhiyun                                      0,  # Reserved
1002*4882a593Smuzhiyun                                      chunk_sz,
1003*4882a593Smuzhiyun                                      total_sz))
1004*4882a593Smuzhiyun        chunk.output_size = num_to_keep
1005*4882a593Smuzhiyun      else:
1006*4882a593Smuzhiyun        # Truncation at trunk boundary.
1007*4882a593Smuzhiyun        truncate_at = chunk.chunk_offset
1008*4882a593Smuzhiyun        chunk_idx_for_update = chunk_idx
1009*4882a593Smuzhiyun
1010*4882a593Smuzhiyun      self._num_total_chunks = chunk_idx_for_update
1011*4882a593Smuzhiyun      self._num_total_blocks = 0
1012*4882a593Smuzhiyun      for i in range(0, chunk_idx_for_update):
1013*4882a593Smuzhiyun        self._num_total_blocks += self._chunks[i].output_size / self.block_size
1014*4882a593Smuzhiyun      self._update_chunks_and_blocks()
1015*4882a593Smuzhiyun      self._image.truncate(truncate_at)
1016*4882a593Smuzhiyun
1017*4882a593Smuzhiyun      # We've modified the file so re-read all data.
1018*4882a593Smuzhiyun      self._read_header()
1019*4882a593Smuzhiyun    else:
1020*4882a593Smuzhiyun      # Truncating to grow - just add a DONT_CARE section.
1021*4882a593Smuzhiyun      self.append_dont_care(size - self.image_size)
1022*4882a593Smuzhiyun
1023*4882a593Smuzhiyun
1024*4882a593Smuzhiyunclass AvbDescriptor(object):
1025*4882a593Smuzhiyun  """Class for AVB descriptor.
1026*4882a593Smuzhiyun
1027*4882a593Smuzhiyun  See the |AvbDescriptor| C struct for more information.
1028*4882a593Smuzhiyun
1029*4882a593Smuzhiyun  Attributes:
1030*4882a593Smuzhiyun    tag: The tag identifying what kind of descriptor this is.
1031*4882a593Smuzhiyun    data: The data in the descriptor.
1032*4882a593Smuzhiyun  """
1033*4882a593Smuzhiyun
1034*4882a593Smuzhiyun  SIZE = 16
1035*4882a593Smuzhiyun  FORMAT_STRING = ('!QQ')  # tag, num_bytes_following (descriptor header)
1036*4882a593Smuzhiyun
1037*4882a593Smuzhiyun  def __init__(self, data):
1038*4882a593Smuzhiyun    """Initializes a new property descriptor.
1039*4882a593Smuzhiyun
1040*4882a593Smuzhiyun    Arguments:
1041*4882a593Smuzhiyun      data: If not None, must be a bytearray().
1042*4882a593Smuzhiyun
1043*4882a593Smuzhiyun    Raises:
1044*4882a593Smuzhiyun      LookupError: If the given descriptor is malformed.
1045*4882a593Smuzhiyun    """
1046*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1047*4882a593Smuzhiyun
1048*4882a593Smuzhiyun    if data:
1049*4882a593Smuzhiyun      (self.tag, num_bytes_following) = (
1050*4882a593Smuzhiyun          struct.unpack(self.FORMAT_STRING, data[0:self.SIZE]))
1051*4882a593Smuzhiyun      self.data = data[self.SIZE:self.SIZE + num_bytes_following]
1052*4882a593Smuzhiyun    else:
1053*4882a593Smuzhiyun      self.tag = None
1054*4882a593Smuzhiyun      self.data = None
1055*4882a593Smuzhiyun
1056*4882a593Smuzhiyun  def print_desc(self, o):
1057*4882a593Smuzhiyun    """Print the descriptor.
1058*4882a593Smuzhiyun
1059*4882a593Smuzhiyun    Arguments:
1060*4882a593Smuzhiyun      o: The object to write the output to.
1061*4882a593Smuzhiyun    """
1062*4882a593Smuzhiyun    o.write('    Unknown descriptor:\n')
1063*4882a593Smuzhiyun    o.write('      Tag:  {}\n'.format(self.tag))
1064*4882a593Smuzhiyun    if len(self.data) < 256:
1065*4882a593Smuzhiyun      o.write('      Data: {} ({} bytes)\n'.format(
1066*4882a593Smuzhiyun          repr(str(self.data)), len(self.data)))
1067*4882a593Smuzhiyun    else:
1068*4882a593Smuzhiyun      o.write('      Data: {} bytes\n'.format(len(self.data)))
1069*4882a593Smuzhiyun
1070*4882a593Smuzhiyun  def encode(self):
1071*4882a593Smuzhiyun    """Serializes the descriptor.
1072*4882a593Smuzhiyun
1073*4882a593Smuzhiyun    Returns:
1074*4882a593Smuzhiyun      A bytearray() with the descriptor data.
1075*4882a593Smuzhiyun    """
1076*4882a593Smuzhiyun    num_bytes_following = len(self.data)
1077*4882a593Smuzhiyun    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
1078*4882a593Smuzhiyun    padding_size = nbf_with_padding - num_bytes_following
1079*4882a593Smuzhiyun    desc = struct.pack(self.FORMAT_STRING, self.tag, nbf_with_padding)
1080*4882a593Smuzhiyun    padding = struct.pack(str(padding_size) + 'x')
1081*4882a593Smuzhiyun    ret = desc + self.data + padding
1082*4882a593Smuzhiyun    return bytearray(ret)
1083*4882a593Smuzhiyun
1084*4882a593Smuzhiyun  def verify(self, image_dir, image_ext, expected_chain_partitions_map):
1085*4882a593Smuzhiyun    """Verifies contents of the descriptor - used in verify_image sub-command.
1086*4882a593Smuzhiyun
1087*4882a593Smuzhiyun    Arguments:
1088*4882a593Smuzhiyun      image_dir: The directory of the file being verified.
1089*4882a593Smuzhiyun      image_ext: The extension of the file being verified (e.g. '.img').
1090*4882a593Smuzhiyun      expected_chain_partitions_map: A map from partition name to the
1091*4882a593Smuzhiyun        tuple (rollback_index_location, key_blob).
1092*4882a593Smuzhiyun
1093*4882a593Smuzhiyun    Returns:
1094*4882a593Smuzhiyun      True if the descriptor verifies, False otherwise.
1095*4882a593Smuzhiyun    """
1096*4882a593Smuzhiyun    # Nothing to do.
1097*4882a593Smuzhiyun    return True
1098*4882a593Smuzhiyun
1099*4882a593Smuzhiyunclass AvbPropertyDescriptor(AvbDescriptor):
1100*4882a593Smuzhiyun  """A class for property descriptors.
1101*4882a593Smuzhiyun
1102*4882a593Smuzhiyun  See the |AvbPropertyDescriptor| C struct for more information.
1103*4882a593Smuzhiyun
1104*4882a593Smuzhiyun  Attributes:
1105*4882a593Smuzhiyun    key: The key.
1106*4882a593Smuzhiyun    value: The key.
1107*4882a593Smuzhiyun  """
1108*4882a593Smuzhiyun
1109*4882a593Smuzhiyun  TAG = 0
1110*4882a593Smuzhiyun  SIZE = 32
1111*4882a593Smuzhiyun  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
1112*4882a593Smuzhiyun                   'Q'  # key size (bytes)
1113*4882a593Smuzhiyun                   'Q')  # value size (bytes)
1114*4882a593Smuzhiyun
1115*4882a593Smuzhiyun  def __init__(self, data=None):
1116*4882a593Smuzhiyun    """Initializes a new property descriptor.
1117*4882a593Smuzhiyun
1118*4882a593Smuzhiyun    Arguments:
1119*4882a593Smuzhiyun      data: If not None, must be a bytearray of size |SIZE|.
1120*4882a593Smuzhiyun
1121*4882a593Smuzhiyun    Raises:
1122*4882a593Smuzhiyun      LookupError: If the given descriptor is malformed.
1123*4882a593Smuzhiyun    """
1124*4882a593Smuzhiyun    AvbDescriptor.__init__(self, None)
1125*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1126*4882a593Smuzhiyun
1127*4882a593Smuzhiyun    if data:
1128*4882a593Smuzhiyun      (tag, num_bytes_following, key_size,
1129*4882a593Smuzhiyun       value_size) = struct.unpack(self.FORMAT_STRING, data[0:self.SIZE])
1130*4882a593Smuzhiyun      expected_size = round_to_multiple(
1131*4882a593Smuzhiyun          self.SIZE - 16 + key_size + 1 + value_size + 1, 8)
1132*4882a593Smuzhiyun      if tag != self.TAG or num_bytes_following != expected_size:
1133*4882a593Smuzhiyun        raise LookupError('Given data does not look like a property '
1134*4882a593Smuzhiyun                          'descriptor.')
1135*4882a593Smuzhiyun      self.key = data[self.SIZE:(self.SIZE + key_size)]
1136*4882a593Smuzhiyun      self.value = data[(self.SIZE + key_size + 1):(self.SIZE + key_size + 1 +
1137*4882a593Smuzhiyun                                                    value_size)]
1138*4882a593Smuzhiyun    else:
1139*4882a593Smuzhiyun      self.key = ''
1140*4882a593Smuzhiyun      self.value = ''
1141*4882a593Smuzhiyun
1142*4882a593Smuzhiyun  def print_desc(self, o):
1143*4882a593Smuzhiyun    """Print the descriptor.
1144*4882a593Smuzhiyun
1145*4882a593Smuzhiyun    Arguments:
1146*4882a593Smuzhiyun      o: The object to write the output to.
1147*4882a593Smuzhiyun    """
1148*4882a593Smuzhiyun    if len(self.value) < 256:
1149*4882a593Smuzhiyun      o.write('    Prop: {} -> {}\n'.format(self.key, repr(str(self.value))))
1150*4882a593Smuzhiyun    else:
1151*4882a593Smuzhiyun      o.write('    Prop: {} -> ({} bytes)\n'.format(self.key, len(self.value)))
1152*4882a593Smuzhiyun
1153*4882a593Smuzhiyun  def encode(self):
1154*4882a593Smuzhiyun    """Serializes the descriptor.
1155*4882a593Smuzhiyun
1156*4882a593Smuzhiyun    Returns:
1157*4882a593Smuzhiyun      A bytearray() with the descriptor data.
1158*4882a593Smuzhiyun    """
1159*4882a593Smuzhiyun    num_bytes_following = self.SIZE + len(self.key) + len(self.value) + 2 - 16
1160*4882a593Smuzhiyun    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
1161*4882a593Smuzhiyun    padding_size = nbf_with_padding - num_bytes_following
1162*4882a593Smuzhiyun    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
1163*4882a593Smuzhiyun                       len(self.key), len(self.value))
1164*4882a593Smuzhiyun    padding = struct.pack(str(padding_size) + 'x')
1165*4882a593Smuzhiyun    ret = desc + self.key + '\0' + self.value + '\0' + padding
1166*4882a593Smuzhiyun    return bytearray(ret)
1167*4882a593Smuzhiyun
1168*4882a593Smuzhiyun  def verify(self, image_dir, image_ext, expected_chain_partitions_map):
1169*4882a593Smuzhiyun    """Verifies contents of the descriptor - used in verify_image sub-command.
1170*4882a593Smuzhiyun
1171*4882a593Smuzhiyun    Arguments:
1172*4882a593Smuzhiyun      image_dir: The directory of the file being verified.
1173*4882a593Smuzhiyun      image_ext: The extension of the file being verified (e.g. '.img').
1174*4882a593Smuzhiyun      expected_chain_partitions_map: A map from partition name to the
1175*4882a593Smuzhiyun        tuple (rollback_index_location, key_blob).
1176*4882a593Smuzhiyun
1177*4882a593Smuzhiyun    Returns:
1178*4882a593Smuzhiyun      True if the descriptor verifies, False otherwise.
1179*4882a593Smuzhiyun    """
1180*4882a593Smuzhiyun    # Nothing to do.
1181*4882a593Smuzhiyun    return True
1182*4882a593Smuzhiyun
1183*4882a593Smuzhiyunclass AvbHashtreeDescriptor(AvbDescriptor):
1184*4882a593Smuzhiyun  """A class for hashtree descriptors.
1185*4882a593Smuzhiyun
1186*4882a593Smuzhiyun  See the |AvbHashtreeDescriptor| C struct for more information.
1187*4882a593Smuzhiyun
1188*4882a593Smuzhiyun  Attributes:
1189*4882a593Smuzhiyun    dm_verity_version: dm-verity version used.
1190*4882a593Smuzhiyun    image_size: Size of the image, after rounding up to |block_size|.
1191*4882a593Smuzhiyun    tree_offset: Offset of the hash tree in the file.
1192*4882a593Smuzhiyun    tree_size: Size of the tree.
1193*4882a593Smuzhiyun    data_block_size: Data block size
1194*4882a593Smuzhiyun    hash_block_size: Hash block size
1195*4882a593Smuzhiyun    fec_num_roots: Number of roots used for FEC (0 if FEC is not used).
1196*4882a593Smuzhiyun    fec_offset: Offset of FEC data (0 if FEC is not used).
1197*4882a593Smuzhiyun    fec_size: Size of FEC data (0 if FEC is not used).
1198*4882a593Smuzhiyun    hash_algorithm: Hash algorithm used.
1199*4882a593Smuzhiyun    partition_name: Partition name.
1200*4882a593Smuzhiyun    salt: Salt used.
1201*4882a593Smuzhiyun    root_digest: Root digest.
1202*4882a593Smuzhiyun    flags: Descriptor flags (see avb_hashtree_descriptor.h).
1203*4882a593Smuzhiyun  """
1204*4882a593Smuzhiyun
1205*4882a593Smuzhiyun  TAG = 1
1206*4882a593Smuzhiyun  RESERVED = 60
1207*4882a593Smuzhiyun  SIZE = 120 + RESERVED
1208*4882a593Smuzhiyun  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
1209*4882a593Smuzhiyun                   'L'  # dm-verity version used
1210*4882a593Smuzhiyun                   'Q'  # image size (bytes)
1211*4882a593Smuzhiyun                   'Q'  # tree offset (bytes)
1212*4882a593Smuzhiyun                   'Q'  # tree size (bytes)
1213*4882a593Smuzhiyun                   'L'  # data block size (bytes)
1214*4882a593Smuzhiyun                   'L'  # hash block size (bytes)
1215*4882a593Smuzhiyun                   'L'  # FEC number of roots
1216*4882a593Smuzhiyun                   'Q'  # FEC offset (bytes)
1217*4882a593Smuzhiyun                   'Q'  # FEC size (bytes)
1218*4882a593Smuzhiyun                   '32s'  # hash algorithm used
1219*4882a593Smuzhiyun                   'L'  # partition name (bytes)
1220*4882a593Smuzhiyun                   'L'  # salt length (bytes)
1221*4882a593Smuzhiyun                   'L'  # root digest length (bytes)
1222*4882a593Smuzhiyun                   'L' +  # flags
1223*4882a593Smuzhiyun                   str(RESERVED) + 's')  # reserved
1224*4882a593Smuzhiyun
1225*4882a593Smuzhiyun  def __init__(self, data=None):
1226*4882a593Smuzhiyun    """Initializes a new hashtree descriptor.
1227*4882a593Smuzhiyun
1228*4882a593Smuzhiyun    Arguments:
1229*4882a593Smuzhiyun      data: If not None, must be a bytearray of size |SIZE|.
1230*4882a593Smuzhiyun
1231*4882a593Smuzhiyun    Raises:
1232*4882a593Smuzhiyun      LookupError: If the given descriptor is malformed.
1233*4882a593Smuzhiyun    """
1234*4882a593Smuzhiyun    AvbDescriptor.__init__(self, None)
1235*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1236*4882a593Smuzhiyun
1237*4882a593Smuzhiyun    if data:
1238*4882a593Smuzhiyun      (tag, num_bytes_following, self.dm_verity_version, self.image_size,
1239*4882a593Smuzhiyun       self.tree_offset, self.tree_size, self.data_block_size,
1240*4882a593Smuzhiyun       self.hash_block_size, self.fec_num_roots, self.fec_offset, self.fec_size,
1241*4882a593Smuzhiyun       self.hash_algorithm, partition_name_len, salt_len,
1242*4882a593Smuzhiyun       root_digest_len, self.flags, _) = struct.unpack(self.FORMAT_STRING,
1243*4882a593Smuzhiyun                                                       data[0:self.SIZE])
1244*4882a593Smuzhiyun      expected_size = round_to_multiple(
1245*4882a593Smuzhiyun          self.SIZE - 16 + partition_name_len + salt_len + root_digest_len, 8)
1246*4882a593Smuzhiyun      if tag != self.TAG or num_bytes_following != expected_size:
1247*4882a593Smuzhiyun        raise LookupError('Given data does not look like a hashtree '
1248*4882a593Smuzhiyun                          'descriptor.')
1249*4882a593Smuzhiyun      # Nuke NUL-bytes at the end.
1250*4882a593Smuzhiyun      self.hash_algorithm = self.hash_algorithm.split('\0', 1)[0]
1251*4882a593Smuzhiyun      o = 0
1252*4882a593Smuzhiyun      self.partition_name = str(data[(self.SIZE + o):(self.SIZE + o +
1253*4882a593Smuzhiyun                                                      partition_name_len)])
1254*4882a593Smuzhiyun      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
1255*4882a593Smuzhiyun      self.partition_name.decode('utf-8')
1256*4882a593Smuzhiyun      o += partition_name_len
1257*4882a593Smuzhiyun      self.salt = data[(self.SIZE + o):(self.SIZE + o + salt_len)]
1258*4882a593Smuzhiyun      o += salt_len
1259*4882a593Smuzhiyun      self.root_digest = data[(self.SIZE + o):(self.SIZE + o + root_digest_len)]
1260*4882a593Smuzhiyun      if root_digest_len != len(hashlib.new(name=self.hash_algorithm).digest()):
1261*4882a593Smuzhiyun        if root_digest_len != 0:
1262*4882a593Smuzhiyun          raise LookupError('root_digest_len doesn\'t match hash algorithm')
1263*4882a593Smuzhiyun
1264*4882a593Smuzhiyun    else:
1265*4882a593Smuzhiyun      self.dm_verity_version = 0
1266*4882a593Smuzhiyun      self.image_size = 0
1267*4882a593Smuzhiyun      self.tree_offset = 0
1268*4882a593Smuzhiyun      self.tree_size = 0
1269*4882a593Smuzhiyun      self.data_block_size = 0
1270*4882a593Smuzhiyun      self.hash_block_size = 0
1271*4882a593Smuzhiyun      self.fec_num_roots = 0
1272*4882a593Smuzhiyun      self.fec_offset = 0
1273*4882a593Smuzhiyun      self.fec_size = 0
1274*4882a593Smuzhiyun      self.hash_algorithm = ''
1275*4882a593Smuzhiyun      self.partition_name = ''
1276*4882a593Smuzhiyun      self.salt = bytearray()
1277*4882a593Smuzhiyun      self.root_digest = bytearray()
1278*4882a593Smuzhiyun      self.flags = 0
1279*4882a593Smuzhiyun
1280*4882a593Smuzhiyun  def print_desc(self, o):
1281*4882a593Smuzhiyun    """Print the descriptor.
1282*4882a593Smuzhiyun
1283*4882a593Smuzhiyun    Arguments:
1284*4882a593Smuzhiyun      o: The object to write the output to.
1285*4882a593Smuzhiyun    """
1286*4882a593Smuzhiyun    o.write('    Hashtree descriptor:\n')
1287*4882a593Smuzhiyun    o.write('      Version of dm-verity:  {}\n'.format(self.dm_verity_version))
1288*4882a593Smuzhiyun    o.write('      Image Size:            {} bytes\n'.format(self.image_size))
1289*4882a593Smuzhiyun    o.write('      Tree Offset:           {}\n'.format(self.tree_offset))
1290*4882a593Smuzhiyun    o.write('      Tree Size:             {} bytes\n'.format(self.tree_size))
1291*4882a593Smuzhiyun    o.write('      Data Block Size:       {} bytes\n'.format(
1292*4882a593Smuzhiyun        self.data_block_size))
1293*4882a593Smuzhiyun    o.write('      Hash Block Size:       {} bytes\n'.format(
1294*4882a593Smuzhiyun        self.hash_block_size))
1295*4882a593Smuzhiyun    o.write('      FEC num roots:         {}\n'.format(self.fec_num_roots))
1296*4882a593Smuzhiyun    o.write('      FEC offset:            {}\n'.format(self.fec_offset))
1297*4882a593Smuzhiyun    o.write('      FEC size:              {} bytes\n'.format(self.fec_size))
1298*4882a593Smuzhiyun    o.write('      Hash Algorithm:        {}\n'.format(self.hash_algorithm))
1299*4882a593Smuzhiyun    o.write('      Partition Name:        {}\n'.format(self.partition_name))
1300*4882a593Smuzhiyun    o.write('      Salt:                  {}\n'.format(str(self.salt).encode(
1301*4882a593Smuzhiyun        'hex')))
1302*4882a593Smuzhiyun    o.write('      Root Digest:           {}\n'.format(str(
1303*4882a593Smuzhiyun        self.root_digest).encode('hex')))
1304*4882a593Smuzhiyun    o.write('      Flags:                 {}\n'.format(self.flags))
1305*4882a593Smuzhiyun
1306*4882a593Smuzhiyun  def encode(self):
1307*4882a593Smuzhiyun    """Serializes the descriptor.
1308*4882a593Smuzhiyun
1309*4882a593Smuzhiyun    Returns:
1310*4882a593Smuzhiyun      A bytearray() with the descriptor data.
1311*4882a593Smuzhiyun    """
1312*4882a593Smuzhiyun    encoded_name = self.partition_name.encode('utf-8')
1313*4882a593Smuzhiyun    num_bytes_following = (self.SIZE + len(encoded_name) + len(self.salt) +
1314*4882a593Smuzhiyun                           len(self.root_digest) - 16)
1315*4882a593Smuzhiyun    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
1316*4882a593Smuzhiyun    padding_size = nbf_with_padding - num_bytes_following
1317*4882a593Smuzhiyun    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
1318*4882a593Smuzhiyun                       self.dm_verity_version, self.image_size,
1319*4882a593Smuzhiyun                       self.tree_offset, self.tree_size, self.data_block_size,
1320*4882a593Smuzhiyun                       self.hash_block_size, self.fec_num_roots,
1321*4882a593Smuzhiyun                       self.fec_offset, self.fec_size, self.hash_algorithm,
1322*4882a593Smuzhiyun                       len(encoded_name), len(self.salt), len(self.root_digest),
1323*4882a593Smuzhiyun                       self.flags, self.RESERVED*'\0')
1324*4882a593Smuzhiyun    padding = struct.pack(str(padding_size) + 'x')
1325*4882a593Smuzhiyun    ret = desc + encoded_name + self.salt + self.root_digest + padding
1326*4882a593Smuzhiyun    return bytearray(ret)
1327*4882a593Smuzhiyun
1328*4882a593Smuzhiyun  def verify(self, image_dir, image_ext, expected_chain_partitions_map):
1329*4882a593Smuzhiyun    """Verifies contents of the descriptor - used in verify_image sub-command.
1330*4882a593Smuzhiyun
1331*4882a593Smuzhiyun    Arguments:
1332*4882a593Smuzhiyun      image_dir: The directory of the file being verified.
1333*4882a593Smuzhiyun      image_ext: The extension of the file being verified (e.g. '.img').
1334*4882a593Smuzhiyun      expected_chain_partitions_map: A map from partition name to the
1335*4882a593Smuzhiyun        tuple (rollback_index_location, key_blob).
1336*4882a593Smuzhiyun
1337*4882a593Smuzhiyun    Returns:
1338*4882a593Smuzhiyun      True if the descriptor verifies, False otherwise.
1339*4882a593Smuzhiyun    """
1340*4882a593Smuzhiyun    image_filename = os.path.join(image_dir, self.partition_name + image_ext)
1341*4882a593Smuzhiyun    image = ImageHandler(image_filename)
1342*4882a593Smuzhiyun    # Generate the hashtree and checks that it matches what's in the file.
1343*4882a593Smuzhiyun    digest_size = len(hashlib.new(name=self.hash_algorithm).digest())
1344*4882a593Smuzhiyun    digest_padding = round_to_pow2(digest_size) - digest_size
1345*4882a593Smuzhiyun    (hash_level_offsets, tree_size) = calc_hash_level_offsets(
1346*4882a593Smuzhiyun      self.image_size, self.data_block_size, digest_size + digest_padding)
1347*4882a593Smuzhiyun    root_digest, hash_tree = generate_hash_tree(image, self.image_size,
1348*4882a593Smuzhiyun                                                self.data_block_size,
1349*4882a593Smuzhiyun                                                self.hash_algorithm, self.salt,
1350*4882a593Smuzhiyun                                                digest_padding,
1351*4882a593Smuzhiyun                                                hash_level_offsets,
1352*4882a593Smuzhiyun                                                tree_size)
1353*4882a593Smuzhiyun    # The root digest must match unless it is not embedded in the descriptor.
1354*4882a593Smuzhiyun    if len(self.root_digest) != 0 and root_digest != self.root_digest:
1355*4882a593Smuzhiyun      sys.stderr.write('hashtree of {} does not match descriptor\n'.
1356*4882a593Smuzhiyun                       format(image_filename))
1357*4882a593Smuzhiyun      return False
1358*4882a593Smuzhiyun    # ... also check that the on-disk hashtree matches
1359*4882a593Smuzhiyun    image.seek(self.tree_offset)
1360*4882a593Smuzhiyun    hash_tree_ondisk = image.read(self.tree_size)
1361*4882a593Smuzhiyun    if hash_tree != hash_tree_ondisk:
1362*4882a593Smuzhiyun      sys.stderr.write('hashtree of {} contains invalid data\n'.
1363*4882a593Smuzhiyun                       format(image_filename))
1364*4882a593Smuzhiyun      return False
1365*4882a593Smuzhiyun    # TODO: we could also verify that the FEC stored in the image is
1366*4882a593Smuzhiyun    # correct but this a) currently requires the 'fec' binary; and b)
1367*4882a593Smuzhiyun    # takes a long time; and c) is not strictly needed for
1368*4882a593Smuzhiyun    # verification purposes as we've already verified the root hash.
1369*4882a593Smuzhiyun    print ('{}: Successfully verified {} hashtree of {} for image of {} bytes'
1370*4882a593Smuzhiyun           .format(self.partition_name, self.hash_algorithm, image_filename,
1371*4882a593Smuzhiyun                   self.image_size))
1372*4882a593Smuzhiyun    return True
1373*4882a593Smuzhiyun
1374*4882a593Smuzhiyun
1375*4882a593Smuzhiyunclass AvbHashDescriptor(AvbDescriptor):
1376*4882a593Smuzhiyun  """A class for hash descriptors.
1377*4882a593Smuzhiyun
1378*4882a593Smuzhiyun  See the |AvbHashDescriptor| C struct for more information.
1379*4882a593Smuzhiyun
1380*4882a593Smuzhiyun  Attributes:
1381*4882a593Smuzhiyun    image_size: Image size, in bytes.
1382*4882a593Smuzhiyun    hash_algorithm: Hash algorithm used.
1383*4882a593Smuzhiyun    partition_name: Partition name.
1384*4882a593Smuzhiyun    salt: Salt used.
1385*4882a593Smuzhiyun    digest: The hash value of salt and data combined.
1386*4882a593Smuzhiyun    flags: The descriptor flags (see avb_hash_descriptor.h).
1387*4882a593Smuzhiyun  """
1388*4882a593Smuzhiyun
1389*4882a593Smuzhiyun  TAG = 2
1390*4882a593Smuzhiyun  RESERVED = 60
1391*4882a593Smuzhiyun  SIZE = 72 + RESERVED
1392*4882a593Smuzhiyun  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
1393*4882a593Smuzhiyun                   'Q'  # image size (bytes)
1394*4882a593Smuzhiyun                   '32s'  # hash algorithm used
1395*4882a593Smuzhiyun                   'L'  # partition name (bytes)
1396*4882a593Smuzhiyun                   'L'  # salt length (bytes)
1397*4882a593Smuzhiyun                   'L'  # digest length (bytes)
1398*4882a593Smuzhiyun                   'L' +  # flags
1399*4882a593Smuzhiyun                   str(RESERVED) + 's')  # reserved
1400*4882a593Smuzhiyun
1401*4882a593Smuzhiyun  def __init__(self, data=None):
1402*4882a593Smuzhiyun    """Initializes a new hash descriptor.
1403*4882a593Smuzhiyun
1404*4882a593Smuzhiyun    Arguments:
1405*4882a593Smuzhiyun      data: If not None, must be a bytearray of size |SIZE|.
1406*4882a593Smuzhiyun
1407*4882a593Smuzhiyun    Raises:
1408*4882a593Smuzhiyun      LookupError: If the given descriptor is malformed.
1409*4882a593Smuzhiyun    """
1410*4882a593Smuzhiyun    AvbDescriptor.__init__(self, None)
1411*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1412*4882a593Smuzhiyun
1413*4882a593Smuzhiyun    if data:
1414*4882a593Smuzhiyun      (tag, num_bytes_following, self.image_size, self.hash_algorithm,
1415*4882a593Smuzhiyun       partition_name_len, salt_len,
1416*4882a593Smuzhiyun       digest_len, self.flags, _) = struct.unpack(self.FORMAT_STRING,
1417*4882a593Smuzhiyun                                                  data[0:self.SIZE])
1418*4882a593Smuzhiyun      expected_size = round_to_multiple(
1419*4882a593Smuzhiyun          self.SIZE - 16 + partition_name_len + salt_len + digest_len, 8)
1420*4882a593Smuzhiyun      if tag != self.TAG or num_bytes_following != expected_size:
1421*4882a593Smuzhiyun        raise LookupError('Given data does not look like a hash ' 'descriptor.')
1422*4882a593Smuzhiyun      # Nuke NUL-bytes at the end.
1423*4882a593Smuzhiyun      self.hash_algorithm = self.hash_algorithm.split('\0', 1)[0]
1424*4882a593Smuzhiyun      o = 0
1425*4882a593Smuzhiyun      self.partition_name = str(data[(self.SIZE + o):(self.SIZE + o +
1426*4882a593Smuzhiyun                                                      partition_name_len)])
1427*4882a593Smuzhiyun      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
1428*4882a593Smuzhiyun      self.partition_name.decode('utf-8')
1429*4882a593Smuzhiyun      o += partition_name_len
1430*4882a593Smuzhiyun      self.salt = data[(self.SIZE + o):(self.SIZE + o + salt_len)]
1431*4882a593Smuzhiyun      o += salt_len
1432*4882a593Smuzhiyun      self.digest = data[(self.SIZE + o):(self.SIZE + o + digest_len)]
1433*4882a593Smuzhiyun      if digest_len != len(hashlib.new(name=self.hash_algorithm).digest()):
1434*4882a593Smuzhiyun        if digest_len != 0:
1435*4882a593Smuzhiyun          raise LookupError('digest_len doesn\'t match hash algorithm')
1436*4882a593Smuzhiyun
1437*4882a593Smuzhiyun    else:
1438*4882a593Smuzhiyun      self.image_size = 0
1439*4882a593Smuzhiyun      self.hash_algorithm = ''
1440*4882a593Smuzhiyun      self.partition_name = ''
1441*4882a593Smuzhiyun      self.salt = bytearray()
1442*4882a593Smuzhiyun      self.digest = bytearray()
1443*4882a593Smuzhiyun      self.flags = 0
1444*4882a593Smuzhiyun
1445*4882a593Smuzhiyun  def print_desc(self, o):
1446*4882a593Smuzhiyun    """Print the descriptor.
1447*4882a593Smuzhiyun
1448*4882a593Smuzhiyun    Arguments:
1449*4882a593Smuzhiyun      o: The object to write the output to.
1450*4882a593Smuzhiyun    """
1451*4882a593Smuzhiyun    o.write('    Hash descriptor:\n')
1452*4882a593Smuzhiyun    o.write('      Image Size:            {} bytes\n'.format(self.image_size))
1453*4882a593Smuzhiyun    o.write('      Hash Algorithm:        {}\n'.format(self.hash_algorithm))
1454*4882a593Smuzhiyun    o.write('      Partition Name:        {}\n'.format(self.partition_name))
1455*4882a593Smuzhiyun    o.write('      Salt:                  {}\n'.format(str(self.salt).encode(
1456*4882a593Smuzhiyun        'hex')))
1457*4882a593Smuzhiyun    o.write('      Digest:                {}\n'.format(str(self.digest).encode(
1458*4882a593Smuzhiyun        'hex')))
1459*4882a593Smuzhiyun    o.write('      Flags:                 {}\n'.format(self.flags))
1460*4882a593Smuzhiyun
1461*4882a593Smuzhiyun  def encode(self):
1462*4882a593Smuzhiyun    """Serializes the descriptor.
1463*4882a593Smuzhiyun
1464*4882a593Smuzhiyun    Returns:
1465*4882a593Smuzhiyun      A bytearray() with the descriptor data.
1466*4882a593Smuzhiyun    """
1467*4882a593Smuzhiyun    encoded_name = self.partition_name.encode('utf-8')
1468*4882a593Smuzhiyun    num_bytes_following = (
1469*4882a593Smuzhiyun        self.SIZE + len(encoded_name) + len(self.salt) + len(self.digest) - 16)
1470*4882a593Smuzhiyun    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
1471*4882a593Smuzhiyun    padding_size = nbf_with_padding - num_bytes_following
1472*4882a593Smuzhiyun    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
1473*4882a593Smuzhiyun                       self.image_size, self.hash_algorithm, len(encoded_name),
1474*4882a593Smuzhiyun                       len(self.salt), len(self.digest), self.flags,
1475*4882a593Smuzhiyun                       self.RESERVED*'\0')
1476*4882a593Smuzhiyun    padding = struct.pack(str(padding_size) + 'x')
1477*4882a593Smuzhiyun    ret = desc + encoded_name + self.salt + self.digest + padding
1478*4882a593Smuzhiyun    return bytearray(ret)
1479*4882a593Smuzhiyun
1480*4882a593Smuzhiyun  def verify(self, image_dir, image_ext, expected_chain_partitions_map):
1481*4882a593Smuzhiyun    """Verifies contents of the descriptor - used in verify_image sub-command.
1482*4882a593Smuzhiyun
1483*4882a593Smuzhiyun    Arguments:
1484*4882a593Smuzhiyun      image_dir: The directory of the file being verified.
1485*4882a593Smuzhiyun      image_ext: The extension of the file being verified (e.g. '.img').
1486*4882a593Smuzhiyun      expected_chain_partitions_map: A map from partition name to the
1487*4882a593Smuzhiyun        tuple (rollback_index_location, key_blob).
1488*4882a593Smuzhiyun
1489*4882a593Smuzhiyun    Returns:
1490*4882a593Smuzhiyun      True if the descriptor verifies, False otherwise.
1491*4882a593Smuzhiyun    """
1492*4882a593Smuzhiyun    image_filename = os.path.join(image_dir, self.partition_name + image_ext)
1493*4882a593Smuzhiyun    image = ImageHandler(image_filename)
1494*4882a593Smuzhiyun    data = image.read(self.image_size)
1495*4882a593Smuzhiyun    ha = hashlib.new(self.hash_algorithm)
1496*4882a593Smuzhiyun    ha.update(self.salt)
1497*4882a593Smuzhiyun    ha.update(data)
1498*4882a593Smuzhiyun    digest = ha.digest()
1499*4882a593Smuzhiyun    # The digest must match unless there is no digest in the descriptor.
1500*4882a593Smuzhiyun    if len(self.digest) != 0 and digest != self.digest:
1501*4882a593Smuzhiyun      sys.stderr.write('{} digest of {} does not match digest in descriptor\n'.
1502*4882a593Smuzhiyun                       format(self.hash_algorithm, image_filename))
1503*4882a593Smuzhiyun      return False
1504*4882a593Smuzhiyun    print ('{}: Successfully verified {} hash of {} for image of {} bytes'
1505*4882a593Smuzhiyun           .format(self.partition_name, self.hash_algorithm, image_filename,
1506*4882a593Smuzhiyun                   self.image_size))
1507*4882a593Smuzhiyun    return True
1508*4882a593Smuzhiyun
1509*4882a593Smuzhiyun
1510*4882a593Smuzhiyunclass AvbKernelCmdlineDescriptor(AvbDescriptor):
1511*4882a593Smuzhiyun  """A class for kernel command-line descriptors.
1512*4882a593Smuzhiyun
1513*4882a593Smuzhiyun  See the |AvbKernelCmdlineDescriptor| C struct for more information.
1514*4882a593Smuzhiyun
1515*4882a593Smuzhiyun  Attributes:
1516*4882a593Smuzhiyun    flags: Flags.
1517*4882a593Smuzhiyun    kernel_cmdline: The kernel command-line.
1518*4882a593Smuzhiyun  """
1519*4882a593Smuzhiyun
1520*4882a593Smuzhiyun  TAG = 3
1521*4882a593Smuzhiyun  SIZE = 24
1522*4882a593Smuzhiyun  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
1523*4882a593Smuzhiyun                   'L'  # flags
1524*4882a593Smuzhiyun                   'L')  # cmdline length (bytes)
1525*4882a593Smuzhiyun
1526*4882a593Smuzhiyun  FLAGS_USE_ONLY_IF_HASHTREE_NOT_DISABLED = (1 << 0)
1527*4882a593Smuzhiyun  FLAGS_USE_ONLY_IF_HASHTREE_DISABLED = (1 << 1)
1528*4882a593Smuzhiyun
1529*4882a593Smuzhiyun  def __init__(self, data=None):
1530*4882a593Smuzhiyun    """Initializes a new kernel cmdline descriptor.
1531*4882a593Smuzhiyun
1532*4882a593Smuzhiyun    Arguments:
1533*4882a593Smuzhiyun      data: If not None, must be a bytearray of size |SIZE|.
1534*4882a593Smuzhiyun
1535*4882a593Smuzhiyun    Raises:
1536*4882a593Smuzhiyun      LookupError: If the given descriptor is malformed.
1537*4882a593Smuzhiyun    """
1538*4882a593Smuzhiyun    AvbDescriptor.__init__(self, None)
1539*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1540*4882a593Smuzhiyun
1541*4882a593Smuzhiyun    if data:
1542*4882a593Smuzhiyun      (tag, num_bytes_following, self.flags, kernel_cmdline_length) = (
1543*4882a593Smuzhiyun          struct.unpack(self.FORMAT_STRING, data[0:self.SIZE]))
1544*4882a593Smuzhiyun      expected_size = round_to_multiple(self.SIZE - 16 + kernel_cmdline_length,
1545*4882a593Smuzhiyun                                        8)
1546*4882a593Smuzhiyun      if tag != self.TAG or num_bytes_following != expected_size:
1547*4882a593Smuzhiyun        raise LookupError('Given data does not look like a kernel cmdline '
1548*4882a593Smuzhiyun                          'descriptor.')
1549*4882a593Smuzhiyun      # Nuke NUL-bytes at the end.
1550*4882a593Smuzhiyun      self.kernel_cmdline = str(data[self.SIZE:(self.SIZE +
1551*4882a593Smuzhiyun                                                kernel_cmdline_length)])
1552*4882a593Smuzhiyun      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
1553*4882a593Smuzhiyun      self.kernel_cmdline.decode('utf-8')
1554*4882a593Smuzhiyun    else:
1555*4882a593Smuzhiyun      self.flags = 0
1556*4882a593Smuzhiyun      self.kernel_cmdline = ''
1557*4882a593Smuzhiyun
1558*4882a593Smuzhiyun  def print_desc(self, o):
1559*4882a593Smuzhiyun    """Print the descriptor.
1560*4882a593Smuzhiyun
1561*4882a593Smuzhiyun    Arguments:
1562*4882a593Smuzhiyun      o: The object to write the output to.
1563*4882a593Smuzhiyun    """
1564*4882a593Smuzhiyun    o.write('    Kernel Cmdline descriptor:\n')
1565*4882a593Smuzhiyun    o.write('      Flags:                 {}\n'.format(self.flags))
1566*4882a593Smuzhiyun    o.write('      Kernel Cmdline:        {}\n'.format(repr(
1567*4882a593Smuzhiyun        self.kernel_cmdline)))
1568*4882a593Smuzhiyun
1569*4882a593Smuzhiyun  def encode(self):
1570*4882a593Smuzhiyun    """Serializes the descriptor.
1571*4882a593Smuzhiyun
1572*4882a593Smuzhiyun    Returns:
1573*4882a593Smuzhiyun      A bytearray() with the descriptor data.
1574*4882a593Smuzhiyun    """
1575*4882a593Smuzhiyun    encoded_str = self.kernel_cmdline.encode('utf-8')
1576*4882a593Smuzhiyun    num_bytes_following = (self.SIZE + len(encoded_str) - 16)
1577*4882a593Smuzhiyun    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
1578*4882a593Smuzhiyun    padding_size = nbf_with_padding - num_bytes_following
1579*4882a593Smuzhiyun    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
1580*4882a593Smuzhiyun                       self.flags, len(encoded_str))
1581*4882a593Smuzhiyun    padding = struct.pack(str(padding_size) + 'x')
1582*4882a593Smuzhiyun    ret = desc + encoded_str + padding
1583*4882a593Smuzhiyun    return bytearray(ret)
1584*4882a593Smuzhiyun
1585*4882a593Smuzhiyun  def verify(self, image_dir, image_ext, expected_chain_partitions_map):
1586*4882a593Smuzhiyun    """Verifies contents of the descriptor - used in verify_image sub-command.
1587*4882a593Smuzhiyun
1588*4882a593Smuzhiyun    Arguments:
1589*4882a593Smuzhiyun      image_dir: The directory of the file being verified.
1590*4882a593Smuzhiyun      image_ext: The extension of the file being verified (e.g. '.img').
1591*4882a593Smuzhiyun      expected_chain_partitions_map: A map from partition name to the
1592*4882a593Smuzhiyun        tuple (rollback_index_location, key_blob).
1593*4882a593Smuzhiyun
1594*4882a593Smuzhiyun    Returns:
1595*4882a593Smuzhiyun      True if the descriptor verifies, False otherwise.
1596*4882a593Smuzhiyun    """
1597*4882a593Smuzhiyun    # Nothing to verify.
1598*4882a593Smuzhiyun    return True
1599*4882a593Smuzhiyun
1600*4882a593Smuzhiyunclass AvbChainPartitionDescriptor(AvbDescriptor):
1601*4882a593Smuzhiyun  """A class for chained partition descriptors.
1602*4882a593Smuzhiyun
1603*4882a593Smuzhiyun  See the |AvbChainPartitionDescriptor| C struct for more information.
1604*4882a593Smuzhiyun
1605*4882a593Smuzhiyun  Attributes:
1606*4882a593Smuzhiyun    rollback_index_location: The rollback index location to use.
1607*4882a593Smuzhiyun    partition_name: Partition name.
1608*4882a593Smuzhiyun    public_key: Bytes for the public key.
1609*4882a593Smuzhiyun  """
1610*4882a593Smuzhiyun
1611*4882a593Smuzhiyun  TAG = 4
1612*4882a593Smuzhiyun  RESERVED = 64
1613*4882a593Smuzhiyun  SIZE = 28 + RESERVED
1614*4882a593Smuzhiyun  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
1615*4882a593Smuzhiyun                   'L'  # rollback_index_location
1616*4882a593Smuzhiyun                   'L'  # partition_name_size (bytes)
1617*4882a593Smuzhiyun                   'L' +  # public_key_size (bytes)
1618*4882a593Smuzhiyun                   str(RESERVED) + 's')  # reserved
1619*4882a593Smuzhiyun
1620*4882a593Smuzhiyun  def __init__(self, data=None):
1621*4882a593Smuzhiyun    """Initializes a new chain partition descriptor.
1622*4882a593Smuzhiyun
1623*4882a593Smuzhiyun    Arguments:
1624*4882a593Smuzhiyun      data: If not None, must be a bytearray of size |SIZE|.
1625*4882a593Smuzhiyun
1626*4882a593Smuzhiyun    Raises:
1627*4882a593Smuzhiyun      LookupError: If the given descriptor is malformed.
1628*4882a593Smuzhiyun    """
1629*4882a593Smuzhiyun    AvbDescriptor.__init__(self, None)
1630*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1631*4882a593Smuzhiyun
1632*4882a593Smuzhiyun    if data:
1633*4882a593Smuzhiyun      (tag, num_bytes_following, self.rollback_index_location,
1634*4882a593Smuzhiyun       partition_name_len,
1635*4882a593Smuzhiyun       public_key_len, _) = struct.unpack(self.FORMAT_STRING, data[0:self.SIZE])
1636*4882a593Smuzhiyun      expected_size = round_to_multiple(
1637*4882a593Smuzhiyun          self.SIZE - 16 + partition_name_len + public_key_len, 8)
1638*4882a593Smuzhiyun      if tag != self.TAG or num_bytes_following != expected_size:
1639*4882a593Smuzhiyun        raise LookupError('Given data does not look like a chain partition '
1640*4882a593Smuzhiyun                          'descriptor.')
1641*4882a593Smuzhiyun      o = 0
1642*4882a593Smuzhiyun      self.partition_name = str(data[(self.SIZE + o):(self.SIZE + o +
1643*4882a593Smuzhiyun                                                      partition_name_len)])
1644*4882a593Smuzhiyun      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
1645*4882a593Smuzhiyun      self.partition_name.decode('utf-8')
1646*4882a593Smuzhiyun      o += partition_name_len
1647*4882a593Smuzhiyun      self.public_key = data[(self.SIZE + o):(self.SIZE + o + public_key_len)]
1648*4882a593Smuzhiyun
1649*4882a593Smuzhiyun    else:
1650*4882a593Smuzhiyun      self.rollback_index_location = 0
1651*4882a593Smuzhiyun      self.partition_name = ''
1652*4882a593Smuzhiyun      self.public_key = bytearray()
1653*4882a593Smuzhiyun
1654*4882a593Smuzhiyun  def print_desc(self, o):
1655*4882a593Smuzhiyun    """Print the descriptor.
1656*4882a593Smuzhiyun
1657*4882a593Smuzhiyun    Arguments:
1658*4882a593Smuzhiyun      o: The object to write the output to.
1659*4882a593Smuzhiyun    """
1660*4882a593Smuzhiyun    o.write('    Chain Partition descriptor:\n')
1661*4882a593Smuzhiyun    o.write('      Partition Name:          {}\n'.format(self.partition_name))
1662*4882a593Smuzhiyun    o.write('      Rollback Index Location: {}\n'.format(
1663*4882a593Smuzhiyun        self.rollback_index_location))
1664*4882a593Smuzhiyun    # Just show the SHA1 of the key, for size reasons.
1665*4882a593Smuzhiyun    hexdig = hashlib.sha1(self.public_key).hexdigest()
1666*4882a593Smuzhiyun    o.write('      Public key (sha1):       {}\n'.format(hexdig))
1667*4882a593Smuzhiyun
1668*4882a593Smuzhiyun  def encode(self):
1669*4882a593Smuzhiyun    """Serializes the descriptor.
1670*4882a593Smuzhiyun
1671*4882a593Smuzhiyun    Returns:
1672*4882a593Smuzhiyun      A bytearray() with the descriptor data.
1673*4882a593Smuzhiyun    """
1674*4882a593Smuzhiyun    encoded_name = self.partition_name.encode('utf-8')
1675*4882a593Smuzhiyun    num_bytes_following = (
1676*4882a593Smuzhiyun        self.SIZE + len(encoded_name) + len(self.public_key) - 16)
1677*4882a593Smuzhiyun    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
1678*4882a593Smuzhiyun    padding_size = nbf_with_padding - num_bytes_following
1679*4882a593Smuzhiyun    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
1680*4882a593Smuzhiyun                       self.rollback_index_location, len(encoded_name),
1681*4882a593Smuzhiyun                       len(self.public_key), self.RESERVED*'\0')
1682*4882a593Smuzhiyun    padding = struct.pack(str(padding_size) + 'x')
1683*4882a593Smuzhiyun    ret = desc + encoded_name + self.public_key + padding
1684*4882a593Smuzhiyun    return bytearray(ret)
1685*4882a593Smuzhiyun
1686*4882a593Smuzhiyun  def verify(self, image_dir, image_ext, expected_chain_partitions_map):
1687*4882a593Smuzhiyun    """Verifies contents of the descriptor - used in verify_image sub-command.
1688*4882a593Smuzhiyun
1689*4882a593Smuzhiyun    Arguments:
1690*4882a593Smuzhiyun      image_dir: The directory of the file being verified.
1691*4882a593Smuzhiyun      image_ext: The extension of the file being verified (e.g. '.img').
1692*4882a593Smuzhiyun      expected_chain_partitions_map: A map from partition name to the
1693*4882a593Smuzhiyun        tuple (rollback_index_location, key_blob).
1694*4882a593Smuzhiyun
1695*4882a593Smuzhiyun    Returns:
1696*4882a593Smuzhiyun      True if the descriptor verifies, False otherwise.
1697*4882a593Smuzhiyun    """
1698*4882a593Smuzhiyun    value = expected_chain_partitions_map.get(self.partition_name)
1699*4882a593Smuzhiyun    if not value:
1700*4882a593Smuzhiyun      sys.stderr.write('No expected chain partition for partition {}. Use '
1701*4882a593Smuzhiyun                       '--expected_chain_partition to specify expected '
1702*4882a593Smuzhiyun                       'contents.\n'.
1703*4882a593Smuzhiyun                       format(self.partition_name))
1704*4882a593Smuzhiyun      return False
1705*4882a593Smuzhiyun    rollback_index_location, pk_blob = value
1706*4882a593Smuzhiyun
1707*4882a593Smuzhiyun    if self.rollback_index_location != rollback_index_location:
1708*4882a593Smuzhiyun      sys.stderr.write('Expected rollback_index_location {} does not '
1709*4882a593Smuzhiyun                       'match {} in descriptor for partition {}\n'.
1710*4882a593Smuzhiyun                       format(rollback_index_location,
1711*4882a593Smuzhiyun                              self.rollback_index_location,
1712*4882a593Smuzhiyun                              self.partition_name))
1713*4882a593Smuzhiyun      return False
1714*4882a593Smuzhiyun
1715*4882a593Smuzhiyun    if self.public_key != pk_blob:
1716*4882a593Smuzhiyun      sys.stderr.write('Expected public key blob does not match public '
1717*4882a593Smuzhiyun                       'key blob in descriptor for partition {}\n'.
1718*4882a593Smuzhiyun                       format(self.partition_name))
1719*4882a593Smuzhiyun      return False
1720*4882a593Smuzhiyun
1721*4882a593Smuzhiyun    print ('{}: Successfully verified chain partition descriptor matches '
1722*4882a593Smuzhiyun           'expected data'.format(self.partition_name))
1723*4882a593Smuzhiyun
1724*4882a593Smuzhiyun    return True
1725*4882a593Smuzhiyun
1726*4882a593SmuzhiyunDESCRIPTOR_CLASSES = [
1727*4882a593Smuzhiyun    AvbPropertyDescriptor, AvbHashtreeDescriptor, AvbHashDescriptor,
1728*4882a593Smuzhiyun    AvbKernelCmdlineDescriptor, AvbChainPartitionDescriptor
1729*4882a593Smuzhiyun]
1730*4882a593Smuzhiyun
1731*4882a593Smuzhiyun
1732*4882a593Smuzhiyundef parse_descriptors(data):
1733*4882a593Smuzhiyun  """Parses a blob of data into descriptors.
1734*4882a593Smuzhiyun
1735*4882a593Smuzhiyun  Arguments:
1736*4882a593Smuzhiyun    data: A bytearray() with encoded descriptors.
1737*4882a593Smuzhiyun
1738*4882a593Smuzhiyun  Returns:
1739*4882a593Smuzhiyun    A list of instances of objects derived from AvbDescriptor. For
1740*4882a593Smuzhiyun    unknown descriptors, the class AvbDescriptor is used.
1741*4882a593Smuzhiyun  """
1742*4882a593Smuzhiyun  o = 0
1743*4882a593Smuzhiyun  ret = []
1744*4882a593Smuzhiyun  while o < len(data):
1745*4882a593Smuzhiyun    tag, nb_following = struct.unpack('!2Q', data[o:o + 16])
1746*4882a593Smuzhiyun    if tag < len(DESCRIPTOR_CLASSES):
1747*4882a593Smuzhiyun      c = DESCRIPTOR_CLASSES[tag]
1748*4882a593Smuzhiyun    else:
1749*4882a593Smuzhiyun      c = AvbDescriptor
1750*4882a593Smuzhiyun    ret.append(c(bytearray(data[o:o + 16 + nb_following])))
1751*4882a593Smuzhiyun    o += 16 + nb_following
1752*4882a593Smuzhiyun  return ret
1753*4882a593Smuzhiyun
1754*4882a593Smuzhiyun
1755*4882a593Smuzhiyunclass AvbFooter(object):
1756*4882a593Smuzhiyun  """A class for parsing and writing footers.
1757*4882a593Smuzhiyun
1758*4882a593Smuzhiyun  Footers are stored at the end of partitions and point to where the
1759*4882a593Smuzhiyun  AvbVBMeta blob is located. They also contain the original size of
1760*4882a593Smuzhiyun  the image before AVB information was added.
1761*4882a593Smuzhiyun
1762*4882a593Smuzhiyun  Attributes:
1763*4882a593Smuzhiyun    magic: Magic for identifying the footer, see |MAGIC|.
1764*4882a593Smuzhiyun    version_major: The major version of avbtool that wrote the footer.
1765*4882a593Smuzhiyun    version_minor: The minor version of avbtool that wrote the footer.
1766*4882a593Smuzhiyun    original_image_size: Original image size.
1767*4882a593Smuzhiyun    vbmeta_offset: Offset of where the AvbVBMeta blob is stored.
1768*4882a593Smuzhiyun    vbmeta_size: Size of the AvbVBMeta blob.
1769*4882a593Smuzhiyun  """
1770*4882a593Smuzhiyun
1771*4882a593Smuzhiyun  MAGIC = 'AVBf'
1772*4882a593Smuzhiyun  SIZE = 64
1773*4882a593Smuzhiyun  RESERVED = 28
1774*4882a593Smuzhiyun  FOOTER_VERSION_MAJOR = AVB_FOOTER_VERSION_MAJOR
1775*4882a593Smuzhiyun  FOOTER_VERSION_MINOR = AVB_FOOTER_VERSION_MINOR
1776*4882a593Smuzhiyun  FORMAT_STRING = ('!4s2L'  # magic, 2 x version.
1777*4882a593Smuzhiyun                   'Q'  # Original image size.
1778*4882a593Smuzhiyun                   'Q'  # Offset of VBMeta blob.
1779*4882a593Smuzhiyun                   'Q' +  # Size of VBMeta blob.
1780*4882a593Smuzhiyun                   str(RESERVED) + 'x')  # padding for reserved bytes
1781*4882a593Smuzhiyun
1782*4882a593Smuzhiyun  def __init__(self, data=None):
1783*4882a593Smuzhiyun    """Initializes a new footer object.
1784*4882a593Smuzhiyun
1785*4882a593Smuzhiyun    Arguments:
1786*4882a593Smuzhiyun      data: If not None, must be a bytearray of size 4096.
1787*4882a593Smuzhiyun
1788*4882a593Smuzhiyun    Raises:
1789*4882a593Smuzhiyun      LookupError: If the given footer is malformed.
1790*4882a593Smuzhiyun      struct.error: If the given data has no footer.
1791*4882a593Smuzhiyun    """
1792*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1793*4882a593Smuzhiyun
1794*4882a593Smuzhiyun    if data:
1795*4882a593Smuzhiyun      (self.magic, self.version_major, self.version_minor,
1796*4882a593Smuzhiyun       self.original_image_size, self.vbmeta_offset,
1797*4882a593Smuzhiyun       self.vbmeta_size) = struct.unpack(self.FORMAT_STRING, data)
1798*4882a593Smuzhiyun      if self.magic != self.MAGIC:
1799*4882a593Smuzhiyun        raise LookupError('Given data does not look like a AVB footer.')
1800*4882a593Smuzhiyun    else:
1801*4882a593Smuzhiyun      self.magic = self.MAGIC
1802*4882a593Smuzhiyun      self.version_major = self.FOOTER_VERSION_MAJOR
1803*4882a593Smuzhiyun      self.version_minor = self.FOOTER_VERSION_MINOR
1804*4882a593Smuzhiyun      self.original_image_size = 0
1805*4882a593Smuzhiyun      self.vbmeta_offset = 0
1806*4882a593Smuzhiyun      self.vbmeta_size = 0
1807*4882a593Smuzhiyun
1808*4882a593Smuzhiyun  def encode(self):
1809*4882a593Smuzhiyun    """Gets a string representing the binary encoding of the footer.
1810*4882a593Smuzhiyun
1811*4882a593Smuzhiyun    Returns:
1812*4882a593Smuzhiyun      A bytearray() with a binary representation of the footer.
1813*4882a593Smuzhiyun    """
1814*4882a593Smuzhiyun    return struct.pack(self.FORMAT_STRING, self.magic, self.version_major,
1815*4882a593Smuzhiyun                       self.version_minor, self.original_image_size,
1816*4882a593Smuzhiyun                       self.vbmeta_offset, self.vbmeta_size)
1817*4882a593Smuzhiyun
1818*4882a593Smuzhiyun
1819*4882a593Smuzhiyunclass AvbVBMetaHeader(object):
1820*4882a593Smuzhiyun  """A class for parsing and writing AVB vbmeta images.
1821*4882a593Smuzhiyun
1822*4882a593Smuzhiyun  Attributes:
1823*4882a593Smuzhiyun    The attributes correspond to the |AvbVBMetaHeader| struct
1824*4882a593Smuzhiyun    defined in avb_vbmeta_header.h.
1825*4882a593Smuzhiyun  """
1826*4882a593Smuzhiyun
1827*4882a593Smuzhiyun  SIZE = 256
1828*4882a593Smuzhiyun
1829*4882a593Smuzhiyun  # Keep in sync with |reserved0| and |reserved| field of
1830*4882a593Smuzhiyun  # |AvbVBMetaImageHeader|.
1831*4882a593Smuzhiyun  RESERVED0 = 4
1832*4882a593Smuzhiyun  RESERVED = 80
1833*4882a593Smuzhiyun
1834*4882a593Smuzhiyun  # Keep in sync with |AvbVBMetaImageHeader|.
1835*4882a593Smuzhiyun  FORMAT_STRING = ('!4s2L'  # magic, 2 x version
1836*4882a593Smuzhiyun                   '2Q'  # 2 x block size
1837*4882a593Smuzhiyun                   'L'  # algorithm type
1838*4882a593Smuzhiyun                   '2Q'  # offset, size (hash)
1839*4882a593Smuzhiyun                   '2Q'  # offset, size (signature)
1840*4882a593Smuzhiyun                   '2Q'  # offset, size (public key)
1841*4882a593Smuzhiyun                   '2Q'  # offset, size (public key metadata)
1842*4882a593Smuzhiyun                   '2Q'  # offset, size (descriptors)
1843*4882a593Smuzhiyun                   'Q'  # rollback_index
1844*4882a593Smuzhiyun                   'L' +  # flags
1845*4882a593Smuzhiyun                   str(RESERVED0) + 'x' +  # padding for reserved bytes
1846*4882a593Smuzhiyun                   '47sx' +  # NUL-terminated release string
1847*4882a593Smuzhiyun                   str(RESERVED) + 'x')  # padding for reserved bytes
1848*4882a593Smuzhiyun
1849*4882a593Smuzhiyun  def __init__(self, data=None):
1850*4882a593Smuzhiyun    """Initializes a new header object.
1851*4882a593Smuzhiyun
1852*4882a593Smuzhiyun    Arguments:
1853*4882a593Smuzhiyun      data: If not None, must be a bytearray of size 8192.
1854*4882a593Smuzhiyun
1855*4882a593Smuzhiyun    Raises:
1856*4882a593Smuzhiyun      Exception: If the given data is malformed.
1857*4882a593Smuzhiyun    """
1858*4882a593Smuzhiyun    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE
1859*4882a593Smuzhiyun
1860*4882a593Smuzhiyun    if data:
1861*4882a593Smuzhiyun      (self.magic, self.required_libavb_version_major,
1862*4882a593Smuzhiyun       self.required_libavb_version_minor,
1863*4882a593Smuzhiyun       self.authentication_data_block_size, self.auxiliary_data_block_size,
1864*4882a593Smuzhiyun       self.algorithm_type, self.hash_offset, self.hash_size,
1865*4882a593Smuzhiyun       self.signature_offset, self.signature_size, self.public_key_offset,
1866*4882a593Smuzhiyun       self.public_key_size, self.public_key_metadata_offset,
1867*4882a593Smuzhiyun       self.public_key_metadata_size, self.descriptors_offset,
1868*4882a593Smuzhiyun       self.descriptors_size,
1869*4882a593Smuzhiyun       self.rollback_index,
1870*4882a593Smuzhiyun       self.flags,
1871*4882a593Smuzhiyun       self.release_string) = struct.unpack(self.FORMAT_STRING, data)
1872*4882a593Smuzhiyun      # Nuke NUL-bytes at the end of the string.
1873*4882a593Smuzhiyun      if self.magic != 'AVB0':
1874*4882a593Smuzhiyun        raise AvbError('Given image does not look like a vbmeta image.')
1875*4882a593Smuzhiyun    else:
1876*4882a593Smuzhiyun      self.magic = 'AVB0'
1877*4882a593Smuzhiyun      # Start by just requiring version 1.0. Code that adds features
1878*4882a593Smuzhiyun      # in a future version can use bump_required_libavb_version_minor() to
1879*4882a593Smuzhiyun      # bump the minor.
1880*4882a593Smuzhiyun      self.required_libavb_version_major = AVB_VERSION_MAJOR
1881*4882a593Smuzhiyun      self.required_libavb_version_minor = 0
1882*4882a593Smuzhiyun      self.authentication_data_block_size = 0
1883*4882a593Smuzhiyun      self.auxiliary_data_block_size = 0
1884*4882a593Smuzhiyun      self.algorithm_type = 0
1885*4882a593Smuzhiyun      self.hash_offset = 0
1886*4882a593Smuzhiyun      self.hash_size = 0
1887*4882a593Smuzhiyun      self.signature_offset = 0
1888*4882a593Smuzhiyun      self.signature_size = 0
1889*4882a593Smuzhiyun      self.public_key_offset = 0
1890*4882a593Smuzhiyun      self.public_key_size = 0
1891*4882a593Smuzhiyun      self.public_key_metadata_offset = 0
1892*4882a593Smuzhiyun      self.public_key_metadata_size = 0
1893*4882a593Smuzhiyun      self.descriptors_offset = 0
1894*4882a593Smuzhiyun      self.descriptors_size = 0
1895*4882a593Smuzhiyun      self.rollback_index = 0
1896*4882a593Smuzhiyun      self.flags = 0
1897*4882a593Smuzhiyun      self.release_string = get_release_string()
1898*4882a593Smuzhiyun
1899*4882a593Smuzhiyun  def bump_required_libavb_version_minor(self, minor):
1900*4882a593Smuzhiyun    """Function to bump required_libavb_version_minor.
1901*4882a593Smuzhiyun
1902*4882a593Smuzhiyun    Call this when writing data that requires a specific libavb
1903*4882a593Smuzhiyun    version to parse it.
1904*4882a593Smuzhiyun
1905*4882a593Smuzhiyun    Arguments:
1906*4882a593Smuzhiyun      minor: The minor version of libavb that has support for the feature.
1907*4882a593Smuzhiyun    """
1908*4882a593Smuzhiyun    self.required_libavb_version_minor = (
1909*4882a593Smuzhiyun        max(self.required_libavb_version_minor, minor))
1910*4882a593Smuzhiyun
1911*4882a593Smuzhiyun  def save(self, output):
1912*4882a593Smuzhiyun    """Serializes the header (256 bytes) to disk.
1913*4882a593Smuzhiyun
1914*4882a593Smuzhiyun    Arguments:
1915*4882a593Smuzhiyun      output: The object to write the output to.
1916*4882a593Smuzhiyun    """
1917*4882a593Smuzhiyun    output.write(struct.pack(
1918*4882a593Smuzhiyun        self.FORMAT_STRING, self.magic, self.required_libavb_version_major,
1919*4882a593Smuzhiyun        self.required_libavb_version_minor, self.authentication_data_block_size,
1920*4882a593Smuzhiyun        self.auxiliary_data_block_size, self.algorithm_type, self.hash_offset,
1921*4882a593Smuzhiyun        self.hash_size, self.signature_offset, self.signature_size,
1922*4882a593Smuzhiyun        self.public_key_offset, self.public_key_size,
1923*4882a593Smuzhiyun        self.public_key_metadata_offset, self.public_key_metadata_size,
1924*4882a593Smuzhiyun        self.descriptors_offset, self.descriptors_size, self.rollback_index,
1925*4882a593Smuzhiyun        self.flags, self.release_string))
1926*4882a593Smuzhiyun
1927*4882a593Smuzhiyun  def encode(self):
1928*4882a593Smuzhiyun    """Serializes the header (256) to a bytearray().
1929*4882a593Smuzhiyun
1930*4882a593Smuzhiyun    Returns:
1931*4882a593Smuzhiyun      A bytearray() with the encoded header.
1932*4882a593Smuzhiyun    """
1933*4882a593Smuzhiyun    return struct.pack(self.FORMAT_STRING, self.magic,
1934*4882a593Smuzhiyun                       self.required_libavb_version_major,
1935*4882a593Smuzhiyun                       self.required_libavb_version_minor,
1936*4882a593Smuzhiyun                       self.authentication_data_block_size,
1937*4882a593Smuzhiyun                       self.auxiliary_data_block_size, self.algorithm_type,
1938*4882a593Smuzhiyun                       self.hash_offset, self.hash_size, self.signature_offset,
1939*4882a593Smuzhiyun                       self.signature_size, self.public_key_offset,
1940*4882a593Smuzhiyun                       self.public_key_size, self.public_key_metadata_offset,
1941*4882a593Smuzhiyun                       self.public_key_metadata_size, self.descriptors_offset,
1942*4882a593Smuzhiyun                       self.descriptors_size, self.rollback_index, self.flags,
1943*4882a593Smuzhiyun                       self.release_string)
1944*4882a593Smuzhiyun
1945*4882a593Smuzhiyun
1946*4882a593Smuzhiyunclass Avb(object):
1947*4882a593Smuzhiyun  """Business logic for avbtool command-line tool."""
1948*4882a593Smuzhiyun
1949*4882a593Smuzhiyun  # Keep in sync with avb_ab_flow.h.
1950*4882a593Smuzhiyun  AB_FORMAT_NO_CRC = '!4sBB2xBBBxBBBx12x'
1951*4882a593Smuzhiyun  AB_MAGIC = '\0AB0'
1952*4882a593Smuzhiyun  AB_MAJOR_VERSION = 1
1953*4882a593Smuzhiyun  AB_MINOR_VERSION = 0
1954*4882a593Smuzhiyun  AB_MISC_METADATA_OFFSET = 2048
1955*4882a593Smuzhiyun
1956*4882a593Smuzhiyun  # Constants for maximum metadata size. These are used to give
1957*4882a593Smuzhiyun  # meaningful errors if the value passed in via --partition_size is
1958*4882a593Smuzhiyun  # too small and when --calc_max_image_size is used. We use
1959*4882a593Smuzhiyun  # conservative figures.
1960*4882a593Smuzhiyun  MAX_VBMETA_SIZE = 64 * 1024
1961*4882a593Smuzhiyun  MAX_FOOTER_SIZE = 4096
1962*4882a593Smuzhiyun
1963*4882a593Smuzhiyun  def erase_footer(self, image_filename, keep_hashtree):
1964*4882a593Smuzhiyun    """Implements the 'erase_footer' command.
1965*4882a593Smuzhiyun
1966*4882a593Smuzhiyun    Arguments:
1967*4882a593Smuzhiyun      image_filename: File to erase a footer from.
1968*4882a593Smuzhiyun      keep_hashtree: If True, keep the hashtree and FEC around.
1969*4882a593Smuzhiyun
1970*4882a593Smuzhiyun    Raises:
1971*4882a593Smuzhiyun      AvbError: If there's no footer in the image.
1972*4882a593Smuzhiyun    """
1973*4882a593Smuzhiyun
1974*4882a593Smuzhiyun    image = ImageHandler(image_filename)
1975*4882a593Smuzhiyun
1976*4882a593Smuzhiyun    (footer, _, descriptors, _) = self._parse_image(image)
1977*4882a593Smuzhiyun
1978*4882a593Smuzhiyun    if not footer:
1979*4882a593Smuzhiyun      raise AvbError('Given image does not have a footer.')
1980*4882a593Smuzhiyun
1981*4882a593Smuzhiyun    new_image_size = None
1982*4882a593Smuzhiyun    if not keep_hashtree:
1983*4882a593Smuzhiyun      new_image_size = footer.original_image_size
1984*4882a593Smuzhiyun    else:
1985*4882a593Smuzhiyun      # If requested to keep the hashtree, search for a hashtree
1986*4882a593Smuzhiyun      # descriptor to figure out the location and size of the hashtree
1987*4882a593Smuzhiyun      # and FEC.
1988*4882a593Smuzhiyun      for desc in descriptors:
1989*4882a593Smuzhiyun        if isinstance(desc, AvbHashtreeDescriptor):
1990*4882a593Smuzhiyun          # The hashtree is always just following the main data so the
1991*4882a593Smuzhiyun          # new size is easily derived.
1992*4882a593Smuzhiyun          new_image_size = desc.tree_offset + desc.tree_size
1993*4882a593Smuzhiyun          # If the image has FEC codes, also keep those.
1994*4882a593Smuzhiyun          if desc.fec_offset > 0:
1995*4882a593Smuzhiyun            fec_end = desc.fec_offset + desc.fec_size
1996*4882a593Smuzhiyun            new_image_size = max(new_image_size, fec_end)
1997*4882a593Smuzhiyun          break
1998*4882a593Smuzhiyun      if not new_image_size:
1999*4882a593Smuzhiyun        raise AvbError('Requested to keep hashtree but no hashtree '
2000*4882a593Smuzhiyun                       'descriptor was found.')
2001*4882a593Smuzhiyun
2002*4882a593Smuzhiyun    # And cut...
2003*4882a593Smuzhiyun    image.truncate(new_image_size)
2004*4882a593Smuzhiyun
2005*4882a593Smuzhiyun  def resize_image(self, image_filename, partition_size):
2006*4882a593Smuzhiyun    """Implements the 'resize_image' command.
2007*4882a593Smuzhiyun
2008*4882a593Smuzhiyun    Arguments:
2009*4882a593Smuzhiyun      image_filename: File with footer to resize.
2010*4882a593Smuzhiyun      partition_size: The new size of the image.
2011*4882a593Smuzhiyun
2012*4882a593Smuzhiyun    Raises:
2013*4882a593Smuzhiyun      AvbError: If there's no footer in the image.
2014*4882a593Smuzhiyun    """
2015*4882a593Smuzhiyun
2016*4882a593Smuzhiyun    image = ImageHandler(image_filename)
2017*4882a593Smuzhiyun
2018*4882a593Smuzhiyun    if partition_size % image.block_size != 0:
2019*4882a593Smuzhiyun      raise AvbError('Partition size of {} is not a multiple of the image '
2020*4882a593Smuzhiyun                     'block size {}.'.format(partition_size,
2021*4882a593Smuzhiyun                                             image.block_size))
2022*4882a593Smuzhiyun
2023*4882a593Smuzhiyun    (footer, vbmeta_header, descriptors, _) = self._parse_image(image)
2024*4882a593Smuzhiyun
2025*4882a593Smuzhiyun    if not footer:
2026*4882a593Smuzhiyun      raise AvbError('Given image does not have a footer.')
2027*4882a593Smuzhiyun
2028*4882a593Smuzhiyun    # The vbmeta blob is always at the end of the data so resizing an
2029*4882a593Smuzhiyun    # image amounts to just moving the footer around.
2030*4882a593Smuzhiyun
2031*4882a593Smuzhiyun    vbmeta_end_offset = footer.vbmeta_offset + footer.vbmeta_size
2032*4882a593Smuzhiyun    if vbmeta_end_offset % image.block_size != 0:
2033*4882a593Smuzhiyun      vbmeta_end_offset += image.block_size - (vbmeta_end_offset % image.block_size)
2034*4882a593Smuzhiyun
2035*4882a593Smuzhiyun    if partition_size < vbmeta_end_offset + 1*image.block_size:
2036*4882a593Smuzhiyun      raise AvbError('Requested size of {} is too small for an image '
2037*4882a593Smuzhiyun                     'of size {}.'
2038*4882a593Smuzhiyun                     .format(partition_size,
2039*4882a593Smuzhiyun                             vbmeta_end_offset + 1*image.block_size))
2040*4882a593Smuzhiyun
2041*4882a593Smuzhiyun    # Cut at the end of the vbmeta blob and insert a DONT_CARE chunk
2042*4882a593Smuzhiyun    # with enough bytes such that the final Footer block is at the end
2043*4882a593Smuzhiyun    # of partition_size.
2044*4882a593Smuzhiyun    image.truncate(vbmeta_end_offset)
2045*4882a593Smuzhiyun    image.append_dont_care(partition_size - vbmeta_end_offset -
2046*4882a593Smuzhiyun                           1*image.block_size)
2047*4882a593Smuzhiyun
2048*4882a593Smuzhiyun    # Just reuse the same footer - only difference is that we're
2049*4882a593Smuzhiyun    # writing it in a different place.
2050*4882a593Smuzhiyun    footer_blob = footer.encode()
2051*4882a593Smuzhiyun    footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
2052*4882a593Smuzhiyun                                footer_blob)
2053*4882a593Smuzhiyun    image.append_raw(footer_blob_with_padding)
2054*4882a593Smuzhiyun
2055*4882a593Smuzhiyun  def set_ab_metadata(self, misc_image, slot_data):
2056*4882a593Smuzhiyun    """Implements the 'set_ab_metadata' command.
2057*4882a593Smuzhiyun
2058*4882a593Smuzhiyun    The |slot_data| argument must be of the form 'A_priority:A_tries_remaining:
2059*4882a593Smuzhiyun    A_successful_boot:B_priority:B_tries_remaining:B_successful_boot'.
2060*4882a593Smuzhiyun
2061*4882a593Smuzhiyun    Arguments:
2062*4882a593Smuzhiyun      misc_image: The misc image to write to.
2063*4882a593Smuzhiyun      slot_data: Slot data as a string
2064*4882a593Smuzhiyun
2065*4882a593Smuzhiyun    Raises:
2066*4882a593Smuzhiyun      AvbError: If slot data is malformed.
2067*4882a593Smuzhiyun    """
2068*4882a593Smuzhiyun    tokens = slot_data.split(':')
2069*4882a593Smuzhiyun    if len(tokens) != 6:
2070*4882a593Smuzhiyun      raise AvbError('Malformed slot data "{}".'.format(slot_data))
2071*4882a593Smuzhiyun    a_priority = int(tokens[0])
2072*4882a593Smuzhiyun    a_tries_remaining = int(tokens[1])
2073*4882a593Smuzhiyun    a_success = True if int(tokens[2]) != 0 else False
2074*4882a593Smuzhiyun    b_priority = int(tokens[3])
2075*4882a593Smuzhiyun    b_tries_remaining = int(tokens[4])
2076*4882a593Smuzhiyun    b_success = True if int(tokens[5]) != 0 else False
2077*4882a593Smuzhiyun
2078*4882a593Smuzhiyun    ab_data_no_crc = struct.pack(self.AB_FORMAT_NO_CRC,
2079*4882a593Smuzhiyun                                 self.AB_MAGIC,
2080*4882a593Smuzhiyun                                 self.AB_MAJOR_VERSION, self.AB_MINOR_VERSION,
2081*4882a593Smuzhiyun                                 a_priority, a_tries_remaining, a_success,
2082*4882a593Smuzhiyun                                 b_priority, b_tries_remaining, b_success)
2083*4882a593Smuzhiyun    # Force CRC to be unsigned, see https://bugs.python.org/issue4903 for why.
2084*4882a593Smuzhiyun    crc_value = binascii.crc32(ab_data_no_crc) & 0xffffffff
2085*4882a593Smuzhiyun    ab_data = ab_data_no_crc + struct.pack('!I', crc_value)
2086*4882a593Smuzhiyun    misc_image.seek(self.AB_MISC_METADATA_OFFSET)
2087*4882a593Smuzhiyun    misc_image.write(ab_data)
2088*4882a593Smuzhiyun
2089*4882a593Smuzhiyun  def info_image(self, image_filename, output):
2090*4882a593Smuzhiyun    """Implements the 'info_image' command.
2091*4882a593Smuzhiyun
2092*4882a593Smuzhiyun    Arguments:
2093*4882a593Smuzhiyun      image_filename: Image file to get information from (file object).
2094*4882a593Smuzhiyun      output: Output file to write human-readable information to (file object).
2095*4882a593Smuzhiyun    """
2096*4882a593Smuzhiyun
2097*4882a593Smuzhiyun    image = ImageHandler(image_filename)
2098*4882a593Smuzhiyun
2099*4882a593Smuzhiyun    o = output
2100*4882a593Smuzhiyun
2101*4882a593Smuzhiyun    (footer, header, descriptors, image_size) = self._parse_image(image)
2102*4882a593Smuzhiyun
2103*4882a593Smuzhiyun    if footer:
2104*4882a593Smuzhiyun      o.write('Footer version:           {}.{}\n'.format(footer.version_major,
2105*4882a593Smuzhiyun                                                         footer.version_minor))
2106*4882a593Smuzhiyun      o.write('Image size:               {} bytes\n'.format(image_size))
2107*4882a593Smuzhiyun      o.write('Original image size:      {} bytes\n'.format(
2108*4882a593Smuzhiyun          footer.original_image_size))
2109*4882a593Smuzhiyun      o.write('VBMeta offset:            {}\n'.format(footer.vbmeta_offset))
2110*4882a593Smuzhiyun      o.write('VBMeta size:              {} bytes\n'.format(footer.vbmeta_size))
2111*4882a593Smuzhiyun      o.write('--\n')
2112*4882a593Smuzhiyun
2113*4882a593Smuzhiyun    (alg_name, _) = lookup_algorithm_by_type(header.algorithm_type)
2114*4882a593Smuzhiyun
2115*4882a593Smuzhiyun    o.write('Minimum libavb version:   {}.{}{}\n'.format(
2116*4882a593Smuzhiyun        header.required_libavb_version_major,
2117*4882a593Smuzhiyun        header.required_libavb_version_minor,
2118*4882a593Smuzhiyun        ' (Sparse)' if image.is_sparse else ''))
2119*4882a593Smuzhiyun    o.write('Header Block:             {} bytes\n'.format(AvbVBMetaHeader.SIZE))
2120*4882a593Smuzhiyun    o.write('Authentication Block:     {} bytes\n'.format(
2121*4882a593Smuzhiyun        header.authentication_data_block_size))
2122*4882a593Smuzhiyun    o.write('Auxiliary Block:          {} bytes\n'.format(
2123*4882a593Smuzhiyun        header.auxiliary_data_block_size))
2124*4882a593Smuzhiyun    o.write('Algorithm:                {}\n'.format(alg_name))
2125*4882a593Smuzhiyun    o.write('Rollback Index:           {}\n'.format(header.rollback_index))
2126*4882a593Smuzhiyun    o.write('Flags:                    {}\n'.format(header.flags))
2127*4882a593Smuzhiyun    o.write('Release String:           \'{}\'\n'.format(
2128*4882a593Smuzhiyun        header.release_string.rstrip('\0')))
2129*4882a593Smuzhiyun
2130*4882a593Smuzhiyun    # Print descriptors.
2131*4882a593Smuzhiyun    num_printed = 0
2132*4882a593Smuzhiyun    o.write('Descriptors:\n')
2133*4882a593Smuzhiyun    for desc in descriptors:
2134*4882a593Smuzhiyun      desc.print_desc(o)
2135*4882a593Smuzhiyun      num_printed += 1
2136*4882a593Smuzhiyun    if num_printed == 0:
2137*4882a593Smuzhiyun      o.write('    (none)\n')
2138*4882a593Smuzhiyun
2139*4882a593Smuzhiyun  def verify_image(self, image_filename, key_path, expected_chain_partitions):
2140*4882a593Smuzhiyun    """Implements the 'verify_image' command.
2141*4882a593Smuzhiyun
2142*4882a593Smuzhiyun    Arguments:
2143*4882a593Smuzhiyun      image_filename: Image file to get information from (file object).
2144*4882a593Smuzhiyun      key_path: None or check that embedded public key matches key at given path.
2145*4882a593Smuzhiyun      expected_chain_partitions: List of chain partitions to check or None.
2146*4882a593Smuzhiyun    """
2147*4882a593Smuzhiyun
2148*4882a593Smuzhiyun    expected_chain_partitions_map = {}
2149*4882a593Smuzhiyun    if expected_chain_partitions:
2150*4882a593Smuzhiyun      used_locations = {}
2151*4882a593Smuzhiyun      for cp in expected_chain_partitions:
2152*4882a593Smuzhiyun        cp_tokens = cp.split(':')
2153*4882a593Smuzhiyun        if len(cp_tokens) != 3:
2154*4882a593Smuzhiyun          raise AvbError('Malformed chained partition "{}".'.format(cp))
2155*4882a593Smuzhiyun        partition_name = cp_tokens[0]
2156*4882a593Smuzhiyun        rollback_index_location = int(cp_tokens[1])
2157*4882a593Smuzhiyun        file_path = cp_tokens[2]
2158*4882a593Smuzhiyun        pk_blob = open(file_path).read()
2159*4882a593Smuzhiyun        expected_chain_partitions_map[partition_name] = (rollback_index_location, pk_blob)
2160*4882a593Smuzhiyun
2161*4882a593Smuzhiyun    image_dir = os.path.dirname(image_filename)
2162*4882a593Smuzhiyun    image_ext = os.path.splitext(image_filename)[1]
2163*4882a593Smuzhiyun
2164*4882a593Smuzhiyun    key_blob = None
2165*4882a593Smuzhiyun    if key_path:
2166*4882a593Smuzhiyun      print 'Verifying image {} using key at {}'.format(image_filename, key_path)
2167*4882a593Smuzhiyun      key_blob = encode_rsa_key(key_path)
2168*4882a593Smuzhiyun    else:
2169*4882a593Smuzhiyun      print 'Verifying image {} using embedded public key'.format(image_filename)
2170*4882a593Smuzhiyun
2171*4882a593Smuzhiyun    image = ImageHandler(image_filename)
2172*4882a593Smuzhiyun    (footer, header, descriptors, image_size) = self._parse_image(image)
2173*4882a593Smuzhiyun    offset = 0
2174*4882a593Smuzhiyun    if footer:
2175*4882a593Smuzhiyun      offset = footer.vbmeta_offset
2176*4882a593Smuzhiyun    size = (header.SIZE + header.authentication_data_block_size +
2177*4882a593Smuzhiyun            header.auxiliary_data_block_size)
2178*4882a593Smuzhiyun    image.seek(offset)
2179*4882a593Smuzhiyun    vbmeta_blob = image.read(size)
2180*4882a593Smuzhiyun    h = AvbVBMetaHeader(vbmeta_blob[0:AvbVBMetaHeader.SIZE])
2181*4882a593Smuzhiyun    alg_name, _ = lookup_algorithm_by_type(header.algorithm_type)
2182*4882a593Smuzhiyun    if not verify_vbmeta_signature(header, vbmeta_blob):
2183*4882a593Smuzhiyun      raise AvbError('Signature check failed for {} vbmeta struct {}'
2184*4882a593Smuzhiyun                     .format(alg_name, image_filename))
2185*4882a593Smuzhiyun
2186*4882a593Smuzhiyun    if key_blob:
2187*4882a593Smuzhiyun      # The embedded public key is in the auxiliary block at an offset.
2188*4882a593Smuzhiyun      key_offset = AvbVBMetaHeader.SIZE
2189*4882a593Smuzhiyun      key_offset += h.authentication_data_block_size
2190*4882a593Smuzhiyun      key_offset += h.public_key_offset
2191*4882a593Smuzhiyun      key_blob_in_vbmeta = vbmeta_blob[key_offset:key_offset + h.public_key_size]
2192*4882a593Smuzhiyun      if key_blob != key_blob_in_vbmeta:
2193*4882a593Smuzhiyun        raise AvbError('Embedded public key does not match given key.')
2194*4882a593Smuzhiyun
2195*4882a593Smuzhiyun    if footer:
2196*4882a593Smuzhiyun      print ('vbmeta: Successfully verified footer and {} vbmeta struct in {}'
2197*4882a593Smuzhiyun             .format(alg_name, image_filename))
2198*4882a593Smuzhiyun    else:
2199*4882a593Smuzhiyun      print ('vbmeta: Successfully verified {} vbmeta struct in {}'
2200*4882a593Smuzhiyun             .format(alg_name, image_filename))
2201*4882a593Smuzhiyun
2202*4882a593Smuzhiyun    for desc in descriptors:
2203*4882a593Smuzhiyun      if not desc.verify(image_dir, image_ext, expected_chain_partitions_map):
2204*4882a593Smuzhiyun        raise AvbError('Error verifying descriptor.')
2205*4882a593Smuzhiyun
2206*4882a593Smuzhiyun
2207*4882a593Smuzhiyun  def _parse_image(self, image):
2208*4882a593Smuzhiyun    """Gets information about an image.
2209*4882a593Smuzhiyun
2210*4882a593Smuzhiyun    The image can either be a vbmeta or an image with a footer.
2211*4882a593Smuzhiyun
2212*4882a593Smuzhiyun    Arguments:
2213*4882a593Smuzhiyun      image: An ImageHandler (vbmeta or footer) with a hashtree descriptor.
2214*4882a593Smuzhiyun
2215*4882a593Smuzhiyun    Returns:
2216*4882a593Smuzhiyun      A tuple where the first argument is a AvbFooter (None if there
2217*4882a593Smuzhiyun      is no footer on the image), the second argument is a
2218*4882a593Smuzhiyun      AvbVBMetaHeader, the third argument is a list of
2219*4882a593Smuzhiyun      AvbDescriptor-derived instances, and the fourth argument is the
2220*4882a593Smuzhiyun      size of |image|.
2221*4882a593Smuzhiyun    """
2222*4882a593Smuzhiyun    assert isinstance(image, ImageHandler)
2223*4882a593Smuzhiyun    footer = None
2224*4882a593Smuzhiyun    image.seek(image.image_size - AvbFooter.SIZE)
2225*4882a593Smuzhiyun    try:
2226*4882a593Smuzhiyun      footer = AvbFooter(image.read(AvbFooter.SIZE))
2227*4882a593Smuzhiyun    except (LookupError, struct.error):
2228*4882a593Smuzhiyun      # Nope, just seek back to the start.
2229*4882a593Smuzhiyun      image.seek(0)
2230*4882a593Smuzhiyun
2231*4882a593Smuzhiyun    vbmeta_offset = 0
2232*4882a593Smuzhiyun    if footer:
2233*4882a593Smuzhiyun      vbmeta_offset = footer.vbmeta_offset
2234*4882a593Smuzhiyun
2235*4882a593Smuzhiyun    image.seek(vbmeta_offset)
2236*4882a593Smuzhiyun    h = AvbVBMetaHeader(image.read(AvbVBMetaHeader.SIZE))
2237*4882a593Smuzhiyun
2238*4882a593Smuzhiyun    auth_block_offset = vbmeta_offset + AvbVBMetaHeader.SIZE
2239*4882a593Smuzhiyun    aux_block_offset = auth_block_offset + h.authentication_data_block_size
2240*4882a593Smuzhiyun    desc_start_offset = aux_block_offset + h.descriptors_offset
2241*4882a593Smuzhiyun    image.seek(desc_start_offset)
2242*4882a593Smuzhiyun    descriptors = parse_descriptors(image.read(h.descriptors_size))
2243*4882a593Smuzhiyun
2244*4882a593Smuzhiyun    return footer, h, descriptors, image.image_size
2245*4882a593Smuzhiyun
2246*4882a593Smuzhiyun  def _load_vbmeta_blob(self, image):
2247*4882a593Smuzhiyun    """Gets the vbmeta struct and associated sections.
2248*4882a593Smuzhiyun
2249*4882a593Smuzhiyun    The image can either be a vbmeta.img or an image with a footer.
2250*4882a593Smuzhiyun
2251*4882a593Smuzhiyun    Arguments:
2252*4882a593Smuzhiyun      image: An ImageHandler (vbmeta or footer).
2253*4882a593Smuzhiyun
2254*4882a593Smuzhiyun    Returns:
2255*4882a593Smuzhiyun      A blob with the vbmeta struct and other sections.
2256*4882a593Smuzhiyun    """
2257*4882a593Smuzhiyun    assert isinstance(image, ImageHandler)
2258*4882a593Smuzhiyun    footer = None
2259*4882a593Smuzhiyun    image.seek(image.image_size - AvbFooter.SIZE)
2260*4882a593Smuzhiyun    try:
2261*4882a593Smuzhiyun      footer = AvbFooter(image.read(AvbFooter.SIZE))
2262*4882a593Smuzhiyun    except (LookupError, struct.error):
2263*4882a593Smuzhiyun      # Nope, just seek back to the start.
2264*4882a593Smuzhiyun      image.seek(0)
2265*4882a593Smuzhiyun
2266*4882a593Smuzhiyun    vbmeta_offset = 0
2267*4882a593Smuzhiyun    if footer:
2268*4882a593Smuzhiyun      vbmeta_offset = footer.vbmeta_offset
2269*4882a593Smuzhiyun
2270*4882a593Smuzhiyun    image.seek(vbmeta_offset)
2271*4882a593Smuzhiyun    h = AvbVBMetaHeader(image.read(AvbVBMetaHeader.SIZE))
2272*4882a593Smuzhiyun
2273*4882a593Smuzhiyun    image.seek(vbmeta_offset)
2274*4882a593Smuzhiyun    data_size = AvbVBMetaHeader.SIZE
2275*4882a593Smuzhiyun    data_size += h.authentication_data_block_size
2276*4882a593Smuzhiyun    data_size += h.auxiliary_data_block_size
2277*4882a593Smuzhiyun    return image.read(data_size)
2278*4882a593Smuzhiyun
2279*4882a593Smuzhiyun  def _get_cmdline_descriptors_for_hashtree_descriptor(self, ht):
2280*4882a593Smuzhiyun    """Generate kernel cmdline descriptors for dm-verity.
2281*4882a593Smuzhiyun
2282*4882a593Smuzhiyun    Arguments:
2283*4882a593Smuzhiyun      ht: A AvbHashtreeDescriptor
2284*4882a593Smuzhiyun
2285*4882a593Smuzhiyun    Returns:
2286*4882a593Smuzhiyun      A list with two AvbKernelCmdlineDescriptor with dm-verity kernel cmdline
2287*4882a593Smuzhiyun      instructions. There is one for when hashtree is not disabled and one for
2288*4882a593Smuzhiyun      when it is.
2289*4882a593Smuzhiyun
2290*4882a593Smuzhiyun    """
2291*4882a593Smuzhiyun
2292*4882a593Smuzhiyun    c = 'dm="1 vroot none ro 1,'
2293*4882a593Smuzhiyun    c += '0'  # start
2294*4882a593Smuzhiyun    c += ' {}'.format((ht.image_size / 512))  # size (# sectors)
2295*4882a593Smuzhiyun    c += ' verity {}'.format(ht.dm_verity_version)  # type and version
2296*4882a593Smuzhiyun    c += ' PARTUUID=$(ANDROID_SYSTEM_PARTUUID)'  # data_dev
2297*4882a593Smuzhiyun    c += ' PARTUUID=$(ANDROID_SYSTEM_PARTUUID)'  # hash_dev
2298*4882a593Smuzhiyun    c += ' {}'.format(ht.data_block_size)  # data_block
2299*4882a593Smuzhiyun    c += ' {}'.format(ht.hash_block_size)  # hash_block
2300*4882a593Smuzhiyun    c += ' {}'.format(ht.image_size / ht.data_block_size)  # #blocks
2301*4882a593Smuzhiyun    c += ' {}'.format(ht.image_size / ht.data_block_size)  # hash_offset
2302*4882a593Smuzhiyun    c += ' {}'.format(ht.hash_algorithm)  # hash_alg
2303*4882a593Smuzhiyun    c += ' {}'.format(str(ht.root_digest).encode('hex'))  # root_digest
2304*4882a593Smuzhiyun    c += ' {}'.format(str(ht.salt).encode('hex'))  # salt
2305*4882a593Smuzhiyun    if ht.fec_num_roots > 0:
2306*4882a593Smuzhiyun      c += ' 10'  # number of optional args
2307*4882a593Smuzhiyun      c += ' restart_on_corruption'
2308*4882a593Smuzhiyun      c += ' ignore_zero_blocks'
2309*4882a593Smuzhiyun      c += ' use_fec_from_device PARTUUID=$(ANDROID_SYSTEM_PARTUUID)'
2310*4882a593Smuzhiyun      c += ' fec_roots {}'.format(ht.fec_num_roots)
2311*4882a593Smuzhiyun      # Note that fec_blocks is the size that FEC covers, *not* the
2312*4882a593Smuzhiyun      # size of the FEC data. Since we use FEC for everything up until
2313*4882a593Smuzhiyun      # the FEC data, it's the same as the offset.
2314*4882a593Smuzhiyun      c += ' fec_blocks {}'.format(ht.fec_offset/ht.data_block_size)
2315*4882a593Smuzhiyun      c += ' fec_start {}'.format(ht.fec_offset/ht.data_block_size)
2316*4882a593Smuzhiyun    else:
2317*4882a593Smuzhiyun      c += ' 2'  # number of optional args
2318*4882a593Smuzhiyun      c += ' restart_on_corruption'
2319*4882a593Smuzhiyun      c += ' ignore_zero_blocks'
2320*4882a593Smuzhiyun    c += '" root=/dev/dm-0'
2321*4882a593Smuzhiyun
2322*4882a593Smuzhiyun    # Now that we have the command-line, generate the descriptor.
2323*4882a593Smuzhiyun    desc = AvbKernelCmdlineDescriptor()
2324*4882a593Smuzhiyun    desc.kernel_cmdline = c
2325*4882a593Smuzhiyun    desc.flags = (
2326*4882a593Smuzhiyun        AvbKernelCmdlineDescriptor.FLAGS_USE_ONLY_IF_HASHTREE_NOT_DISABLED)
2327*4882a593Smuzhiyun
2328*4882a593Smuzhiyun    # The descriptor for when hashtree verification is disabled is a lot
2329*4882a593Smuzhiyun    # simpler - we just set the root to the partition.
2330*4882a593Smuzhiyun    desc_no_ht = AvbKernelCmdlineDescriptor()
2331*4882a593Smuzhiyun    desc_no_ht.kernel_cmdline = 'root=PARTUUID=$(ANDROID_SYSTEM_PARTUUID)'
2332*4882a593Smuzhiyun    desc_no_ht.flags = (
2333*4882a593Smuzhiyun        AvbKernelCmdlineDescriptor.FLAGS_USE_ONLY_IF_HASHTREE_DISABLED)
2334*4882a593Smuzhiyun
2335*4882a593Smuzhiyun    return [desc, desc_no_ht]
2336*4882a593Smuzhiyun
2337*4882a593Smuzhiyun  def _get_cmdline_descriptors_for_dm_verity(self, image):
2338*4882a593Smuzhiyun    """Generate kernel cmdline descriptors for dm-verity.
2339*4882a593Smuzhiyun
2340*4882a593Smuzhiyun    Arguments:
2341*4882a593Smuzhiyun      image: An ImageHandler (vbmeta or footer) with a hashtree descriptor.
2342*4882a593Smuzhiyun
2343*4882a593Smuzhiyun    Returns:
2344*4882a593Smuzhiyun      A list with two AvbKernelCmdlineDescriptor with dm-verity kernel cmdline
2345*4882a593Smuzhiyun      instructions. There is one for when hashtree is not disabled and one for
2346*4882a593Smuzhiyun      when it is.
2347*4882a593Smuzhiyun
2348*4882a593Smuzhiyun    Raises:
2349*4882a593Smuzhiyun      AvbError: If  |image| doesn't have a hashtree descriptor.
2350*4882a593Smuzhiyun
2351*4882a593Smuzhiyun    """
2352*4882a593Smuzhiyun
2353*4882a593Smuzhiyun    (_, _, descriptors, _) = self._parse_image(image)
2354*4882a593Smuzhiyun
2355*4882a593Smuzhiyun    ht = None
2356*4882a593Smuzhiyun    for desc in descriptors:
2357*4882a593Smuzhiyun      if isinstance(desc, AvbHashtreeDescriptor):
2358*4882a593Smuzhiyun        ht = desc
2359*4882a593Smuzhiyun        break
2360*4882a593Smuzhiyun
2361*4882a593Smuzhiyun    if not ht:
2362*4882a593Smuzhiyun      raise AvbError('No hashtree descriptor in given image')
2363*4882a593Smuzhiyun
2364*4882a593Smuzhiyun    return self._get_cmdline_descriptors_for_hashtree_descriptor(ht)
2365*4882a593Smuzhiyun
2366*4882a593Smuzhiyun  def make_vbmeta_image(self, output, chain_partitions, algorithm_name,
2367*4882a593Smuzhiyun                        key_path, public_key_metadata_path, rollback_index,
2368*4882a593Smuzhiyun                        flags, props, props_from_file, kernel_cmdlines,
2369*4882a593Smuzhiyun                        setup_rootfs_from_kernel,
2370*4882a593Smuzhiyun                        include_descriptors_from_image,
2371*4882a593Smuzhiyun                        signing_helper,
2372*4882a593Smuzhiyun                        signing_helper_with_files,
2373*4882a593Smuzhiyun                        release_string,
2374*4882a593Smuzhiyun                        append_to_release_string,
2375*4882a593Smuzhiyun                        print_required_libavb_version,
2376*4882a593Smuzhiyun                        padding_size):
2377*4882a593Smuzhiyun    """Implements the 'make_vbmeta_image' command.
2378*4882a593Smuzhiyun
2379*4882a593Smuzhiyun    Arguments:
2380*4882a593Smuzhiyun      output: File to write the image to.
2381*4882a593Smuzhiyun      chain_partitions: List of partitions to chain or None.
2382*4882a593Smuzhiyun      algorithm_name: Name of algorithm to use.
2383*4882a593Smuzhiyun      key_path: Path to key to use or None.
2384*4882a593Smuzhiyun      public_key_metadata_path: Path to public key metadata or None.
2385*4882a593Smuzhiyun      rollback_index: The rollback index to use.
2386*4882a593Smuzhiyun      flags: Flags value to use in the image.
2387*4882a593Smuzhiyun      props: Properties to insert (list of strings of the form 'key:value').
2388*4882a593Smuzhiyun      props_from_file: Properties to insert (list of strings 'key:<path>').
2389*4882a593Smuzhiyun      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
2390*4882a593Smuzhiyun      setup_rootfs_from_kernel: None or file to generate from.
2391*4882a593Smuzhiyun      include_descriptors_from_image: List of file objects with descriptors.
2392*4882a593Smuzhiyun      signing_helper: Program which signs a hash and return signature.
2393*4882a593Smuzhiyun      signing_helper_with_files: Same as signing_helper but uses files instead.
2394*4882a593Smuzhiyun      release_string: None or avbtool release string to use instead of default.
2395*4882a593Smuzhiyun      append_to_release_string: None or string to append.
2396*4882a593Smuzhiyun      print_required_libavb_version: True to only print required libavb version.
2397*4882a593Smuzhiyun      padding_size: If not 0, pads output so size is a multiple of the number.
2398*4882a593Smuzhiyun
2399*4882a593Smuzhiyun    Raises:
2400*4882a593Smuzhiyun      AvbError: If a chained partition is malformed.
2401*4882a593Smuzhiyun    """
2402*4882a593Smuzhiyun
2403*4882a593Smuzhiyun    # If we're asked to calculate minimum required libavb version, we're done.
2404*4882a593Smuzhiyun    if print_required_libavb_version:
2405*4882a593Smuzhiyun      if include_descriptors_from_image:
2406*4882a593Smuzhiyun        # Use the bump logic in AvbVBMetaHeader to calculate the max required
2407*4882a593Smuzhiyun        # version of all included descriptors.
2408*4882a593Smuzhiyun        tmp_header = AvbVBMetaHeader()
2409*4882a593Smuzhiyun        for image in include_descriptors_from_image:
2410*4882a593Smuzhiyun          (_, image_header, _, _) = self._parse_image(ImageHandler(image.name))
2411*4882a593Smuzhiyun          tmp_header.bump_required_libavb_version_minor(
2412*4882a593Smuzhiyun              image_header.required_libavb_version_minor)
2413*4882a593Smuzhiyun        print '1.{}'.format(tmp_header.required_libavb_version_minor)
2414*4882a593Smuzhiyun      else:
2415*4882a593Smuzhiyun        # Descriptors aside, all vbmeta features are supported in 1.0.
2416*4882a593Smuzhiyun        print '1.0'
2417*4882a593Smuzhiyun      return
2418*4882a593Smuzhiyun
2419*4882a593Smuzhiyun    if not output:
2420*4882a593Smuzhiyun      raise AvbError('No output file given')
2421*4882a593Smuzhiyun
2422*4882a593Smuzhiyun    descriptors = []
2423*4882a593Smuzhiyun    ht_desc_to_setup = None
2424*4882a593Smuzhiyun    vbmeta_blob = self._generate_vbmeta_blob(
2425*4882a593Smuzhiyun        algorithm_name, key_path, public_key_metadata_path, descriptors,
2426*4882a593Smuzhiyun        chain_partitions, rollback_index, flags, props, props_from_file,
2427*4882a593Smuzhiyun        kernel_cmdlines, setup_rootfs_from_kernel, ht_desc_to_setup,
2428*4882a593Smuzhiyun        include_descriptors_from_image, signing_helper,
2429*4882a593Smuzhiyun        signing_helper_with_files, release_string,
2430*4882a593Smuzhiyun        append_to_release_string, 0)
2431*4882a593Smuzhiyun
2432*4882a593Smuzhiyun    # Write entire vbmeta blob (header, authentication, auxiliary).
2433*4882a593Smuzhiyun    output.seek(0)
2434*4882a593Smuzhiyun    output.write(vbmeta_blob)
2435*4882a593Smuzhiyun
2436*4882a593Smuzhiyun    if padding_size > 0:
2437*4882a593Smuzhiyun      padded_size = round_to_multiple(len(vbmeta_blob), padding_size)
2438*4882a593Smuzhiyun      padding_needed = padded_size - len(vbmeta_blob)
2439*4882a593Smuzhiyun      output.write('\0' * padding_needed)
2440*4882a593Smuzhiyun
2441*4882a593Smuzhiyun  def _generate_vbmeta_blob(self, algorithm_name, key_path,
2442*4882a593Smuzhiyun                            public_key_metadata_path, descriptors,
2443*4882a593Smuzhiyun                            chain_partitions,
2444*4882a593Smuzhiyun                            rollback_index, flags, props, props_from_file,
2445*4882a593Smuzhiyun                            kernel_cmdlines,
2446*4882a593Smuzhiyun                            setup_rootfs_from_kernel,
2447*4882a593Smuzhiyun                            ht_desc_to_setup,
2448*4882a593Smuzhiyun                            include_descriptors_from_image, signing_helper,
2449*4882a593Smuzhiyun                            signing_helper_with_files,
2450*4882a593Smuzhiyun                            release_string, append_to_release_string,
2451*4882a593Smuzhiyun                            required_libavb_version_minor):
2452*4882a593Smuzhiyun    """Generates a VBMeta blob.
2453*4882a593Smuzhiyun
2454*4882a593Smuzhiyun    This blob contains the header (struct AvbVBMetaHeader), the
2455*4882a593Smuzhiyun    authentication data block (which contains the hash and signature
2456*4882a593Smuzhiyun    for the header and auxiliary block), and the auxiliary block
2457*4882a593Smuzhiyun    (which contains descriptors, the public key used, and other data).
2458*4882a593Smuzhiyun
2459*4882a593Smuzhiyun    The |key| parameter can |None| only if the |algorithm_name| is
2460*4882a593Smuzhiyun    'NONE'.
2461*4882a593Smuzhiyun
2462*4882a593Smuzhiyun    Arguments:
2463*4882a593Smuzhiyun      algorithm_name: The algorithm name as per the ALGORITHMS dict.
2464*4882a593Smuzhiyun      key_path: The path to the .pem file used to sign the blob.
2465*4882a593Smuzhiyun      public_key_metadata_path: Path to public key metadata or None.
2466*4882a593Smuzhiyun      descriptors: A list of descriptors to insert or None.
2467*4882a593Smuzhiyun      chain_partitions: List of partitions to chain or None.
2468*4882a593Smuzhiyun      rollback_index: The rollback index to use.
2469*4882a593Smuzhiyun      flags: Flags to use in the image.
2470*4882a593Smuzhiyun      props: Properties to insert (List of strings of the form 'key:value').
2471*4882a593Smuzhiyun      props_from_file: Properties to insert (List of strings 'key:<path>').
2472*4882a593Smuzhiyun      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
2473*4882a593Smuzhiyun      setup_rootfs_from_kernel: None or file to generate
2474*4882a593Smuzhiyun        dm-verity kernel cmdline from.
2475*4882a593Smuzhiyun      ht_desc_to_setup: If not None, an AvbHashtreeDescriptor to
2476*4882a593Smuzhiyun        generate dm-verity kernel cmdline descriptors from.
2477*4882a593Smuzhiyun      include_descriptors_from_image: List of file objects for which
2478*4882a593Smuzhiyun        to insert descriptors from.
2479*4882a593Smuzhiyun      signing_helper: Program which signs a hash and return signature.
2480*4882a593Smuzhiyun      signing_helper_with_files: Same as signing_helper but uses files instead.
2481*4882a593Smuzhiyun      release_string: None or avbtool release string.
2482*4882a593Smuzhiyun      append_to_release_string: None or string to append.
2483*4882a593Smuzhiyun      required_libavb_version_minor: Use at least this required minor version.
2484*4882a593Smuzhiyun
2485*4882a593Smuzhiyun    Returns:
2486*4882a593Smuzhiyun      A bytearray() with the VBMeta blob.
2487*4882a593Smuzhiyun
2488*4882a593Smuzhiyun    Raises:
2489*4882a593Smuzhiyun      Exception: If the |algorithm_name| is not found, if no key has
2490*4882a593Smuzhiyun        been given and the given algorithm requires one, or the key is
2491*4882a593Smuzhiyun        of the wrong size.
2492*4882a593Smuzhiyun
2493*4882a593Smuzhiyun    """
2494*4882a593Smuzhiyun    try:
2495*4882a593Smuzhiyun      alg = ALGORITHMS[algorithm_name]
2496*4882a593Smuzhiyun    except KeyError:
2497*4882a593Smuzhiyun      raise AvbError('Unknown algorithm with name {}'.format(algorithm_name))
2498*4882a593Smuzhiyun
2499*4882a593Smuzhiyun    if not descriptors:
2500*4882a593Smuzhiyun      descriptors = []
2501*4882a593Smuzhiyun
2502*4882a593Smuzhiyun    h = AvbVBMetaHeader()
2503*4882a593Smuzhiyun    h.bump_required_libavb_version_minor(required_libavb_version_minor)
2504*4882a593Smuzhiyun
2505*4882a593Smuzhiyun    # Insert chained partition descriptors, if any
2506*4882a593Smuzhiyun    if chain_partitions:
2507*4882a593Smuzhiyun      used_locations = {}
2508*4882a593Smuzhiyun      for cp in chain_partitions:
2509*4882a593Smuzhiyun        cp_tokens = cp.split(':')
2510*4882a593Smuzhiyun        if len(cp_tokens) != 3:
2511*4882a593Smuzhiyun          raise AvbError('Malformed chained partition "{}".'.format(cp))
2512*4882a593Smuzhiyun        partition_name = cp_tokens[0]
2513*4882a593Smuzhiyun        rollback_index_location = int(cp_tokens[1])
2514*4882a593Smuzhiyun        file_path = cp_tokens[2]
2515*4882a593Smuzhiyun        # Check that the same rollback location isn't being used by
2516*4882a593Smuzhiyun        # multiple chained partitions.
2517*4882a593Smuzhiyun        if used_locations.get(rollback_index_location):
2518*4882a593Smuzhiyun          raise AvbError('Rollback Index Location {} is already in use.'.format(
2519*4882a593Smuzhiyun              rollback_index_location))
2520*4882a593Smuzhiyun        used_locations[rollback_index_location] = True
2521*4882a593Smuzhiyun        desc = AvbChainPartitionDescriptor()
2522*4882a593Smuzhiyun        desc.partition_name = partition_name
2523*4882a593Smuzhiyun        desc.rollback_index_location = rollback_index_location
2524*4882a593Smuzhiyun        if desc.rollback_index_location < 1:
2525*4882a593Smuzhiyun          raise AvbError('Rollback index location must be 1 or larger.')
2526*4882a593Smuzhiyun        desc.public_key = open(file_path, 'rb').read()
2527*4882a593Smuzhiyun        descriptors.append(desc)
2528*4882a593Smuzhiyun
2529*4882a593Smuzhiyun    # Descriptors.
2530*4882a593Smuzhiyun    encoded_descriptors = bytearray()
2531*4882a593Smuzhiyun    for desc in descriptors:
2532*4882a593Smuzhiyun      encoded_descriptors.extend(desc.encode())
2533*4882a593Smuzhiyun
2534*4882a593Smuzhiyun    # Add properties.
2535*4882a593Smuzhiyun    if props:
2536*4882a593Smuzhiyun      for prop in props:
2537*4882a593Smuzhiyun        idx = prop.find(':')
2538*4882a593Smuzhiyun        if idx == -1:
2539*4882a593Smuzhiyun          raise AvbError('Malformed property "{}".'.format(prop))
2540*4882a593Smuzhiyun        desc = AvbPropertyDescriptor()
2541*4882a593Smuzhiyun        desc.key = prop[0:idx]
2542*4882a593Smuzhiyun        desc.value = prop[(idx + 1):]
2543*4882a593Smuzhiyun        encoded_descriptors.extend(desc.encode())
2544*4882a593Smuzhiyun    if props_from_file:
2545*4882a593Smuzhiyun      for prop in props_from_file:
2546*4882a593Smuzhiyun        idx = prop.find(':')
2547*4882a593Smuzhiyun        if idx == -1:
2548*4882a593Smuzhiyun          raise AvbError('Malformed property "{}".'.format(prop))
2549*4882a593Smuzhiyun        desc = AvbPropertyDescriptor()
2550*4882a593Smuzhiyun        desc.key = prop[0:idx]
2551*4882a593Smuzhiyun        desc.value = prop[(idx + 1):]
2552*4882a593Smuzhiyun        file_path = prop[(idx + 1):]
2553*4882a593Smuzhiyun        desc.value = open(file_path, 'rb').read()
2554*4882a593Smuzhiyun        encoded_descriptors.extend(desc.encode())
2555*4882a593Smuzhiyun
2556*4882a593Smuzhiyun    # Add AvbKernelCmdline descriptor for dm-verity from an image, if requested.
2557*4882a593Smuzhiyun    if setup_rootfs_from_kernel:
2558*4882a593Smuzhiyun      image_handler = ImageHandler(
2559*4882a593Smuzhiyun          setup_rootfs_from_kernel.name)
2560*4882a593Smuzhiyun      cmdline_desc = self._get_cmdline_descriptors_for_dm_verity(image_handler)
2561*4882a593Smuzhiyun      encoded_descriptors.extend(cmdline_desc[0].encode())
2562*4882a593Smuzhiyun      encoded_descriptors.extend(cmdline_desc[1].encode())
2563*4882a593Smuzhiyun
2564*4882a593Smuzhiyun    # Add AvbKernelCmdline descriptor for dm-verity from desc, if requested.
2565*4882a593Smuzhiyun    if ht_desc_to_setup:
2566*4882a593Smuzhiyun      cmdline_desc = self._get_cmdline_descriptors_for_hashtree_descriptor(
2567*4882a593Smuzhiyun          ht_desc_to_setup)
2568*4882a593Smuzhiyun      encoded_descriptors.extend(cmdline_desc[0].encode())
2569*4882a593Smuzhiyun      encoded_descriptors.extend(cmdline_desc[1].encode())
2570*4882a593Smuzhiyun
2571*4882a593Smuzhiyun    # Add kernel command-lines.
2572*4882a593Smuzhiyun    if kernel_cmdlines:
2573*4882a593Smuzhiyun      for i in kernel_cmdlines:
2574*4882a593Smuzhiyun        desc = AvbKernelCmdlineDescriptor()
2575*4882a593Smuzhiyun        desc.kernel_cmdline = i
2576*4882a593Smuzhiyun        encoded_descriptors.extend(desc.encode())
2577*4882a593Smuzhiyun
2578*4882a593Smuzhiyun    # Add descriptors from other images.
2579*4882a593Smuzhiyun    if include_descriptors_from_image:
2580*4882a593Smuzhiyun      descriptors_dict = dict()
2581*4882a593Smuzhiyun      for image in include_descriptors_from_image:
2582*4882a593Smuzhiyun        image_handler = ImageHandler(image.name)
2583*4882a593Smuzhiyun        (_, image_vbmeta_header, image_descriptors, _) = self._parse_image(
2584*4882a593Smuzhiyun            image_handler)
2585*4882a593Smuzhiyun        # Bump the required libavb version to support all included descriptors.
2586*4882a593Smuzhiyun        h.bump_required_libavb_version_minor(
2587*4882a593Smuzhiyun            image_vbmeta_header.required_libavb_version_minor)
2588*4882a593Smuzhiyun        for desc in image_descriptors:
2589*4882a593Smuzhiyun          # The --include_descriptors_from_image option is used in some setups
2590*4882a593Smuzhiyun          # with images A and B where both A and B contain a descriptor
2591*4882a593Smuzhiyun          # for a partition with the same name. Since it's not meaningful
2592*4882a593Smuzhiyun          # to include both descriptors, only include the last seen descriptor.
2593*4882a593Smuzhiyun          # See bug 76386656 for details.
2594*4882a593Smuzhiyun          if hasattr(desc, 'partition_name'):
2595*4882a593Smuzhiyun            key = type(desc).__name__ + '_' + desc.partition_name
2596*4882a593Smuzhiyun            descriptors_dict[key] = desc.encode()
2597*4882a593Smuzhiyun          else:
2598*4882a593Smuzhiyun            encoded_descriptors.extend(desc.encode())
2599*4882a593Smuzhiyun      for key in sorted(descriptors_dict.keys()):
2600*4882a593Smuzhiyun        encoded_descriptors.extend(descriptors_dict[key])
2601*4882a593Smuzhiyun
2602*4882a593Smuzhiyun    # Load public key metadata blob, if requested.
2603*4882a593Smuzhiyun    pkmd_blob = []
2604*4882a593Smuzhiyun    if public_key_metadata_path:
2605*4882a593Smuzhiyun      with open(public_key_metadata_path) as f:
2606*4882a593Smuzhiyun        pkmd_blob = f.read()
2607*4882a593Smuzhiyun
2608*4882a593Smuzhiyun    key = None
2609*4882a593Smuzhiyun    encoded_key = bytearray()
2610*4882a593Smuzhiyun    if alg.public_key_num_bytes > 0:
2611*4882a593Smuzhiyun      if not key_path:
2612*4882a593Smuzhiyun        raise AvbError('Key is required for algorithm {}'.format(
2613*4882a593Smuzhiyun            algorithm_name))
2614*4882a593Smuzhiyun      encoded_key = encode_rsa_key(key_path)
2615*4882a593Smuzhiyun      if len(encoded_key) != alg.public_key_num_bytes:
2616*4882a593Smuzhiyun        raise AvbError('Key is wrong size for algorithm {}'.format(
2617*4882a593Smuzhiyun            algorithm_name))
2618*4882a593Smuzhiyun
2619*4882a593Smuzhiyun    # Override release string, if requested.
2620*4882a593Smuzhiyun    if isinstance(release_string, (str, unicode)):
2621*4882a593Smuzhiyun      h.release_string = release_string
2622*4882a593Smuzhiyun
2623*4882a593Smuzhiyun    # Append to release string, if requested. Also insert a space before.
2624*4882a593Smuzhiyun    if isinstance(append_to_release_string, (str, unicode)):
2625*4882a593Smuzhiyun      h.release_string += ' ' + append_to_release_string
2626*4882a593Smuzhiyun
2627*4882a593Smuzhiyun    # For the Auxiliary data block, descriptors are stored at offset 0,
2628*4882a593Smuzhiyun    # followed by the public key, followed by the public key metadata blob.
2629*4882a593Smuzhiyun    h.auxiliary_data_block_size = round_to_multiple(
2630*4882a593Smuzhiyun        len(encoded_descriptors) + len(encoded_key) + len(pkmd_blob), 64)
2631*4882a593Smuzhiyun    h.descriptors_offset = 0
2632*4882a593Smuzhiyun    h.descriptors_size = len(encoded_descriptors)
2633*4882a593Smuzhiyun    h.public_key_offset = h.descriptors_size
2634*4882a593Smuzhiyun    h.public_key_size = len(encoded_key)
2635*4882a593Smuzhiyun    h.public_key_metadata_offset = h.public_key_offset + h.public_key_size
2636*4882a593Smuzhiyun    h.public_key_metadata_size = len(pkmd_blob)
2637*4882a593Smuzhiyun
2638*4882a593Smuzhiyun    # For the Authentication data block, the hash is first and then
2639*4882a593Smuzhiyun    # the signature.
2640*4882a593Smuzhiyun    h.authentication_data_block_size = round_to_multiple(
2641*4882a593Smuzhiyun        alg.hash_num_bytes + alg.signature_num_bytes, 64)
2642*4882a593Smuzhiyun    h.algorithm_type = alg.algorithm_type
2643*4882a593Smuzhiyun    h.hash_offset = 0
2644*4882a593Smuzhiyun    h.hash_size = alg.hash_num_bytes
2645*4882a593Smuzhiyun    # Signature offset and size - it's stored right after the hash
2646*4882a593Smuzhiyun    # (in Authentication data block).
2647*4882a593Smuzhiyun    h.signature_offset = alg.hash_num_bytes
2648*4882a593Smuzhiyun    h.signature_size = alg.signature_num_bytes
2649*4882a593Smuzhiyun
2650*4882a593Smuzhiyun    h.rollback_index = rollback_index
2651*4882a593Smuzhiyun    h.flags = flags
2652*4882a593Smuzhiyun
2653*4882a593Smuzhiyun    # Generate Header data block.
2654*4882a593Smuzhiyun    header_data_blob = h.encode()
2655*4882a593Smuzhiyun
2656*4882a593Smuzhiyun    # Generate Auxiliary data block.
2657*4882a593Smuzhiyun    aux_data_blob = bytearray()
2658*4882a593Smuzhiyun    aux_data_blob.extend(encoded_descriptors)
2659*4882a593Smuzhiyun    aux_data_blob.extend(encoded_key)
2660*4882a593Smuzhiyun    aux_data_blob.extend(pkmd_blob)
2661*4882a593Smuzhiyun    padding_bytes = h.auxiliary_data_block_size - len(aux_data_blob)
2662*4882a593Smuzhiyun    aux_data_blob.extend('\0' * padding_bytes)
2663*4882a593Smuzhiyun
2664*4882a593Smuzhiyun    # Calculate the hash.
2665*4882a593Smuzhiyun    binary_hash = bytearray()
2666*4882a593Smuzhiyun    binary_signature = bytearray()
2667*4882a593Smuzhiyun    if algorithm_name != 'NONE':
2668*4882a593Smuzhiyun      ha = hashlib.new(alg.hash_name)
2669*4882a593Smuzhiyun      ha.update(header_data_blob)
2670*4882a593Smuzhiyun      ha.update(aux_data_blob)
2671*4882a593Smuzhiyun      binary_hash.extend(ha.digest())
2672*4882a593Smuzhiyun
2673*4882a593Smuzhiyun      # Calculate the signature.
2674*4882a593Smuzhiyun      padding_and_hash = str(bytearray(alg.padding)) + binary_hash
2675*4882a593Smuzhiyun      binary_signature.extend(raw_sign(signing_helper,
2676*4882a593Smuzhiyun                                       signing_helper_with_files,
2677*4882a593Smuzhiyun                                       algorithm_name,
2678*4882a593Smuzhiyun                                       alg.signature_num_bytes, key_path,
2679*4882a593Smuzhiyun                                       padding_and_hash))
2680*4882a593Smuzhiyun
2681*4882a593Smuzhiyun    # Generate Authentication data block.
2682*4882a593Smuzhiyun    auth_data_blob = bytearray()
2683*4882a593Smuzhiyun    auth_data_blob.extend(binary_hash)
2684*4882a593Smuzhiyun    auth_data_blob.extend(binary_signature)
2685*4882a593Smuzhiyun    padding_bytes = h.authentication_data_block_size - len(auth_data_blob)
2686*4882a593Smuzhiyun    auth_data_blob.extend('\0' * padding_bytes)
2687*4882a593Smuzhiyun
2688*4882a593Smuzhiyun    return header_data_blob + auth_data_blob + aux_data_blob
2689*4882a593Smuzhiyun
2690*4882a593Smuzhiyun  def extract_public_key(self, key_path, output):
2691*4882a593Smuzhiyun    """Implements the 'extract_public_key' command.
2692*4882a593Smuzhiyun
2693*4882a593Smuzhiyun    Arguments:
2694*4882a593Smuzhiyun      key_path: The path to a RSA private key file.
2695*4882a593Smuzhiyun      output: The file to write to.
2696*4882a593Smuzhiyun    """
2697*4882a593Smuzhiyun    output.write(encode_rsa_key(key_path))
2698*4882a593Smuzhiyun
2699*4882a593Smuzhiyun  def append_vbmeta_image(self, image_filename, vbmeta_image_filename,
2700*4882a593Smuzhiyun                          partition_size):
2701*4882a593Smuzhiyun    """Implementation of the append_vbmeta_image command.
2702*4882a593Smuzhiyun
2703*4882a593Smuzhiyun    Arguments:
2704*4882a593Smuzhiyun      image_filename: File to add the footer to.
2705*4882a593Smuzhiyun      vbmeta_image_filename: File to get vbmeta struct from.
2706*4882a593Smuzhiyun      partition_size: Size of partition.
2707*4882a593Smuzhiyun
2708*4882a593Smuzhiyun    Raises:
2709*4882a593Smuzhiyun      AvbError: If an argument is incorrect.
2710*4882a593Smuzhiyun    """
2711*4882a593Smuzhiyun    image = ImageHandler(image_filename)
2712*4882a593Smuzhiyun
2713*4882a593Smuzhiyun    if partition_size % image.block_size != 0:
2714*4882a593Smuzhiyun      raise AvbError('Partition size of {} is not a multiple of the image '
2715*4882a593Smuzhiyun                     'block size {}.'.format(partition_size,
2716*4882a593Smuzhiyun                                             image.block_size))
2717*4882a593Smuzhiyun
2718*4882a593Smuzhiyun    # If there's already a footer, truncate the image to its original
2719*4882a593Smuzhiyun    # size. This way 'avbtool append_vbmeta_image' is idempotent.
2720*4882a593Smuzhiyun    if image.image_size >= AvbFooter.SIZE:
2721*4882a593Smuzhiyun      image.seek(image.image_size - AvbFooter.SIZE)
2722*4882a593Smuzhiyun      try:
2723*4882a593Smuzhiyun        footer = AvbFooter(image.read(AvbFooter.SIZE))
2724*4882a593Smuzhiyun        # Existing footer found. Just truncate.
2725*4882a593Smuzhiyun        original_image_size = footer.original_image_size
2726*4882a593Smuzhiyun        image.truncate(footer.original_image_size)
2727*4882a593Smuzhiyun      except (LookupError, struct.error):
2728*4882a593Smuzhiyun        original_image_size = image.image_size
2729*4882a593Smuzhiyun    else:
2730*4882a593Smuzhiyun      # Image size is too small to possibly contain a footer.
2731*4882a593Smuzhiyun      original_image_size = image.image_size
2732*4882a593Smuzhiyun
2733*4882a593Smuzhiyun    # If anything goes wrong from here-on, restore the image back to
2734*4882a593Smuzhiyun    # its original size.
2735*4882a593Smuzhiyun    try:
2736*4882a593Smuzhiyun      vbmeta_image_handler = ImageHandler(vbmeta_image_filename)
2737*4882a593Smuzhiyun      vbmeta_blob = self._load_vbmeta_blob(vbmeta_image_handler)
2738*4882a593Smuzhiyun
2739*4882a593Smuzhiyun      # If the image isn't sparse, its size might not be a multiple of
2740*4882a593Smuzhiyun      # the block size. This will screw up padding later so just grow it.
2741*4882a593Smuzhiyun      if image.image_size % image.block_size != 0:
2742*4882a593Smuzhiyun        assert not image.is_sparse
2743*4882a593Smuzhiyun        padding_needed = image.block_size - (image.image_size%image.block_size)
2744*4882a593Smuzhiyun        image.truncate(image.image_size + padding_needed)
2745*4882a593Smuzhiyun
2746*4882a593Smuzhiyun      # The append_raw() method requires content with size being a
2747*4882a593Smuzhiyun      # multiple of |block_size| so add padding as needed. Also record
2748*4882a593Smuzhiyun      # where this is written to since we'll need to put that in the
2749*4882a593Smuzhiyun      # footer.
2750*4882a593Smuzhiyun      vbmeta_offset = image.image_size
2751*4882a593Smuzhiyun      padding_needed = (round_to_multiple(len(vbmeta_blob), image.block_size) -
2752*4882a593Smuzhiyun                        len(vbmeta_blob))
2753*4882a593Smuzhiyun      vbmeta_blob_with_padding = vbmeta_blob + '\0'*padding_needed
2754*4882a593Smuzhiyun
2755*4882a593Smuzhiyun      # Append vbmeta blob and footer
2756*4882a593Smuzhiyun      image.append_raw(vbmeta_blob_with_padding)
2757*4882a593Smuzhiyun      vbmeta_end_offset = vbmeta_offset + len(vbmeta_blob_with_padding)
2758*4882a593Smuzhiyun
2759*4882a593Smuzhiyun      # Now insert a DONT_CARE chunk with enough bytes such that the
2760*4882a593Smuzhiyun      # final Footer block is at the end of partition_size..
2761*4882a593Smuzhiyun      image.append_dont_care(partition_size - vbmeta_end_offset -
2762*4882a593Smuzhiyun                             1*image.block_size)
2763*4882a593Smuzhiyun
2764*4882a593Smuzhiyun      # Generate the Footer that tells where the VBMeta footer
2765*4882a593Smuzhiyun      # is. Also put enough padding in the front of the footer since
2766*4882a593Smuzhiyun      # we'll write out an entire block.
2767*4882a593Smuzhiyun      footer = AvbFooter()
2768*4882a593Smuzhiyun      footer.original_image_size = original_image_size
2769*4882a593Smuzhiyun      footer.vbmeta_offset = vbmeta_offset
2770*4882a593Smuzhiyun      footer.vbmeta_size = len(vbmeta_blob)
2771*4882a593Smuzhiyun      footer_blob = footer.encode()
2772*4882a593Smuzhiyun      footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
2773*4882a593Smuzhiyun                                  footer_blob)
2774*4882a593Smuzhiyun      image.append_raw(footer_blob_with_padding)
2775*4882a593Smuzhiyun
2776*4882a593Smuzhiyun    except:
2777*4882a593Smuzhiyun      # Truncate back to original size, then re-raise
2778*4882a593Smuzhiyun      image.truncate(original_image_size)
2779*4882a593Smuzhiyun      raise
2780*4882a593Smuzhiyun
2781*4882a593Smuzhiyun  def add_hash_footer(self, image_filename, partition_size, partition_name,
2782*4882a593Smuzhiyun                      hash_algorithm, salt, chain_partitions, algorithm_name,
2783*4882a593Smuzhiyun                      key_path,
2784*4882a593Smuzhiyun                      public_key_metadata_path, rollback_index, flags, props,
2785*4882a593Smuzhiyun                      props_from_file, kernel_cmdlines,
2786*4882a593Smuzhiyun                      setup_rootfs_from_kernel,
2787*4882a593Smuzhiyun                      include_descriptors_from_image, calc_max_image_size,
2788*4882a593Smuzhiyun                      signing_helper, signing_helper_with_files,
2789*4882a593Smuzhiyun                      release_string, append_to_release_string,
2790*4882a593Smuzhiyun                      output_vbmeta_image, do_not_append_vbmeta_image,
2791*4882a593Smuzhiyun                      print_required_libavb_version, use_persistent_digest,
2792*4882a593Smuzhiyun                      do_not_use_ab):
2793*4882a593Smuzhiyun    """Implementation of the add_hash_footer on unsparse images.
2794*4882a593Smuzhiyun
2795*4882a593Smuzhiyun    Arguments:
2796*4882a593Smuzhiyun      image_filename: File to add the footer to.
2797*4882a593Smuzhiyun      partition_size: Size of partition.
2798*4882a593Smuzhiyun      partition_name: Name of partition (without A/B suffix).
2799*4882a593Smuzhiyun      hash_algorithm: Hash algorithm to use.
2800*4882a593Smuzhiyun      salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
2801*4882a593Smuzhiyun      chain_partitions: List of partitions to chain.
2802*4882a593Smuzhiyun      algorithm_name: Name of algorithm to use.
2803*4882a593Smuzhiyun      key_path: Path to key to use or None.
2804*4882a593Smuzhiyun      public_key_metadata_path: Path to public key metadata or None.
2805*4882a593Smuzhiyun      rollback_index: Rollback index.
2806*4882a593Smuzhiyun      flags: Flags value to use in the image.
2807*4882a593Smuzhiyun      props: Properties to insert (List of strings of the form 'key:value').
2808*4882a593Smuzhiyun      props_from_file: Properties to insert (List of strings 'key:<path>').
2809*4882a593Smuzhiyun      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
2810*4882a593Smuzhiyun      setup_rootfs_from_kernel: None or file to generate
2811*4882a593Smuzhiyun        dm-verity kernel cmdline from.
2812*4882a593Smuzhiyun      include_descriptors_from_image: List of file objects for which
2813*4882a593Smuzhiyun        to insert descriptors from.
2814*4882a593Smuzhiyun      calc_max_image_size: Don't store the footer - instead calculate the
2815*4882a593Smuzhiyun        maximum image size leaving enough room for metadata with the
2816*4882a593Smuzhiyun        given |partition_size|.
2817*4882a593Smuzhiyun      signing_helper: Program which signs a hash and return signature.
2818*4882a593Smuzhiyun      signing_helper_with_files: Same as signing_helper but uses files instead.
2819*4882a593Smuzhiyun      release_string: None or avbtool release string.
2820*4882a593Smuzhiyun      append_to_release_string: None or string to append.
2821*4882a593Smuzhiyun      output_vbmeta_image: If not None, also write vbmeta struct to this file.
2822*4882a593Smuzhiyun      do_not_append_vbmeta_image: If True, don't append vbmeta struct.
2823*4882a593Smuzhiyun      print_required_libavb_version: True to only print required libavb version.
2824*4882a593Smuzhiyun      use_persistent_digest: Use a persistent digest on device.
2825*4882a593Smuzhiyun      do_not_use_ab: This partition does not use A/B.
2826*4882a593Smuzhiyun
2827*4882a593Smuzhiyun    Raises:
2828*4882a593Smuzhiyun      AvbError: If an argument is incorrect.
2829*4882a593Smuzhiyun    """
2830*4882a593Smuzhiyun
2831*4882a593Smuzhiyun    required_libavb_version_minor = 0
2832*4882a593Smuzhiyun    if use_persistent_digest or do_not_use_ab:
2833*4882a593Smuzhiyun      required_libavb_version_minor = 1
2834*4882a593Smuzhiyun
2835*4882a593Smuzhiyun    # If we're asked to calculate minimum required libavb version, we're done.
2836*4882a593Smuzhiyun    if print_required_libavb_version:
2837*4882a593Smuzhiyun      print '1.{}'.format(required_libavb_version_minor)
2838*4882a593Smuzhiyun      return
2839*4882a593Smuzhiyun
2840*4882a593Smuzhiyun    # First, calculate the maximum image size such that an image
2841*4882a593Smuzhiyun    # this size + metadata (footer + vbmeta struct) fits in
2842*4882a593Smuzhiyun    # |partition_size|.
2843*4882a593Smuzhiyun    max_metadata_size = self.MAX_VBMETA_SIZE + self.MAX_FOOTER_SIZE
2844*4882a593Smuzhiyun    if partition_size < max_metadata_size:
2845*4882a593Smuzhiyun      raise AvbError('Parition size of {} is too small. '
2846*4882a593Smuzhiyun                     'Needs to be at least {}'.format(
2847*4882a593Smuzhiyun                         partition_size, max_metadata_size))
2848*4882a593Smuzhiyun    max_image_size = partition_size - max_metadata_size
2849*4882a593Smuzhiyun
2850*4882a593Smuzhiyun    # If we're asked to only calculate the maximum image size, we're done.
2851*4882a593Smuzhiyun    if calc_max_image_size:
2852*4882a593Smuzhiyun      print '{}'.format(max_image_size)
2853*4882a593Smuzhiyun      return
2854*4882a593Smuzhiyun
2855*4882a593Smuzhiyun    image = ImageHandler(image_filename)
2856*4882a593Smuzhiyun
2857*4882a593Smuzhiyun    if partition_size % image.block_size != 0:
2858*4882a593Smuzhiyun      raise AvbError('Partition size of {} is not a multiple of the image '
2859*4882a593Smuzhiyun                     'block size {}.'.format(partition_size,
2860*4882a593Smuzhiyun                                             image.block_size))
2861*4882a593Smuzhiyun
2862*4882a593Smuzhiyun    # If there's already a footer, truncate the image to its original
2863*4882a593Smuzhiyun    # size. This way 'avbtool add_hash_footer' is idempotent (modulo
2864*4882a593Smuzhiyun    # salts).
2865*4882a593Smuzhiyun    if image.image_size >= AvbFooter.SIZE:
2866*4882a593Smuzhiyun      image.seek(image.image_size - AvbFooter.SIZE)
2867*4882a593Smuzhiyun      try:
2868*4882a593Smuzhiyun        footer = AvbFooter(image.read(AvbFooter.SIZE))
2869*4882a593Smuzhiyun        # Existing footer found. Just truncate.
2870*4882a593Smuzhiyun        original_image_size = footer.original_image_size
2871*4882a593Smuzhiyun        image.truncate(footer.original_image_size)
2872*4882a593Smuzhiyun      except (LookupError, struct.error):
2873*4882a593Smuzhiyun        original_image_size = image.image_size
2874*4882a593Smuzhiyun    else:
2875*4882a593Smuzhiyun      # Image size is too small to possibly contain a footer.
2876*4882a593Smuzhiyun      original_image_size = image.image_size
2877*4882a593Smuzhiyun
2878*4882a593Smuzhiyun    # If anything goes wrong from here-on, restore the image back to
2879*4882a593Smuzhiyun    # its original size.
2880*4882a593Smuzhiyun    try:
2881*4882a593Smuzhiyun      # If image size exceeds the maximum image size, fail.
2882*4882a593Smuzhiyun      if image.image_size > max_image_size:
2883*4882a593Smuzhiyun        raise AvbError('Image size of {} exceeds maximum image '
2884*4882a593Smuzhiyun                       'size of {} in order to fit in a partition '
2885*4882a593Smuzhiyun                       'size of {}.'.format(image.image_size, max_image_size,
2886*4882a593Smuzhiyun                                            partition_size))
2887*4882a593Smuzhiyun
2888*4882a593Smuzhiyun      digest_size = len(hashlib.new(name=hash_algorithm).digest())
2889*4882a593Smuzhiyun      if salt:
2890*4882a593Smuzhiyun        salt = salt.decode('hex')
2891*4882a593Smuzhiyun      else:
2892*4882a593Smuzhiyun        if salt is None:
2893*4882a593Smuzhiyun          # If salt is not explicitly specified, choose a hash
2894*4882a593Smuzhiyun          # that's the same size as the hash size.
2895*4882a593Smuzhiyun          hash_size = digest_size
2896*4882a593Smuzhiyun          salt = open('/dev/urandom').read(hash_size)
2897*4882a593Smuzhiyun        else:
2898*4882a593Smuzhiyun          salt = ''
2899*4882a593Smuzhiyun
2900*4882a593Smuzhiyun      hasher = hashlib.new(name=hash_algorithm, string=salt)
2901*4882a593Smuzhiyun      # TODO(zeuthen): might want to read this in chunks to avoid
2902*4882a593Smuzhiyun      # memory pressure, then again, this is only supposed to be used
2903*4882a593Smuzhiyun      # on kernel/initramfs partitions. Possible optimization.
2904*4882a593Smuzhiyun      image.seek(0)
2905*4882a593Smuzhiyun      hasher.update(image.read(image.image_size))
2906*4882a593Smuzhiyun      digest = hasher.digest()
2907*4882a593Smuzhiyun
2908*4882a593Smuzhiyun      h_desc = AvbHashDescriptor()
2909*4882a593Smuzhiyun      h_desc.image_size = image.image_size
2910*4882a593Smuzhiyun      h_desc.hash_algorithm = hash_algorithm
2911*4882a593Smuzhiyun      h_desc.partition_name = partition_name
2912*4882a593Smuzhiyun      h_desc.salt = salt
2913*4882a593Smuzhiyun      h_desc.flags = 0
2914*4882a593Smuzhiyun      if do_not_use_ab:
2915*4882a593Smuzhiyun        h_desc.flags |= 1  # AVB_HASH_DESCRIPTOR_FLAGS_DO_NOT_USE_AB
2916*4882a593Smuzhiyun      if not use_persistent_digest:
2917*4882a593Smuzhiyun        h_desc.digest = digest
2918*4882a593Smuzhiyun
2919*4882a593Smuzhiyun      # Generate the VBMeta footer.
2920*4882a593Smuzhiyun      ht_desc_to_setup = None
2921*4882a593Smuzhiyun      vbmeta_blob = self._generate_vbmeta_blob(
2922*4882a593Smuzhiyun          algorithm_name, key_path, public_key_metadata_path, [h_desc],
2923*4882a593Smuzhiyun          chain_partitions, rollback_index, flags, props, props_from_file,
2924*4882a593Smuzhiyun          kernel_cmdlines, setup_rootfs_from_kernel, ht_desc_to_setup,
2925*4882a593Smuzhiyun          include_descriptors_from_image, signing_helper,
2926*4882a593Smuzhiyun          signing_helper_with_files, release_string,
2927*4882a593Smuzhiyun          append_to_release_string, required_libavb_version_minor)
2928*4882a593Smuzhiyun
2929*4882a593Smuzhiyun      # Write vbmeta blob, if requested.
2930*4882a593Smuzhiyun      if output_vbmeta_image:
2931*4882a593Smuzhiyun        output_vbmeta_image.write(vbmeta_blob)
2932*4882a593Smuzhiyun
2933*4882a593Smuzhiyun      # Append vbmeta blob and footer, unless requested not to.
2934*4882a593Smuzhiyun      if not do_not_append_vbmeta_image:
2935*4882a593Smuzhiyun        # If the image isn't sparse, its size might not be a multiple of
2936*4882a593Smuzhiyun        # the block size. This will screw up padding later so just grow it.
2937*4882a593Smuzhiyun        if image.image_size % image.block_size != 0:
2938*4882a593Smuzhiyun          assert not image.is_sparse
2939*4882a593Smuzhiyun          padding_needed = image.block_size - (
2940*4882a593Smuzhiyun              image.image_size % image.block_size)
2941*4882a593Smuzhiyun          image.truncate(image.image_size + padding_needed)
2942*4882a593Smuzhiyun
2943*4882a593Smuzhiyun        # The append_raw() method requires content with size being a
2944*4882a593Smuzhiyun        # multiple of |block_size| so add padding as needed. Also record
2945*4882a593Smuzhiyun        # where this is written to since we'll need to put that in the
2946*4882a593Smuzhiyun        # footer.
2947*4882a593Smuzhiyun        vbmeta_offset = image.image_size
2948*4882a593Smuzhiyun        padding_needed = (
2949*4882a593Smuzhiyun            round_to_multiple(len(vbmeta_blob), image.block_size) -
2950*4882a593Smuzhiyun            len(vbmeta_blob))
2951*4882a593Smuzhiyun        vbmeta_blob_with_padding = vbmeta_blob + '\0' * padding_needed
2952*4882a593Smuzhiyun
2953*4882a593Smuzhiyun        image.append_raw(vbmeta_blob_with_padding)
2954*4882a593Smuzhiyun        vbmeta_end_offset = vbmeta_offset + len(vbmeta_blob_with_padding)
2955*4882a593Smuzhiyun
2956*4882a593Smuzhiyun        # Now insert a DONT_CARE chunk with enough bytes such that the
2957*4882a593Smuzhiyun        # final Footer block is at the end of partition_size..
2958*4882a593Smuzhiyun        image.append_dont_care(partition_size - vbmeta_end_offset -
2959*4882a593Smuzhiyun                               1*image.block_size)
2960*4882a593Smuzhiyun
2961*4882a593Smuzhiyun        # Generate the Footer that tells where the VBMeta footer
2962*4882a593Smuzhiyun        # is. Also put enough padding in the front of the footer since
2963*4882a593Smuzhiyun        # we'll write out an entire block.
2964*4882a593Smuzhiyun        footer = AvbFooter()
2965*4882a593Smuzhiyun        footer.original_image_size = original_image_size
2966*4882a593Smuzhiyun        footer.vbmeta_offset = vbmeta_offset
2967*4882a593Smuzhiyun        footer.vbmeta_size = len(vbmeta_blob)
2968*4882a593Smuzhiyun        footer_blob = footer.encode()
2969*4882a593Smuzhiyun        footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
2970*4882a593Smuzhiyun                                    footer_blob)
2971*4882a593Smuzhiyun        image.append_raw(footer_blob_with_padding)
2972*4882a593Smuzhiyun
2973*4882a593Smuzhiyun    except:
2974*4882a593Smuzhiyun      # Truncate back to original size, then re-raise
2975*4882a593Smuzhiyun      image.truncate(original_image_size)
2976*4882a593Smuzhiyun      raise
2977*4882a593Smuzhiyun
2978*4882a593Smuzhiyun  def add_hashtree_footer(self, image_filename, partition_size, partition_name,
2979*4882a593Smuzhiyun                          generate_fec, fec_num_roots, hash_algorithm,
2980*4882a593Smuzhiyun                          block_size, salt, chain_partitions, algorithm_name,
2981*4882a593Smuzhiyun                          key_path,
2982*4882a593Smuzhiyun                          public_key_metadata_path, rollback_index, flags,
2983*4882a593Smuzhiyun                          props, props_from_file, kernel_cmdlines,
2984*4882a593Smuzhiyun                          setup_rootfs_from_kernel,
2985*4882a593Smuzhiyun                          setup_as_rootfs_from_kernel,
2986*4882a593Smuzhiyun                          include_descriptors_from_image,
2987*4882a593Smuzhiyun                          calc_max_image_size, signing_helper,
2988*4882a593Smuzhiyun                          signing_helper_with_files,
2989*4882a593Smuzhiyun                          release_string, append_to_release_string,
2990*4882a593Smuzhiyun                          output_vbmeta_image, do_not_append_vbmeta_image,
2991*4882a593Smuzhiyun                          print_required_libavb_version,
2992*4882a593Smuzhiyun                          use_persistent_root_digest, do_not_use_ab):
2993*4882a593Smuzhiyun    """Implements the 'add_hashtree_footer' command.
2994*4882a593Smuzhiyun
2995*4882a593Smuzhiyun    See https://gitlab.com/cryptsetup/cryptsetup/wikis/DMVerity for
2996*4882a593Smuzhiyun    more information about dm-verity and these hashes.
2997*4882a593Smuzhiyun
2998*4882a593Smuzhiyun    Arguments:
2999*4882a593Smuzhiyun      image_filename: File to add the footer to.
3000*4882a593Smuzhiyun      partition_size: Size of partition.
3001*4882a593Smuzhiyun      partition_name: Name of partition (without A/B suffix).
3002*4882a593Smuzhiyun      generate_fec: If True, generate FEC codes.
3003*4882a593Smuzhiyun      fec_num_roots: Number of roots for FEC.
3004*4882a593Smuzhiyun      hash_algorithm: Hash algorithm to use.
3005*4882a593Smuzhiyun      block_size: Block size to use.
3006*4882a593Smuzhiyun      salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
3007*4882a593Smuzhiyun      chain_partitions: List of partitions to chain.
3008*4882a593Smuzhiyun      algorithm_name: Name of algorithm to use.
3009*4882a593Smuzhiyun      key_path: Path to key to use or None.
3010*4882a593Smuzhiyun      public_key_metadata_path: Path to public key metadata or None.
3011*4882a593Smuzhiyun      rollback_index: Rollback index.
3012*4882a593Smuzhiyun      flags: Flags value to use in the image.
3013*4882a593Smuzhiyun      props: Properties to insert (List of strings of the form 'key:value').
3014*4882a593Smuzhiyun      props_from_file: Properties to insert (List of strings 'key:<path>').
3015*4882a593Smuzhiyun      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
3016*4882a593Smuzhiyun      setup_rootfs_from_kernel: None or file to generate
3017*4882a593Smuzhiyun        dm-verity kernel cmdline from.
3018*4882a593Smuzhiyun      setup_as_rootfs_from_kernel: If True, generate dm-verity kernel
3019*4882a593Smuzhiyun        cmdline to set up rootfs.
3020*4882a593Smuzhiyun      include_descriptors_from_image: List of file objects for which
3021*4882a593Smuzhiyun        to insert descriptors from.
3022*4882a593Smuzhiyun      calc_max_image_size: Don't store the hashtree or footer - instead
3023*4882a593Smuzhiyun        calculate the maximum image size leaving enough room for hashtree
3024*4882a593Smuzhiyun        and metadata with the given |partition_size|.
3025*4882a593Smuzhiyun      signing_helper: Program which signs a hash and return signature.
3026*4882a593Smuzhiyun      signing_helper_with_files: Same as signing_helper but uses files instead.
3027*4882a593Smuzhiyun      release_string: None or avbtool release string.
3028*4882a593Smuzhiyun      append_to_release_string: None or string to append.
3029*4882a593Smuzhiyun      output_vbmeta_image: If not None, also write vbmeta struct to this file.
3030*4882a593Smuzhiyun      do_not_append_vbmeta_image: If True, don't append vbmeta struct.
3031*4882a593Smuzhiyun      print_required_libavb_version: True to only print required libavb version.
3032*4882a593Smuzhiyun      use_persistent_root_digest: Use a persistent root digest on device.
3033*4882a593Smuzhiyun      do_not_use_ab: The partition does not use A/B.
3034*4882a593Smuzhiyun
3035*4882a593Smuzhiyun    Raises:
3036*4882a593Smuzhiyun      AvbError: If an argument is incorrect.
3037*4882a593Smuzhiyun    """
3038*4882a593Smuzhiyun
3039*4882a593Smuzhiyun    required_libavb_version_minor = 0
3040*4882a593Smuzhiyun    if use_persistent_root_digest or do_not_use_ab:
3041*4882a593Smuzhiyun      required_libavb_version_minor = 1
3042*4882a593Smuzhiyun
3043*4882a593Smuzhiyun    # If we're asked to calculate minimum required libavb version, we're done.
3044*4882a593Smuzhiyun    if print_required_libavb_version:
3045*4882a593Smuzhiyun      print '1.{}'.format(required_libavb_version_minor)
3046*4882a593Smuzhiyun      return
3047*4882a593Smuzhiyun
3048*4882a593Smuzhiyun    digest_size = len(hashlib.new(name=hash_algorithm).digest())
3049*4882a593Smuzhiyun    digest_padding = round_to_pow2(digest_size) - digest_size
3050*4882a593Smuzhiyun
3051*4882a593Smuzhiyun    # First, calculate the maximum image size such that an image
3052*4882a593Smuzhiyun    # this size + the hashtree + metadata (footer + vbmeta struct)
3053*4882a593Smuzhiyun    # fits in |partition_size|. We use very conservative figures for
3054*4882a593Smuzhiyun    # metadata.
3055*4882a593Smuzhiyun    (_, max_tree_size) = calc_hash_level_offsets(
3056*4882a593Smuzhiyun        partition_size, block_size, digest_size + digest_padding)
3057*4882a593Smuzhiyun    max_fec_size = 0
3058*4882a593Smuzhiyun    if generate_fec:
3059*4882a593Smuzhiyun      max_fec_size = calc_fec_data_size(partition_size, fec_num_roots)
3060*4882a593Smuzhiyun    max_metadata_size = (max_fec_size + max_tree_size +
3061*4882a593Smuzhiyun                         self.MAX_VBMETA_SIZE +
3062*4882a593Smuzhiyun                         self.MAX_FOOTER_SIZE)
3063*4882a593Smuzhiyun    max_image_size = partition_size - max_metadata_size
3064*4882a593Smuzhiyun
3065*4882a593Smuzhiyun    # If we're asked to only calculate the maximum image size, we're done.
3066*4882a593Smuzhiyun    if calc_max_image_size:
3067*4882a593Smuzhiyun      print '{}'.format(max_image_size)
3068*4882a593Smuzhiyun      return
3069*4882a593Smuzhiyun
3070*4882a593Smuzhiyun    image = ImageHandler(image_filename)
3071*4882a593Smuzhiyun
3072*4882a593Smuzhiyun    if partition_size % image.block_size != 0:
3073*4882a593Smuzhiyun      raise AvbError('Partition size of {} is not a multiple of the image '
3074*4882a593Smuzhiyun                     'block size {}.'.format(partition_size,
3075*4882a593Smuzhiyun                                             image.block_size))
3076*4882a593Smuzhiyun
3077*4882a593Smuzhiyun    # If there's already a footer, truncate the image to its original
3078*4882a593Smuzhiyun    # size. This way 'avbtool add_hashtree_footer' is idempotent
3079*4882a593Smuzhiyun    # (modulo salts).
3080*4882a593Smuzhiyun    if image.image_size >= AvbFooter.SIZE:
3081*4882a593Smuzhiyun      image.seek(image.image_size - AvbFooter.SIZE)
3082*4882a593Smuzhiyun      try:
3083*4882a593Smuzhiyun        footer = AvbFooter(image.read(AvbFooter.SIZE))
3084*4882a593Smuzhiyun        # Existing footer found. Just truncate.
3085*4882a593Smuzhiyun        original_image_size = footer.original_image_size
3086*4882a593Smuzhiyun        image.truncate(footer.original_image_size)
3087*4882a593Smuzhiyun      except (LookupError, struct.error):
3088*4882a593Smuzhiyun        original_image_size = image.image_size
3089*4882a593Smuzhiyun    else:
3090*4882a593Smuzhiyun      # Image size is too small to possibly contain a footer.
3091*4882a593Smuzhiyun      original_image_size = image.image_size
3092*4882a593Smuzhiyun
3093*4882a593Smuzhiyun    # If anything goes wrong from here-on, restore the image back to
3094*4882a593Smuzhiyun    # its original size.
3095*4882a593Smuzhiyun    try:
3096*4882a593Smuzhiyun      # Ensure image is multiple of block_size.
3097*4882a593Smuzhiyun      rounded_image_size = round_to_multiple(image.image_size, block_size)
3098*4882a593Smuzhiyun      if rounded_image_size > image.image_size:
3099*4882a593Smuzhiyun        image.append_raw('\0' * (rounded_image_size - image.image_size))
3100*4882a593Smuzhiyun
3101*4882a593Smuzhiyun      # If image size exceeds the maximum image size, fail.
3102*4882a593Smuzhiyun      if image.image_size > max_image_size:
3103*4882a593Smuzhiyun        raise AvbError('Image size of {} exceeds maximum image '
3104*4882a593Smuzhiyun                       'size of {} in order to fit in a partition '
3105*4882a593Smuzhiyun                       'size of {}.'.format(image.image_size, max_image_size,
3106*4882a593Smuzhiyun                                            partition_size))
3107*4882a593Smuzhiyun
3108*4882a593Smuzhiyun      if salt:
3109*4882a593Smuzhiyun        salt = salt.decode('hex')
3110*4882a593Smuzhiyun      else:
3111*4882a593Smuzhiyun        if salt is None:
3112*4882a593Smuzhiyun          # If salt is not explicitly specified, choose a hash
3113*4882a593Smuzhiyun          # that's the same size as the hash size.
3114*4882a593Smuzhiyun          hash_size = digest_size
3115*4882a593Smuzhiyun          salt = open('/dev/urandom').read(hash_size)
3116*4882a593Smuzhiyun        else:
3117*4882a593Smuzhiyun          salt = ''
3118*4882a593Smuzhiyun
3119*4882a593Smuzhiyun      # Hashes are stored upside down so we need to calculate hash
3120*4882a593Smuzhiyun      # offsets in advance.
3121*4882a593Smuzhiyun      (hash_level_offsets, tree_size) = calc_hash_level_offsets(
3122*4882a593Smuzhiyun          image.image_size, block_size, digest_size + digest_padding)
3123*4882a593Smuzhiyun
3124*4882a593Smuzhiyun      # If the image isn't sparse, its size might not be a multiple of
3125*4882a593Smuzhiyun      # the block size. This will screw up padding later so just grow it.
3126*4882a593Smuzhiyun      if image.image_size % image.block_size != 0:
3127*4882a593Smuzhiyun        assert not image.is_sparse
3128*4882a593Smuzhiyun        padding_needed = image.block_size - (image.image_size%image.block_size)
3129*4882a593Smuzhiyun        image.truncate(image.image_size + padding_needed)
3130*4882a593Smuzhiyun
3131*4882a593Smuzhiyun      # Generate the tree and add padding as needed.
3132*4882a593Smuzhiyun      tree_offset = image.image_size
3133*4882a593Smuzhiyun      root_digest, hash_tree = generate_hash_tree(image, image.image_size,
3134*4882a593Smuzhiyun                                                  block_size,
3135*4882a593Smuzhiyun                                                  hash_algorithm, salt,
3136*4882a593Smuzhiyun                                                  digest_padding,
3137*4882a593Smuzhiyun                                                  hash_level_offsets,
3138*4882a593Smuzhiyun                                                  tree_size)
3139*4882a593Smuzhiyun
3140*4882a593Smuzhiyun      # Generate HashtreeDescriptor with details about the tree we
3141*4882a593Smuzhiyun      # just generated.
3142*4882a593Smuzhiyun      ht_desc = AvbHashtreeDescriptor()
3143*4882a593Smuzhiyun      ht_desc.dm_verity_version = 1
3144*4882a593Smuzhiyun      ht_desc.image_size = image.image_size
3145*4882a593Smuzhiyun      ht_desc.tree_offset = tree_offset
3146*4882a593Smuzhiyun      ht_desc.tree_size = tree_size
3147*4882a593Smuzhiyun      ht_desc.data_block_size = block_size
3148*4882a593Smuzhiyun      ht_desc.hash_block_size = block_size
3149*4882a593Smuzhiyun      ht_desc.hash_algorithm = hash_algorithm
3150*4882a593Smuzhiyun      ht_desc.partition_name = partition_name
3151*4882a593Smuzhiyun      ht_desc.salt = salt
3152*4882a593Smuzhiyun      if do_not_use_ab:
3153*4882a593Smuzhiyun        ht_desc.flags |= 1  # AVB_HASHTREE_DESCRIPTOR_FLAGS_DO_NOT_USE_AB
3154*4882a593Smuzhiyun      if not use_persistent_root_digest:
3155*4882a593Smuzhiyun        ht_desc.root_digest = root_digest
3156*4882a593Smuzhiyun
3157*4882a593Smuzhiyun      # Write the hash tree
3158*4882a593Smuzhiyun      padding_needed = (round_to_multiple(len(hash_tree), image.block_size) -
3159*4882a593Smuzhiyun                        len(hash_tree))
3160*4882a593Smuzhiyun      hash_tree_with_padding = hash_tree + '\0'*padding_needed
3161*4882a593Smuzhiyun      image.append_raw(hash_tree_with_padding)
3162*4882a593Smuzhiyun      len_hashtree_and_fec = len(hash_tree_with_padding)
3163*4882a593Smuzhiyun
3164*4882a593Smuzhiyun      # Generate FEC codes, if requested.
3165*4882a593Smuzhiyun      if generate_fec:
3166*4882a593Smuzhiyun        fec_data = generate_fec_data(image_filename, fec_num_roots)
3167*4882a593Smuzhiyun        padding_needed = (round_to_multiple(len(fec_data), image.block_size) -
3168*4882a593Smuzhiyun                          len(fec_data))
3169*4882a593Smuzhiyun        fec_data_with_padding = fec_data + '\0'*padding_needed
3170*4882a593Smuzhiyun        fec_offset = image.image_size
3171*4882a593Smuzhiyun        image.append_raw(fec_data_with_padding)
3172*4882a593Smuzhiyun        len_hashtree_and_fec += len(fec_data_with_padding)
3173*4882a593Smuzhiyun        # Update the hashtree descriptor.
3174*4882a593Smuzhiyun        ht_desc.fec_num_roots = fec_num_roots
3175*4882a593Smuzhiyun        ht_desc.fec_offset = fec_offset
3176*4882a593Smuzhiyun        ht_desc.fec_size = len(fec_data)
3177*4882a593Smuzhiyun
3178*4882a593Smuzhiyun      ht_desc_to_setup = None
3179*4882a593Smuzhiyun      if setup_as_rootfs_from_kernel:
3180*4882a593Smuzhiyun        ht_desc_to_setup = ht_desc
3181*4882a593Smuzhiyun
3182*4882a593Smuzhiyun      # Generate the VBMeta footer and add padding as needed.
3183*4882a593Smuzhiyun      vbmeta_offset = tree_offset + len_hashtree_and_fec
3184*4882a593Smuzhiyun      vbmeta_blob = self._generate_vbmeta_blob(
3185*4882a593Smuzhiyun          algorithm_name, key_path, public_key_metadata_path, [ht_desc],
3186*4882a593Smuzhiyun          chain_partitions, rollback_index, flags, props, props_from_file,
3187*4882a593Smuzhiyun          kernel_cmdlines, setup_rootfs_from_kernel, ht_desc_to_setup,
3188*4882a593Smuzhiyun          include_descriptors_from_image, signing_helper,
3189*4882a593Smuzhiyun          signing_helper_with_files, release_string,
3190*4882a593Smuzhiyun          append_to_release_string, required_libavb_version_minor)
3191*4882a593Smuzhiyun      padding_needed = (round_to_multiple(len(vbmeta_blob), image.block_size) -
3192*4882a593Smuzhiyun                        len(vbmeta_blob))
3193*4882a593Smuzhiyun      vbmeta_blob_with_padding = vbmeta_blob + '\0'*padding_needed
3194*4882a593Smuzhiyun
3195*4882a593Smuzhiyun      # Write vbmeta blob, if requested.
3196*4882a593Smuzhiyun      if output_vbmeta_image:
3197*4882a593Smuzhiyun        output_vbmeta_image.write(vbmeta_blob)
3198*4882a593Smuzhiyun
3199*4882a593Smuzhiyun      # Append vbmeta blob and footer, unless requested not to.
3200*4882a593Smuzhiyun      if not do_not_append_vbmeta_image:
3201*4882a593Smuzhiyun        image.append_raw(vbmeta_blob_with_padding)
3202*4882a593Smuzhiyun
3203*4882a593Smuzhiyun        # Now insert a DONT_CARE chunk with enough bytes such that the
3204*4882a593Smuzhiyun        # final Footer block is at the end of partition_size..
3205*4882a593Smuzhiyun        image.append_dont_care(partition_size - image.image_size -
3206*4882a593Smuzhiyun                               1*image.block_size)
3207*4882a593Smuzhiyun
3208*4882a593Smuzhiyun        # Generate the Footer that tells where the VBMeta footer
3209*4882a593Smuzhiyun        # is. Also put enough padding in the front of the footer since
3210*4882a593Smuzhiyun        # we'll write out an entire block.
3211*4882a593Smuzhiyun        footer = AvbFooter()
3212*4882a593Smuzhiyun        footer.original_image_size = original_image_size
3213*4882a593Smuzhiyun        footer.vbmeta_offset = vbmeta_offset
3214*4882a593Smuzhiyun        footer.vbmeta_size = len(vbmeta_blob)
3215*4882a593Smuzhiyun        footer_blob = footer.encode()
3216*4882a593Smuzhiyun        footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
3217*4882a593Smuzhiyun                                    footer_blob)
3218*4882a593Smuzhiyun        image.append_raw(footer_blob_with_padding)
3219*4882a593Smuzhiyun
3220*4882a593Smuzhiyun    except:
3221*4882a593Smuzhiyun      # Truncate back to original size, then re-raise.
3222*4882a593Smuzhiyun      image.truncate(original_image_size)
3223*4882a593Smuzhiyun      raise
3224*4882a593Smuzhiyun
3225*4882a593Smuzhiyun  def make_atx_certificate(self, output, authority_key_path, subject_key_path,
3226*4882a593Smuzhiyun                           subject_key_version, subject,
3227*4882a593Smuzhiyun                           is_intermediate_authority, usage, signing_helper,
3228*4882a593Smuzhiyun                           signing_helper_with_files):
3229*4882a593Smuzhiyun    """Implements the 'make_atx_certificate' command.
3230*4882a593Smuzhiyun
3231*4882a593Smuzhiyun    Android Things certificates are required for Android Things public key
3232*4882a593Smuzhiyun    metadata. They chain the vbmeta signing key for a particular product back to
3233*4882a593Smuzhiyun    a fused, permanent root key. These certificates are fixed-length and fixed-
3234*4882a593Smuzhiyun    format with the explicit goal of not parsing ASN.1 in bootloader code.
3235*4882a593Smuzhiyun
3236*4882a593Smuzhiyun    Arguments:
3237*4882a593Smuzhiyun      output: Certificate will be written to this file on success.
3238*4882a593Smuzhiyun      authority_key_path: A PEM file path with the authority private key.
3239*4882a593Smuzhiyun                          If None, then a certificate will be created without a
3240*4882a593Smuzhiyun                          signature. The signature can be created out-of-band
3241*4882a593Smuzhiyun                          and appended.
3242*4882a593Smuzhiyun      subject_key_path: Path to a PEM or DER subject public key.
3243*4882a593Smuzhiyun      subject_key_version: A 64-bit version value. If this is None, the number
3244*4882a593Smuzhiyun                           of seconds since the epoch is used.
3245*4882a593Smuzhiyun      subject: A subject identifier. For Product Signing Key certificates this
3246*4882a593Smuzhiyun               should be the same Product ID found in the permanent attributes.
3247*4882a593Smuzhiyun      is_intermediate_authority: True if the certificate is for an intermediate
3248*4882a593Smuzhiyun                                 authority.
3249*4882a593Smuzhiyun      usage: If not empty, overrides the cert usage with a hash of this value.
3250*4882a593Smuzhiyun      signing_helper: Program which signs a hash and returns the signature.
3251*4882a593Smuzhiyun      signing_helper_with_files: Same as signing_helper but uses files instead.
3252*4882a593Smuzhiyun    """
3253*4882a593Smuzhiyun    signed_data = bytearray()
3254*4882a593Smuzhiyun    signed_data.extend(struct.pack('<I', 1))  # Format Version
3255*4882a593Smuzhiyun    signed_data.extend(encode_rsa_key(subject_key_path))
3256*4882a593Smuzhiyun    hasher = hashlib.sha256()
3257*4882a593Smuzhiyun    hasher.update(subject)
3258*4882a593Smuzhiyun    signed_data.extend(hasher.digest())
3259*4882a593Smuzhiyun    if not usage:
3260*4882a593Smuzhiyun      usage = 'com.google.android.things.vboot'
3261*4882a593Smuzhiyun      if is_intermediate_authority:
3262*4882a593Smuzhiyun        usage += '.ca'
3263*4882a593Smuzhiyun    hasher = hashlib.sha256()
3264*4882a593Smuzhiyun    hasher.update(usage)
3265*4882a593Smuzhiyun    signed_data.extend(hasher.digest())
3266*4882a593Smuzhiyun    if not subject_key_version:
3267*4882a593Smuzhiyun      subject_key_version = int(time.time())
3268*4882a593Smuzhiyun    signed_data.extend(struct.pack('<Q', subject_key_version))
3269*4882a593Smuzhiyun    signature = bytearray()
3270*4882a593Smuzhiyun    if authority_key_path:
3271*4882a593Smuzhiyun      padding_and_hash = bytearray()
3272*4882a593Smuzhiyun      algorithm_name = 'SHA512_RSA4096'
3273*4882a593Smuzhiyun      alg = ALGORITHMS[algorithm_name]
3274*4882a593Smuzhiyun      hasher = hashlib.sha512()
3275*4882a593Smuzhiyun      padding_and_hash.extend(alg.padding)
3276*4882a593Smuzhiyun      hasher.update(signed_data)
3277*4882a593Smuzhiyun      padding_and_hash.extend(hasher.digest())
3278*4882a593Smuzhiyun      signature.extend(raw_sign(signing_helper, signing_helper_with_files,
3279*4882a593Smuzhiyun                                algorithm_name,
3280*4882a593Smuzhiyun                                alg.signature_num_bytes, authority_key_path,
3281*4882a593Smuzhiyun                                padding_and_hash))
3282*4882a593Smuzhiyun    output.write(signed_data)
3283*4882a593Smuzhiyun    output.write(signature)
3284*4882a593Smuzhiyun
3285*4882a593Smuzhiyun  def make_atx_permanent_attributes(self, output, root_authority_key_path,
3286*4882a593Smuzhiyun                                    product_id):
3287*4882a593Smuzhiyun    """Implements the 'make_atx_permanent_attributes' command.
3288*4882a593Smuzhiyun
3289*4882a593Smuzhiyun    Android Things permanent attributes are designed to be permanent for a
3290*4882a593Smuzhiyun    particular product and a hash of these attributes should be fused into
3291*4882a593Smuzhiyun    hardware to enforce this.
3292*4882a593Smuzhiyun
3293*4882a593Smuzhiyun    Arguments:
3294*4882a593Smuzhiyun      output: Attributes will be written to this file on success.
3295*4882a593Smuzhiyun      root_authority_key_path: Path to a PEM or DER public key for
3296*4882a593Smuzhiyun        the root authority.
3297*4882a593Smuzhiyun      product_id: A 16-byte Product ID.
3298*4882a593Smuzhiyun
3299*4882a593Smuzhiyun    Raises:
3300*4882a593Smuzhiyun      AvbError: If an argument is incorrect.
3301*4882a593Smuzhiyun    """
3302*4882a593Smuzhiyun    EXPECTED_PRODUCT_ID_SIZE = 16
3303*4882a593Smuzhiyun    if len(product_id) != EXPECTED_PRODUCT_ID_SIZE:
3304*4882a593Smuzhiyun      raise AvbError('Invalid Product ID length.')
3305*4882a593Smuzhiyun    output.write(struct.pack('<I', 1))  # Format Version
3306*4882a593Smuzhiyun    output.write(encode_rsa_key(root_authority_key_path))
3307*4882a593Smuzhiyun    output.write(product_id)
3308*4882a593Smuzhiyun
3309*4882a593Smuzhiyun  def make_atx_metadata(self, output, intermediate_key_certificate,
3310*4882a593Smuzhiyun                        product_key_certificate):
3311*4882a593Smuzhiyun    """Implements the 'make_atx_metadata' command.
3312*4882a593Smuzhiyun
3313*4882a593Smuzhiyun    Android Things metadata are included in vbmeta images to facilitate
3314*4882a593Smuzhiyun    verification. The output of this command can be used as the
3315*4882a593Smuzhiyun    public_key_metadata argument to other commands.
3316*4882a593Smuzhiyun
3317*4882a593Smuzhiyun    Arguments:
3318*4882a593Smuzhiyun      output: Metadata will be written to this file on success.
3319*4882a593Smuzhiyun      intermediate_key_certificate: A certificate file as output by
3320*4882a593Smuzhiyun                                    make_atx_certificate with
3321*4882a593Smuzhiyun                                    is_intermediate_authority set to true.
3322*4882a593Smuzhiyun      product_key_certificate: A certificate file as output by
3323*4882a593Smuzhiyun                               make_atx_certificate with
3324*4882a593Smuzhiyun                               is_intermediate_authority set to false.
3325*4882a593Smuzhiyun
3326*4882a593Smuzhiyun    Raises:
3327*4882a593Smuzhiyun      AvbError: If an argument is incorrect.
3328*4882a593Smuzhiyun    """
3329*4882a593Smuzhiyun    EXPECTED_CERTIFICATE_SIZE = 1620
3330*4882a593Smuzhiyun    if len(intermediate_key_certificate) != EXPECTED_CERTIFICATE_SIZE:
3331*4882a593Smuzhiyun      raise AvbError('Invalid intermediate key certificate length.')
3332*4882a593Smuzhiyun    if len(product_key_certificate) != EXPECTED_CERTIFICATE_SIZE:
3333*4882a593Smuzhiyun      raise AvbError('Invalid product key certificate length.')
3334*4882a593Smuzhiyun    output.write(struct.pack('<I', 1))  # Format Version
3335*4882a593Smuzhiyun    output.write(intermediate_key_certificate)
3336*4882a593Smuzhiyun    output.write(product_key_certificate)
3337*4882a593Smuzhiyun
3338*4882a593Smuzhiyun  def make_atx_unlock_credential(self, output, intermediate_key_certificate,
3339*4882a593Smuzhiyun                                 unlock_key_certificate, challenge_path,
3340*4882a593Smuzhiyun                                 unlock_key_path, signing_helper,
3341*4882a593Smuzhiyun                                 signing_helper_with_files):
3342*4882a593Smuzhiyun    """Implements the 'make_atx_unlock_credential' command.
3343*4882a593Smuzhiyun
3344*4882a593Smuzhiyun    Android Things unlock credentials can be used to authorize the unlock of AVB
3345*4882a593Smuzhiyun    on a device. These credentials are presented to an Android Things bootloader
3346*4882a593Smuzhiyun    via the fastboot interface in response to a 16-byte challenge. This method
3347*4882a593Smuzhiyun    creates all fields of the credential except the challenge signature field
3348*4882a593Smuzhiyun    (which is the last field) and can optionally create the challenge signature
3349*4882a593Smuzhiyun    field as well if a challenge and the unlock_key_path is provided.
3350*4882a593Smuzhiyun
3351*4882a593Smuzhiyun    Arguments:
3352*4882a593Smuzhiyun      output: The credential will be written to this file on success.
3353*4882a593Smuzhiyun      intermediate_key_certificate: A certificate file as output by
3354*4882a593Smuzhiyun                                    make_atx_certificate with
3355*4882a593Smuzhiyun                                    is_intermediate_authority set to true.
3356*4882a593Smuzhiyun      unlock_key_certificate: A certificate file as output by
3357*4882a593Smuzhiyun                              make_atx_certificate with
3358*4882a593Smuzhiyun                              is_intermediate_authority set to false and the
3359*4882a593Smuzhiyun                              usage set to
3360*4882a593Smuzhiyun                              'com.google.android.things.vboot.unlock'.
3361*4882a593Smuzhiyun      challenge_path: [optional] A path to the challenge to sign.
3362*4882a593Smuzhiyun      unlock_key_path: [optional] A PEM file path with the unlock private key.
3363*4882a593Smuzhiyun      signing_helper: Program which signs a hash and returns the signature.
3364*4882a593Smuzhiyun      signing_helper_with_files: Same as signing_helper but uses files instead.
3365*4882a593Smuzhiyun
3366*4882a593Smuzhiyun    Raises:
3367*4882a593Smuzhiyun      AvbError: If an argument is incorrect.
3368*4882a593Smuzhiyun    """
3369*4882a593Smuzhiyun    EXPECTED_CERTIFICATE_SIZE = 1620
3370*4882a593Smuzhiyun    EXPECTED_CHALLENGE_SIZE = 16
3371*4882a593Smuzhiyun    if len(intermediate_key_certificate) != EXPECTED_CERTIFICATE_SIZE:
3372*4882a593Smuzhiyun      raise AvbError('Invalid intermediate key certificate length.')
3373*4882a593Smuzhiyun    if len(unlock_key_certificate) != EXPECTED_CERTIFICATE_SIZE:
3374*4882a593Smuzhiyun      raise AvbError('Invalid product key certificate length.')
3375*4882a593Smuzhiyun    challenge = bytearray()
3376*4882a593Smuzhiyun    if challenge_path:
3377*4882a593Smuzhiyun      with open(challenge_path, 'r') as f:
3378*4882a593Smuzhiyun        challenge = f.read()
3379*4882a593Smuzhiyun      if len(challenge) != EXPECTED_CHALLENGE_SIZE:
3380*4882a593Smuzhiyun        raise AvbError('Invalid unlock challenge length.')
3381*4882a593Smuzhiyun    output.write(struct.pack('<I', 1))  # Format Version
3382*4882a593Smuzhiyun    output.write(intermediate_key_certificate)
3383*4882a593Smuzhiyun    output.write(unlock_key_certificate)
3384*4882a593Smuzhiyun    if challenge_path and unlock_key_path:
3385*4882a593Smuzhiyun      signature = bytearray()
3386*4882a593Smuzhiyun      padding_and_hash = bytearray()
3387*4882a593Smuzhiyun      algorithm_name = 'SHA512_RSA4096'
3388*4882a593Smuzhiyun      alg = ALGORITHMS[algorithm_name]
3389*4882a593Smuzhiyun      hasher = hashlib.sha512()
3390*4882a593Smuzhiyun      padding_and_hash.extend(alg.padding)
3391*4882a593Smuzhiyun      hasher.update(challenge)
3392*4882a593Smuzhiyun      padding_and_hash.extend(hasher.digest())
3393*4882a593Smuzhiyun      signature.extend(raw_sign(signing_helper, signing_helper_with_files,
3394*4882a593Smuzhiyun                                algorithm_name,
3395*4882a593Smuzhiyun                                alg.signature_num_bytes, unlock_key_path,
3396*4882a593Smuzhiyun                                padding_and_hash))
3397*4882a593Smuzhiyun      output.write(signature)
3398*4882a593Smuzhiyun
3399*4882a593Smuzhiyun
3400*4882a593Smuzhiyundef calc_hash_level_offsets(image_size, block_size, digest_size):
3401*4882a593Smuzhiyun  """Calculate the offsets of all the hash-levels in a Merkle-tree.
3402*4882a593Smuzhiyun
3403*4882a593Smuzhiyun  Arguments:
3404*4882a593Smuzhiyun    image_size: The size of the image to calculate a Merkle-tree for.
3405*4882a593Smuzhiyun    block_size: The block size, e.g. 4096.
3406*4882a593Smuzhiyun    digest_size: The size of each hash, e.g. 32 for SHA-256.
3407*4882a593Smuzhiyun
3408*4882a593Smuzhiyun  Returns:
3409*4882a593Smuzhiyun    A tuple where the first argument is an array of offsets and the
3410*4882a593Smuzhiyun    second is size of the tree, in bytes.
3411*4882a593Smuzhiyun  """
3412*4882a593Smuzhiyun  level_offsets = []
3413*4882a593Smuzhiyun  level_sizes = []
3414*4882a593Smuzhiyun  tree_size = 0
3415*4882a593Smuzhiyun
3416*4882a593Smuzhiyun  num_levels = 0
3417*4882a593Smuzhiyun  size = image_size
3418*4882a593Smuzhiyun  while size > block_size:
3419*4882a593Smuzhiyun    num_blocks = (size + block_size - 1) / block_size
3420*4882a593Smuzhiyun    level_size = round_to_multiple(num_blocks * digest_size, block_size)
3421*4882a593Smuzhiyun
3422*4882a593Smuzhiyun    level_sizes.append(level_size)
3423*4882a593Smuzhiyun    tree_size += level_size
3424*4882a593Smuzhiyun    num_levels += 1
3425*4882a593Smuzhiyun
3426*4882a593Smuzhiyun    size = level_size
3427*4882a593Smuzhiyun
3428*4882a593Smuzhiyun  for n in range(0, num_levels):
3429*4882a593Smuzhiyun    offset = 0
3430*4882a593Smuzhiyun    for m in range(n + 1, num_levels):
3431*4882a593Smuzhiyun      offset += level_sizes[m]
3432*4882a593Smuzhiyun    level_offsets.append(offset)
3433*4882a593Smuzhiyun
3434*4882a593Smuzhiyun  return level_offsets, tree_size
3435*4882a593Smuzhiyun
3436*4882a593Smuzhiyun
3437*4882a593Smuzhiyun# See system/extras/libfec/include/fec/io.h for these definitions.
3438*4882a593SmuzhiyunFEC_FOOTER_FORMAT = '<LLLLLQ32s'
3439*4882a593SmuzhiyunFEC_MAGIC = 0xfecfecfe
3440*4882a593Smuzhiyun
3441*4882a593Smuzhiyun
3442*4882a593Smuzhiyundef calc_fec_data_size(image_size, num_roots):
3443*4882a593Smuzhiyun  """Calculates how much space FEC data will take.
3444*4882a593Smuzhiyun
3445*4882a593Smuzhiyun  Args:
3446*4882a593Smuzhiyun    image_size: The size of the image.
3447*4882a593Smuzhiyun    num_roots: Number of roots.
3448*4882a593Smuzhiyun
3449*4882a593Smuzhiyun  Returns:
3450*4882a593Smuzhiyun    The number of bytes needed for FEC for an image of the given size
3451*4882a593Smuzhiyun    and with the requested number of FEC roots.
3452*4882a593Smuzhiyun
3453*4882a593Smuzhiyun  Raises:
3454*4882a593Smuzhiyun    ValueError: If output from the 'fec' tool is invalid.
3455*4882a593Smuzhiyun
3456*4882a593Smuzhiyun  """
3457*4882a593Smuzhiyun  p = subprocess.Popen(
3458*4882a593Smuzhiyun      ['fec', '--print-fec-size', str(image_size), '--roots', str(num_roots)],
3459*4882a593Smuzhiyun      stdout=subprocess.PIPE,
3460*4882a593Smuzhiyun      stderr=subprocess.PIPE)
3461*4882a593Smuzhiyun  (pout, perr) = p.communicate()
3462*4882a593Smuzhiyun  retcode = p.wait()
3463*4882a593Smuzhiyun  if retcode != 0:
3464*4882a593Smuzhiyun    raise ValueError('Error invoking fec: {}'.format(perr))
3465*4882a593Smuzhiyun  return int(pout)
3466*4882a593Smuzhiyun
3467*4882a593Smuzhiyun
3468*4882a593Smuzhiyundef generate_fec_data(image_filename, num_roots):
3469*4882a593Smuzhiyun  """Generate FEC codes for an image.
3470*4882a593Smuzhiyun
3471*4882a593Smuzhiyun  Args:
3472*4882a593Smuzhiyun    image_filename: The filename of the image.
3473*4882a593Smuzhiyun    num_roots: Number of roots.
3474*4882a593Smuzhiyun
3475*4882a593Smuzhiyun  Returns:
3476*4882a593Smuzhiyun    The FEC data blob.
3477*4882a593Smuzhiyun
3478*4882a593Smuzhiyun  Raises:
3479*4882a593Smuzhiyun    ValueError: If output from the 'fec' tool is invalid.
3480*4882a593Smuzhiyun  """
3481*4882a593Smuzhiyun  fec_tmpfile = tempfile.NamedTemporaryFile()
3482*4882a593Smuzhiyun  subprocess.check_call(
3483*4882a593Smuzhiyun      ['fec', '--encode', '--roots', str(num_roots), image_filename,
3484*4882a593Smuzhiyun       fec_tmpfile.name],
3485*4882a593Smuzhiyun      stderr=open(os.devnull))
3486*4882a593Smuzhiyun  fec_data = fec_tmpfile.read()
3487*4882a593Smuzhiyun  footer_size = struct.calcsize(FEC_FOOTER_FORMAT)
3488*4882a593Smuzhiyun  footer_data = fec_data[-footer_size:]
3489*4882a593Smuzhiyun  (magic, _, _, num_roots, fec_size, _, _) = struct.unpack(FEC_FOOTER_FORMAT,
3490*4882a593Smuzhiyun                                                           footer_data)
3491*4882a593Smuzhiyun  if magic != FEC_MAGIC:
3492*4882a593Smuzhiyun    raise ValueError('Unexpected magic in FEC footer')
3493*4882a593Smuzhiyun  return fec_data[0:fec_size]
3494*4882a593Smuzhiyun
3495*4882a593Smuzhiyun
3496*4882a593Smuzhiyundef generate_hash_tree(image, image_size, block_size, hash_alg_name, salt,
3497*4882a593Smuzhiyun                       digest_padding, hash_level_offsets, tree_size):
3498*4882a593Smuzhiyun  """Generates a Merkle-tree for a file.
3499*4882a593Smuzhiyun
3500*4882a593Smuzhiyun  Args:
3501*4882a593Smuzhiyun    image: The image, as a file.
3502*4882a593Smuzhiyun    image_size: The size of the image.
3503*4882a593Smuzhiyun    block_size: The block size, e.g. 4096.
3504*4882a593Smuzhiyun    hash_alg_name: The hash algorithm, e.g. 'sha256' or 'sha1'.
3505*4882a593Smuzhiyun    salt: The salt to use.
3506*4882a593Smuzhiyun    digest_padding: The padding for each digest.
3507*4882a593Smuzhiyun    hash_level_offsets: The offsets from calc_hash_level_offsets().
3508*4882a593Smuzhiyun    tree_size: The size of the tree, in number of bytes.
3509*4882a593Smuzhiyun
3510*4882a593Smuzhiyun  Returns:
3511*4882a593Smuzhiyun    A tuple where the first element is the top-level hash and the
3512*4882a593Smuzhiyun    second element is the hash-tree.
3513*4882a593Smuzhiyun  """
3514*4882a593Smuzhiyun  hash_ret = bytearray(tree_size)
3515*4882a593Smuzhiyun  hash_src_offset = 0
3516*4882a593Smuzhiyun  hash_src_size = image_size
3517*4882a593Smuzhiyun  level_num = 0
3518*4882a593Smuzhiyun  while hash_src_size > block_size:
3519*4882a593Smuzhiyun    level_output = ''
3520*4882a593Smuzhiyun    remaining = hash_src_size
3521*4882a593Smuzhiyun    while remaining > 0:
3522*4882a593Smuzhiyun      hasher = hashlib.new(name=hash_alg_name, string=salt)
3523*4882a593Smuzhiyun      # Only read from the file for the first level - for subsequent
3524*4882a593Smuzhiyun      # levels, access the array we're building.
3525*4882a593Smuzhiyun      if level_num == 0:
3526*4882a593Smuzhiyun        image.seek(hash_src_offset + hash_src_size - remaining)
3527*4882a593Smuzhiyun        data = image.read(min(remaining, block_size))
3528*4882a593Smuzhiyun      else:
3529*4882a593Smuzhiyun        offset = hash_level_offsets[level_num - 1] + hash_src_size - remaining
3530*4882a593Smuzhiyun        data = hash_ret[offset:offset + block_size]
3531*4882a593Smuzhiyun      hasher.update(data)
3532*4882a593Smuzhiyun
3533*4882a593Smuzhiyun      remaining -= len(data)
3534*4882a593Smuzhiyun      if len(data) < block_size:
3535*4882a593Smuzhiyun        hasher.update('\0' * (block_size - len(data)))
3536*4882a593Smuzhiyun      level_output += hasher.digest()
3537*4882a593Smuzhiyun      if digest_padding > 0:
3538*4882a593Smuzhiyun        level_output += '\0' * digest_padding
3539*4882a593Smuzhiyun
3540*4882a593Smuzhiyun    padding_needed = (round_to_multiple(
3541*4882a593Smuzhiyun        len(level_output), block_size) - len(level_output))
3542*4882a593Smuzhiyun    level_output += '\0' * padding_needed
3543*4882a593Smuzhiyun
3544*4882a593Smuzhiyun    # Copy level-output into resulting tree.
3545*4882a593Smuzhiyun    offset = hash_level_offsets[level_num]
3546*4882a593Smuzhiyun    hash_ret[offset:offset + len(level_output)] = level_output
3547*4882a593Smuzhiyun
3548*4882a593Smuzhiyun    # Continue on to the next level.
3549*4882a593Smuzhiyun    hash_src_size = len(level_output)
3550*4882a593Smuzhiyun    level_num += 1
3551*4882a593Smuzhiyun
3552*4882a593Smuzhiyun  hasher = hashlib.new(name=hash_alg_name, string=salt)
3553*4882a593Smuzhiyun  hasher.update(level_output)
3554*4882a593Smuzhiyun  return hasher.digest(), hash_ret
3555*4882a593Smuzhiyun
3556*4882a593Smuzhiyun
3557*4882a593Smuzhiyunclass AvbTool(object):
3558*4882a593Smuzhiyun  """Object for avbtool command-line tool."""
3559*4882a593Smuzhiyun
3560*4882a593Smuzhiyun  def __init__(self):
3561*4882a593Smuzhiyun    """Initializer method."""
3562*4882a593Smuzhiyun    self.avb = Avb()
3563*4882a593Smuzhiyun
3564*4882a593Smuzhiyun  def _add_common_args(self, sub_parser):
3565*4882a593Smuzhiyun    """Adds arguments used by several sub-commands.
3566*4882a593Smuzhiyun
3567*4882a593Smuzhiyun    Arguments:
3568*4882a593Smuzhiyun      sub_parser: The parser to add arguments to.
3569*4882a593Smuzhiyun    """
3570*4882a593Smuzhiyun    sub_parser.add_argument('--algorithm',
3571*4882a593Smuzhiyun                            help='Algorithm to use (default: NONE)',
3572*4882a593Smuzhiyun                            metavar='ALGORITHM',
3573*4882a593Smuzhiyun                            default='NONE')
3574*4882a593Smuzhiyun    sub_parser.add_argument('--key',
3575*4882a593Smuzhiyun                            help='Path to RSA private key file',
3576*4882a593Smuzhiyun                            metavar='KEY',
3577*4882a593Smuzhiyun                            required=False)
3578*4882a593Smuzhiyun    sub_parser.add_argument('--signing_helper',
3579*4882a593Smuzhiyun                            help='Path to helper used for signing',
3580*4882a593Smuzhiyun                            metavar='APP',
3581*4882a593Smuzhiyun                            default=None,
3582*4882a593Smuzhiyun                            required=False)
3583*4882a593Smuzhiyun    sub_parser.add_argument('--signing_helper_with_files',
3584*4882a593Smuzhiyun                            help='Path to helper used for signing using files',
3585*4882a593Smuzhiyun                            metavar='APP',
3586*4882a593Smuzhiyun                            default=None,
3587*4882a593Smuzhiyun                            required=False)
3588*4882a593Smuzhiyun    sub_parser.add_argument('--public_key_metadata',
3589*4882a593Smuzhiyun                            help='Path to public key metadata file',
3590*4882a593Smuzhiyun                            metavar='KEY_METADATA',
3591*4882a593Smuzhiyun                            required=False)
3592*4882a593Smuzhiyun    sub_parser.add_argument('--rollback_index',
3593*4882a593Smuzhiyun                            help='Rollback Index',
3594*4882a593Smuzhiyun                            type=parse_number,
3595*4882a593Smuzhiyun                            default=0)
3596*4882a593Smuzhiyun    # This is used internally for unit tests. Do not include in --help output.
3597*4882a593Smuzhiyun    sub_parser.add_argument('--internal_release_string',
3598*4882a593Smuzhiyun                            help=argparse.SUPPRESS)
3599*4882a593Smuzhiyun    sub_parser.add_argument('--append_to_release_string',
3600*4882a593Smuzhiyun                            help='Text to append to release string',
3601*4882a593Smuzhiyun                            metavar='STR')
3602*4882a593Smuzhiyun    sub_parser.add_argument('--prop',
3603*4882a593Smuzhiyun                            help='Add property',
3604*4882a593Smuzhiyun                            metavar='KEY:VALUE',
3605*4882a593Smuzhiyun                            action='append')
3606*4882a593Smuzhiyun    sub_parser.add_argument('--prop_from_file',
3607*4882a593Smuzhiyun                            help='Add property from file',
3608*4882a593Smuzhiyun                            metavar='KEY:PATH',
3609*4882a593Smuzhiyun                            action='append')
3610*4882a593Smuzhiyun    sub_parser.add_argument('--kernel_cmdline',
3611*4882a593Smuzhiyun                            help='Add kernel cmdline',
3612*4882a593Smuzhiyun                            metavar='CMDLINE',
3613*4882a593Smuzhiyun                            action='append')
3614*4882a593Smuzhiyun    # TODO(zeuthen): the --setup_rootfs_from_kernel option used to be called
3615*4882a593Smuzhiyun    # --generate_dm_verity_cmdline_from_hashtree. Remove support for the latter
3616*4882a593Smuzhiyun    # at some future point.
3617*4882a593Smuzhiyun    sub_parser.add_argument('--setup_rootfs_from_kernel',
3618*4882a593Smuzhiyun                            '--generate_dm_verity_cmdline_from_hashtree',
3619*4882a593Smuzhiyun                            metavar='IMAGE',
3620*4882a593Smuzhiyun                            help='Adds kernel cmdline to set up IMAGE',
3621*4882a593Smuzhiyun                            type=argparse.FileType('rb'))
3622*4882a593Smuzhiyun    sub_parser.add_argument('--include_descriptors_from_image',
3623*4882a593Smuzhiyun                            help='Include descriptors from image',
3624*4882a593Smuzhiyun                            metavar='IMAGE',
3625*4882a593Smuzhiyun                            action='append',
3626*4882a593Smuzhiyun                            type=argparse.FileType('rb'))
3627*4882a593Smuzhiyun    sub_parser.add_argument('--print_required_libavb_version',
3628*4882a593Smuzhiyun                            help=('Don\'t store the footer - '
3629*4882a593Smuzhiyun                                  'instead calculate the required libavb '
3630*4882a593Smuzhiyun                                  'version for the given options.'),
3631*4882a593Smuzhiyun                            action='store_true')
3632*4882a593Smuzhiyun    # These are only allowed from top-level vbmeta and boot-in-lieu-of-vbmeta.
3633*4882a593Smuzhiyun    sub_parser.add_argument('--chain_partition',
3634*4882a593Smuzhiyun                            help='Allow signed integrity-data for partition',
3635*4882a593Smuzhiyun                            metavar='PART_NAME:ROLLBACK_SLOT:KEY_PATH',
3636*4882a593Smuzhiyun                            action='append')
3637*4882a593Smuzhiyun    sub_parser.add_argument('--flags',
3638*4882a593Smuzhiyun                            help='VBMeta flags',
3639*4882a593Smuzhiyun                            type=parse_number,
3640*4882a593Smuzhiyun                            default=0)
3641*4882a593Smuzhiyun    sub_parser.add_argument('--set_hashtree_disabled_flag',
3642*4882a593Smuzhiyun                            help='Set the HASHTREE_DISABLED flag',
3643*4882a593Smuzhiyun                            action='store_true')
3644*4882a593Smuzhiyun
3645*4882a593Smuzhiyun  def _add_common_footer_args(self, sub_parser):
3646*4882a593Smuzhiyun    """Adds arguments used by add_*_footer sub-commands.
3647*4882a593Smuzhiyun
3648*4882a593Smuzhiyun    Arguments:
3649*4882a593Smuzhiyun      sub_parser: The parser to add arguments to.
3650*4882a593Smuzhiyun    """
3651*4882a593Smuzhiyun    sub_parser.add_argument('--use_persistent_digest',
3652*4882a593Smuzhiyun                            help='Use a persistent digest on device instead of '
3653*4882a593Smuzhiyun                                 'storing the digest in the descriptor. This '
3654*4882a593Smuzhiyun                                 'cannot be used with A/B so must be combined '
3655*4882a593Smuzhiyun                                 'with --do_not_use_ab when an A/B suffix is '
3656*4882a593Smuzhiyun                                 'expected at runtime.',
3657*4882a593Smuzhiyun                            action='store_true')
3658*4882a593Smuzhiyun    sub_parser.add_argument('--do_not_use_ab',
3659*4882a593Smuzhiyun                            help='The partition does not use A/B even when an '
3660*4882a593Smuzhiyun                                 'A/B suffix is present. This must not be used '
3661*4882a593Smuzhiyun                                 'for vbmeta or chained partitions.',
3662*4882a593Smuzhiyun                            action='store_true')
3663*4882a593Smuzhiyun
3664*4882a593Smuzhiyun  def _fixup_common_args(self, args):
3665*4882a593Smuzhiyun    """Common fixups needed by subcommands.
3666*4882a593Smuzhiyun
3667*4882a593Smuzhiyun    Arguments:
3668*4882a593Smuzhiyun      args: Arguments to modify.
3669*4882a593Smuzhiyun
3670*4882a593Smuzhiyun    Returns:
3671*4882a593Smuzhiyun      The modified arguments.
3672*4882a593Smuzhiyun    """
3673*4882a593Smuzhiyun    if args.set_hashtree_disabled_flag:
3674*4882a593Smuzhiyun      args.flags |= AVB_VBMETA_IMAGE_FLAGS_HASHTREE_DISABLED
3675*4882a593Smuzhiyun    return args
3676*4882a593Smuzhiyun
3677*4882a593Smuzhiyun  def run(self, argv):
3678*4882a593Smuzhiyun    """Command-line processor.
3679*4882a593Smuzhiyun
3680*4882a593Smuzhiyun    Arguments:
3681*4882a593Smuzhiyun      argv: Pass sys.argv from main.
3682*4882a593Smuzhiyun    """
3683*4882a593Smuzhiyun    parser = argparse.ArgumentParser()
3684*4882a593Smuzhiyun    subparsers = parser.add_subparsers(title='subcommands')
3685*4882a593Smuzhiyun
3686*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('version',
3687*4882a593Smuzhiyun                                       help='Prints version of avbtool.')
3688*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.version)
3689*4882a593Smuzhiyun
3690*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('extract_public_key',
3691*4882a593Smuzhiyun                                       help='Extract public key.')
3692*4882a593Smuzhiyun    sub_parser.add_argument('--key',
3693*4882a593Smuzhiyun                            help='Path to RSA private key file',
3694*4882a593Smuzhiyun                            required=True)
3695*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3696*4882a593Smuzhiyun                            help='Output file name',
3697*4882a593Smuzhiyun                            type=argparse.FileType('wb'),
3698*4882a593Smuzhiyun                            required=True)
3699*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.extract_public_key)
3700*4882a593Smuzhiyun
3701*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('make_vbmeta_image',
3702*4882a593Smuzhiyun                                       help='Makes a vbmeta image.')
3703*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3704*4882a593Smuzhiyun                            help='Output file name',
3705*4882a593Smuzhiyun                            type=argparse.FileType('wb'))
3706*4882a593Smuzhiyun    sub_parser.add_argument('--padding_size',
3707*4882a593Smuzhiyun                            metavar='NUMBER',
3708*4882a593Smuzhiyun                            help='If non-zero, pads output with NUL bytes so '
3709*4882a593Smuzhiyun                                 'its size is a multiple of NUMBER (default: 0)',
3710*4882a593Smuzhiyun                            type=parse_number,
3711*4882a593Smuzhiyun                            default=0)
3712*4882a593Smuzhiyun    self._add_common_args(sub_parser)
3713*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.make_vbmeta_image)
3714*4882a593Smuzhiyun
3715*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('add_hash_footer',
3716*4882a593Smuzhiyun                                       help='Add hashes and footer to image.')
3717*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3718*4882a593Smuzhiyun                            help='Image to add hashes to',
3719*4882a593Smuzhiyun                            type=argparse.FileType('rab+'))
3720*4882a593Smuzhiyun    sub_parser.add_argument('--partition_size',
3721*4882a593Smuzhiyun                            help='Partition size',
3722*4882a593Smuzhiyun                            type=parse_number)
3723*4882a593Smuzhiyun    sub_parser.add_argument('--partition_name',
3724*4882a593Smuzhiyun                            help='Partition name',
3725*4882a593Smuzhiyun                            default=None)
3726*4882a593Smuzhiyun    sub_parser.add_argument('--hash_algorithm',
3727*4882a593Smuzhiyun                            help='Hash algorithm to use (default: sha256)',
3728*4882a593Smuzhiyun                            default='sha256')
3729*4882a593Smuzhiyun    sub_parser.add_argument('--salt',
3730*4882a593Smuzhiyun                            help='Salt in hex (default: /dev/urandom)')
3731*4882a593Smuzhiyun    sub_parser.add_argument('--calc_max_image_size',
3732*4882a593Smuzhiyun                            help=('Don\'t store the footer - '
3733*4882a593Smuzhiyun                                  'instead calculate the maximum image size '
3734*4882a593Smuzhiyun                                  'leaving enough room for metadata with '
3735*4882a593Smuzhiyun                                  'the given partition size.'),
3736*4882a593Smuzhiyun                            action='store_true')
3737*4882a593Smuzhiyun    sub_parser.add_argument('--output_vbmeta_image',
3738*4882a593Smuzhiyun                            help='Also write vbmeta struct to file',
3739*4882a593Smuzhiyun                            type=argparse.FileType('wb'))
3740*4882a593Smuzhiyun    sub_parser.add_argument('--do_not_append_vbmeta_image',
3741*4882a593Smuzhiyun                            help=('Do not append vbmeta struct or footer '
3742*4882a593Smuzhiyun                                  'to the image'),
3743*4882a593Smuzhiyun                            action='store_true')
3744*4882a593Smuzhiyun    self._add_common_args(sub_parser)
3745*4882a593Smuzhiyun    self._add_common_footer_args(sub_parser)
3746*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.add_hash_footer)
3747*4882a593Smuzhiyun
3748*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('append_vbmeta_image',
3749*4882a593Smuzhiyun                                       help='Append vbmeta image to image.')
3750*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3751*4882a593Smuzhiyun                            help='Image to append vbmeta blob to',
3752*4882a593Smuzhiyun                            type=argparse.FileType('rab+'))
3753*4882a593Smuzhiyun    sub_parser.add_argument('--partition_size',
3754*4882a593Smuzhiyun                            help='Partition size',
3755*4882a593Smuzhiyun                            type=parse_number,
3756*4882a593Smuzhiyun                            required=True)
3757*4882a593Smuzhiyun    sub_parser.add_argument('--vbmeta_image',
3758*4882a593Smuzhiyun                            help='Image with vbmeta blob to append',
3759*4882a593Smuzhiyun                            type=argparse.FileType('rb'))
3760*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.append_vbmeta_image)
3761*4882a593Smuzhiyun
3762*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('add_hashtree_footer',
3763*4882a593Smuzhiyun                                       help='Add hashtree and footer to image.')
3764*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3765*4882a593Smuzhiyun                            help='Image to add hashtree to',
3766*4882a593Smuzhiyun                            type=argparse.FileType('rab+'))
3767*4882a593Smuzhiyun    sub_parser.add_argument('--partition_size',
3768*4882a593Smuzhiyun                            help='Partition size',
3769*4882a593Smuzhiyun                            type=parse_number)
3770*4882a593Smuzhiyun    sub_parser.add_argument('--partition_name',
3771*4882a593Smuzhiyun                            help='Partition name',
3772*4882a593Smuzhiyun                            default=None)
3773*4882a593Smuzhiyun    sub_parser.add_argument('--hash_algorithm',
3774*4882a593Smuzhiyun                            help='Hash algorithm to use (default: sha1)',
3775*4882a593Smuzhiyun                            default='sha1')
3776*4882a593Smuzhiyun    sub_parser.add_argument('--salt',
3777*4882a593Smuzhiyun                            help='Salt in hex (default: /dev/urandom)')
3778*4882a593Smuzhiyun    sub_parser.add_argument('--block_size',
3779*4882a593Smuzhiyun                            help='Block size (default: 4096)',
3780*4882a593Smuzhiyun                            type=parse_number,
3781*4882a593Smuzhiyun                            default=4096)
3782*4882a593Smuzhiyun    # TODO(zeuthen): The --generate_fec option was removed when we
3783*4882a593Smuzhiyun    # moved to generating FEC by default. To avoid breaking existing
3784*4882a593Smuzhiyun    # users needing to transition we simply just print a warning below
3785*4882a593Smuzhiyun    # in add_hashtree_footer(). Remove this option and the warning at
3786*4882a593Smuzhiyun    # some point in the future.
3787*4882a593Smuzhiyun    sub_parser.add_argument('--generate_fec',
3788*4882a593Smuzhiyun                            help=argparse.SUPPRESS,
3789*4882a593Smuzhiyun                            action='store_true')
3790*4882a593Smuzhiyun    sub_parser.add_argument('--do_not_generate_fec',
3791*4882a593Smuzhiyun                            help='Do not generate forward-error-correction codes',
3792*4882a593Smuzhiyun                            action='store_true')
3793*4882a593Smuzhiyun    sub_parser.add_argument('--fec_num_roots',
3794*4882a593Smuzhiyun                            help='Number of roots for FEC (default: 2)',
3795*4882a593Smuzhiyun                            type=parse_number,
3796*4882a593Smuzhiyun                            default=2)
3797*4882a593Smuzhiyun    sub_parser.add_argument('--calc_max_image_size',
3798*4882a593Smuzhiyun                            help=('Don\'t store the hashtree or footer - '
3799*4882a593Smuzhiyun                                  'instead calculate the maximum image size '
3800*4882a593Smuzhiyun                                  'leaving enough room for hashtree '
3801*4882a593Smuzhiyun                                  'and metadata with the given partition '
3802*4882a593Smuzhiyun                                  'size.'),
3803*4882a593Smuzhiyun                            action='store_true')
3804*4882a593Smuzhiyun    sub_parser.add_argument('--output_vbmeta_image',
3805*4882a593Smuzhiyun                            help='Also write vbmeta struct to file',
3806*4882a593Smuzhiyun                            type=argparse.FileType('wb'))
3807*4882a593Smuzhiyun    sub_parser.add_argument('--do_not_append_vbmeta_image',
3808*4882a593Smuzhiyun                            help=('Do not append vbmeta struct or footer '
3809*4882a593Smuzhiyun                                  'to the image'),
3810*4882a593Smuzhiyun                            action='store_true')
3811*4882a593Smuzhiyun    # This is different from --setup_rootfs_from_kernel insofar that
3812*4882a593Smuzhiyun    # it doesn't take an IMAGE, the generated cmdline will be for the
3813*4882a593Smuzhiyun    # hashtree we're adding.
3814*4882a593Smuzhiyun    sub_parser.add_argument('--setup_as_rootfs_from_kernel',
3815*4882a593Smuzhiyun                            action='store_true',
3816*4882a593Smuzhiyun                            help='Adds kernel cmdline for setting up rootfs')
3817*4882a593Smuzhiyun    self._add_common_args(sub_parser)
3818*4882a593Smuzhiyun    self._add_common_footer_args(sub_parser)
3819*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.add_hashtree_footer)
3820*4882a593Smuzhiyun
3821*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('erase_footer',
3822*4882a593Smuzhiyun                                       help='Erase footer from an image.')
3823*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3824*4882a593Smuzhiyun                            help='Image with a footer',
3825*4882a593Smuzhiyun                            type=argparse.FileType('rwb+'),
3826*4882a593Smuzhiyun                            required=True)
3827*4882a593Smuzhiyun    sub_parser.add_argument('--keep_hashtree',
3828*4882a593Smuzhiyun                            help='Keep the hashtree and FEC in the image',
3829*4882a593Smuzhiyun                            action='store_true')
3830*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.erase_footer)
3831*4882a593Smuzhiyun
3832*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('resize_image',
3833*4882a593Smuzhiyun                                       help='Resize image with a footer.')
3834*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3835*4882a593Smuzhiyun                            help='Image with a footer',
3836*4882a593Smuzhiyun                            type=argparse.FileType('rwb+'),
3837*4882a593Smuzhiyun                            required=True)
3838*4882a593Smuzhiyun    sub_parser.add_argument('--partition_size',
3839*4882a593Smuzhiyun                            help='New partition size',
3840*4882a593Smuzhiyun                            type=parse_number)
3841*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.resize_image)
3842*4882a593Smuzhiyun
3843*4882a593Smuzhiyun    sub_parser = subparsers.add_parser(
3844*4882a593Smuzhiyun        'info_image',
3845*4882a593Smuzhiyun        help='Show information about vbmeta or footer.')
3846*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3847*4882a593Smuzhiyun                            help='Image to show information about',
3848*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3849*4882a593Smuzhiyun                            required=True)
3850*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3851*4882a593Smuzhiyun                            help='Write info to file',
3852*4882a593Smuzhiyun                            type=argparse.FileType('wt'),
3853*4882a593Smuzhiyun                            default=sys.stdout)
3854*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.info_image)
3855*4882a593Smuzhiyun
3856*4882a593Smuzhiyun    sub_parser = subparsers.add_parser(
3857*4882a593Smuzhiyun        'verify_image',
3858*4882a593Smuzhiyun        help='Verify an image.')
3859*4882a593Smuzhiyun    sub_parser.add_argument('--image',
3860*4882a593Smuzhiyun                            help='Image to verify',
3861*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3862*4882a593Smuzhiyun                            required=True)
3863*4882a593Smuzhiyun    sub_parser.add_argument('--key',
3864*4882a593Smuzhiyun                            help='Check embedded public key matches KEY',
3865*4882a593Smuzhiyun                            metavar='KEY',
3866*4882a593Smuzhiyun                            required=False)
3867*4882a593Smuzhiyun    sub_parser.add_argument('--expected_chain_partition',
3868*4882a593Smuzhiyun                            help='Expected chain partition',
3869*4882a593Smuzhiyun                            metavar='PART_NAME:ROLLBACK_SLOT:KEY_PATH',
3870*4882a593Smuzhiyun                            action='append')
3871*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.verify_image)
3872*4882a593Smuzhiyun
3873*4882a593Smuzhiyun    sub_parser = subparsers.add_parser('set_ab_metadata',
3874*4882a593Smuzhiyun                                       help='Set A/B metadata.')
3875*4882a593Smuzhiyun    sub_parser.add_argument('--misc_image',
3876*4882a593Smuzhiyun                            help=('The misc image to modify. If the image does '
3877*4882a593Smuzhiyun                                  'not exist, it will be created.'),
3878*4882a593Smuzhiyun                            type=argparse.FileType('r+b'),
3879*4882a593Smuzhiyun                            required=True)
3880*4882a593Smuzhiyun    sub_parser.add_argument('--slot_data',
3881*4882a593Smuzhiyun                            help=('Slot data of the form "priority", '
3882*4882a593Smuzhiyun                                  '"tries_remaining", "sucessful_boot" for '
3883*4882a593Smuzhiyun                                  'slot A followed by the same for slot B, '
3884*4882a593Smuzhiyun                                  'separated by colons. The default value '
3885*4882a593Smuzhiyun                                  'is 15:7:0:14:7:0.'),
3886*4882a593Smuzhiyun                            default='15:7:0:14:7:0')
3887*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.set_ab_metadata)
3888*4882a593Smuzhiyun
3889*4882a593Smuzhiyun    sub_parser = subparsers.add_parser(
3890*4882a593Smuzhiyun        'make_atx_certificate',
3891*4882a593Smuzhiyun        help='Create an Android Things eXtension (ATX) certificate.')
3892*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3893*4882a593Smuzhiyun                            help='Write certificate to file',
3894*4882a593Smuzhiyun                            type=argparse.FileType('wb'),
3895*4882a593Smuzhiyun                            default=sys.stdout)
3896*4882a593Smuzhiyun    sub_parser.add_argument('--subject',
3897*4882a593Smuzhiyun                            help=('Path to subject file'),
3898*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3899*4882a593Smuzhiyun                            required=True)
3900*4882a593Smuzhiyun    sub_parser.add_argument('--subject_key',
3901*4882a593Smuzhiyun                            help=('Path to subject RSA public key file'),
3902*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3903*4882a593Smuzhiyun                            required=True)
3904*4882a593Smuzhiyun    sub_parser.add_argument('--subject_key_version',
3905*4882a593Smuzhiyun                            help=('Version of the subject key'),
3906*4882a593Smuzhiyun                            type=parse_number,
3907*4882a593Smuzhiyun                            required=False)
3908*4882a593Smuzhiyun    sub_parser.add_argument('--subject_is_intermediate_authority',
3909*4882a593Smuzhiyun                            help=('Generate an intermediate authority '
3910*4882a593Smuzhiyun                                  'certificate'),
3911*4882a593Smuzhiyun                            action='store_true')
3912*4882a593Smuzhiyun    sub_parser.add_argument('--usage',
3913*4882a593Smuzhiyun                            help=('Override usage with a hash of the provided'
3914*4882a593Smuzhiyun                                  'string'),
3915*4882a593Smuzhiyun                            required=False)
3916*4882a593Smuzhiyun    sub_parser.add_argument('--authority_key',
3917*4882a593Smuzhiyun                            help='Path to authority RSA private key file',
3918*4882a593Smuzhiyun                            required=False)
3919*4882a593Smuzhiyun    sub_parser.add_argument('--signing_helper',
3920*4882a593Smuzhiyun                            help='Path to helper used for signing',
3921*4882a593Smuzhiyun                            metavar='APP',
3922*4882a593Smuzhiyun                            default=None,
3923*4882a593Smuzhiyun                            required=False)
3924*4882a593Smuzhiyun    sub_parser.add_argument('--signing_helper_with_files',
3925*4882a593Smuzhiyun                            help='Path to helper used for signing using files',
3926*4882a593Smuzhiyun                            metavar='APP',
3927*4882a593Smuzhiyun                            default=None,
3928*4882a593Smuzhiyun                            required=False)
3929*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.make_atx_certificate)
3930*4882a593Smuzhiyun
3931*4882a593Smuzhiyun    sub_parser = subparsers.add_parser(
3932*4882a593Smuzhiyun        'make_atx_permanent_attributes',
3933*4882a593Smuzhiyun        help='Create Android Things eXtension (ATX) permanent attributes.')
3934*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3935*4882a593Smuzhiyun                            help='Write attributes to file',
3936*4882a593Smuzhiyun                            type=argparse.FileType('wb'),
3937*4882a593Smuzhiyun                            default=sys.stdout)
3938*4882a593Smuzhiyun    sub_parser.add_argument('--root_authority_key',
3939*4882a593Smuzhiyun                            help='Path to authority RSA public key file',
3940*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3941*4882a593Smuzhiyun                            required=True)
3942*4882a593Smuzhiyun    sub_parser.add_argument('--product_id',
3943*4882a593Smuzhiyun                            help=('Path to Product ID file'),
3944*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3945*4882a593Smuzhiyun                            required=True)
3946*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.make_atx_permanent_attributes)
3947*4882a593Smuzhiyun
3948*4882a593Smuzhiyun    sub_parser = subparsers.add_parser(
3949*4882a593Smuzhiyun        'make_atx_metadata',
3950*4882a593Smuzhiyun        help='Create Android Things eXtension (ATX) metadata.')
3951*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3952*4882a593Smuzhiyun                            help='Write metadata to file',
3953*4882a593Smuzhiyun                            type=argparse.FileType('wb'),
3954*4882a593Smuzhiyun                            default=sys.stdout)
3955*4882a593Smuzhiyun    sub_parser.add_argument('--intermediate_key_certificate',
3956*4882a593Smuzhiyun                            help='Path to intermediate key certificate file',
3957*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3958*4882a593Smuzhiyun                            required=True)
3959*4882a593Smuzhiyun    sub_parser.add_argument('--product_key_certificate',
3960*4882a593Smuzhiyun                            help='Path to product key certificate file',
3961*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3962*4882a593Smuzhiyun                            required=True)
3963*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.make_atx_metadata)
3964*4882a593Smuzhiyun
3965*4882a593Smuzhiyun    sub_parser = subparsers.add_parser(
3966*4882a593Smuzhiyun        'make_atx_unlock_credential',
3967*4882a593Smuzhiyun        help='Create an Android Things eXtension (ATX) unlock credential.')
3968*4882a593Smuzhiyun    sub_parser.add_argument('--output',
3969*4882a593Smuzhiyun                            help='Write credential to file',
3970*4882a593Smuzhiyun                            type=argparse.FileType('wb'),
3971*4882a593Smuzhiyun                            default=sys.stdout)
3972*4882a593Smuzhiyun    sub_parser.add_argument('--intermediate_key_certificate',
3973*4882a593Smuzhiyun                            help='Path to intermediate key certificate file',
3974*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3975*4882a593Smuzhiyun                            required=True)
3976*4882a593Smuzhiyun    sub_parser.add_argument('--unlock_key_certificate',
3977*4882a593Smuzhiyun                            help='Path to unlock key certificate file',
3978*4882a593Smuzhiyun                            type=argparse.FileType('rb'),
3979*4882a593Smuzhiyun                            required=True)
3980*4882a593Smuzhiyun    sub_parser.add_argument('--challenge',
3981*4882a593Smuzhiyun                            help='Path to the challenge to sign (optional). If '
3982*4882a593Smuzhiyun                                 'this is not provided the challenge signature '
3983*4882a593Smuzhiyun                                 'field is omitted and can be concatenated '
3984*4882a593Smuzhiyun                                 'later.',
3985*4882a593Smuzhiyun                            required=False)
3986*4882a593Smuzhiyun    sub_parser.add_argument('--unlock_key',
3987*4882a593Smuzhiyun                            help='Path to unlock key (optional). Must be '
3988*4882a593Smuzhiyun                                 'provided if using --challenge.',
3989*4882a593Smuzhiyun                            required=False)
3990*4882a593Smuzhiyun    sub_parser.add_argument('--signing_helper',
3991*4882a593Smuzhiyun                            help='Path to helper used for signing',
3992*4882a593Smuzhiyun                            metavar='APP',
3993*4882a593Smuzhiyun                            default=None,
3994*4882a593Smuzhiyun                            required=False)
3995*4882a593Smuzhiyun    sub_parser.add_argument('--signing_helper_with_files',
3996*4882a593Smuzhiyun                            help='Path to helper used for signing using files',
3997*4882a593Smuzhiyun                            metavar='APP',
3998*4882a593Smuzhiyun                            default=None,
3999*4882a593Smuzhiyun                            required=False)
4000*4882a593Smuzhiyun    sub_parser.set_defaults(func=self.make_atx_unlock_credential)
4001*4882a593Smuzhiyun
4002*4882a593Smuzhiyun    args = parser.parse_args(argv[1:])
4003*4882a593Smuzhiyun    try:
4004*4882a593Smuzhiyun      args.func(args)
4005*4882a593Smuzhiyun    except AvbError as e:
4006*4882a593Smuzhiyun      sys.stderr.write('{}: {}\n'.format(argv[0], e.message))
4007*4882a593Smuzhiyun      sys.exit(1)
4008*4882a593Smuzhiyun
4009*4882a593Smuzhiyun  def version(self, _):
4010*4882a593Smuzhiyun    """Implements the 'version' sub-command."""
4011*4882a593Smuzhiyun    print get_release_string()
4012*4882a593Smuzhiyun
4013*4882a593Smuzhiyun  def extract_public_key(self, args):
4014*4882a593Smuzhiyun    """Implements the 'extract_public_key' sub-command."""
4015*4882a593Smuzhiyun    self.avb.extract_public_key(args.key, args.output)
4016*4882a593Smuzhiyun
4017*4882a593Smuzhiyun  def make_vbmeta_image(self, args):
4018*4882a593Smuzhiyun    """Implements the 'make_vbmeta_image' sub-command."""
4019*4882a593Smuzhiyun    args = self._fixup_common_args(args)
4020*4882a593Smuzhiyun    self.avb.make_vbmeta_image(args.output, args.chain_partition,
4021*4882a593Smuzhiyun                               args.algorithm, args.key,
4022*4882a593Smuzhiyun                               args.public_key_metadata, args.rollback_index,
4023*4882a593Smuzhiyun                               args.flags, args.prop, args.prop_from_file,
4024*4882a593Smuzhiyun                               args.kernel_cmdline,
4025*4882a593Smuzhiyun                               args.setup_rootfs_from_kernel,
4026*4882a593Smuzhiyun                               args.include_descriptors_from_image,
4027*4882a593Smuzhiyun                               args.signing_helper,
4028*4882a593Smuzhiyun                               args.signing_helper_with_files,
4029*4882a593Smuzhiyun                               args.internal_release_string,
4030*4882a593Smuzhiyun                               args.append_to_release_string,
4031*4882a593Smuzhiyun                               args.print_required_libavb_version,
4032*4882a593Smuzhiyun                               args.padding_size)
4033*4882a593Smuzhiyun
4034*4882a593Smuzhiyun  def append_vbmeta_image(self, args):
4035*4882a593Smuzhiyun    """Implements the 'append_vbmeta_image' sub-command."""
4036*4882a593Smuzhiyun    self.avb.append_vbmeta_image(args.image.name, args.vbmeta_image.name,
4037*4882a593Smuzhiyun                                 args.partition_size)
4038*4882a593Smuzhiyun
4039*4882a593Smuzhiyun  def add_hash_footer(self, args):
4040*4882a593Smuzhiyun    """Implements the 'add_hash_footer' sub-command."""
4041*4882a593Smuzhiyun    args = self._fixup_common_args(args)
4042*4882a593Smuzhiyun    self.avb.add_hash_footer(args.image.name if args.image else None,
4043*4882a593Smuzhiyun                             args.partition_size,
4044*4882a593Smuzhiyun                             args.partition_name, args.hash_algorithm,
4045*4882a593Smuzhiyun                             args.salt, args.chain_partition, args.algorithm,
4046*4882a593Smuzhiyun                             args.key,
4047*4882a593Smuzhiyun                             args.public_key_metadata, args.rollback_index,
4048*4882a593Smuzhiyun                             args.flags, args.prop, args.prop_from_file,
4049*4882a593Smuzhiyun                             args.kernel_cmdline,
4050*4882a593Smuzhiyun                             args.setup_rootfs_from_kernel,
4051*4882a593Smuzhiyun                             args.include_descriptors_from_image,
4052*4882a593Smuzhiyun                             args.calc_max_image_size,
4053*4882a593Smuzhiyun                             args.signing_helper,
4054*4882a593Smuzhiyun                             args.signing_helper_with_files,
4055*4882a593Smuzhiyun                             args.internal_release_string,
4056*4882a593Smuzhiyun                             args.append_to_release_string,
4057*4882a593Smuzhiyun                             args.output_vbmeta_image,
4058*4882a593Smuzhiyun                             args.do_not_append_vbmeta_image,
4059*4882a593Smuzhiyun                             args.print_required_libavb_version,
4060*4882a593Smuzhiyun                             args.use_persistent_digest,
4061*4882a593Smuzhiyun                             args.do_not_use_ab)
4062*4882a593Smuzhiyun
4063*4882a593Smuzhiyun  def add_hashtree_footer(self, args):
4064*4882a593Smuzhiyun    """Implements the 'add_hashtree_footer' sub-command."""
4065*4882a593Smuzhiyun    args = self._fixup_common_args(args)
4066*4882a593Smuzhiyun    # TODO(zeuthen): Remove when removing support for the
4067*4882a593Smuzhiyun    # '--generate_fec' option above.
4068*4882a593Smuzhiyun    if args.generate_fec:
4069*4882a593Smuzhiyun      sys.stderr.write('The --generate_fec option is deprecated since FEC '
4070*4882a593Smuzhiyun                       'is now generated by default. Use the option '
4071*4882a593Smuzhiyun                       '--do_not_generate_fec to not generate FEC.\n')
4072*4882a593Smuzhiyun    self.avb.add_hashtree_footer(args.image.name if args.image else None,
4073*4882a593Smuzhiyun                                 args.partition_size,
4074*4882a593Smuzhiyun                                 args.partition_name,
4075*4882a593Smuzhiyun                                 not args.do_not_generate_fec, args.fec_num_roots,
4076*4882a593Smuzhiyun                                 args.hash_algorithm, args.block_size,
4077*4882a593Smuzhiyun                                 args.salt, args.chain_partition, args.algorithm,
4078*4882a593Smuzhiyun                                 args.key, args.public_key_metadata,
4079*4882a593Smuzhiyun                                 args.rollback_index, args.flags, args.prop,
4080*4882a593Smuzhiyun                                 args.prop_from_file,
4081*4882a593Smuzhiyun                                 args.kernel_cmdline,
4082*4882a593Smuzhiyun                                 args.setup_rootfs_from_kernel,
4083*4882a593Smuzhiyun                                 args.setup_as_rootfs_from_kernel,
4084*4882a593Smuzhiyun                                 args.include_descriptors_from_image,
4085*4882a593Smuzhiyun                                 args.calc_max_image_size,
4086*4882a593Smuzhiyun                                 args.signing_helper,
4087*4882a593Smuzhiyun                                 args.signing_helper_with_files,
4088*4882a593Smuzhiyun                                 args.internal_release_string,
4089*4882a593Smuzhiyun                                 args.append_to_release_string,
4090*4882a593Smuzhiyun                                 args.output_vbmeta_image,
4091*4882a593Smuzhiyun                                 args.do_not_append_vbmeta_image,
4092*4882a593Smuzhiyun                                 args.print_required_libavb_version,
4093*4882a593Smuzhiyun                                 args.use_persistent_digest,
4094*4882a593Smuzhiyun                                 args.do_not_use_ab)
4095*4882a593Smuzhiyun
4096*4882a593Smuzhiyun  def erase_footer(self, args):
4097*4882a593Smuzhiyun    """Implements the 'erase_footer' sub-command."""
4098*4882a593Smuzhiyun    self.avb.erase_footer(args.image.name, args.keep_hashtree)
4099*4882a593Smuzhiyun
4100*4882a593Smuzhiyun  def resize_image(self, args):
4101*4882a593Smuzhiyun    """Implements the 'resize_image' sub-command."""
4102*4882a593Smuzhiyun    self.avb.resize_image(args.image.name, args.partition_size)
4103*4882a593Smuzhiyun
4104*4882a593Smuzhiyun  def set_ab_metadata(self, args):
4105*4882a593Smuzhiyun    """Implements the 'set_ab_metadata' sub-command."""
4106*4882a593Smuzhiyun    self.avb.set_ab_metadata(args.misc_image, args.slot_data)
4107*4882a593Smuzhiyun
4108*4882a593Smuzhiyun  def info_image(self, args):
4109*4882a593Smuzhiyun    """Implements the 'info_image' sub-command."""
4110*4882a593Smuzhiyun    self.avb.info_image(args.image.name, args.output)
4111*4882a593Smuzhiyun
4112*4882a593Smuzhiyun  def verify_image(self, args):
4113*4882a593Smuzhiyun    """Implements the 'verify_image' sub-command."""
4114*4882a593Smuzhiyun    self.avb.verify_image(args.image.name, args.key,
4115*4882a593Smuzhiyun                          args.expected_chain_partition)
4116*4882a593Smuzhiyun
4117*4882a593Smuzhiyun  def make_atx_certificate(self, args):
4118*4882a593Smuzhiyun    """Implements the 'make_atx_certificate' sub-command."""
4119*4882a593Smuzhiyun    self.avb.make_atx_certificate(args.output, args.authority_key,
4120*4882a593Smuzhiyun                                  args.subject_key.name,
4121*4882a593Smuzhiyun                                  args.subject_key_version,
4122*4882a593Smuzhiyun                                  args.subject.read(),
4123*4882a593Smuzhiyun                                  args.subject_is_intermediate_authority,
4124*4882a593Smuzhiyun                                  args.usage,
4125*4882a593Smuzhiyun                                  args.signing_helper,
4126*4882a593Smuzhiyun                                  args.signing_helper_with_files)
4127*4882a593Smuzhiyun
4128*4882a593Smuzhiyun  def make_atx_permanent_attributes(self, args):
4129*4882a593Smuzhiyun    """Implements the 'make_atx_permanent_attributes' sub-command."""
4130*4882a593Smuzhiyun    self.avb.make_atx_permanent_attributes(args.output,
4131*4882a593Smuzhiyun                                           args.root_authority_key.name,
4132*4882a593Smuzhiyun                                           args.product_id.read())
4133*4882a593Smuzhiyun
4134*4882a593Smuzhiyun  def make_atx_metadata(self, args):
4135*4882a593Smuzhiyun    """Implements the 'make_atx_metadata' sub-command."""
4136*4882a593Smuzhiyun    self.avb.make_atx_metadata(args.output,
4137*4882a593Smuzhiyun                               args.intermediate_key_certificate.read(),
4138*4882a593Smuzhiyun                               args.product_key_certificate.read())
4139*4882a593Smuzhiyun
4140*4882a593Smuzhiyun  def make_atx_unlock_credential(self, args):
4141*4882a593Smuzhiyun    """Implements the 'make_atx_unlock_credential' sub-command."""
4142*4882a593Smuzhiyun    self.avb.make_atx_unlock_credential(
4143*4882a593Smuzhiyun        args.output,
4144*4882a593Smuzhiyun        args.intermediate_key_certificate.read(),
4145*4882a593Smuzhiyun        args.unlock_key_certificate.read(),
4146*4882a593Smuzhiyun        args.challenge,
4147*4882a593Smuzhiyun        args.unlock_key,
4148*4882a593Smuzhiyun        args.signing_helper,
4149*4882a593Smuzhiyun        args.signing_helper_with_files)
4150*4882a593Smuzhiyun
4151*4882a593Smuzhiyun
4152*4882a593Smuzhiyunif __name__ == '__main__':
4153*4882a593Smuzhiyun  tool = AvbTool()
4154*4882a593Smuzhiyun  tool.run(sys.argv)
4155