#!/usr/bin/env python

# Copyright 2016, The Android Open Source Project
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use, copy,
# modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
"""Command-line tool for working with Android Verified Boot images."""

import argparse
import binascii
import bisect
import hashlib
import os
import struct
import subprocess
import sys

import Crypto.PublicKey.RSA

# Keep in sync with avb_vbmeta_header.h.
AVB_VERSION_MAJOR = 1
AVB_VERSION_MINOR = 0


class AvbError(Exception):
  """Application-specific errors.

  These errors represent issues for which a stack-trace should not be
  presented.

  Attributes:
    message: Error message.
  """

  def __init__(self, message):
    Exception.__init__(self, message)


class Algorithm(object):
  """Contains details about an algorithm.

  See the avb_vbmeta_header.h file for more details about
  algorithms.

  The constant |ALGORITHMS| is a dictionary from human-readable
  names (e.g 'SHA256_RSA2048') to instances of this class.

  Attributes:
    algorithm_type: Integer code corresponding to |AvbAlgorithmType|.
    hash_num_bytes: Number of bytes used to store the hash.
    signature_num_bytes: Number of bytes used to store the signature.
    public_key_num_bytes: Number of bytes used to store the public key.
    padding: Padding used for signature, if any.
  """

  def __init__(self, algorithm_type, hash_num_bytes, signature_num_bytes,
               public_key_num_bytes, padding):
    self.algorithm_type = algorithm_type
    self.hash_num_bytes = hash_num_bytes
    self.signature_num_bytes = signature_num_bytes
    self.public_key_num_bytes = public_key_num_bytes
    self.padding = padding

# This must be kept in sync with the avb_crypto.h file.
#
# The PKC1-v1.5 padding is a blob of binary DER of ASN.1 and is
# obtained from section 5.2.2 of RFC 4880.
ALGORITHMS = {
    'NONE': Algorithm(
        algorithm_type=0,        # AVB_ALGORITHM_TYPE_NONE
        hash_num_bytes=0,
        signature_num_bytes=0,
        public_key_num_bytes=0,
        padding=[]),
    'SHA256_RSA2048': Algorithm(
        algorithm_type=1,        # AVB_ALGORITHM_TYPE_SHA256_RSA2048
        hash_num_bytes=32,
        signature_num_bytes=256,
        public_key_num_bytes=8 + 2*2048/8,
        padding=[
            # PKCS1-v1_5 padding
            0x00, 0x01] + [0xff]*202 + [0x00] + [
                # ASN.1 header
                0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05,
                0x00, 0x04, 0x20,
            ]),
    'SHA256_RSA4096': Algorithm(
        algorithm_type=2,        # AVB_ALGORITHM_TYPE_SHA256_RSA4096
        hash_num_bytes=32,
        signature_num_bytes=512,
        public_key_num_bytes=8 + 2*4096/8,
        padding=[
            # PKCS1-v1_5 padding
            0x00, 0x01] + [0xff]*458 + [0x00] + [
                # ASN.1 header
                0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05,
                0x00, 0x04, 0x20,
            ]),
    'SHA256_RSA8192': Algorithm(
        algorithm_type=3,        # AVB_ALGORITHM_TYPE_SHA256_RSA8192
        hash_num_bytes=32,
        signature_num_bytes=1024,
        public_key_num_bytes=8 + 2*8192/8,
        padding=[
            # PKCS1-v1_5 padding
            0x00, 0x01] + [0xff]*970 + [0x00] + [
                # ASN.1 header
                0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05,
                0x00, 0x04, 0x20,
            ]),
    'SHA512_RSA2048': Algorithm(
        algorithm_type=4,        # AVB_ALGORITHM_TYPE_SHA512_RSA2048
        hash_num_bytes=64,
        signature_num_bytes=256,
        public_key_num_bytes=8 + 2*2048/8,
        padding=[
            # PKCS1-v1_5 padding
            0x00, 0x01] + [0xff]*170 + [0x00] + [
                # ASN.1 header
                0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05,
                0x00, 0x04, 0x40
            ]),
    'SHA512_RSA4096': Algorithm(
        algorithm_type=5,        # AVB_ALGORITHM_TYPE_SHA512_RSA4096
        hash_num_bytes=64,
        signature_num_bytes=512,
        public_key_num_bytes=8 + 2*4096/8,
        padding=[
            # PKCS1-v1_5 padding
            0x00, 0x01] + [0xff]*426 + [0x00] + [
                # ASN.1 header
                0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05,
                0x00, 0x04, 0x40
            ]),
    'SHA512_RSA8192': Algorithm(
        algorithm_type=6,        # AVB_ALGORITHM_TYPE_SHA512_RSA8192
        hash_num_bytes=64,
        signature_num_bytes=1024,
        public_key_num_bytes=8 + 2*8192/8,
        padding=[
            # PKCS1-v1_5 padding
            0x00, 0x01] + [0xff]*938 + [0x00] + [
                # ASN.1 header
                0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
                0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05,
                0x00, 0x04, 0x40
            ]),
}


def round_to_multiple(number, size):
  """Rounds a number up to nearest multiple of another number.

  Args:
    number: The number to round up.
    size: The multiple to round up to.

  Returns:
    If |number| is a multiple of |size|, returns |number|, otherwise
    returns |number| + |size|.
  """
  remainder = number % size
  if remainder == 0:
    return number
  return number + size - remainder


def round_to_pow2(number):
  """Rounds a number up to the next power of 2.

  Args:
    number: The number to round up.

  Returns:
    If |number| is already a power of 2 then |number| is
    returned. Otherwise the smallest power of 2 greater than |number|
    is returned.
  """
  return 2**((number - 1).bit_length())


def write_long(output, num_bits, value):
  """Writes a long to an output stream using a given amount of bits.

  This number is written big-endian, e.g. with the most significant
  bit first.

  Arguments:
    output: The object to write the output to.
    num_bits: The number of bits to write, e.g. 2048.
    value: The value to write.
  """
  for bit_pos in range(num_bits, 0, -8):
    octet = (value >> (bit_pos - 8)) & 0xff
    output.write(struct.pack('!B', octet))


def encode_long(num_bits, value):
  """Encodes a long to a bytearray() using a given amount of bits.

  This number is written big-endian, e.g. with the most significant
  bit first.

  Arguments:
    num_bits: The number of bits to write, e.g. 2048.
    value: The value to write.

  Returns:
    A bytearray() with the encoded long.
  """
  ret = bytearray()
  for bit_pos in range(num_bits, 0, -8):
    octet = (value >> (bit_pos - 8)) & 0xff
    ret.extend(struct.pack('!B', octet))
  return ret


def egcd(a, b):
  """Calculate greatest common divisor of two numbers.

  This implementation uses a recursive version of the extended
  Euclidian algorithm.

  Arguments:
    a: First number.
    b: Second number.

  Returns:
    A tuple (gcd, x, y) that where |gcd| is the greatest common
    divisor of |a| and |b| and |a|*|x| + |b|*|y| = |gcd|.
  """
  if a == 0:
    return (b, 0, 1)
  else:
    g, y, x = egcd(b % a, a)
    return (g, x - (b // a) * y, y)


def modinv(a, m):
  """Calculate modular multiplicative inverse of |a| modulo |m|.

  This calculates the number |x| such that |a| * |x| == 1 (modulo
  |m|). This number only exists if |a| and |m| are co-prime - |None|
  is returned if this isn't true.

  Arguments:
    a: The number to calculate a modular inverse of.
    m: The modulo to use.

  Returns:
    The modular multiplicative inverse of |a| and |m| or |None| if
    these numbers are not co-prime.
  """
  gcd, x, _ = egcd(a, m)
  if gcd != 1:
    return None  # modular inverse does not exist
  else:
    return x % m


def parse_number(string):
  """Parse a string as a number.

  This is just a short-hand for int(string, 0) suitable for use in the
  |type| parameter of |ArgumentParser|'s add_argument() function. An
  improvement to just using type=int is that this function supports
  numbers in other bases, e.g. "0x1234".

  Arguments:
    string: The string to parse.

  Returns:
    The parsed integer.

  Raises:
    ValueError: If the number could not be parsed.
  """
  return int(string, 0)


def write_rsa_key(output, key):
  """Writes a public RSA key in |AvbRSAPublicKeyHeader| format.

  This writes the |AvbRSAPublicKeyHeader| as well as the two large
  numbers (|key_num_bits| bits long) following it.

  Arguments:
    output: The object to write the output to.
    key: A Crypto.PublicKey.RSA object.
  """
  # key.e is exponent
  # key.n is modulus
  key_num_bits = key.size() + 1
  # Calculate n0inv = -1/n[0] (mod 2^32)
  b = 2L**32
  n0inv = b - modinv(key.n, b)
  # Calculate rr = r^2 (mod N), where r = 2^(# of key bits)
  r = 2L**key.n.bit_length()
  rrmodn = r * r % key.n
  output.write(struct.pack('!II', key_num_bits, n0inv))
  write_long(output, key_num_bits, key.n)
  write_long(output, key_num_bits, rrmodn)


def encode_rsa_key(key):
  """Encodes a public RSA key in |AvbRSAPublicKeyHeader| format.

  This creates a |AvbRSAPublicKeyHeader| as well as the two large
  numbers (|key_num_bits| bits long) following it.

  Arguments:
    key: A Crypto.PublicKey.RSA object.

  Returns:
    A bytearray() with the |AvbRSAPublicKeyHeader|.
  """
  ret = bytearray()
  # key.e is exponent
  # key.n is modulus
  key_num_bits = key.size() + 1
  # Calculate n0inv = -1/n[0] (mod 2^32)
  b = 2L**32
  n0inv = b - modinv(key.n, b)
  # Calculate rr = r^2 (mod N), where r = 2^(# of key bits)
  r = 2L**key.n.bit_length()
  rrmodn = r * r % key.n
  ret.extend(struct.pack('!II', key_num_bits, n0inv))
  ret.extend(encode_long(key_num_bits, key.n))
  ret.extend(encode_long(key_num_bits, rrmodn))
  return ret


def lookup_algorithm_by_type(alg_type):
  """Looks up algorithm by type.

  Arguments:
    alg_type: The integer representing the type.

  Returns:
    A tuple with the algorithm name and an |Algorithm| instance.

  Raises:
    Exception: If the algorithm cannot be found
  """
  for alg_name in ALGORITHMS:
    alg_data = ALGORITHMS[alg_name]
    if alg_data.algorithm_type == alg_type:
      return (alg_name, alg_data)
  raise AvbError('Unknown algorithm type {}'.format(alg_type))


class ImageChunk(object):
  """Data structure used for representing chunks in Android sparse files.

  Attributes:
    chunk_type: One of TYPE_RAW, TYPE_FILL, or TYPE_DONT_CARE.
    chunk_offset: Offset in the sparse file where this chunk begins.
    output_offset: Offset in de-sparsified file where output begins.
    output_size: Number of bytes in output.
    input_offset: Offset in sparse file for data if TYPE_RAW otherwise None.
    fill_data: Blob with data to fill if TYPE_FILL otherwise None.
  """

  FORMAT = '<2H2I'
  TYPE_RAW = 0xcac1
  TYPE_FILL = 0xcac2
  TYPE_DONT_CARE = 0xcac3
  TYPE_CRC32 = 0xcac4

  def __init__(self, chunk_type, chunk_offset, output_offset, output_size,
               input_offset, fill_data):
    """Initializes an ImageChunk object.

    Arguments:
      chunk_type: One of TYPE_RAW, TYPE_FILL, or TYPE_DONT_CARE.
      chunk_offset: Offset in the sparse file where this chunk begins.
      output_offset: Offset in de-sparsified file.
      output_size: Number of bytes in output.
      input_offset: Offset in sparse file if TYPE_RAW otherwise None.
      fill_data: Blob with data to fill if TYPE_FILL otherwise None.

    Raises:
      ValueError: If data is not well-formed.
    """
    self.chunk_type = chunk_type
    self.chunk_offset = chunk_offset
    self.output_offset = output_offset
    self.output_size = output_size
    self.input_offset = input_offset
    self.fill_data = fill_data
    # Check invariants.
    if self.chunk_type == self.TYPE_RAW:
      if self.fill_data is not None:
        raise ValueError('RAW chunk cannot have fill_data set.')
      if not self.input_offset:
        raise ValueError('RAW chunk must have input_offset set.')
    elif self.chunk_type == self.TYPE_FILL:
      if self.fill_data is None:
        raise ValueError('FILL chunk must have fill_data set.')
      if self.input_offset:
        raise ValueError('FILL chunk cannot have input_offset set.')
    elif self.chunk_type == self.TYPE_DONT_CARE:
      if self.fill_data is not None:
        raise ValueError('DONT_CARE chunk cannot have fill_data set.')
      if self.input_offset:
        raise ValueError('DONT_CARE chunk cannot have input_offset set.')
    else:
      raise ValueError('Invalid chunk type')


class ImageHandler(object):
  """Abstraction for image I/O with support for Android sparse images.

  This class provides an interface for working with image files that
  may be using the Android Sparse Image format. When an instance is
  constructed, we test whether it's an Android sparse file. If so,
  operations will be on the sparse file by interpreting the sparse
  format, otherwise they will be directly on the file. Either way the
  operations do the same.

  For reading, this interface mimics a file object - it has seek(),
  tell(), and read() methods. For writing, only truncation
  (truncate()) and appending is supported (append_raw() and
  append_dont_care()). Additionally, data can only be written in units
  of the block size.

  Attributes:
    is_sparse: Whether the file being operated on is sparse.
    block_size: The block size, typically 4096.
    image_size: The size of the unsparsified file.
    care_size: Position in the unsparsified file where only
      DONT_CARE data follows.
  """
  # See system/core/libsparse/sparse_format.h for details.
  MAGIC = 0xed26ff3a
  HEADER_FORMAT = '<I4H4I'

  # These are formats and offset of just the |total_chunks| and
  # |total_blocks| fields.
  NUM_CHUNKS_AND_BLOCKS_FORMAT = '<II'
  NUM_CHUNKS_AND_BLOCKS_OFFSET = 16

  def __init__(self, image_filename):
    """Initializes an image handler.

    Arguments:
      image_filename: The name of the file to operate on.

    Raises:
      ValueError: If data in the file is invalid.
    """
    self._image_filename = image_filename
    self._read_header()

  def _read_header(self):
    """Initializes internal data structures used for reading file.

    This may be called multiple times and is typically called after
    modifying the file (e.g. appending, truncation).

    Raises:
      ValueError: If data in the file is invalid.
    """
    self.is_sparse = False
    self.block_size = 4096
    self._file_pos = 0
    self._image = open(self._image_filename, 'r+b')
    self._image.seek(0, os.SEEK_END)
    self.care_size = self._image.tell()
    self.image_size = self._image.tell()

    self._image.seek(0, os.SEEK_SET)
    header_bin = self._image.read(struct.calcsize(self.HEADER_FORMAT))
    (magic, major_version, minor_version, file_hdr_sz, chunk_hdr_sz,
     block_size, self._num_total_blocks, self._num_total_chunks,
     _) = struct.unpack(self.HEADER_FORMAT, header_bin)
    if magic != self.MAGIC:
      # Not a sparse image, our job here is done.
      return
    if not (major_version == 1 and minor_version == 0):
      raise ValueError('Encountered sparse image format version {}.{} but '
                       'only 1.0 is supported'.format(major_version,
                                                      minor_version))
    if file_hdr_sz != struct.calcsize(self.HEADER_FORMAT):
      raise ValueError('Unexpected file_hdr_sz value {}.'.
                       format(file_hdr_sz))
    if chunk_hdr_sz != struct.calcsize(ImageChunk.FORMAT):
      raise ValueError('Unexpected chunk_hdr_sz value {}.'.
                       format(chunk_hdr_sz))

    self.block_size = block_size

    # Build an list of chunks by parsing the file.
    self._chunks = []

    # Find the smallest offset where only "Don't care" chunks
    # follow. This will be the size of the content in the sparse
    # image.
    offset = 0
    output_offset = 0
    last_dont_care_section_output_offset = None
    last_section_was_dont_care = False
    for _ in xrange(1, self._num_total_chunks + 1):
      chunk_offset = self._image.tell()

      header_bin = self._image.read(struct.calcsize(ImageChunk.FORMAT))
      (chunk_type, _, chunk_sz, total_sz) = struct.unpack(ImageChunk.FORMAT,
                                                          header_bin)
      data_sz = total_sz - struct.calcsize(ImageChunk.FORMAT)

      last_section_was_dont_care = False

      if chunk_type == ImageChunk.TYPE_RAW:
        if data_sz != (chunk_sz * self.block_size):
          raise ValueError('Raw chunk input size ({}) does not match output '
                           'size ({})'.
                           format(data_sz, chunk_sz*self.block_size))
        self._chunks.append(ImageChunk(ImageChunk.TYPE_RAW,
                                       chunk_offset,
                                       output_offset,
                                       chunk_sz*self.block_size,
                                       self._image.tell(),
                                       None))
        self._image.read(data_sz)

      elif chunk_type == ImageChunk.TYPE_FILL:
        if data_sz != 4:
          raise ValueError('Fill chunk should have 4 bytes of fill, but this '
                           'has {}'.format(data_sz))
        fill_data = self._image.read(4)
        self._chunks.append(ImageChunk(ImageChunk.TYPE_FILL,
                                       chunk_offset,
                                       output_offset,
                                       chunk_sz*self.block_size,
                                       None,
                                       fill_data))
      elif chunk_type == ImageChunk.TYPE_DONT_CARE:
        if data_sz != 0:
          raise ValueError('Don\'t care chunk input size is non-zero ({})'.
                           format(data_sz))
        else:
          if not last_section_was_dont_care:
            last_dont_care_section_output_offset = output_offset
            last_section_was_dont_care = True
        self._chunks.append(ImageChunk(ImageChunk.TYPE_DONT_CARE,
                                       chunk_offset,
                                       output_offset,
                                       chunk_sz*self.block_size,
                                       None,
                                       None))
      elif chunk_type == ImageChunk.TYPE_CRC32:
        if data_sz != 4:
          raise ValueError('CRC32 chunk should have 4 bytes of CRC, but '
                           'this has {}'.format(data_sz))
        self._image.read(4)
      else:
        raise ValueError('Unknown chunk type {}'.format(chunk_type))

      offset += chunk_sz
      output_offset += chunk_sz*self.block_size

    # Record where sparse data end.
    self._sparse_end = self._image.tell()

    # Now that we've traversed all chunks, sanity check.
    if self._num_total_blocks != offset:
      raise ValueError('The header said we should have {} output blocks, '
                       'but we saw {}'.format(self._num_total_blocks, offset))
    junk_len = len(self._image.read())
    if junk_len > 0:
      raise ValueError('There were {} bytes of extra data at the end of the '
                       'file.'.format(junk_len))

    # Assign |image_size| and |care_size| attributes.
    self.image_size = output_offset
    if last_section_was_dont_care:
      self.care_size = last_dont_care_section_output_offset
    else:
      self.care_size = output_offset

    # This is used when bisecting in read() to find the initial slice.
    self._chunk_output_offsets = [i.output_offset for i in self._chunks]

    self.is_sparse = True

  def _update_chunks_and_blocks(self):
    """Helper function to update the image header.

    The the |total_chunks| and |total_blocks| fields in the header
    will be set to value of the |_num_total_blocks| and
    |_num_total_chunks| attributes.

    """
    self._image.seek(self.NUM_CHUNKS_AND_BLOCKS_OFFSET, os.SEEK_SET)
    self._image.write(struct.pack(self.NUM_CHUNKS_AND_BLOCKS_FORMAT,
                                  self._num_total_blocks,
                                  self._num_total_chunks))

  def append_dont_care(self, num_bytes):
    """Appends a DONT_CARE chunk to the sparse file.

    The given number of bytes must be a multiple of the block size.

    Arguments:
      num_bytes: Size in number of bytes of the DONT_CARE chunk.
    """
    assert num_bytes % self.block_size == 0

    if not self.is_sparse:
      self._image.seek(0, os.SEEK_END)
      # This is more efficient that writing NUL bytes since it'll add
      # a hole on file systems that support sparse files (native
      # sparse, not Android sparse).
      self._image.truncate(self._image.tell() + num_bytes)
      self._read_header()
      return

    self._num_total_chunks += 1
    self._num_total_blocks += num_bytes / self.block_size
    self._update_chunks_and_blocks()

    self._image.seek(self._sparse_end, os.SEEK_SET)
    self._image.write(struct.pack(ImageChunk.FORMAT,
                                  ImageChunk.TYPE_DONT_CARE,
                                  0,  # Reserved
                                  num_bytes / self.block_size,
                                  struct.calcsize(ImageChunk.FORMAT)))
    self._read_header()

  def append_raw(self, data):
    """Appends a RAW chunk to the sparse file.

    The length of the given data must be a multiple of the block size.

    Arguments:
      data: Data to append.
    """
    assert len(data) % self.block_size == 0

    if not self.is_sparse:
      self._image.seek(0, os.SEEK_END)
      self._image.write(data)
      self._read_header()
      return

    self._num_total_chunks += 1
    self._num_total_blocks += len(data) / self.block_size
    self._update_chunks_and_blocks()

    self._image.seek(self._sparse_end, os.SEEK_SET)
    self._image.write(struct.pack(ImageChunk.FORMAT,
                                  ImageChunk.TYPE_RAW,
                                  0,  # Reserved
                                  len(data) / self.block_size,
                                  len(data) +
                                  struct.calcsize(ImageChunk.FORMAT)))
    self._image.write(data)
    self._read_header()

  def append_fill(self, fill_data, size):
    """Appends a fill chunk to the sparse file.

    The total length of the fill data must be a multiple of the block size.

    Arguments:
      fill_data: Fill data to append - must be four bytes.
      size: Number of chunk - must be a multiple of four and the block size.
    """
    assert len(fill_data) == 4
    assert size % 4 == 0
    assert size % self.block_size == 0

    if not self.is_sparse:
      self._image.seek(0, os.SEEK_END)
      self._image.write(fill_data * (size/4))
      self._read_header()
      return

    self._num_total_chunks += 1
    self._num_total_blocks += size / self.block_size
    self._update_chunks_and_blocks()

    self._image.seek(self._sparse_end, os.SEEK_SET)
    self._image.write(struct.pack(ImageChunk.FORMAT,
                                  ImageChunk.TYPE_FILL,
                                  0,  # Reserved
                                  size / self.block_size,
                                  4 + struct.calcsize(ImageChunk.FORMAT)))
    self._image.write(fill_data)
    self._read_header()

  def seek(self, offset):
    """Sets the cursor position for reading from unsparsified file.

    Arguments:
      offset: Offset to seek to from the beginning of the file.
    """
    self._file_pos = offset

  def read(self, size):
    """Reads data from the unsparsified file.

    This method may return fewer than |size| bytes of data if the end
    of the file was encountered.

    The file cursor for reading is advanced by the number of bytes
    read.

    Arguments:
      size: Number of bytes to read.

    Returns:
      The data.

    """
    if not self.is_sparse:
      self._image.seek(self._file_pos)
      data = self._image.read(size)
      self._file_pos += len(data)
      return data

    # Iterate over all chunks.
    chunk_idx = bisect.bisect_right(self._chunk_output_offsets,
                                    self._file_pos) - 1
    data = bytearray()
    to_go = size
    while to_go > 0:
      chunk = self._chunks[chunk_idx]
      chunk_pos_offset = self._file_pos - chunk.output_offset
      chunk_pos_to_go = min(chunk.output_size - chunk_pos_offset, to_go)

      if chunk.chunk_type == ImageChunk.TYPE_RAW:
        self._image.seek(chunk.input_offset + chunk_pos_offset)
        data.extend(self._image.read(chunk_pos_to_go))
      elif chunk.chunk_type == ImageChunk.TYPE_FILL:
        all_data = chunk.fill_data*(chunk_pos_to_go/len(chunk.fill_data) + 2)
        offset_mod = chunk_pos_offset % len(chunk.fill_data)
        data.extend(all_data[offset_mod:(offset_mod + chunk_pos_to_go)])
      else:
        assert chunk.chunk_type == ImageChunk.TYPE_DONT_CARE
        data.extend('\0' * chunk_pos_to_go)

      to_go -= chunk_pos_to_go
      self._file_pos += chunk_pos_to_go
      chunk_idx += 1
      # Generate partial read in case of EOF.
      if chunk_idx >= len(self._chunks):
        break

    return data

  def tell(self):
    """Returns the file cursor position for reading from unsparsified file.

    Returns:
      The file cursor position for reading.
    """
    return self._file_pos

  def truncate(self, size):
    """Truncates the unsparsified file.

    Arguments:
      size: Desired size of unsparsified file.

    Raises:
      ValueError: If desired size isn't a multiple of the block size.
    """
    if not self.is_sparse:
      self._image.truncate(size)
      self._read_header()
      return

    if size % self.block_size != 0:
      raise ValueError('Cannot truncate to a size which is not a multiple '
                       'of the block size')

    if size == self.image_size:
      # Trivial where there's nothing to do.
      return
    elif size < self.image_size:
      chunk_idx = bisect.bisect_right(self._chunk_output_offsets, size) - 1
      chunk = self._chunks[chunk_idx]
      if chunk.output_offset != size:
        # Truncation in the middle of a trunk - need to keep the chunk
        # and modify it.
        chunk_idx_for_update = chunk_idx + 1
        num_to_keep = size - chunk.output_offset
        assert num_to_keep % self.block_size == 0
        if chunk.chunk_type == ImageChunk.TYPE_RAW:
          truncate_at = (chunk.chunk_offset +
                         struct.calcsize(ImageChunk.FORMAT) + num_to_keep)
          data_sz = num_to_keep
        elif chunk.chunk_type == ImageChunk.TYPE_FILL:
          truncate_at = (chunk.chunk_offset +
                         struct.calcsize(ImageChunk.FORMAT) + 4)
          data_sz = 4
        else:
          assert chunk.chunk_type == ImageChunk.TYPE_DONT_CARE
          truncate_at = chunk.chunk_offset + struct.calcsize(ImageChunk.FORMAT)
          data_sz = 0
        chunk_sz = num_to_keep/self.block_size
        total_sz = data_sz + struct.calcsize(ImageChunk.FORMAT)
        self._image.seek(chunk.chunk_offset)
        self._image.write(struct.pack(ImageChunk.FORMAT,
                                      chunk.chunk_type,
                                      0,  # Reserved
                                      chunk_sz,
                                      total_sz))
        chunk.output_size = num_to_keep
      else:
        # Truncation at trunk boundary.
        truncate_at = chunk.chunk_offset
        chunk_idx_for_update = chunk_idx

      self._num_total_chunks = chunk_idx_for_update
      self._num_total_blocks = 0
      for i in range(0, chunk_idx_for_update):
        self._num_total_blocks += self._chunks[i].output_size / self.block_size
      self._update_chunks_and_blocks()
      self._image.truncate(truncate_at)

      # We've modified the file so re-read all data.
      self._read_header()
    else:
      # Truncating to grow - just add a DONT_CARE section.
      self.append_dont_care(size - self.image_size)


class AvbDescriptor(object):
  """Class for AVB descriptor.

  See the |AvbDescriptor| C struct for more information.

  Attributes:
    tag: The tag identifying what kind of descriptor this is.
    data: The data in the descriptor.
  """

  SIZE = 16
  FORMAT_STRING = ('!QQ')  # tag, num_bytes_following (descriptor header)

  def __init__(self, data):
    """Initializes a new property descriptor.

    Arguments:
      data: If not None, must be a bytearray().

    Raises:
      LookupError: If the given descriptor is malformed.
    """
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (self.tag, num_bytes_following) = (
          struct.unpack(self.FORMAT_STRING, data[0:self.SIZE]))
      self.data = data[self.SIZE:self.SIZE + num_bytes_following]
    else:
      self.tag = None
      self.data = None

  def print_desc(self, o):
    """Print the descriptor.

    Arguments:
      o: The object to write the output to.
    """
    o.write('    Unknown descriptor:\n')
    o.write('      Tag:  {}\n'.format(self.tag))
    if len(self.data) < 256:
      o.write('      Data: {} ({} bytes)\n'.format(
          repr(str(self.data)), len(self.data)))
    else:
      o.write('      Data: {} bytes\n'.format(len(self.data)))

  def encode(self):
    """Serializes the descriptor.

    Returns:
      A bytearray() with the descriptor data.
    """
    num_bytes_following = len(self.data)
    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
    padding_size = nbf_with_padding - num_bytes_following
    desc = struct.pack(self.FORMAT_STRING, self.tag, nbf_with_padding)
    padding = struct.pack(str(padding_size) + 'x')
    ret = desc + self.data + padding
    return bytearray(ret)


class AvbPropertyDescriptor(AvbDescriptor):
  """A class for property descriptors.

  See the |AvbPropertyDescriptor| C struct for more information.

  Attributes:
    key: The key.
    value: The key.
  """

  TAG = 0
  SIZE = 32
  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
                   'Q'  # key size (bytes)
                   'Q')  # value size (bytes)

  def __init__(self, data=None):
    """Initializes a new property descriptor.

    Arguments:
      data: If not None, must be a bytearray of size |SIZE|.

    Raises:
      LookupError: If the given descriptor is malformed.
    """
    AvbDescriptor.__init__(self, None)
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (tag, num_bytes_following, key_size,
       value_size) = struct.unpack(self.FORMAT_STRING, data[0:self.SIZE])
      expected_size = round_to_multiple(
          self.SIZE - 16 + key_size + 1 + value_size + 1, 8)
      if tag != self.TAG or num_bytes_following != expected_size:
        raise LookupError('Given data does not look like a property '
                          'descriptor.')
      self.key = data[self.SIZE:(self.SIZE + key_size)]
      self.value = data[(self.SIZE + key_size + 1):(self.SIZE + key_size + 1 +
                                                    value_size)]
    else:
      self.key = ''
      self.value = ''

  def print_desc(self, o):
    """Print the descriptor.

    Arguments:
      o: The object to write the output to.
    """
    if len(self.value) < 256:
      o.write('    Prop: {} -> {}\n'.format(self.key, repr(str(self.value))))
    else:
      o.write('    Prop: {} -> ({} bytes)\n'.format(self.key, len(self.value)))

  def encode(self):
    """Serializes the descriptor.

    Returns:
      A bytearray() with the descriptor data.
    """
    num_bytes_following = self.SIZE + len(self.key) + len(self.value) + 2 - 16
    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
    padding_size = nbf_with_padding - num_bytes_following
    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
                       len(self.key), len(self.value))
    padding = struct.pack(str(padding_size) + 'x')
    ret = desc + self.key + '\0' + self.value + '\0' + padding
    return bytearray(ret)


class AvbHashtreeDescriptor(AvbDescriptor):
  """A class for hashtree descriptors.

  See the |AvbHashtreeDescriptor| C struct for more information.

  Attributes:
    dm_verity_version: dm-verity version used.
    image_size: Size of the image, after rounding up to |block_size|.
    tree_offset: Offset of the hash tree in the file.
    tree_size: Size of the tree.
    data_block_size: Data block size
    hash_block_size: Hash block size
    hash_algorithm: Hash algorithm used.
    partition_name: Partition name.
    salt: Salt used.
    root_digest: Root digest.
  """

  TAG = 1
  SIZE = 96
  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
                   'L'  # dm-verity version used
                   'Q'  # image size (bytes)
                   'Q'  # tree offset (bytes)
                   'Q'  # tree size (bytes)
                   'L'  # data block size (bytes)
                   'L'  # hash block size (bytes)
                   '32s'  # hash algorithm used
                   'L'  # partition name (bytes)
                   'L'  # salt length (bytes)
                   'L')  # root digest length (bytes)

  def __init__(self, data=None):
    """Initializes a new hashtree descriptor.

    Arguments:
      data: If not None, must be a bytearray of size |SIZE|.

    Raises:
      LookupError: If the given descriptor is malformed.
    """
    AvbDescriptor.__init__(self, None)
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (tag, num_bytes_following, self.dm_verity_version, self.image_size,
       self.tree_offset, self.tree_size, self.data_block_size,
       self.hash_block_size, self.hash_algorithm, partition_name_len, salt_len,
       root_digest_len) = struct.unpack(self.FORMAT_STRING, data[0:self.SIZE])
      expected_size = round_to_multiple(
          self.SIZE - 16 + partition_name_len + salt_len + root_digest_len, 8)
      if tag != self.TAG or num_bytes_following != expected_size:
        raise LookupError('Given data does not look like a hashtree '
                          'descriptor.')
      # Nuke NUL-bytes at the end.
      self.hash_algorithm = self.hash_algorithm.split('\0', 1)[0]
      o = 0
      self.partition_name = str(data[(self.SIZE + o):(self.SIZE + o +
                                                      partition_name_len)])
      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
      self.partition_name.decode('utf-8')
      o += partition_name_len
      self.salt = data[(self.SIZE + o):(self.SIZE + o + salt_len)]
      o += salt_len
      self.root_digest = data[(self.SIZE + o):(self.SIZE + o + root_digest_len)]
      if root_digest_len != len(hashlib.new(name=self.hash_algorithm).digest()):
        raise LookupError('root_digest_len doesn\'t match hash algorithm')

    else:
      self.dm_verity_version = 0
      self.image_size = 0
      self.tree_offset = 0
      self.tree_size = 0
      self.data_block_size = 0
      self.hash_block_size = 0
      self.hash_algorithm = ''
      self.partition_name = ''
      self.salt = bytearray()
      self.root_digest = bytearray()

  def print_desc(self, o):
    """Print the descriptor.

    Arguments:
      o: The object to write the output to.
    """
    o.write('    Hashtree descriptor:\n')
    o.write('      Version of dm-verity:  {}\n'.format(self.dm_verity_version))
    o.write('      Image Size:            {} bytes\n'.format(self.image_size))
    o.write('      Tree Offset:           {}\n'.format(self.tree_offset))
    o.write('      Tree Size:             {} bytes\n'.format(self.tree_size))
    o.write('      Data Block Size:       {} bytes\n'.format(
        self.data_block_size))
    o.write('      Hash Block Size:       {} bytes\n'.format(
        self.hash_block_size))
    o.write('      Hash Algorithm:        {}\n'.format(self.hash_algorithm))
    o.write('      Partition Name:        {}\n'.format(self.partition_name))
    o.write('      Salt:                  {}\n'.format(str(self.salt).encode(
        'hex')))
    o.write('      Root Digest:           {}\n'.format(str(
        self.root_digest).encode('hex')))

  def encode(self):
    """Serializes the descriptor.

    Returns:
      A bytearray() with the descriptor data.
    """
    encoded_name = self.partition_name.encode('utf-8')
    num_bytes_following = (self.SIZE + len(encoded_name) + len(self.salt) +
                           len(self.root_digest) - 16)
    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
    padding_size = nbf_with_padding - num_bytes_following
    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
                       self.dm_verity_version, self.image_size,
                       self.tree_offset, self.tree_size, self.data_block_size,
                       self.hash_block_size, self.hash_algorithm,
                       len(encoded_name), len(self.salt), len(self.root_digest))
    padding = struct.pack(str(padding_size) + 'x')
    ret = desc + encoded_name + self.salt + self.root_digest + padding
    return bytearray(ret)


class AvbHashDescriptor(AvbDescriptor):
  """A class for hash descriptors.

  See the |AvbHashDescriptor| C struct for more information.

  Attributes:
    image_size: Image size, in bytes.
    hash_algorithm: Hash algorithm used.
    partition_name: Partition name.
    salt: Salt used.
    digest: The hash value of salt and data combined.
  """

  TAG = 2
  SIZE = 68
  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
                   'Q'  # image size (bytes)
                   '32s'  # hash algorithm used
                   'L'  # partition name (bytes)
                   'L'  # salt length (bytes)
                   'L')  # digest length (bytes)

  def __init__(self, data=None):
    """Initializes a new hash descriptor.

    Arguments:
      data: If not None, must be a bytearray of size |SIZE|.

    Raises:
      LookupError: If the given descriptor is malformed.
    """
    AvbDescriptor.__init__(self, None)
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (tag, num_bytes_following, self.image_size, self.hash_algorithm,
       partition_name_len, salt_len,
       digest_len) = struct.unpack(self.FORMAT_STRING, data[0:self.SIZE])
      expected_size = round_to_multiple(
          self.SIZE - 16 + partition_name_len + salt_len + digest_len, 8)
      if tag != self.TAG or num_bytes_following != expected_size:
        raise LookupError('Given data does not look like a hash ' 'descriptor.')
      # Nuke NUL-bytes at the end.
      self.hash_algorithm = self.hash_algorithm.split('\0', 1)[0]
      o = 0
      self.partition_name = str(data[(self.SIZE + o):(self.SIZE + o +
                                                      partition_name_len)])
      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
      self.partition_name.decode('utf-8')
      o += partition_name_len
      self.salt = data[(self.SIZE + o):(self.SIZE + o + salt_len)]
      o += salt_len
      self.digest = data[(self.SIZE + o):(self.SIZE + o + digest_len)]
      if digest_len != len(hashlib.new(name=self.hash_algorithm).digest()):
        raise LookupError('digest_len doesn\'t match hash algorithm')

    else:
      self.image_size = 0
      self.hash_algorithm = ''
      self.partition_name = ''
      self.salt = bytearray()
      self.digest = bytearray()

  def print_desc(self, o):
    """Print the descriptor.

    Arguments:
      o: The object to write the output to.
    """
    o.write('    Hash descriptor:\n')
    o.write('      Image Size:            {} bytes\n'.format(self.image_size))
    o.write('      Hash Algorithm:        {}\n'.format(self.hash_algorithm))
    o.write('      Partition Name:        {}\n'.format(self.partition_name))
    o.write('      Salt:                  {}\n'.format(str(self.salt).encode(
        'hex')))
    o.write('      Digest:                {}\n'.format(str(self.digest).encode(
        'hex')))

  def encode(self):
    """Serializes the descriptor.

    Returns:
      A bytearray() with the descriptor data.
    """
    encoded_name = self.partition_name.encode('utf-8')
    num_bytes_following = (
        self.SIZE + len(encoded_name) + len(self.salt) + len(self.digest) - 16)
    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
    padding_size = nbf_with_padding - num_bytes_following
    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
                       self.image_size, self.hash_algorithm, len(encoded_name),
                       len(self.salt), len(self.digest))
    padding = struct.pack(str(padding_size) + 'x')
    ret = desc + encoded_name + self.salt + self.digest + padding
    return bytearray(ret)


class AvbKernelCmdlineDescriptor(AvbDescriptor):
  """A class for kernel command-line descriptors.

  See the |AvbKernelCmdlineDescriptor| C struct for more information.

  Attributes:
    kernel_cmdline: The kernel command-line.
  """

  TAG = 3
  SIZE = 20
  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
                   'L')  # cmdline length (bytes)

  def __init__(self, data=None):
    """Initializes a new kernel cmdline descriptor.

    Arguments:
      data: If not None, must be a bytearray of size |SIZE|.

    Raises:
      LookupError: If the given descriptor is malformed.
    """
    AvbDescriptor.__init__(self, None)
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (tag, num_bytes_following, kernel_cmdline_length) = (
          struct.unpack(self.FORMAT_STRING, data[0:self.SIZE]))
      expected_size = round_to_multiple(self.SIZE - 16 + kernel_cmdline_length,
                                        8)
      if tag != self.TAG or num_bytes_following != expected_size:
        raise LookupError('Given data does not look like a kernel cmdline '
                          'descriptor.')
      # Nuke NUL-bytes at the end.
      self.kernel_cmdline = str(data[self.SIZE:(self.SIZE +
                                                kernel_cmdline_length)])
      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
      self.kernel_cmdline.decode('utf-8')
    else:
      self.kernel_cmdline = ''

  def print_desc(self, o):
    """Print the descriptor.

    Arguments:
      o: The object to write the output to.
    """
    o.write('    Kernel Cmdline descriptor:\n')
    o.write('      Kernel Cmdline:        {}\n'.format(repr(
        self.kernel_cmdline)))

  def encode(self):
    """Serializes the descriptor.

    Returns:
      A bytearray() with the descriptor data.
    """
    encoded_str = self.kernel_cmdline.encode('utf-8')
    num_bytes_following = (self.SIZE + len(encoded_str) - 16)
    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
    padding_size = nbf_with_padding - num_bytes_following
    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
                       len(encoded_str))
    padding = struct.pack(str(padding_size) + 'x')
    ret = desc + encoded_str + padding
    return bytearray(ret)


class AvbChainPartitionDescriptor(AvbDescriptor):
  """A class for chained partition descriptors.

  See the |AvbChainPartitionDescriptor| C struct for more information.

  Attributes:
    rollback_index_slot: The rollback index slot to use.
    partition_name: Partition name.
    public_key: Bytes for the public key.
  """

  TAG = 4
  SIZE = 28
  FORMAT_STRING = ('!QQ'  # tag, num_bytes_following (descriptor header)
                   'L'  # rollback_index_slot
                   'L'  # partition_name_size (bytes)
                   'L')  # public_key_size (bytes)

  def __init__(self, data=None):
    """Initializes a new chain partition descriptor.

    Arguments:
      data: If not None, must be a bytearray of size |SIZE|.

    Raises:
      LookupError: If the given descriptor is malformed.
    """
    AvbDescriptor.__init__(self, None)
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (tag, num_bytes_following, self.rollback_index_slot, partition_name_len,
       public_key_len) = struct.unpack(self.FORMAT_STRING, data[0:self.SIZE])
      expected_size = round_to_multiple(
          self.SIZE - 16 + partition_name_len + public_key_len, 8)
      if tag != self.TAG or num_bytes_following != expected_size:
        raise LookupError('Given data does not look like a chain partition '
                          'descriptor.')
      o = 0
      self.partition_name = str(data[(self.SIZE + o):(self.SIZE + o +
                                                      partition_name_len)])
      # Validate UTF-8 - decode() raises UnicodeDecodeError if not valid UTF-8.
      self.partition_name.decode('utf-8')
      o += partition_name_len
      self.public_key = data[(self.SIZE + o):(self.SIZE + o + public_key_len)]

    else:
      self.rollback_index_slot = 0
      self.partition_name = ''
      self.public_key = bytearray()

  def print_desc(self, o):
    """Print the descriptor.

    Arguments:
      o: The object to write the output to.
    """
    o.write('    Chain Partition descriptor:\n')
    o.write('      Partition Name:        {}\n'.format(self.partition_name))
    o.write('      Rollback Index Slot:   {}\n'.format(
        self.rollback_index_slot))
    # Just show the SHA1 of the key, for size reasons.
    hexdig = hashlib.sha1(self.public_key).hexdigest()
    o.write('      Public key (sha1):     {}\n'.format(hexdig))

  def encode(self):
    """Serializes the descriptor.

    Returns:
      A bytearray() with the descriptor data.
    """
    encoded_name = self.partition_name.encode('utf-8')
    num_bytes_following = (
        self.SIZE + len(encoded_name) + len(self.public_key) - 16)
    nbf_with_padding = round_to_multiple(num_bytes_following, 8)
    padding_size = nbf_with_padding - num_bytes_following
    desc = struct.pack(self.FORMAT_STRING, self.TAG, nbf_with_padding,
                       self.rollback_index_slot, len(encoded_name),
                       len(self.public_key))
    padding = struct.pack(str(padding_size) + 'x')
    ret = desc + encoded_name + self.public_key + padding
    return bytearray(ret)


DESCRIPTOR_CLASSES = [
    AvbPropertyDescriptor, AvbHashtreeDescriptor, AvbHashDescriptor,
    AvbKernelCmdlineDescriptor, AvbChainPartitionDescriptor
]


def parse_descriptors(data):
  """Parses a blob of data into descriptors.

  Arguments:
    data: A bytearray() with encoded descriptors.

  Returns:
    A list of instances of objects derived from AvbDescriptor. For
    unknown descriptors, the class AvbDescriptor is used.
  """
  o = 0
  ret = []
  while o < len(data):
    tag, nb_following = struct.unpack('!2Q', data[o:o + 16])
    if tag < len(DESCRIPTOR_CLASSES):
      c = DESCRIPTOR_CLASSES[tag]
    else:
      c = AvbDescriptor
    ret.append(c(bytearray(data[o:o + 16 + nb_following])))
    o += 16 + nb_following
  return ret


class AvbFooter(object):
  """A class for parsing and writing footers.

  Footers are stored at the end of partitions and point to where the
  AvbVBMeta blob is located. They also contain the original size of
  the image before AVB information was added.

  Attributes:
    magic: Magic for identifying the footer, see |MAGIC|.
    version_major: The major version of avbtool that wrote the footer.
    version_minor: The minor version of avbtool that wrote the footer.
    original_image_size: Original image size.
    vbmeta_offset: Offset of where the AvbVBMeta blob is stored.
    vbmeta_size: Size of the AvbVBMeta blob.
  """

  MAGIC = 'AVBf'
  SIZE = 64
  RESERVED = 28
  FORMAT_STRING = ('!4s2L'  # magic, 2 x version.
                   'Q'  # Original image size.
                   'Q'  # Offset of VBMeta blob.
                   'Q' +  # Size of VBMeta blob.
                   str(RESERVED) + 'x')  # padding for reserved bytes

  def __init__(self, data=None):
    """Initializes a new footer object.

    Arguments:
      data: If not None, must be a bytearray of size 4096.

    Raises:
      LookupError: If the given footer is malformed.
      struct.error: If the given data has no footer.
    """
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (self.magic, self.version_major, self.version_minor,
       self.original_image_size, self.vbmeta_offset,
       self.vbmeta_size) = struct.unpack(self.FORMAT_STRING, data)
      if self.magic != self.MAGIC:
        raise LookupError('Given data does not look like a AVB footer.')
    else:
      self.magic = self.MAGIC
      self.version_major = AVB_VERSION_MAJOR
      self.version_minor = AVB_VERSION_MINOR
      self.original_image_size = 0
      self.vbmeta_offset = 0
      self.vbmeta_size = 0

  def encode(self):
    """Gets a string representing the binary encoding of the footer.

    Returns:
      A bytearray() with a binary representation of the footer.
    """
    return struct.pack(self.FORMAT_STRING, self.magic, self.version_major,
                       self.version_minor, self.original_image_size,
                       self.vbmeta_offset, self.vbmeta_size)


class AvbVBMetaHeader(object):
  """A class for parsing and writing AVB vbmeta images.

  Attributes:
    The attributes correspond to the |AvbVBMetaHeader| struct
    defined in avb_vbmeta_header.h.
  """

  SIZE = 256

  # Keep in sync with |reserved| field of |AvbVBMetaImageHeader|.
  RESERVED = 152

  # Keep in sync with |AvbVBMetaImageHeader|.
  FORMAT_STRING = ('!4s2L'  # magic, 2 x version
                   '2Q'  # 2 x block size
                   'L'  # algorithm type
                   '2Q'  # offset, size (hash)
                   '2Q'  # offset, size (signature)
                   '2Q'  # offset, size (public key)
                   '2Q'  # offset, size (descriptors)
                   'Q' +  # rollback_index
                   str(RESERVED) + 'x')  # padding for reserved bytes

  def __init__(self, data=None):
    """Initializes a new header object.

    Arguments:
      data: If not None, must be a bytearray of size 8192.

    Raises:
      Exception: If the given data is malformed.
    """
    assert struct.calcsize(self.FORMAT_STRING) == self.SIZE

    if data:
      (self.magic, self.header_version_major, self.header_version_minor,
       self.authentication_data_block_size, self.auxiliary_data_block_size,
       self.algorithm_type, self.hash_offset, self.hash_size,
       self.signature_offset, self.signature_size, self.public_key_offset,
       self.public_key_size, self.descriptors_offset, self.descriptors_size,
       self.rollback_index) = struct.unpack(self.FORMAT_STRING, data)
      # Nuke NUL-bytes at the end of the string.
      if self.magic != 'AVB0':
        raise AvbError('Given image does not look like a vbmeta image.')
    else:
      self.magic = 'AVB0'
      self.header_version_major = AVB_VERSION_MAJOR
      self.header_version_minor = AVB_VERSION_MINOR
      self.authentication_data_block_size = 0
      self.auxiliary_data_block_size = 0
      self.algorithm_type = 0
      self.hash_offset = 0
      self.hash_size = 0
      self.signature_offset = 0
      self.signature_size = 0
      self.public_key_offset = 0
      self.public_key_size = 0
      self.descriptors_offset = 0
      self.descriptors_size = 0
      self.rollback_index = 0

  def save(self, output):
    """Serializes the header (256 bytes) to disk.

    Arguments:
      output: The object to write the output to.
    """
    output.write(struct.pack(
        self.FORMAT_STRING, self.magic, self.header_version_major,
        self.header_version_minor, self.authentication_data_block_size,
        self.auxiliary_data_block_size, self.algorithm_type, self.hash_offset,
        self.hash_size, self.signature_offset, self.signature_size,
        self.public_key_offset, self.public_key_size, self.descriptors_offset,
        self.descriptors_size, self.rollback_index))

  def encode(self):
    """Serializes the header (256) to a bytearray().

    Returns:
      A bytearray() with the encoded header.
    """
    return struct.pack(self.FORMAT_STRING, self.magic,
                       self.header_version_major, self.header_version_minor,
                       self.authentication_data_block_size,
                       self.auxiliary_data_block_size, self.algorithm_type,
                       self.hash_offset, self.hash_size, self.signature_offset,
                       self.signature_size, self.public_key_offset,
                       self.public_key_size, self.descriptors_offset,
                       self.descriptors_size, self.rollback_index)


class Avb(object):
  """Business logic for avbtool command-line tool."""

  # Keep in sync with avb_ab_flow.h.
  AB_FORMAT_NO_CRC = '!4sBB2xBBBxBBBx12x'
  AB_MAGIC = '\0AB0'
  AB_MAJOR_VERSION = 1
  AB_MINOR_VERSION = 0
  AB_MISC_METADATA_OFFSET = 2048

  def erase_footer(self, image_filename, keep_hashtree):
    """Implements the 'erase_footer' command.

    Arguments:
      image_filename: File to erase a footer from.
      keep_hashtree: If True, keep the hashtree around.

    Raises:
      AvbError: If there's no footer in the image.
    """

    image = ImageHandler(image_filename)

    (footer, _, descriptors, _) = self._parse_image(image)

    if not footer:
      raise AvbError('Given image does not have a footer.')

    new_image_size = None
    if not keep_hashtree:
      new_image_size = footer.original_image_size
    else:
      # If requested to keep the hashtree, search for a hashtree
      # descriptor to figure out the location and size of the hashtree.
      for desc in descriptors:
        if isinstance(desc, AvbHashtreeDescriptor):
          # The hashtree is always just following the main data so the
          # new size is easily derived.
          new_image_size = desc.tree_offset + desc.tree_size
          break
      if not new_image_size:
        raise AvbError('Requested to keep hashtree but no hashtree '
                       'descriptor was found.')

    # And cut...
    image.truncate(new_image_size)

  def set_ab_metadata(self, misc_image, slot_data):
    """Implements the 'set_ab_metadata' command.

    The |slot_data| argument must be of the form 'A_priority:A_tries_remaining:
    A_successful_boot:B_priority:B_tries_remaining:B_successful_boot'.

    Arguments:
      misc_image: The misc image to write to.
      slot_data: Slot data as a string

    Raises:
      AvbError: If slot data is malformed.
    """
    tokens = slot_data.split(':')
    if len(tokens) != 6:
      raise AvbError('Malformed slot data "{}".'.format(slot_data))
    a_priority = int(tokens[0])
    a_tries_remaining = int(tokens[1])
    a_success = True if int(tokens[2]) != 0 else False
    b_priority = int(tokens[3])
    b_tries_remaining = int(tokens[4])
    b_success = True if int(tokens[5]) != 0 else False

    ab_data_no_crc = struct.pack(self.AB_FORMAT_NO_CRC,
                                 self.AB_MAGIC,
                                 self.AB_MAJOR_VERSION, self.AB_MINOR_VERSION,
                                 a_priority, a_tries_remaining, a_success,
                                 b_priority, b_tries_remaining, b_success)
    # Force CRC to be unsigned, see https://bugs.python.org/issue4903 for why.
    crc_value = binascii.crc32(ab_data_no_crc) & 0xffffffff
    ab_data = ab_data_no_crc + struct.pack('!I', crc_value)
    misc_image.seek(self.AB_MISC_METADATA_OFFSET)
    misc_image.write(ab_data)

  def info_image(self, image_filename, output):
    """Implements the 'info_image' command.

    Arguments:
      image_filename: Image file to get information from (file object).
      output: Output file to write human-readable information to (file object).
    """

    image = ImageHandler(image_filename)

    o = output

    (footer, header, descriptors, image_size) = self._parse_image(image)

    if footer:
      o.write('Footer version:           {}.{}\n'.format(footer.version_major,
                                                         footer.version_minor))
      o.write('Image size:               {} bytes\n'.format(image_size))
      o.write('Original image size:      {} bytes\n'.format(
          footer.original_image_size))
      o.write('VBMeta offset:            {}\n'.format(footer.vbmeta_offset))
      o.write('VBMeta size:              {} bytes\n'.format(footer.vbmeta_size))
      o.write('--\n')

    (alg_name, _) = lookup_algorithm_by_type(header.algorithm_type)

    o.write('VBMeta image version:     {}.{}{}\n'.format(
        header.header_version_major, header.header_version_minor,
        ' (Sparse)' if image.is_sparse else ''))
    o.write('Header Block:             {} bytes\n'.format(AvbVBMetaHeader.SIZE))
    o.write('Authentication Block:     {} bytes\n'.format(
        header.authentication_data_block_size))
    o.write('Auxiliary Block:          {} bytes\n'.format(
        header.auxiliary_data_block_size))
    o.write('Algorithm:                {}\n'.format(alg_name))
    o.write('Rollback Index:           {}\n'.format(header.rollback_index))

    # Print descriptors.
    num_printed = 0
    o.write('Descriptors:\n')
    for desc in descriptors:
      desc.print_desc(o)
      num_printed += 1
    if num_printed == 0:
      o.write('    (none)\n')

  def _parse_image(self, image):
    """Gets information about an image.

    The image can either be a vbmeta or an image with a footer.

    Arguments:
      image: An ImageHandler (vbmeta or footer) with a hashtree descriptor.

    Returns:
      A tuple where the first argument is a AvbFooter (None if there
      is no footer on the image), the second argument is a
      AvbVBMetaHeader, the third argument is a list of
      AvbDescriptor-derived instances, and the fourth argument is the
      size of |image|.
    """
    assert isinstance(image, ImageHandler)
    footer = None
    image_size = image.care_size
    image.seek(image_size - AvbFooter.SIZE)
    try:
      footer = AvbFooter(image.read(AvbFooter.SIZE))
    except (LookupError, struct.error):
      # Nope, just seek back to the start.
      image.seek(0)

    vbmeta_offset = 0
    if footer:
      vbmeta_offset = footer.vbmeta_offset

    image.seek(vbmeta_offset)
    h = AvbVBMetaHeader(image.read(AvbVBMetaHeader.SIZE))

    auth_block_offset = vbmeta_offset + AvbVBMetaHeader.SIZE
    aux_block_offset = auth_block_offset + h.authentication_data_block_size
    desc_start_offset = aux_block_offset + h.descriptors_offset
    image.seek(desc_start_offset)
    descriptors = parse_descriptors(image.read(h.descriptors_size))

    return footer, h, descriptors, image_size

  def _get_cmdline_descriptor_for_dm_verity(self, image):
    """Generate kernel cmdline descriptor for dm-verity.

    Arguments:
      image: An ImageHandler (vbmeta or footer) with a hashtree descriptor.

    Returns:
      A AvbKernelCmdlineDescriptor with dm-verity kernel cmdline
      instructions for the hashtree.

    Raises:
      AvbError: If  |image| doesn't have a hashtree descriptor.

    """

    (_, _, descriptors, _) = self._parse_image(image)

    ht = None
    for desc in descriptors:
      if isinstance(desc, AvbHashtreeDescriptor):
        ht = desc
        break

    if not ht:
      raise AvbError('No hashtree descriptor in given image')

    c = 'dm="1 vroot none ro 1,'
    c += '0 '  # start
    c += '{} '.format((ht.image_size / 512))  # size (# sectors)
    c += 'verity {} '.format(ht.dm_verity_version)  # type and version
    c += 'PARTUUID=$(ANDROID_SYSTEM_PARTUUID) '  # data_dev
    c += 'PARTUUID=$(ANDROID_SYSTEM_PARTUUID) '  # hash_dev
    c += '{} '.format(ht.data_block_size)  # data_block
    c += '{} '.format(ht.hash_block_size)  # hash_block
    c += '{} '.format(ht.image_size / ht.data_block_size)  # #blocks
    c += '{} '.format(ht.image_size / ht.data_block_size)  # hash_offset
    c += '{} '.format(ht.hash_algorithm)  # hash_alg
    c += '{} '.format(str(ht.root_digest).encode('hex'))  # root_digest
    c += '{}'.format(str(ht.salt).encode('hex'))  # salt
    c += '"'

    desc = AvbKernelCmdlineDescriptor()
    desc.kernel_cmdline = c
    return desc

  def make_vbmeta_image(self, output, chain_partitions, algorithm_name,
                        key_path, rollback_index, props, props_from_file,
                        kernel_cmdlines,
                        generate_dm_verity_cmdline_from_hashtree,
                        include_descriptors_from_image):
    """Implements the 'make_vbmeta_image' command.

    Arguments:
      output: File to write the image to.
      chain_partitions: List of partitions to chain.
      algorithm_name: Name of algorithm to use.
      key_path: Path to key to use or None.
      rollback_index: The rollback index to use.
      props: Properties to insert (list of strings of the form 'key:value').
      props_from_file: Properties to insert (list of strings 'key:<path>').
      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
      generate_dm_verity_cmdline_from_hashtree: None or file to generate from.
      include_descriptors_from_image: List of file objects with descriptors.

    Raises:
      AvbError: If a chained partition is malformed.
    """

    descriptors = []

    # Insert chained partition descriptors.
    if chain_partitions:
      for cp in chain_partitions:
        cp_tokens = cp.split(':')
        if len(cp_tokens) != 3:
          raise AvbError('Malformed chained partition "{}".'.format(cp))
        desc = AvbChainPartitionDescriptor()
        desc.partition_name = cp_tokens[0]
        desc.rollback_index_slot = int(cp_tokens[1])
        if desc.rollback_index_slot < 1:
          raise AvbError('Rollback index slot must be 1 or larger.')
        file_path = cp_tokens[2]
        desc.public_key = open(file_path, 'rb').read()
        descriptors.append(desc)

    vbmeta_blob = self._generate_vbmeta_blob(
        algorithm_name, key_path, descriptors, rollback_index, props,
        props_from_file, kernel_cmdlines,
        generate_dm_verity_cmdline_from_hashtree,
        include_descriptors_from_image)

    # Write entire vbmeta blob (header, authentication, auxiliary).
    output.seek(0)
    output.write(vbmeta_blob)

  def _generate_vbmeta_blob(self, algorithm_name, key_path, descriptors,
                            rollback_index, props, props_from_file,
                            kernel_cmdlines,
                            generate_dm_verity_cmdline_from_hashtree,
                            include_descriptors_from_image):
    """Generates a VBMeta blob.

    This blob contains the header (struct AvbVBMetaHeader), the
    authentication data block (which contains the hash and signature
    for the header and auxiliary block), and the auxiliary block
    (which contains descriptors, the public key used, and other data).

    The |key| parameter can |None| only if the |algorithm_name| is
    'NONE'.

    Arguments:
      algorithm_name: The algorithm name as per the ALGORITHMS dict.
      key_path: The path to the .pem file used to sign the blob.
      descriptors: A list of descriptors to insert or None.
      rollback_index: The rollback index to use.
      props: Properties to insert (List of strings of the form 'key:value').
      props_from_file: Properties to insert (List of strings 'key:<path>').
      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
      generate_dm_verity_cmdline_from_hashtree: None or file to generate
        dm-verity kernel cmdline from.
      include_descriptors_from_image: List of file objects for which
        to insert descriptors from.

    Returns:
      A bytearray() with the VBMeta blob.

    Raises:
      Exception: If the |algorithm_name| is not found, if no key has
        been given and the given algorithm requires one, or the key is
        of the wrong size.

    """
    try:
      alg = ALGORITHMS[algorithm_name]
    except KeyError:
      raise AvbError('Unknown algorithm with name {}'.format(algorithm_name))

    # Descriptors.
    encoded_descriptors = bytearray()
    if descriptors:
      for desc in descriptors:
        encoded_descriptors.extend(desc.encode())

    # Add properties.
    if props:
      for prop in props:
        idx = prop.find(':')
        if idx == -1:
          raise AvbError('Malformed property "{}".'.format(prop))
        desc = AvbPropertyDescriptor()
        desc.key = prop[0:idx]
        desc.value = prop[(idx + 1):]
        encoded_descriptors.extend(desc.encode())
    if props_from_file:
      for prop in props_from_file:
        idx = prop.find(':')
        if idx == -1:
          raise AvbError('Malformed property "{}".'.format(prop))
        desc = AvbPropertyDescriptor()
        desc.key = prop[0:idx]
        desc.value = prop[(idx + 1):]
        file_path = prop[(idx + 1):]
        desc.value = open(file_path, 'rb').read()
        encoded_descriptors.extend(desc.encode())

    # Add AvbKernelCmdline descriptor for dm-verity, if requested.
    if generate_dm_verity_cmdline_from_hashtree:
      image_handler = ImageHandler(
          generate_dm_verity_cmdline_from_hashtree.name)
      encoded_descriptors.extend(self._get_cmdline_descriptor_for_dm_verity(
          image_handler).encode())

    # Add kernel command-lines.
    if kernel_cmdlines:
      for i in kernel_cmdlines:
        desc = AvbKernelCmdlineDescriptor()
        desc.kernel_cmdline = i
        encoded_descriptors.extend(desc.encode())

    # Add descriptors from other images.
    if include_descriptors_from_image:
      for image in include_descriptors_from_image:
        image_handler = ImageHandler(image.name)
        (_, _, image_descriptors, _) = self._parse_image(image_handler)
        for desc in image_descriptors:
          encoded_descriptors.extend(desc.encode())

    key = None
    encoded_key = bytearray()
    if alg.public_key_num_bytes > 0:
      if not key_path:
        raise AvbError('Key is required for algorithm {}'.format(
            algorithm_name))
      key = Crypto.PublicKey.RSA.importKey(open(key_path).read())
      encoded_key = encode_rsa_key(key)
      if len(encoded_key) != alg.public_key_num_bytes:
        raise AvbError('Key is wrong size for algorithm {}'.format(
            algorithm_name))

    h = AvbVBMetaHeader()

    # For the Auxiliary data block, descriptors are stored at offset 0
    # and the public key is immediately after that.
    h.auxiliary_data_block_size = round_to_multiple(
        len(encoded_descriptors) + len(encoded_key), 64)
    h.descriptors_offset = 0
    h.descriptors_size = len(encoded_descriptors)
    h.public_key_offset = h.descriptors_size
    h.public_key_size = len(encoded_key)

    # For the Authentication data block, the hash is first and then
    # the signature.
    h.authentication_data_block_size = round_to_multiple(
        alg.hash_num_bytes + alg.public_key_num_bytes, 64)
    h.algorithm_type = alg.algorithm_type
    h.hash_offset = 0
    h.hash_size = alg.hash_num_bytes
    # Signature offset and size - it's stored right after the hash
    # (in Authentication data block).
    h.signature_offset = alg.hash_num_bytes
    h.signature_size = alg.signature_num_bytes

    h.rollback_index = rollback_index

    # Generate Header data block.
    header_data_blob = h.encode()

    # Generate Auxiliary data block.
    aux_data_blob = bytearray()
    aux_data_blob.extend(encoded_descriptors)
    aux_data_blob.extend(encoded_key)
    padding_bytes = h.auxiliary_data_block_size - len(aux_data_blob)
    aux_data_blob.extend('\0' * padding_bytes)

    # Calculate the hash.
    binary_hash = bytearray()
    binary_signature = bytearray()
    if algorithm_name != 'NONE':
      if algorithm_name[0:6] == 'SHA256':
        ha = hashlib.sha256()
      elif algorithm_name[0:6] == 'SHA512':
        ha = hashlib.sha512()
      else:
        raise AvbError('Unsupported algorithm {}.'.format(algorithm_name))
      ha.update(header_data_blob)
      ha.update(aux_data_blob)
      binary_hash.extend(ha.digest())

      # Calculate the signature.
      p = subprocess.Popen(
          ['openssl', 'rsautl', '-sign', '-inkey', key_path, '-raw'],
          stdin=subprocess.PIPE,
          stdout=subprocess.PIPE,
          stderr=subprocess.PIPE)
      padding_and_hash = str(bytearray(alg.padding)) + binary_hash
      (pout, perr) = p.communicate(padding_and_hash)
      retcode = p.wait()
      if retcode != 0:
        raise AvbError('Error signing: {}'.format(perr))
      binary_signature.extend(pout)

    # Generate Authentication data block.
    auth_data_blob = bytearray()
    auth_data_blob.extend(binary_hash)
    auth_data_blob.extend(binary_signature)
    padding_bytes = h.authentication_data_block_size - len(auth_data_blob)
    auth_data_blob.extend('\0' * padding_bytes)

    return header_data_blob + auth_data_blob + aux_data_blob

  def extract_public_key(self, key_path, output):
    """Implements the 'extract_public_key' command.

    Arguments:
      key_path: The path to a RSA private key file.
      output: The file to write to.
    """
    key = Crypto.PublicKey.RSA.importKey(open(key_path).read())
    write_rsa_key(output, key)

  def add_hash_footer(self, image_filename, partition_size, partition_name,
                      hash_algorithm, salt, algorithm_name, key_path,
                      rollback_index, props, props_from_file, kernel_cmdlines,
                      generate_dm_verity_cmdline_from_hashtree,
                      include_descriptors_from_image):
    """Implementation of the add_hash_footer on unsparse images.

    Arguments:
      image_filename: File to add the footer to.
      partition_size: Size of partition.
      partition_name: Name of partition (without A/B suffix).
      hash_algorithm: Hash algorithm to use.
      salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
      algorithm_name: Name of algorithm to use.
      key_path: Path to key to use or None.
      rollback_index: Rollback index.
      props: Properties to insert (List of strings of the form 'key:value').
      props_from_file: Properties to insert (List of strings 'key:<path>').
      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
      generate_dm_verity_cmdline_from_hashtree: None or file to generate
        dm-verity kernel cmdline from.
      include_descriptors_from_image: List of file objects for which
        to insert descriptors from.

    Raises:
      AvbError: If an argument is incorrect.
    """
    image = ImageHandler(image_filename)

    if partition_size % image.block_size != 0:
      raise AvbError('Partition size of {} is not a multiple of the image '
                     'block size {}.'.format(partition_size,
                                             image.block_size))

    # If there's already a footer, truncate the image to its original
    # size. This way 'avbtool add_hash_footer' is idempotent (modulo
    # salts).
    image_size = image.care_size
    image.seek(image_size - AvbFooter.SIZE)
    try:
      footer = AvbFooter(image.read(AvbFooter.SIZE))
      # Existing footer found. Just truncate.
      original_image_size = footer.original_image_size
      image_size = footer.original_image_size
      image.truncate(image_size)
    except (LookupError, struct.error):
      original_image_size = image_size

    # If anything goes wrong from here-on, restore the image back to
    # its original size.
    try:
      digest_size = len(hashlib.new(name=hash_algorithm).digest())
      if salt:
        salt = salt.decode('hex')
      else:
        if salt is None:
          # If salt is not explicitly specified, choose a hash
          # that's the same size as the hash size.
          hash_size = digest_size
          salt = open('/dev/urandom').read(hash_size)
        else:
          salt = ''

      hasher = hashlib.new(name=hash_algorithm, string=salt)
      # TODO(zeuthen): might want to read this in chunks to avoid
      # memory pressure, then again, this is only supposed to be used
      # on kernel/initramfs partitions. Possible optimization.
      image.seek(0)
      hasher.update(image.read(image_size))
      digest = hasher.digest()

      h_desc = AvbHashDescriptor()
      h_desc.image_size = image_size
      h_desc.hash_algorithm = hash_algorithm
      h_desc.partition_name = partition_name
      h_desc.salt = salt
      h_desc.digest = digest

      # Generate the VBMeta footer.
      vbmeta_blob = self._generate_vbmeta_blob(
          algorithm_name, key_path, [h_desc], rollback_index, props,
          props_from_file, kernel_cmdlines,
          generate_dm_verity_cmdline_from_hashtree,
          include_descriptors_from_image)

      # We might have a DONT_CARE hole at the end (in which case
      # |image.care_size| < |image.image_size|) so truncate here.
      image.truncate(image.care_size)

      # If the image isn't sparse, its size might not be a multiple of
      # the block size. This will screw up padding later so just grow it.
      if image.care_size % image.block_size != 0:
        assert not image.is_sparse
        padding_needed = image.block_size - (image.care_size%image.block_size)
        image.truncate(image.care_size + padding_needed)

      # The append_raw() method requires content with size being a
      # multiple of |block_size| so add padding as needed. Also record
      # where this is written to since we'll need to put that in the
      # footer.
      vbmeta_offset = image.care_size
      padding_needed = (round_to_multiple(len(vbmeta_blob), image.block_size) -
                        len(vbmeta_blob))
      vbmeta_blob_with_padding = vbmeta_blob + '\0'*padding_needed
      image.append_raw(vbmeta_blob_with_padding)
      vbmeta_end_offset = vbmeta_offset + len(vbmeta_blob_with_padding)

      # Now insert a DONT_CARE chunk with enough bytes such that the
      # final Footer block is at the end of partition_size..
      image.append_dont_care(partition_size - vbmeta_end_offset -
                             1*image.block_size)

      # Generate the Footer that tells where the VBMeta footer
      # is. Also put enough padding in the front of the footer since
      # we'll write out an entire block.
      footer = AvbFooter()
      footer.original_image_size = original_image_size
      footer.vbmeta_offset = vbmeta_offset
      footer.vbmeta_size = len(vbmeta_blob)
      footer_blob = footer.encode()
      footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
                                  footer_blob)
      image.append_raw(footer_blob_with_padding)

    except:
      # Truncate back to original size, then re-raise
      image.truncate(original_image_size)
      raise

  def add_hashtree_footer(self, image_filename, partition_size, partition_name,
                          hash_algorithm, block_size, salt, algorithm_name,
                          key_path, rollback_index, props, props_from_file,
                          kernel_cmdlines,
                          generate_dm_verity_cmdline_from_hashtree,
                          include_descriptors_from_image):
    """Implements the 'add_hashtree_footer' command.

    See https://gitlab.com/cryptsetup/cryptsetup/wikis/DMVerity for
    more information about dm-verity and these hashes.

    Arguments:
      image_filename: File to add the footer to.
      partition_size: Size of partition.
      partition_name: Name of partition (without A/B suffix).
      hash_algorithm: Hash algorithm to use.
      block_size: Block size to use.
      salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
      algorithm_name: Name of algorithm to use.
      key_path: Path to key to use or None.
      rollback_index: Rollback index.
      props: Properties to insert (List of strings of the form 'key:value').
      props_from_file: Properties to insert (List of strings 'key:<path>').
      kernel_cmdlines: Kernel cmdlines to insert (list of strings).
      generate_dm_verity_cmdline_from_hashtree: None or file to generate
        dm-verity kernel cmdline from.
      include_descriptors_from_image: List of file objects for which
        to insert descriptors from.

    Raises:
      AvbError: If an argument is incorrect.
    """
    image = ImageHandler(image_filename)

    if partition_size % image.block_size != 0:
      raise AvbError('Partition size of {} is not a multiple of the image '
                     'block size {}.'.format(partition_size,
                                             image.block_size))

    # If there's already a footer, truncate the image to its original
    # size. This way 'avbtool add_hashtree_footer' is idempotent
    # (modulo salts).
    image_size = image.care_size
    image.seek(image_size - AvbFooter.SIZE)
    try:
      footer = AvbFooter(image.read(AvbFooter.SIZE))
      # Existing footer found. Just truncate.
      original_image_size = footer.original_image_size
      image_size = footer.original_image_size
      image.truncate(image_size)
    except (LookupError, struct.error):
      original_image_size = image_size

    # If anything goes wrong from here-on, restore the image back to
    # its original size.
    try:
      # Ensure image is multiple of block_size.
      rounded_image_size = round_to_multiple(image_size, block_size)
      if rounded_image_size > image_size:
        image.append_raw('\0' * (rounded_image_size - image_size))
        image_size = rounded_image_size

      digest_size = len(hashlib.new(name=hash_algorithm).digest())
      digest_padding = round_to_pow2(digest_size) - digest_size

      if salt:
        salt = salt.decode('hex')
      else:
        if salt is None:
          # If salt is not explicitly specified, choose a hash
          # that's the same size as the hash size.
          hash_size = digest_size
          salt = open('/dev/urandom').read(hash_size)
        else:
          salt = ''

      # Hashes are stored upside down so we need to calculate hash
      # offsets in advance.
      (hash_level_offsets, tree_size) = calc_hash_level_offsets(
          image_size, block_size, digest_size + digest_padding)

      # We might have a DONT_CARE hole at the end (in which case
      # |image.care_size| < |image.image_size|) so truncate here.
      image.truncate(image.care_size)

      # If the image isn't sparse, its size might not be a multiple of
      # the block size. This will screw up padding later so just grow it.
      if image.care_size % image.block_size != 0:
        assert not image.is_sparse
        padding_needed = image.block_size - (image.care_size%image.block_size)
        image.truncate(image.care_size + padding_needed)

      # Generate the tree and add padding as needed.
      tree_offset = image.care_size
      root_digest, hash_tree = generate_hash_tree(image, image_size,
                                                  block_size,
                                                  hash_algorithm, salt,
                                                  digest_padding,
                                                  hash_level_offsets,
                                                  tree_size)
      padding_needed = (round_to_multiple(len(hash_tree), image.block_size) -
                        len(hash_tree))
      hash_tree_with_padding = hash_tree + '\0'*padding_needed
      image.append_raw(hash_tree_with_padding)

      # Generate HashtreeDescriptor with details about the tree we
      # just generated.
      ht_desc = AvbHashtreeDescriptor()
      ht_desc.dm_verity_version = 1
      ht_desc.image_size = image_size
      ht_desc.tree_offset = tree_offset
      ht_desc.tree_size = tree_size
      ht_desc.data_block_size = block_size
      ht_desc.hash_block_size = block_size
      ht_desc.hash_algorithm = hash_algorithm
      ht_desc.partition_name = partition_name
      ht_desc.salt = salt
      ht_desc.root_digest = root_digest

      # Generate the VBMeta footer and add padding as needed.
      vbmeta_offset = tree_offset + len(hash_tree_with_padding)
      vbmeta_blob = self._generate_vbmeta_blob(
          algorithm_name, key_path, [ht_desc], rollback_index, props,
          props_from_file, kernel_cmdlines,
          generate_dm_verity_cmdline_from_hashtree,
          include_descriptors_from_image)
      padding_needed = (round_to_multiple(len(vbmeta_blob), image.block_size) -
                        len(vbmeta_blob))
      vbmeta_blob_with_padding = vbmeta_blob + '\0'*padding_needed
      image.append_raw(vbmeta_blob_with_padding)

      # Now insert a DONT_CARE chunk with enough bytes such that the
      # final Footer block is at the end of partition_size..
      image.append_dont_care(partition_size - image.care_size -
                             1*image.block_size)

      # Generate the Footer that tells where the VBMeta footer
      # is. Also put enough padding in the front of the footer since
      # we'll write out an entire block.
      footer = AvbFooter()
      footer.original_image_size = original_image_size
      footer.vbmeta_offset = vbmeta_offset
      footer.vbmeta_size = len(vbmeta_blob)
      footer_blob = footer.encode()
      footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
                                  footer_blob)
      image.append_raw(footer_blob_with_padding)

    except:
      # Truncate back to original size, then re-raise
      image.truncate(original_image_size)
      raise


def calc_hash_level_offsets(image_size, block_size, digest_size):
  """Calculate the offsets of all the hash-levels in a Merkle-tree.

  Arguments:
    image_size: The size of the image to calculate a Merkle-tree for.
    block_size: The block size, e.g. 4096.
    digest_size: The size of each hash, e.g. 32 for SHA-256.

  Returns:
    A tuple where the first argument is an array of offsets and the
    second is size of the tree, in bytes.
  """
  level_offsets = []
  level_sizes = []
  tree_size = 0

  num_levels = 0
  size = image_size
  while size > block_size:
    num_blocks = (size + block_size - 1) / block_size
    level_size = round_to_multiple(num_blocks * digest_size, block_size)

    level_sizes.append(level_size)
    tree_size += level_size
    num_levels += 1

    size = level_size

  for n in range(0, num_levels):
    offset = 0
    for m in range(n + 1, num_levels):
      offset += level_sizes[m]
    level_offsets.append(offset)

  return level_offsets, tree_size


def generate_hash_tree(image, image_size, block_size, hash_alg_name, salt,
                       digest_padding, hash_level_offsets, tree_size):
  """Generates a Merkle-tree for a file.

  Args:
    image: The image, as a file.
    image_size: The size of the image.
    block_size: The block size, e.g. 4096.
    hash_alg_name: The hash algorithm, e.g. 'sha256' or 'sha1'.
    salt: The salt to use.
    digest_padding: The padding for each digest.
    hash_level_offsets: The offsets from calc_hash_level_offsets().
    tree_size: The size of the tree, in number of bytes.

  Returns:
    A tuple where the first element is the top-level hash and the
    second element is the hash-tree.
  """
  hash_ret = bytearray(tree_size)
  hash_src_offset = 0
  hash_src_size = image_size
  level_num = 0
  while hash_src_size > block_size:
    level_output = ''
    remaining = hash_src_size
    while remaining > 0:
      hasher = hashlib.new(name=hash_alg_name, string=salt)
      # Only read from the file for the first level - for subsequent
      # levels, access the array we're building.
      if level_num == 0:
        image.seek(hash_src_offset + hash_src_size - remaining)
        data = image.read(min(remaining, block_size))
      else:
        offset = hash_level_offsets[level_num - 1] + hash_src_size - remaining
        data = hash_ret[offset:offset + block_size]
      hasher.update(data)

      remaining -= len(data)
      if len(data) < block_size:
        hasher.update('\0' * (block_size - len(data)))
      level_output += hasher.digest()
      if digest_padding > 0:
        level_output += '\0' * digest_padding

    padding_needed = (round_to_multiple(
        len(level_output), block_size) - len(level_output))
    level_output += '\0' * padding_needed

    # Copy level-output into resulting tree.
    offset = hash_level_offsets[level_num]
    hash_ret[offset:offset + len(level_output)] = level_output

    # Continue on to the next level.
    hash_src_size = len(level_output)
    level_num += 1

  hasher = hashlib.new(name=hash_alg_name, string=salt)
  hasher.update(level_output)
  return hasher.digest(), hash_ret


class AvbTool(object):
  """Object for avbtool command-line tool."""

  def __init__(self):
    """Initializer method."""
    self.avb = Avb()

  def _add_common_args(self, sub_parser):
    """Adds arguments used by several sub-commands.

    Arguments:
      sub_parser: The parser to add arguments to.
    """
    sub_parser.add_argument('--algorithm',
                            help='Algorithm to use (default: NONE)',
                            metavar='ALGORITHM',
                            default='NONE')
    sub_parser.add_argument('--key',
                            help='Path to RSA private key file',
                            metavar='KEY',
                            required=False)
    sub_parser.add_argument('--rollback_index',
                            help='Rollback Index',
                            type=parse_number,
                            default=0)
    sub_parser.add_argument('--prop',
                            help='Add property',
                            metavar='KEY:VALUE',
                            action='append')
    sub_parser.add_argument('--prop_from_file',
                            help='Add property from file',
                            metavar='KEY:PATH',
                            action='append')
    sub_parser.add_argument('--kernel_cmdline',
                            help='Add kernel cmdline',
                            metavar='CMDLINE',
                            action='append')
    sub_parser.add_argument('--generate_dm_verity_cmdline_from_hashtree',
                            metavar='IMAGE',
                            help='Generate kernel cmdline for dm-verity',
                            type=argparse.FileType('rb'))
    sub_parser.add_argument('--include_descriptors_from_image',
                            help='Include descriptors from image',
                            metavar='IMAGE',
                            action='append',
                            type=argparse.FileType('rb'))

  def run(self, argv):
    """Command-line processor.

    Arguments:
      argv: Pass sys.argv from main.
    """
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(title='subcommands')

    sub_parser = subparsers.add_parser('version',
                                       help='Prints version of avbtool.')
    sub_parser.set_defaults(func=self.version)

    sub_parser = subparsers.add_parser('extract_public_key',
                                       help='Extract public key.')
    sub_parser.add_argument('--key',
                            help='Path to RSA private key file',
                            required=True)
    sub_parser.add_argument('--output',
                            help='Output file name',
                            type=argparse.FileType('wb'),
                            required=True)
    sub_parser.set_defaults(func=self.extract_public_key)

    sub_parser = subparsers.add_parser('make_vbmeta_image',
                                       help='Makes a vbmeta image.')
    sub_parser.add_argument('--output',
                            help='Output file name',
                            type=argparse.FileType('wb'),
                            required=True)
    self._add_common_args(sub_parser)
    sub_parser.add_argument('--chain_partition',
                            help='Allow signed integrity-data for partition',
                            metavar='PART_NAME:ROLLBACK_SLOT:KEY_PATH',
                            action='append')
    sub_parser.set_defaults(func=self.make_vbmeta_image)

    sub_parser = subparsers.add_parser('add_hash_footer',
                                       help='Add hashes and footer to image.')
    sub_parser.add_argument('--image',
                            help='Image to add hashes to',
                            type=argparse.FileType('rab+'))
    sub_parser.add_argument('--partition_size',
                            help='Partition size',
                            type=parse_number,
                            required=True)
    sub_parser.add_argument('--partition_name',
                            help='Partition name',
                            required=True)
    sub_parser.add_argument('--hash_algorithm',
                            help='Hash algorithm to use (default: sha256)',
                            default='sha256')
    sub_parser.add_argument('--salt',
                            help='Salt in hex (default: /dev/urandom)')
    self._add_common_args(sub_parser)
    sub_parser.set_defaults(func=self.add_hash_footer)

    sub_parser = subparsers.add_parser('add_hashtree_footer',
                                       help='Add hashtree and footer to image.')
    sub_parser.add_argument('--image',
                            help='Image to add hashtree to',
                            type=argparse.FileType('rab+'))
    sub_parser.add_argument('--partition_size',
                            help='Partition size',
                            type=parse_number,
                            required=True)
    sub_parser.add_argument('--partition_name',
                            help='Partition name',
                            required=True)
    sub_parser.add_argument('--hash_algorithm',
                            help='Hash algorithm to use (default: sha1)',
                            default='sha1')
    sub_parser.add_argument('--salt',
                            help='Salt in hex (default: /dev/urandom)')
    sub_parser.add_argument('--block_size',
                            help='Block size (default: 4096)',
                            type=parse_number,
                            default=4096)
    self._add_common_args(sub_parser)
    sub_parser.set_defaults(func=self.add_hashtree_footer)

    sub_parser = subparsers.add_parser('erase_footer',
                                       help='Erase footer from an image.')
    sub_parser.add_argument('--image',
                            help='Image with a footer',
                            type=argparse.FileType('rwb+'),
                            required=True)
    sub_parser.add_argument('--keep_hashtree',
                            help='Keep the hashtree in the image',
                            action='store_true')
    sub_parser.set_defaults(func=self.erase_footer)

    sub_parser = subparsers.add_parser(
        'info_image',
        help='Show information about vbmeta or footer.')
    sub_parser.add_argument('--image',
                            help='Image to show information about',
                            type=argparse.FileType('rb'),
                            required=True)
    sub_parser.add_argument('--output',
                            help='Write info to file',
                            type=argparse.FileType('wt'),
                            default=sys.stdout)
    sub_parser.set_defaults(func=self.info_image)

    sub_parser = subparsers.add_parser('set_ab_metadata',
                                       help='Set A/B metadata.')
    sub_parser.add_argument('--misc_image',
                            help=('The misc image to modify. If the image does '
                                  'not exist, it will be created.'),
                            type=argparse.FileType('r+b'),
                            required=True)
    sub_parser.add_argument('--slot_data',
                            help=('Slot data of the form "priority", '
                                  '"tries_remaining", "sucessful_boot" for '
                                  'slot A followed by the same for slot B, '
                                  'separated by colons. The default value '
                                  'is 15:7:0:14:7:0.'),
                            default='15:7:0:14:7:0')
    sub_parser.set_defaults(func=self.set_ab_metadata)

    args = parser.parse_args(argv[1:])
    try:
      args.func(args)
    except AvbError as e:
      sys.stderr.write('{}: {}\n'.format(argv[0], e.message))
      sys.exit(1)

  def version(self, _):
    """Implements the 'version' sub-command."""
    print '{}.{}'.format(AVB_VERSION_MAJOR, AVB_VERSION_MINOR)

  def extract_public_key(self, args):
    """Implements the 'extract_public_key' sub-command."""
    self.avb.extract_public_key(args.key, args.output)

  def make_vbmeta_image(self, args):
    """Implements the 'make_vbmeta_image' sub-command."""
    self.avb.make_vbmeta_image(args.output, args.chain_partition,
                               args.algorithm, args.key, args.rollback_index,
                               args.prop, args.prop_from_file,
                               args.kernel_cmdline,
                               args.generate_dm_verity_cmdline_from_hashtree,
                               args.include_descriptors_from_image)

  def add_hash_footer(self, args):
    """Implements the 'add_hash_footer' sub-command."""
    self.avb.add_hash_footer(args.image.name, args.partition_size,
                             args.partition_name, args.hash_algorithm,
                             args.salt, args.algorithm, args.key,
                             args.rollback_index, args.prop,
                             args.prop_from_file, args.kernel_cmdline,
                             args.generate_dm_verity_cmdline_from_hashtree,
                             args.include_descriptors_from_image)

  def add_hashtree_footer(self, args):
    """Implements the 'add_hashtree_footer' sub-command."""
    self.avb.add_hashtree_footer(args.image.name, args.partition_size,
                                 args.partition_name, args.hash_algorithm,
                                 args.block_size, args.salt, args.algorithm,
                                 args.key, args.rollback_index, args.prop,
                                 args.prop_from_file, args.kernel_cmdline,
                                 args.generate_dm_verity_cmdline_from_hashtree,
                                 args.include_descriptors_from_image)

  def erase_footer(self, args):
    """Implements the 'erase_footer' sub-command."""
    self.avb.erase_footer(args.image.name, args.keep_hashtree)

  def set_ab_metadata(self, args):
    """Implements the 'set_ab_metadata' sub-command."""
    self.avb.set_ab_metadata(args.misc_image, args.slot_data)

  def info_image(self, args):
    """Implements the 'info_image' sub-command."""
    self.avb.info_image(args.image.name, args.output)


if __name__ == '__main__':
  tool = AvbTool()
  tool.run(sys.argv)