avbtool: Add new verify_image command.

This can be used to verify that a custom signer works correctly.

Since this is only expected to be used on developer machines (e.g. not
in the build system), just use PyCrypto instead of shelling out to
openssl(1).

Bug: 36590704
Test: New unit tests and unit tests pass
Change-Id: Ifd26c7fe0e06cff102d35a46399f9d251cd779c6
diff --git a/README.md b/README.md
index 8c0d550..56d792a 100644
--- a/README.md
+++ b/README.md
@@ -274,6 +274,11 @@
 
     /path/to/my_signing_program SHA256_RSA2048 /path/to/publickey.pem
 
+The `verify_image` command verifies that the signature on the vbmeta
+struct is valid and that it was made with the embedded public
+key. This can be used to check that a custom signing helper works as
+intended.
+
 The `append_vbmeta_image` command can be used to append an entire
 vbmeta blob to the end of another image. This is useful for cases when
 not using any vbmeta partitions, for example:
diff --git a/avbtool b/avbtool
index 10ff39a..574139c 100755
--- a/avbtool
+++ b/avbtool
@@ -69,20 +69,23 @@
 
   Attributes:
     algorithm_type: Integer code corresponding to |AvbAlgorithmType|.
+    hash_name: Empty or a name from |hashlib.algorithms|.
     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):
+  def __init__(self, algorithm_type, hash_name, hash_num_bytes,
+               signature_num_bytes, public_key_num_bytes, padding):
     self.algorithm_type = algorithm_type
+    self.hash_name = hash_name
     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
@@ -90,12 +93,14 @@
 ALGORITHMS = {
     'NONE': Algorithm(
         algorithm_type=0,        # AVB_ALGORITHM_TYPE_NONE
+        hash_name='',
         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_name='sha256',
         hash_num_bytes=32,
         signature_num_bytes=256,
         public_key_num_bytes=8 + 2*2048/8,
@@ -109,6 +114,7 @@
             ]),
     'SHA256_RSA4096': Algorithm(
         algorithm_type=2,        # AVB_ALGORITHM_TYPE_SHA256_RSA4096
+        hash_name='sha256',
         hash_num_bytes=32,
         signature_num_bytes=512,
         public_key_num_bytes=8 + 2*4096/8,
@@ -122,6 +128,7 @@
             ]),
     'SHA256_RSA8192': Algorithm(
         algorithm_type=3,        # AVB_ALGORITHM_TYPE_SHA256_RSA8192
+        hash_name='sha256',
         hash_num_bytes=32,
         signature_num_bytes=1024,
         public_key_num_bytes=8 + 2*8192/8,
@@ -135,6 +142,7 @@
             ]),
     'SHA512_RSA2048': Algorithm(
         algorithm_type=4,        # AVB_ALGORITHM_TYPE_SHA512_RSA2048
+        hash_name='sha512',
         hash_num_bytes=64,
         signature_num_bytes=256,
         public_key_num_bytes=8 + 2*2048/8,
@@ -148,6 +156,7 @@
             ]),
     'SHA512_RSA4096': Algorithm(
         algorithm_type=5,        # AVB_ALGORITHM_TYPE_SHA512_RSA4096
+        hash_name='sha512',
         hash_num_bytes=64,
         signature_num_bytes=512,
         public_key_num_bytes=8 + 2*4096/8,
@@ -161,6 +170,7 @@
             ]),
     'SHA512_RSA8192': Algorithm(
         algorithm_type=6,        # AVB_ALGORITHM_TYPE_SHA512_RSA8192
+        hash_name='sha512',
         hash_num_bytes=64,
         signature_num_bytes=1024,
         public_key_num_bytes=8 + 2*8192/8,
@@ -220,6 +230,8 @@
   This number is written big-endian, e.g. with the most significant
   bit first.
 
+  This is the reverse of decode_long().
+
   Arguments:
     num_bits: The number of bits to write, e.g. 2048.
     value: The value to write.
@@ -234,6 +246,27 @@
   return ret
 
 
+def decode_long(blob):
+  """Decodes a long from a bytearray() using a given amount of bits.
+
+  This number is expected to be in big-endian, e.g. with the most
+  significant bit first.
+
+  This is the reverse of encode_long().
+
+  Arguments:
+    value: A bytearray() with the encoded long.
+
+  Returns:
+    The decoded value.
+  """
+  ret = 0
+  for b in bytearray(blob):
+    ret *= 256
+    ret += b
+  return ret
+
+
 def egcd(a, b):
   """Calculate greatest common divisor of two numbers.
 
@@ -441,6 +474,70 @@
   return signature
 
 
+def verify_vbmeta_signature(vbmeta_header, vbmeta_blob):
+  """Checks that the signature in a vbmeta blob was made by
+     the embedded public key.
+
+  Arguments:
+    vbmeta_header: A AvbVBMetaHeader.
+    vbmeta_blob: The whole vbmeta blob, including the header.
+
+  Returns:
+    True if the signature is valid and corresponds to the embedded
+    public key. Also returns True if the vbmeta blob is not signed.
+  """
+  (_, alg) = lookup_algorithm_by_type(vbmeta_header.algorithm_type)
+  if alg.hash_name == '':
+    return True
+  header_blob = vbmeta_blob[0:256]
+  auth_offset = 256
+  aux_offset = auth_offset + vbmeta_header.authentication_data_block_size
+  aux_size = vbmeta_header.auxiliary_data_block_size
+  aux_blob = vbmeta_blob[aux_offset:aux_offset + aux_size]
+  pubkey_offset = aux_offset + vbmeta_header.public_key_offset
+  pubkey_size = vbmeta_header.public_key_size
+  pubkey_blob = vbmeta_blob[pubkey_offset:pubkey_offset + pubkey_size]
+
+  digest_offset = auth_offset + vbmeta_header.hash_offset
+  digest_size = vbmeta_header.hash_size
+  digest_blob = vbmeta_blob[digest_offset:digest_offset + digest_size]
+
+  sig_offset = auth_offset + vbmeta_header.signature_offset
+  sig_size = vbmeta_header.signature_size
+  sig_blob = vbmeta_blob[sig_offset:sig_offset + sig_size]
+
+  # Now that we've got the stored digest, public key, and signature
+  # all we need to do is to verify. This is the exactly the same
+  # steps as performed in the avb_vbmeta_image_verify() function in
+  # libavb/avb_vbmeta_image.c.
+
+  ha = hashlib.new(alg.hash_name)
+  ha.update(header_blob)
+  ha.update(aux_blob)
+  computed_digest = ha.digest()
+
+  if computed_digest != digest_blob:
+    return False
+
+  padding_and_digest = bytearray(alg.padding)
+  padding_and_digest.extend(computed_digest)
+
+  (num_bits,) = struct.unpack('!I', pubkey_blob[0:4])
+  modulus_blob = pubkey_blob[8:8 + num_bits/8]
+  modulus = decode_long(modulus_blob)
+  exponent = 65537
+
+  # For now, just use Crypto.PublicKey.RSA to verify the signature. This
+  # is OK since 'avbtool verify_image' is not expected to run on the
+  # Android builders (see bug #36809096).
+  import Crypto.PublicKey.RSA
+  key = Crypto.PublicKey.RSA.construct((modulus, long(exponent)))
+  if not key.verify(decode_long(padding_and_digest),
+                    (decode_long(sig_blob), None)):
+    return False
+  return True
+
+
 class ImageChunk(object):
   """Data structure used for representing chunks in Android sparse files.
 
@@ -1801,6 +1898,25 @@
     if num_printed == 0:
       o.write('    (none)\n')
 
+  def verify_image(self, image_filename):
+    """Implements the 'verify_image' command.
+
+    Arguments:
+      image_filename: Image file to get information from (file object).
+    """
+
+    image = ImageHandler(image_filename)
+    (footer, header, descriptors, image_size) = self._parse_image(image)
+    offset = 0
+    if footer:
+      offset = footer.vbmeta_offset
+    size = (header.SIZE + header.authentication_data_block_size +
+            header.auxiliary_data_block_size)
+    image.seek(offset)
+    vbmeta_blob = image.read(size)
+    if not verify_vbmeta_signature(header, vbmeta_blob):
+      raise AvbError('Signature check failed.')
+
   def _parse_image(self, image):
     """Gets information about an image.
 
@@ -2174,12 +2290,7 @@
     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 = hashlib.new(alg.hash_name)
       ha.update(header_data_blob)
       ha.update(aux_data_blob)
       binary_hash.extend(ha.digest())
@@ -3157,6 +3268,15 @@
                             default=sys.stdout)
     sub_parser.set_defaults(func=self.info_image)
 
+    sub_parser = subparsers.add_parser(
+        'verify_image',
+        help='Verify an image.')
+    sub_parser.add_argument('--image',
+                            help='Image to verify',
+                            type=argparse.FileType('rb'),
+                            required=True)
+    sub_parser.set_defaults(func=self.verify_image)
+
     sub_parser = subparsers.add_parser('set_ab_metadata',
                                        help='Set A/B metadata.')
     sub_parser.add_argument('--misc_image',
@@ -3325,6 +3445,10 @@
     """Implements the 'info_image' sub-command."""
     self.avb.info_image(args.image.name, args.output)
 
+  def verify_image(self, args):
+    """Implements the 'verify_image' sub-command."""
+    self.avb.verify_image(args.image.name)
+
   def make_atx_certificate(self, args):
     """Implements the 'make_atx_certificate' sub-command."""
     self.avb.make_atx_certificate(args.output, args.authority_key,
diff --git a/test/avbtool_signing_helper_test.py b/test/avbtool_signing_helper_test.py
index 7291797..8705403 100755
--- a/test/avbtool_signing_helper_test.py
+++ b/test/avbtool_signing_helper_test.py
@@ -44,6 +44,12 @@
     sys.stderr.write("There is not input data\n")
     return errno.EINVAL
 
+  if os.environ.get('SIGNING_HELPER_GENERATE_WRONG_SIGNATURE'):
+    # We're only called with this algorithm which signature size is 256.
+    assert sys.argv[1] == 'SHA256_RSA2048'
+    sys.stdout.write('X'*256)
+    return 0
+
   if 'SIGNING_HELPER_TEST' not in os.environ or os.environ['SIGNING_HELPER_TEST'] == "":
     sys.stderr.write("env SIGNING_HELPER_TEST is not set or empty\n")
     return errno.EINVAL
diff --git a/test/avbtool_unittest.cc b/test/avbtool_unittest.cc
index 8df5f26..e3af844 100644
--- a/test/avbtool_unittest.cc
+++ b/test/avbtool_unittest.cc
@@ -1486,6 +1486,51 @@
       vbmeta_path.value().c_str());
 }
 
+TEST_F(AvbToolTest, VerifyImageNoSignature) {
+  GenerateVBMetaImage("vbmeta.img",
+                      "",  // NONE
+                      0,
+                      base::FilePath());
+
+  EXPECT_COMMAND(0,
+                 "./avbtool verify_image "
+                 "--image %s ",
+                 vbmeta_image_path_.value().c_str());
+}
+
+TEST_F(AvbToolTest, VerifyImageValidSignature) {
+  GenerateVBMetaImage("vbmeta.img",
+                      "SHA256_RSA2048",
+                      0,
+                      base::FilePath("test/data/testkey_rsa2048.pem"));
+
+  EXPECT_COMMAND(0,
+                 "./avbtool verify_image "
+                 "--image %s ",
+                 vbmeta_image_path_.value().c_str());
+}
+
+TEST_F(AvbToolTest, VerifyImageBrokenSignature) {
+  base::FilePath vbmeta_path = testdir_.Append("vbmeta.bin");
+  base::FilePath signing_helper_test_path =
+      testdir_.Append("signing_helper_test");
+
+  // Intentionally make the signer generate a wrong signature.
+  EXPECT_COMMAND(
+      0,
+      "SIGNING_HELPER_GENERATE_WRONG_SIGNATURE=1 ./avbtool make_vbmeta_image "
+      "--output %s "
+      "--algorithm SHA256_RSA2048 --key test/data/testkey_rsa2048.pem "
+      "--signing_helper test/avbtool_signing_helper_test.py "
+      "--internal_release_string \"\"",
+      vbmeta_path.value().c_str());
+
+  EXPECT_COMMAND(1,
+                 "./avbtool verify_image "
+                 "--image %s ",
+                 vbmeta_path.value().c_str());
+}
+
 TEST_F(AvbToolTest, MakeAtxPikCertificate) {
   base::FilePath subject_path = testdir_.Append("tmp_subject");
   ASSERT_TRUE(base::WriteFile(subject_path, "fake PIK subject", 16));