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