1 files changed, 42 insertions, 31 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)