tpm_test: drbg_test: add support for using lab vectors

This adds support for running drbg_test with the lab vectors.

BUG=b:189376694
TEST=./tpmtest

Change-Id: I44f3671f55e1befcac36006568bf1a4deb2d4685
Signed-off-by: Mary Ruthven <mruthven@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2924406
Reviewed-by: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Reviewed-by: Namyoon Woo <namyoon@chromium.org>

4 files changed, 245 insertions, 33 deletions

diff --git a/test/tpm_test/drbg_test.py b/test/tpm_test/drbg_test.py
index 3a7f086aee..1ada131600 100644
--- a/test/tpm_test/drbg_test.py
+++ b/test/tpm_test/drbg_test.py
@@ -9,10 +9,10 @@ from __future__ import print_function
 
 from binascii import a2b_hex as a2b
 
+import lab_vectors
 import subcmd
 import utils
 
-
 # A standard empty response to DRBG extended commands.
 EMPTY_DRBG_RESPONSE = bytes([0x80, 0x01,
                              0x00, 0x00, 0x00, 0x0c,
@@ -25,34 +25,34 @@ DRBG_GENERATE = 2
 DRBG_GROUP_INIT = 3
 
 TEST_INPUTS = (
-  (DRBG_GROUP_INIT, 32),
-  (DRBG_INIT,
-   ('C40894D0C37712140924115BF8A3110C7258532365BB598F81B127A5E4CB8EB0',
-    'FBB1EDAF92D0C2699F5C0A7418D308B09AC679FFBB0D8918C8E62D35091DD2B9',
-    '2B18535D739F7E75AF4FF0C0C713DD4C9B0A6803D2E0DB2BDE3C4F3650ABF750')),
-  (DRBG_RESEED,
-   ('4D58A621857706450338CCA8A1AF5CD2BD9305F3475CF1A8752518DD8E8267B6',
-    '0153A0A1D7487E2EE9915E2CAA8488F97239C67595F418D9503D0B11CC07044E', '')),
-  (DRBG_GENERATE,
-   ('39AE66C2939D1D73EF21AE22988B04CC7E8EA2D790C75E1FC6ACC7FEEEF90F98',
-    '',
-    False)),
-  (DRBG_GENERATE,
-   ('B8031829E07B09EEEADEBA149D0AC9F08B110197CD8BBDDC32744BCD66FCF3C4',
-    'A1307377F6B472661BC3C6D44C035FB20A13CCB04D6601B2425FC4DDA3B6D7DF',
-    True)),
-  (DRBG_INIT,
-   ('3A2D261884010CCB4C2C4D7B323CCB7BD4515089BEB749C565A7492710922164',
-    '9E4D22471A4546F516099DD4D737967562D1BB77D774B67B7FE4ED893AE336CF',
-    '5837CAA74345CC2D316555EF820E9F3B0FD454D8C5B7BDE68E4A176D52EE7D1C')),
-  (DRBG_GENERATE,
-   ('4D87985505D779F1AD98455E04199FE8F2FE8E550E6FEB1D26177A2C5B744B9F',
-    '',
-    False)),
-  (DRBG_GENERATE,
-   ('85D011A3B36AC6B25A792F213A1C22C80BFD1C5B47BCA04CD0D9834BB466447B',
-    'B03863C42C9396B4936D83A551871A424C5A8EDBDC9D1E0E8E89710D58B5CA1E',
-    True)),
+    (DRBG_GROUP_INIT, 32),
+    (DRBG_INIT,
+     ('C40894D0C37712140924115BF8A3110C7258532365BB598F81B127A5E4CB8EB0',
+      'FBB1EDAF92D0C2699F5C0A7418D308B09AC679FFBB0D8918C8E62D35091DD2B9',
+      '2B18535D739F7E75AF4FF0C0C713DD4C9B0A6803D2E0DB2BDE3C4F3650ABF750')),
+    (DRBG_RESEED,
+     ('4D58A621857706450338CCA8A1AF5CD2BD9305F3475CF1A8752518DD8E8267B6',
+      '0153A0A1D7487E2EE9915E2CAA8488F97239C67595F418D9503D0B11CC07044E', '')),
+    (DRBG_GENERATE,
+     ('39AE66C2939D1D73EF21AE22988B04CC7E8EA2D790C75E1FC6ACC7FEEEF90F98',
+      '',
+      False)),
+    (DRBG_GENERATE,
+     ('B8031829E07B09EEEADEBA149D0AC9F08B110197CD8BBDDC32744BCD66FCF3C4',
+      'A1307377F6B472661BC3C6D44C035FB20A13CCB04D6601B2425FC4DDA3B6D7DF',
+      True)),
+    (DRBG_INIT,
+     ('3A2D261884010CCB4C2C4D7B323CCB7BD4515089BEB749C565A7492710922164',
+      '9E4D22471A4546F516099DD4D737967562D1BB77D774B67B7FE4ED893AE336CF',
+      '5837CAA74345CC2D316555EF820E9F3B0FD454D8C5B7BDE68E4A176D52EE7D1C')),
+    (DRBG_GENERATE,
+     ('4D87985505D779F1AD98455E04199FE8F2FE8E550E6FEB1D26177A2C5B744B9F',
+      '',
+      False)),
+    (DRBG_GENERATE,
+     ('85D011A3B36AC6B25A792F213A1C22C80BFD1C5B47BCA04CD0D9834BB466447B',
+      'B03863C42C9396B4936D83A551871A424C5A8EDBDC9D1E0E8E89710D58B5CA1E',
+      True)),
 )
 
 # DRBG_TEST command structure:
@@ -158,9 +158,10 @@ def drbg_test_inputs(tpm, test_inputs):
 
     Args:
         tpm: a tpm object used to communicate with the device
+        test_inputs: a list of tuples (drbg_op, drbg_params)
 
     Returns:
-        a list of tuples with the generate responses (tgid, tcid, result_str)
+        a list of tuples with the generated responses (tgid, tcid, result_str)
 
     Raises:
         subcmd.TpmTestError: on unexpected target responses
@@ -173,6 +174,7 @@ def drbg_test_inputs(tpm, test_inputs):
         drbg_op, drbg_params = test
         if drbg_op == DRBG_GROUP_INIT:
             tgid += 1
+            print('Start Test Group', tgid)
             outlen = drbg_params
         elif drbg_op == DRBG_INIT:
             drbg_init(tpm, drbg_params)
@@ -191,7 +193,7 @@ def drbg_test_inputs(tpm, test_inputs):
     print('%sSUCCESS: %s' % (utils.cursor_back(), 'DRBG test'))
     return test_results
 
-def drbg_test(tpm):
+def drbg_test(tpm, request, expected, result_dir):
     """Runs DRBG test cases.
 
     Args:
@@ -201,3 +203,12 @@ def drbg_test(tpm):
         subcmd.TpmTestError: on unexpected target responses
     """
     drbg_test_inputs(tpm, TEST_INPUTS)
+
+    if not request:
+        return
+    print("Running through lab inputs from", request)
+    # Use lab inputs
+    lab = lab_vectors.DRBGLabTest(request, expected, result_dir)
+    test_inputs = lab.get_test_inputs()
+    results = drbg_test_inputs(tpm, test_inputs)
+    lab.save_test_results(results)
diff --git a/test/tpm_test/lab_vectors.py b/test/tpm_test/lab_vectors.py
new file mode 100644
index 0000000000..dfa9f6a62f
--- /dev/null
+++ b/test/tpm_test/lab_vectors.py
@@ -0,0 +1,172 @@
+# -*- coding: utf-8 -*-
+# Copyright 2021 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Module for using lab vectors."""
+
+from __future__ import print_function
+
+import os
+
+import drbg_test
+import utils
+
+class LabTest(object):
+    """Base class implementing the lab vector interface.
+
+    This is used to convert the vectors from the lab to the format required
+    by the tpm tests. It also converts the results from the test to the correct
+    format expected by the lab.
+    """
+
+    ALGORITHM = "algorithm"
+    GROUPS = "testGroups"
+    TEST_CASES = "tests"
+    GROUP_ID = "tgId"
+    TEST_CASE_ID = "tcId"
+
+    def __init__(self, request_file, expected_file, result_dir=""):
+        """Initialize the lab test object.
+
+        Args:
+            request_file: the test input vector json filename
+            expected_file: the expected result json filename
+            result_dir: directory to store the actual result vector in
+        """
+        self._request_file = request_file
+        self._expected_file = expected_file
+        self._request_vector = utils.read_vectors(self._request_file)
+        if self._expected_file:
+            self._expected_vector = utils.read_vectors(self._expected_file)
+        else:
+            self._expected_vector = None
+
+        out_file = (os.path.basename(request_file).strip('.json') +
+                    '-output.json')
+        self._result_file = os.path.join(result_dir,
+            os.path.basename(request_file).strip('.json') + '-output.json')
+        self._result_json = None
+        self._test_inputs = []
+
+    def __str__(self):
+        """Return the algorithm name from the input vector."""
+        return self._request_vector[1][self.ALGORITHM]
+
+    def get_test_inputs(self):
+        """Convert the lab vectors into the format required for the test."""
+        raise NotImplementedError('Algorithm needs to provide the vector processing')
+
+    def _algo_get_formatted_results(self, results):
+        """Convert the results list into the lab format."""
+        raise NotImplementedError('Algorithm needs to process results list')
+
+    def save_test_results(self, results):
+        """Convert the results to the lab format and save them to a file."""
+        print('Saving results in', self._result_file)
+        self._result_json = utils.read_vectors(self._request_file)
+        formatted_results = self._algo_get_formatted_results(results)
+        self._result_json[1][self.GROUPS] = formatted_results
+        utils.write_test_result_json(self._result_file, self._result_json)
+
+
+class DRBGLabTest(LabTest):
+    """Class implementing the lab vector interface for drbg_test.
+
+    Convert the request vector to the test_input format from the drbg test.
+    Convert the response list from the drbg test to the same format as the
+    expected vectors.
+    """
+    RESPONSE_KEY = "returnedBits"
+    RESPONSE_BITS = RESPONSE_KEY + "Len"
+    NONCE = "nonce"
+    PERSO = "persoString"
+    INPUT_1 = "additionalInput"
+    ENTROPY = "entropyInput"
+    MODE = "intendedUse"
+    RESEED = "reSeed"
+    GENERATE = "generate"
+    OTHER_INPUT = "otherInput"
+
+    def _get_expected_response(self, group_id, case_id):
+        """Return the response for the given group and test case."""
+        if not self._expected_vector:
+            return ""
+        group = self._expected_vector[1][self.GROUPS][group_id]
+        return group[self.TEST_CASES][case_id][self.RESPONSE_KEY]
+
+    def _add_test_input(self, test_input):
+        """Append the test item to the test input list."""
+        self._test_inputs.append(test_input)
+
+    def _process_test_case(self, test, response):
+        """Add steps from the test case into the test_inputs list."""
+        drbg_op = drbg_test.DRBG_INIT
+        drbg_params = (test[self.ENTROPY], test[self.NONCE], test[self.PERSO])
+        self._add_test_input((drbg_op, drbg_params))
+        generate_calls = 0
+        for step in test[self.OTHER_INPUT]:
+            mode = step[self.MODE]
+            input1 = step[self.INPUT_1]
+            entropy = step[self.ENTROPY]
+            if mode == self.RESEED:
+                drbg_op = drbg_test.DRBG_RESEED
+                drbg_params = (entropy, input1, "")
+            elif mode == self.GENERATE:
+                drbg_op = drbg_test.DRBG_GENERATE
+                generate_calls += 1
+                if entropy:
+                    raise ValueError('Got entropy during generate %r' % step)
+                # The vectors only verify the second generate command. Only pass in
+                # the result if it will match.
+                check_result = generate_calls == 2
+                expected_response = response if check_result else ''
+                drbg_params = (input1, expected_response, check_result)
+            else:
+                raise ValueError("Invalid mode %r" % mode)
+            self._add_test_input((drbg_op, drbg_params))
+
+    def get_test_inputs(self):
+        """Convert the lab input to the format required by drbg_test.
+
+        Returns:
+            a list of tuples (drbg_op, tuple of drbg_params)
+        """
+        for i, request_group in enumerate(self._request_vector[1][self.GROUPS]):
+            response_bytes = request_group[self.RESPONSE_BITS] >> 3
+            # The test expects each group to specify the response size in bytes.
+            self._add_test_input((drbg_test.DRBG_GROUP_INIT, response_bytes))
+            for j, test in enumerate(request_group[self.TEST_CASES]):
+                response = self._get_expected_response(i, j)
+                self._process_test_case(test, response)
+        return self._test_inputs
+
+
+    def _algo_get_formatted_results(self, results):
+        """Format the results into the list the lab expects.
+
+        Args:
+            results: a list of tuples with the generated responses (tgid, tcid,
+                     result_str)
+
+        Returns:
+            a list of dictionaries. There's a dictionary for each test group.
+            Those contain the group id and a list of tests with the test case
+            id and result string
+        """
+        formatted_results = []
+        last_group = -1
+        for result in results:
+            group_id, test_id, response = result
+            if group_id != last_group:
+                last_group = group_id
+                new_group = {}
+                new_group[self.GROUP_ID] = group_id
+                new_group[self.TEST_CASES] = []
+                formatted_results.append(new_group)
+            test_dict = {}
+            test_dict[self.TEST_CASE_ID] = test_id
+            test_dict[self.RESPONSE_KEY] = response
+            # The group id counts from 1. Offset it to get the list index.
+            formatted_results[group_id - 1][self.TEST_CASES].append(test_dict)
+        return formatted_results
diff --git a/test/tpm_test/tpmtest.py b/test/tpm_test/tpmtest.py
index 04567d67cb..20e368e337 100755
--- a/test/tpm_test/tpmtest.py
+++ b/test/tpm_test/tpmtest.py
@@ -154,12 +154,15 @@ def usage():
           '        0 - raw TRNG'
           '        [-o file] - set output file, default /tmp/trng_output\n'
           '        [-s bits] - TRNG sample size in bit, default = 1\n'
+          '     -l path to output lab result vectors.\n'
+          '     -r path for the drbg input vector.\n'
+          '        [-e file] - expected results file\n'
           '     -h - this help\n')
 
 def main():
     """Run TPM tests"""
     try:
-        opts, _ = getopt.getopt(sys.argv[1:], 'dt:hs:o:', 'help')
+        opts, _ = getopt.getopt(sys.argv[1:], 'dt:hs:o:r:e:l:', 'help')
     except getopt.GetoptError as err:
         print(str(err))
         usage()
@@ -169,6 +172,9 @@ def main():
     trng_output = '/tmp/trng_output'
     trng_sample_bits = 1
     trng_mode = 0
+    lab_output = '/tmp/lab_output'
+    drbg_request = ''
+    drbg_expected = ''
 
     for option, arg in opts:
         if option == '-d':
@@ -178,6 +184,12 @@ def main():
             trng_mode = int(arg)
         elif option == '-o':
             trng_output = arg
+        elif option == '-l':
+            lab_output = arg
+        elif option == '-r':
+            drbg_request = arg
+        elif option == '-e':
+            drbg_expected = arg
         elif option == '-s':
             trng_sample_bits = int(arg)
         elif option in ('-h', '--help'):
@@ -189,9 +201,14 @@ def main():
             trng_test.trng_test(tpm_object, trng_output,
                                 trng_mode, trng_sample_bits)
             sys.exit(0)
+        if drbg_request:
+            drbg_test.drbg_test(tpm_object, drbg_request, drbg_expected,
+                                lab_output)
+            sys.exit(0)
         crypto_test.crypto_tests(tpm_object, os.path.join(ROOT_DIR,
                                                           'crypto_test.xml'))
-        drbg_test.drbg_test(tpm_object)
+        drbg_test.drbg_test(tpm_object, drbg_request, drbg_expected,
+                            lab_output)
         ecc_test.ecc_test(tpm_object)
         ecies_test.ecies_test(tpm_object)
         hash_test.hash_test(tpm_object)
diff --git a/test/tpm_test/utils.py b/test/tpm_test/utils.py
index 5a58dc3b38..b4e1748779 100644
--- a/test/tpm_test/utils.py
+++ b/test/tpm_test/utils.py
@@ -5,6 +5,7 @@
 
 """Support functions for extended command based testing."""
 
+import json
 import sys
 
 if hasattr(sys.stdout, 'isatty') and sys.stdout.isatty():
@@ -35,3 +36,14 @@ def hex_dump(binstr):
         i += strsize
     dump_lines.append('')
     return '\n'.join(dump_lines)
+
+def write_test_result_json(filename, results):
+    """Write the test results to the given file."""
+    with open(filename, 'w') as json_file:
+        json.dump(results, json_file, indent=4)
+
+def read_vectors(filename):
+    """Read the test vectors from the given json file."""
+    with open(filename, 'r') as json_file:
+        contents = json.load(json_file)
+    return contents