at_auth_unlock.py: Choose matching unlock credentials if given multiple

This adds support for passing in multiple unlock credential archives
and/or directories containing such credential archives. The tool will
automatically search all provided archives and automatically use any
the first match found. Matches are found based on the SHA256 hash of the
product ID, which is included in the unlock certificate (in the archive)
and in the challenge from the device.

Fixes: 118340336
Test: Verified using unit tests (`atest at_auth_unlock_unittest`) and
against real device
Change-Id: If4980d535039b3da0b0f08482cf2a22ccf8bbfb6
diff --git a/test/at_auth_unlock_unittest.py b/test/at_auth_unlock_unittest.py
index c966000..7683565 100644
--- a/test/at_auth_unlock_unittest.py
+++ b/test/at_auth_unlock_unittest.py
@@ -59,14 +59,14 @@
 
   def testFromValidZipArchive(self):
     with validUnlockCredsZip() as zip:
-      with UnlockCredentials.from_credential_archive(zip) as creds:
-        self.assertIsNotNone(creds.intermediate_cert)
-        self.assertIsNotNone(creds.unlock_cert)
-        self.assertIsNotNone(creds.unlock_key)
+      creds = UnlockCredentials.from_credential_archive(zip)
+      self.assertIsNotNone(creds.intermediate_cert)
+      self.assertIsNotNone(creds.unlock_cert)
+      self.assertIsNotNone(creds.unlock_key)
 
   def testFromInvalidZipArchive(self):
     with self.assertRaises(zipfile.BadZipfile):
-      UnlockCredentials.from_credential_archive(DATA_FILE_PUK_KEY).__enter__()
+      UnlockCredentials.from_credential_archive(DATA_FILE_PUK_KEY)
 
   def testFromArchiveMissingPikCertificate(self):
     with createTempZip({
@@ -74,7 +74,7 @@
         'puk_v1.pem': DATA_FILE_PUK_KEY
     }) as zip:
       with self.assertRaises(ValueError):
-        UnlockCredentials.from_credential_archive(zip).__enter__()
+        UnlockCredentials.from_credential_archive(zip)
 
   def testFromArchiveMissingPukCertificate(self):
     with createTempZip({
@@ -82,7 +82,7 @@
         'puk_v1.pem': DATA_FILE_PUK_KEY
     }) as zip:
       with self.assertRaises(ValueError):
-        UnlockCredentials.from_credential_archive(zip).__enter__()
+        UnlockCredentials.from_credential_archive(zip)
 
   def testFromArchiveMissingPuk(self):
     with createTempZip({
@@ -90,7 +90,7 @@
         'puk_certificate_v1.bin': DATA_FILE_PUK_CERTIFICATE,
     }) as zip:
       with self.assertRaises(ValueError):
-        UnlockCredentials.from_credential_archive(zip).__enter__()
+        UnlockCredentials.from_credential_archive(zip)
 
   def testFromArchiveMultiplePikCertificates(self):
     with createTempZip({
@@ -100,7 +100,7 @@
         'puk_v1.pem': DATA_FILE_PUK_KEY
     }) as zip:
       with self.assertRaises(ValueError):
-        UnlockCredentials.from_credential_archive(zip).__enter__()
+        UnlockCredentials.from_credential_archive(zip)
 
   def testFromArchiveMultiplePukCertificates(self):
     with createTempZip({
@@ -110,7 +110,7 @@
         'puk_v1.pem': DATA_FILE_PUK_KEY
     }) as zip:
       with self.assertRaises(ValueError):
-        UnlockCredentials.from_credential_archive(zip).__enter__()
+        UnlockCredentials.from_credential_archive(zip)
 
   def testFromArchiveMultiplePuks(self):
     with createTempZip({
@@ -120,7 +120,7 @@
         'puk_v2.pem': DATA_FILE_PUK_KEY
     }) as zip:
       with self.assertRaises(ValueError):
-        UnlockCredentials.from_credential_archive(zip).__enter__()
+        UnlockCredentials.from_credential_archive(zip)
 
   def testFromFiles(self):
     creds = UnlockCredentials(
@@ -196,13 +196,17 @@
       shutil.rmtree(tempdir)
 
 
-def writeFullUnlockChallenge(out_file, product_id_hash=b'\x00' * 32):
+def writeFullUnlockChallenge(out_file, product_id_hash=None):
   """Helper function to create a file with a full AvbAtxUnlockChallenge struct.
 
   Arguments:
     product_id_hash: [optional] 32 byte value to include in the challenge as the
-      SHA256 hash of the product ID
+      SHA256 hash of the product ID. If not provided, will default to the
+      product ID hash from the subject of DATA_FILE_PUK_CERTIFICATE.
   """
+  if product_id_hash is None:
+    with open(DATA_FILE_PUK_CERTIFICATE, 'rb') as f:
+      product_id_hash = GetAtxCertificateSubject(f.read())
   assert len(product_id_hash) == 32
 
   with open(out_file, 'wb') as out:
@@ -216,38 +220,43 @@
 
   def testCredentialIsCorrect(self):
     with validUnlockCredsZip() as zip:
-      with UnlockCredentials.from_credential_archive(zip) as creds:
-        tempdir = tempfile.mkdtemp()
-        try:
-          challenge_file = os.path.join(tempdir, 'challenge')
-          writeFullUnlockChallenge(challenge_file)
-          out_cred = os.path.join(tempdir, 'credential')
+      creds = UnlockCredentials.from_credential_archive(zip)
 
-          # Compare unlock credential generated by function with one generated
-          # using 'avbtool make_atx_unlock_credential', to check correctness.
-          MakeAtxUnlockCredential(creds, challenge_file, out_cred)
-          self.assertTrue(filecmp.cmp(out_cred, DATA_FILE_UNLOCK_CREDENTIAL))
-        finally:
-          shutil.rmtree(tempdir)
+      tempdir = tempfile.mkdtemp()
+      try:
+        challenge_file = os.path.join(tempdir, 'challenge')
+        writeFullUnlockChallenge(challenge_file)
+        challenge = UnlockChallenge(challenge_file)
+        out_cred = os.path.join(tempdir, 'credential')
+
+        # Compare unlock credential generated by function with one generated
+        # using 'avbtool make_atx_unlock_credential', to check correctness.
+        MakeAtxUnlockCredential(creds, challenge, out_cred)
+        self.assertTrue(filecmp.cmp(out_cred, DATA_FILE_UNLOCK_CREDENTIAL))
+      finally:
+        shutil.rmtree(tempdir)
 
   def testWrongChallengeSize(self):
     with validUnlockCredsZip() as zip:
-      with UnlockCredentials.from_credential_archive(zip) as creds:
-        tempdir = tempfile.mkdtemp()
-        try:
-          out_cred = os.path.join(tempdir, 'credential')
+      creds = UnlockCredentials.from_credential_archive(zip)
 
-          # The bundled unlock challenge is just the 16 byte challenge, not the
-          # full AvbAtxUnlockChallenge like this expects.
-          with self.assertRaises(ValueError):
-            MakeAtxUnlockCredential(creds, DATA_FILE_UNLOCK_CHALLENGE, out_cred)
-        finally:
-          shutil.rmtree(tempdir)
+      tempdir = tempfile.mkdtemp()
+      try:
+        out_cred = os.path.join(tempdir, 'credential')
+
+        # The bundled unlock challenge is just the 16 byte challenge, not the
+        # full AvbAtxUnlockChallenge like this expects.
+        with self.assertRaises(ValueError):
+          challenge = UnlockChallenge(DATA_FILE_UNLOCK_CHALLENGE)
+          MakeAtxUnlockCredential(creds, challenge, out_cred)
+      finally:
+        shutil.rmtree(tempdir)
 
 
 def makeFastbootCommandFake(testcase,
                             expect_serial=None,
                             error_on_command_number=None,
+                            product_id_hash=None,
                             stay_locked=False):
   """Construct a fake fastboot command handler, to be used with unitttest.mock.Mock.side_effect.
 
@@ -288,7 +297,7 @@
         raise subprocess.CalledProcessError(
             returncode=1, cmd=args, output=b'Fake: No data staged')
 
-      writeFullUnlockChallenge(args[1])
+      writeFullUnlockChallenge(args[1], product_id_hash=product_id_hash)
       handler.challenge_staged = False
     elif args[0] == 'stage':
       handler.staged_file = args[1]
@@ -320,38 +329,11 @@
 class AuthenticatedUnlockTest(unittest.TestCase):
 
   @patch('subprocess.check_output')
-  def testSuccessfulUnlock(self, mock_subp_check_output):
-    with validUnlockCredsZip() as zip:
-      with UnlockCredentials.from_credential_archive(zip) as creds:
-        SERIAL = 'abcde12345'
-        mock_subp_check_output.side_effect = makeFastbootCommandFake(
-            self, expect_serial=SERIAL)
-        self.assertTrue(AuthenticatedUnlock(creds, serial=SERIAL))
-
-  @patch('subprocess.check_output')
-  def testFastbootError(self, mock_subp_check_output):
-    with validUnlockCredsZip() as zip:
-      with UnlockCredentials.from_credential_archive(zip) as creds:
-        mock_subp_check_output.side_effect = makeFastbootCommandFake(
-            self, error_on_command_number=0)
-        self.assertFalse(AuthenticatedUnlock(creds))
-
-  @patch('subprocess.check_output')
-  def testDoesntActuallyUnlock(self, mock_subp_check_output):
-    with validUnlockCredsZip() as zip:
-      with UnlockCredentials.from_credential_archive(zip) as creds:
-        mock_subp_check_output.side_effect = makeFastbootCommandFake(
-            self, stay_locked=True)
-        self.assertFalse(AuthenticatedUnlock(creds))
-
-
-class AuthenticatedUnlockArgParserTest(unittest.TestCase):
-
-  @patch('subprocess.check_output')
   def testUnlockWithZipArchive(self, mock_subp_check_output):
     with validUnlockCredsZip() as zip:
       mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
       self.assertEqual(main([zip.name]), 0)
+      self.assertNotEqual(mock_subp_check_output.call_count, 0)
 
   @patch('subprocess.check_output')
   def testUnlockDeviceBySerial(self, mock_subp_check_output):
@@ -360,6 +342,7 @@
       mock_subp_check_output.side_effect = makeFastbootCommandFake(
           self, expect_serial=SERIAL)
       self.assertEqual(main([zip.name, '-s', SERIAL]), 0)
+      self.assertNotEqual(mock_subp_check_output.call_count, 0)
 
   @patch('subprocess.check_output')
   def testUnlockWithIndividualFiles(self, mock_subp_check_output):
@@ -369,16 +352,146 @@
             '--pik_cert', DATA_FILE_PIK_CERTIFICATE, '--puk_cert',
             DATA_FILE_PUK_CERTIFICATE, '--puk', DATA_FILE_PUK_KEY
         ]), 0)
+    self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+  @patch('subprocess.check_output')
+  def testFastbootError(self, mock_subp_check_output):
+    """Verify that errors are handled properly if fastboot commands error out."""
+    with validUnlockCredsZip() as zip:
+      for n in range(5):
+        mock_subp_check_output.reset_mock()
+        mock_subp_check_output.side_effect = makeFastbootCommandFake(
+            self, error_on_command_number=n)
+        self.assertNotEqual(main([zip.name]), 0)
+        self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+  @patch('subprocess.check_output')
+  def testDoesntActuallyUnlock(self, mock_subp_check_output):
+    """Verify fails if fake set to not actually unlock."""
+    with validUnlockCredsZip() as zip:
+      mock_subp_check_output.side_effect = makeFastbootCommandFake(
+          self, stay_locked=True)
+      self.assertNotEqual(main([zip.name]), 0)
+      self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+  @patch('subprocess.check_output')
+  def testNoCredentialsMatchDeviceProductID(self, mock_subp_check_output):
+    """Test two cases where fake responds with a challenge that has a product ID hash which doesn't match the credentials used."""
+    # Case 1: Change the product ID hash that the fake responds with.
+    with validUnlockCredsZip() as zip:
+      mock_subp_check_output.side_effect = makeFastbootCommandFake(
+          self, product_id_hash=b'\x00' * 32)
+      self.assertNotEqual(main([zip.name]), 0)
+      self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+    # Case 2: Use credentials with a different product ID.
+    with createTempZip({
+        'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        # Note: PIK cert used as PUK cert so subject (i.e. product ID hash) is
+        # different
+        'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        'puk_v1.pem': DATA_FILE_PUK_KEY
+    }) as zip:
+      mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
+      self.assertNotEqual(main([zip.name]), 0)
+      self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+  @patch('subprocess.check_output')
+  def testMatchingCredentialSelectedFromZipArchives(self,
+                                                    mock_subp_check_output):
+    """Test correct credential based on product ID hash used if multiple provided directly through arguments."""
+    with validUnlockCredsZip() as correctCreds, createTempZip({
+        'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        # Note: PIK cert used as PUK cert so subject (i.e. product ID hash)
+        # doesn't match
+        'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        'puk_v1.pem': DATA_FILE_PUK_KEY
+    }) as wrongCreds:
+      mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
+      self.assertEqual(main([wrongCreds.name, correctCreds.name]), 0)
+      self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+  @patch('subprocess.check_output')
+  def testMatchingCredentialSelectedFromDirectory(self, mock_subp_check_output):
+    """Test correct credential based on product ID hash used if multiple provided indirectly through a directory argument."""
+    with validUnlockCredsZip() as correctCreds, createTempZip({
+        'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        # Note: PIK cert used as PUK cert so subject (i.e. product ID hash)
+        # doesn't match
+        'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        'puk_v1.pem': DATA_FILE_PUK_KEY
+    }) as wrongCreds:
+      tempdir = tempfile.mkdtemp()
+      try:
+        shutil.copy2(correctCreds.name, tempdir)
+        shutil.copy2(wrongCreds.name, tempdir)
+
+        mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
+        self.assertEqual(main([tempdir]), 0)
+        self.assertNotEqual(mock_subp_check_output.call_count, 0)
+      finally:
+        shutil.rmtree(tempdir)
+
+  @patch('subprocess.check_output')
+  def testMatchingCredentialSelectedFromEither(self, mock_subp_check_output):
+    """Test correct credential based on product ID hash used if arguments give some combination of file and directory arguments."""
+    with validUnlockCredsZip() as correctCreds, createTempZip({
+        'pik_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        # Note: PIK cert used as PUK cert so subject (i.e. product ID hash)
+        # doesn't match
+        'puk_certificate_v1.bin': DATA_FILE_PIK_CERTIFICATE,
+        'puk_v1.pem': DATA_FILE_PUK_KEY
+    }) as wrongCreds:
+      # Case 1: Correct creds in directory, wrong in file arg
+      tempdir = tempfile.mkdtemp()
+      try:
+        shutil.copy2(correctCreds.name, tempdir)
+
+        mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
+        self.assertEqual(main([wrongCreds.name, tempdir]), 0)
+        self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+      finally:
+        shutil.rmtree(tempdir)
+
+      # Case 2: Correct creds in file arg, wrong in directory
+      tempdir = tempfile.mkdtemp()
+      try:
+        shutil.copy2(wrongCreds.name, tempdir)
+
+        mock_subp_check_output.side_effect = makeFastbootCommandFake(self)
+        self.assertEqual(main([tempdir, correctCreds.name]), 0)
+        self.assertNotEqual(mock_subp_check_output.call_count, 0)
+
+        # Case 2: Correct creds in file arg, wrong in directory
+      finally:
+        shutil.rmtree(tempdir)
 
   @patch('argparse.ArgumentParser.error')
-  def testMutualExclusionArchiveAndFiles(self, mock_parser_error):
+  def testArgparseDirectoryWithNoCredentials(self, mock_parser_error):
+    """Test """
+    tempdir = tempfile.mkdtemp()
+    try:
+      # Make sure random files are ignored.
+      with open(os.path.join(tempdir, 'so_random'), 'w') as f:
+        f.write("I'm a random file")
+
+      mock_parser_error.side_effect = ValueError('ArgumentParser.error')
+      with self.assertRaises(ValueError):
+        main([tempdir])
+      self.assertEqual(mock_parser_error.call_count, 1)
+    finally:
+      shutil.rmtree(tempdir)
+
+  @patch('argparse.ArgumentParser.error')
+  def testArgparseMutualExclusionArchiveAndFiles(self, mock_parser_error):
     mock_parser_error.side_effect = ValueError('ArgumentParser.error')
     with self.assertRaises(ValueError):
       main(['dummy.zip', '--pik_cert', DATA_FILE_PIK_CERTIFICATE])
     self.assertEqual(mock_parser_error.call_count, 1)
 
   @patch('argparse.ArgumentParser.error')
-  def testMutualInclusionOfFileArgs(self, mock_parser_error):
+  def testArgparseMutualInclusionOfFileArgs(self, mock_parser_error):
     mock_parser_error.side_effect = ValueError('ArgumentParser.error')
     with self.assertRaises(ValueError):
       main(['--pik_cert', 'pik_cert.bin', '--puk_cert', 'puk_cert.bin'])
@@ -395,7 +508,7 @@
     self.assertEqual(mock_parser_error.call_count, 1)
 
   @patch('argparse.ArgumentParser.error')
-  def testMissingBundleAndFiles(self, mock_parser_error):
+  def testArgparseMissingBundleAndFiles(self, mock_parser_error):
     mock_parser_error.side_effect = ValueError('ArgumentParser.error')
     with self.assertRaises(ValueError):
       main(['-s', '1234abcd'])
diff --git a/tools/at_auth_unlock.py b/tools/at_auth_unlock.py
index 6c949d4..f49b4a6 100755
--- a/tools/at_auth_unlock.py
+++ b/tools/at_auth_unlock.py
@@ -26,12 +26,18 @@
 
   1) by providing paths to the individual credential files using the
      '--pik_cert', '--puk_cert', and '--puk' command line swtiches, or
+
   2) by providing a path to a zip archive containing the three credential files,
      named as follows:
        - Product Intermediate Key (PIK) certificate: 'pik_certificate.*\.bin'
        - Product Unlock Key (PUK) certificate: 'puk_certificate.*\.bin'
        - PUK private key: 'puk.*\.pem'
 
+     You can also provide one or more archives and/or one or more directories
+     containing such zip archives. In either scenario, the tool will search all
+     of the provided credential archives for a match against the product ID of
+     the device being unlocked and automatically use the first match.
+
 Dependencies:
   - Python 2.7.x, 3.2.x, or newer (for argparse)
   - PyCrypto 2.5 or newer (for PKCS1_v1_5 and RSA PKCS#8 PEM key import)
@@ -43,12 +49,13 @@
 fastboot, given valid unlock credentials for the device."""
 
 HELP_USAGE = """
-  %(prog)s [-h] [-v] [-s SERIAL] unlock_creds.zip
+  %(prog)s [-h] [-v] [-s SERIAL] unlock_creds.zip [unlock_creds_2.zip ...]
   %(prog)s --pik_cert pik_cert.bin --puk_cert puk_cert.bin --puk puk.pem"""
 
 HELP_EPILOG = """examples:
   %(prog)s unlock_creds.zip
-  %(prog)s unlock_creds.zip -s SERIAL
+  %(prog)s unlock_creds.zip unlock_creds_2.zip -s SERIAL
+  %(prog)s path_to_dir_with_multiple_unlock_creds/
   %(prog)s --pik_cert pik_cert.bin --puk_cert puk_cert.bin --puk puk.pem"""
 
 import sys
@@ -59,7 +66,7 @@
   sys.exit(1)
 
 import argparse
-import contextlib
+import binascii
 import os
 import re
 import shutil
@@ -78,25 +85,16 @@
   print('PyCrypto 2.5 or newer required, missing or too old: ' + str(e))
 
 
-class NullContextManager(object):
-  """Local implementation of contextlib.nullcontext, which is Python 3-only."""
-
-  def __init__(self, enter_result=None):
-    self.enter_result = enter_result
-
-  def __enter__(self):
-    return self.enter_result
-
-  def __exit__(self, *args):
-    pass
-
-
 class UnlockCredentials(object):
   """Helper data container class for the 3 unlock credentials involved in an AVB authenticated unlock operation.
 
   """
 
-  def __init__(self, intermediate_cert_file, unlock_cert_file, unlock_key_file):
+  def __init__(self,
+               intermediate_cert_file,
+               unlock_cert_file,
+               unlock_key_file,
+               source_file=None):
     # The certificates are AvbAtxCertificate structs as defined in libavb_atx,
     # not an X.509 certificate. Do a basic length sanity check when reading
     # them.
@@ -117,6 +115,8 @@
       if not self._unlock_key.has_private():
         raise ValueError('Unlock key was not an RSA private key.')
 
+    self._source_file = source_file
+
   @property
   def intermediate_cert(self):
     return self._intermediate_cert
@@ -129,12 +129,11 @@
   def unlock_key(self):
     return self._unlock_key
 
-  @classmethod
-  def from_files(cls, pik_cert, puk_cert, puk):
-    return NullContextManager(cls(pik_cert, puk_cert, puk))
+  @property
+  def source_file(self):
+    return self._source_file
 
   @classmethod
-  @contextlib.contextmanager
   def from_credential_archive(cls, archive):
     """Create UnlockCredentials from an unlock credential zip archive.
 
@@ -186,15 +185,73 @@
 
         zip.extractall(path=tempdir, members=[pik_cert, puk_cert, puk])
 
-        yield cls(
+        return cls(
             intermediate_cert_file=os.path.join(tempdir, pik_cert),
             unlock_cert_file=os.path.join(tempdir, puk_cert),
-            unlock_key_file=os.path.join(tempdir, puk))
+            unlock_key_file=os.path.join(tempdir, puk),
+            source_file=archive)
     finally:
       shutil.rmtree(tempdir)
 
 
-def MakeAtxUnlockCredential(creds, challenge_file, out_file):
+class UnlockChallenge(object):
+  """Helper class for parsing the AvbAtxUnlockChallenge struct returned from 'fastboot oem at-get-vboot-unlock-challenge'.
+
+     The file provided to the constructor should be the full 52-byte
+     AvbAtxUnlockChallenge struct, not just the challenge itself.
+  """
+
+  def __init__(self, challenge_file):
+    CHALLENGE_STRUCT_SIZE = 52
+    PRODUCT_ID_HASH_SIZE = 32
+    CHALLENGE_DATA_SIZE = 16
+    with open(challenge_file, 'rb') as f:
+      data = f.read()
+      if len(data) != CHALLENGE_STRUCT_SIZE:
+        raise ValueError('Invalid unlock challenge length.')
+
+      self._version, self._product_id_hash, self._challenge_data = struct.unpack(
+          '<I{}s{}s'.format(PRODUCT_ID_HASH_SIZE, CHALLENGE_DATA_SIZE), data)
+
+  @property
+  def version(self):
+    return self._version
+
+  @property
+  def product_id_hash(self):
+    return self._product_id_hash
+
+  @property
+  def challenge_data(self):
+    return self._challenge_data
+
+
+def GetAtxCertificateSubject(cert):
+  """Parses and returns the subject field from the given AvbAtxCertificate struct."""
+  CERT_SUBJECT_OFFSET = 4 + 1032  # Format version and public key come before subject
+  CERT_SUBJECT_LENGTH = 32
+  return cert[CERT_SUBJECT_OFFSET:CERT_SUBJECT_OFFSET + CERT_SUBJECT_LENGTH]
+
+
+def SelectMatchingUnlockCredential(all_creds, challenge):
+  """Find and return the first UnlockCredentials object whose product ID matches that of the unlock challenge.
+
+  The Product Unlock Key (PUK) certificate's subject field contains the
+  SHA256 hash of the product ID that it can be used to unlock. This same
+  value (SHA256 hash of the product ID) is contained in the unlock challenge.
+
+  Arguments:
+    all_creds: List of UnlockCredentials objects to be searched for a match
+      against the given challenge.
+    challenge: UnlockChallenge object created from challenge obtained via
+      'fastboot oem at-get-vboot-unlock-challenge'.
+  """
+  for creds in all_creds:
+    if GetAtxCertificateSubject(creds.unlock_cert) == challenge.product_id_hash:
+      return creds
+
+
+def MakeAtxUnlockCredential(creds, challenge, out_file):
   """Simple reimplementation of 'avbtool make_atx_unlock_credential'.
 
   Generates an Android Things authenticated unlock credential to authorize
@@ -208,28 +265,14 @@
   Arguments:
     creds: UnlockCredentials object wrapping the PIK certificate, PUK
       certificate, and PUK private key.
-    challenge_file: Challenge obtained via 'oem at-get-vboot-unlock-challenge'.
-      This should be the full 52-byte AvbAtxUnlockChallenge struct, not just the
-      challenge itself.
+    challenge: UnlockChallenge object created from challenge obtained via
+      'fastboot oem at-get-vboot-unlock-challenge'.
     out_file: Output filename to write the AvbAtxUnlockCredential struct to.
 
   Raises:
     ValueError: If challenge has wrong length.
   """
-  # The 16-byte challenge from the bootloader, which needs to be signed with the
-  # PUK and included in the AvbAtxUnlockCredential response, is located at the
-  # end of the 52-byte AvbAtxUnlockChallenge struct
-  CHALLENGE_STRUCT_SIZE = 52
-  CHALLENGE_FIELD_SIZE = 16
-  with open(challenge_file, 'rb') as f:
-    f.seek(0, os.SEEK_END)
-    if f.tell() != CHALLENGE_STRUCT_SIZE:
-      raise ValueError('Invalid unlock challenge length.')
-
-    f.seek(-CHALLENGE_FIELD_SIZE, os.SEEK_END)
-    challenge = f.read()
-
-  hash = SHA512.new(challenge)
+  hash = SHA512.new(challenge.challenge_data)
   signer = PKCS1_v1_5.new(creds.unlock_key)
   signature = signer.sign(hash)
 
@@ -240,12 +283,13 @@
     out.write(signature)
 
 
-def AuthenticatedUnlock(creds, serial=None, verbose=False):
+def AuthenticatedUnlock(all_creds, serial=None, verbose=False):
   """Performs an authenticated AVB unlock of a device over fastboot.
 
   Arguments:
-    creds: UnlockCredentials object wrapping the PIK certificate, PUK
-      certificate, and PUK private key.
+    all_creds: List of UnlockCredentials objects wrapping the PIK certificate,
+      PUK certificate, and PUK private key. The list will be searched to find
+      matching credentials for the device being unlocked.
     serial: [optional] A device serial number or other valid value to be passed
       to fastboot's '-s' switch to select the device to unlock.
     verbose: [optional] Enable verbose output, which prints the fastboot
@@ -260,7 +304,7 @@
     def fastboot_cmd(args):
       args = ['fastboot'] + (['-s', serial] if serial else []) + args
       if verbose:
-        print('$ ' + ' '.join(args))
+        print('\n$ ' + ' '.join(args))
 
       out = subprocess.check_output(
           args, stderr=subprocess.STDOUT).decode('utf-8')
@@ -272,7 +316,21 @@
     try:
       fastboot_cmd(['oem', 'at-get-vboot-unlock-challenge'])
       fastboot_cmd(['get_staged', challenge_file])
-      MakeAtxUnlockCredential(creds, challenge_file, credential_file)
+
+      challenge = UnlockChallenge(challenge_file)
+      print('Product ID SHA256 hash = {}'.format(
+          binascii.hexlify(challenge.product_id_hash)))
+
+      selected_cred = SelectMatchingUnlockCredential(all_creds, challenge)
+      if not selected_cred:
+        print(
+            'ERROR: None of the provided unlock credentials match this device.')
+        return False
+      if selected_cred.source_file:
+        print('Found matching unlock credentials: {}'.format(
+            selected_cred.source_file))
+      MakeAtxUnlockCredential(selected_cred, challenge, credential_file)
+
       fastboot_cmd(['stage', credential_file])
       fastboot_cmd(['oem', 'at-unlock-vboot'])
 
@@ -288,11 +346,30 @@
       print("Command '{}' returned non-zero exit status {}".format(
           ' '.join(e.cmd), e.returncode))
       return False
-
   finally:
     shutil.rmtree(tempdir)
 
 
+def FindUnlockCredentialsInDirectory(dir, verbose=False):
+  if not os.path.isdir(dir):
+    raise ValueError('Not a directory: ' + dir)
+
+  creds = []
+  for file in os.listdir(dir):
+    path = os.path.join(dir, file)
+    if os.path.isfile(path):
+      try:
+        creds.append(UnlockCredentials.from_credential_archive(path))
+        if verbose:
+          print('Found valid unlock credential bundle: ' + path)
+      except (IOError, ValueError, zipfile.BadZipFile) as e:
+        if verbose:
+          print(
+              "Ignoring file which isn't a valid unlock credential zip bundle: "
+              + path)
+  return creds
+
+
 def main(in_args):
   parser = argparse.ArgumentParser(
       description=HELP_DESCRIPTION,
@@ -305,7 +382,8 @@
       '-v',
       '--verbose',
       action='store_true',
-      help='verbose; prints fastboot commands and their output')
+      help=
+      'enable verbose output, e.g. prints fastboot commands and their output')
   parser.add_argument(
       '-s',
       '--serial',
@@ -320,13 +398,18 @@
   # group to be specified - so we define them as optional arguments and do the
   # validation ourselves below.
 
-  # Argument group #1 - Unlock credential zip bundle/archive
+  # Argument group #1 - Unlock credential zip archive(s) (or directory
+  # containing multiple such archives)
   parser.add_argument(
       'bundle',
       metavar='unlock_creds.zip',
-      nargs='?',
+      nargs='*',
       help=
-      'Unlock using a zip bundle of credentials (e.g. from Developer Console).')
+      'Unlock using a zip bundle/archive of credentials (e.g. from Developer '
+      'Console). You can optionally provide multiple archives and/or a  '
+      'directory of such bundles and the tool will automatically select the '
+      'correct one to use based on matching the product ID against the device '
+      'being unlocked.')
 
   # Argument group #2 - Individual credential files
   parser.add_argument(
@@ -348,7 +431,7 @@
   # Do the custom validation described above.
   if args.pik_cert is not None or args.puk_cert is not None or args.puk is not None:
     # Check mutual exclusion with bundle positional argument
-    if args.bundle is not None:
+    if len(args.bundle):
       parser.error(
           'bundle argument is mutually exclusive with --pik_cert, --puk_cert, and --puk'
       )
@@ -360,17 +443,29 @@
       parser.error("--puk_cert is required if --pik_cert or --puk' is given")
     if args.puk is None:
       parser.error("--puk is required if --pik_cert or --puk_cert' is given")
-  elif args.bundle is None:
+  elif not len(args.bundle):
     parser.error(
         'must provide either credentials bundle or individual credential files')
 
-  if args.bundle is not None:
-    creds = UnlockCredentials.from_credential_archive(args.bundle)
+  # Parse arguments into UnlockCredentials objects
+  if len(args.bundle):
+    creds = []
+    for path in args.bundle:
+      if os.path.isfile(path):
+        creds.append(UnlockCredentials.from_credential_archive(path))
+      elif os.path.isdir(path):
+        creds.extend(
+            FindUnlockCredentialsInDirectory(path, verbose=args.verbose))
+      else:
+        parser.error("path argument '{}' does not exist".format(path))
+
+    if len(creds) == 0:
+      parser.error('No unlock credentials were found in any of the given paths')
   else:
-    creds = UnlockCredentials.from_files(args.pik_cert, args.puk_cert, args.puk)
-  with creds as creds:
-    ret = AuthenticatedUnlock(creds, serial=args.serial, verbose=args.verbose)
-    return 0 if ret else 1
+    creds = [UnlockCredentials(args.pik_cert, args.puk_cert, args.puk)]
+
+  ret = AuthenticatedUnlock(creds, serial=args.serial, verbose=args.verbose)
+  return 0 if ret else 1
 
 
 if __name__ == '__main__':