From a0ee706819037df1107ba7498fe3bfbf8b59540b Mon Sep 17 00:00:00 2001
From: Vadim Bendebury <vbendeb@chromium.org>
Date: Wed, 10 Feb 2016 10:33:21 -0800
Subject: [PATCH] cr50: test: consolidate test exceptions

There is no point in defining tpm test exception classes per test
type, one common class is enough, especially if the source module of
the exception is reported.

BRANCH=none
BUG=none
TEST=tried running the test without the USB FTDI cable plugged in, got
     the following error message:

     $ ./test/tpm_test/tpmtest.py
     Starting MPSSE at 800 kHz

     Error in tpmtest.py:54:  Failed to connect
     $

Change-Id: I5642aa70c8a581099887b58e3a436d7f8d7608a1
Signed-off-by: Vadim Bendebury <vbendeb@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/327300
Reviewed-by: Nagendra Modadugu <ngm@google.com>
---
 test/tpm_test/crypto_test.py | 44 ++++++++++++++++++------------------
 test/tpm_test/ecc_test.py    | 10 +++-----
 test/tpm_test/hash_test.py   | 12 ++++------
 test/tpm_test/rsa_test.py    |  8 ++-----
 test/tpm_test/subcmd.py      |  5 ++++
 test/tpm_test/tpmtest.py     | 37 ++++++++++++++----------------
 6 files changed, 53 insertions(+), 63 deletions(-)

diff --git a/test/tpm_test/crypto_test.py b/test/tpm_test/crypto_test.py
index 7fa4ea6146..929df7c48c 100644
--- a/test/tpm_test/crypto_test.py
+++ b/test/tpm_test/crypto_test.py
@@ -17,9 +17,6 @@ import utils
 DECRYPT = 0
 ENCRYPT = 1
 
-class CryptoError(Exception):
-  pass
-
 def get_attribute(tdesc, attr_name, required=True):
   """Retrieve an attribute value from an XML node.
 
@@ -33,8 +30,8 @@ def get_attribute(tdesc, attr_name, required=True):
   Returns:
     The attribute value as a string (ascii or binary)
   Raises:
-    CryptoError: on various format errors, or in case a required attribute is
-      not found, the error message describes the problem.
+    subcmd.TpmTestError: on various format errors, or in case a required
+      attribute is not found, the error message describes the problem.
 
   """
   # Fields stored in hex format by default.
@@ -43,7 +40,7 @@ def get_attribute(tdesc, attr_name, required=True):
   data = tdesc.find(attr_name)
   if data is None:
     if required:
-      raise CryptoError('node "%s" does not have attribute "%s"' %
+      raise subcmd.TpmTestError('node "%s" does not have attribute "%s"' %
                         (tdesc.get('name'), attr_name))
     return ''
 
@@ -55,7 +52,7 @@ def get_attribute(tdesc, attr_name, required=True):
     else:
       cell_format = 'ascii'
   elif cell_format not in ('hex', 'ascii'):
-    raise CryptoError('%s:%s, unrecognizable format "%s"' %
+    raise subcmd.TpmTestError('%s:%s, unrecognizable format "%s"' %
                       (tdesc.get('name'), attr_name, cell_format))
 
   text = ' '.join(x.strip() for x in data.text.splitlines() if x)
@@ -65,7 +62,7 @@ def get_attribute(tdesc, attr_name, required=True):
   # Drop spaces from hex representation.
   text = text.replace(' ', '')
   if len(text) & 3:
-    raise CryptoError('%s:%s %swrong hex number size' %
+    raise subcmd.TpmTestError('%s:%s %swrong hex number size' %
                       (tdesc.get('name'), attr_name, utils.hex_dump(text)))
   # Convert text to binary
   value = ''
@@ -73,7 +70,7 @@ def get_attribute(tdesc, attr_name, required=True):
     try:
       value += struct.pack('<I', int('0x%s' % text[8*x:8*(x+1)], 16))
     except ValueError:
-      raise CryptoError('%s:%s %swrong hex value' %
+      raise subcmd.TpmTestError('%s:%s %swrong hex value' %
                         (tdesc.get('name'), attr_name, utils.hex_dump(text)))
   return value
 
@@ -134,15 +131,16 @@ def crypto_run(node_name, op_type, key, iv, in_text, out_text, tpm):
     comparison with the expected value was successful.
 
   Raises:
-    CryptoError: in case there were problems parsing the node name, or verifying
-      the operation results.
+    subcmd.TpmTestError: in case there were problems parsing the node name, or
+      verifying the operation results.
+
   """
   mode_name, submode_name = node_name.split(':')
   submode_name = submode_name[:3].upper()
 
   mode = SUPPORTED_MODES.get(mode_name.upper())
   if not mode:
-    raise CryptoError('unrecognizable mode in node "%s"' % node_name)
+    raise subcmd.TpmTestError('unrecognizable mode in node "%s"' % node_name)
 
   submode = mode.submodes.get(submode_name, 0)
   cmd = '%c' % op_type    # Encrypt or decrypt
@@ -166,12 +164,13 @@ def crypto_run(node_name, op_type, key, iv, in_text, out_text, tpm):
       if tpm.debug_enabled():
         print('Out text mismatch in node %s:\n' % node_name)
       else:
-        raise CryptoError('Out text mismatch in node %s, operation %d:\n'
-                          'In text:%sExpected out text:%sReal out text:%s' % (
-                              node_name, op_type,
-                              utils.hex_dump(in_text),
-                              utils.hex_dump(out_text),
-                              utils.hex_dump(real_out_text)))
+        raise subcmd.TpmTestError(
+          'Out text mismatch in node %s, operation %d:\n'
+          'In text:%sExpected out text:%sReal out text:%s' % (
+            node_name, op_type,
+            utils.hex_dump(in_text),
+            utils.hex_dump(out_text),
+            utils.hex_dump(real_out_text)))
   return real_out_text
 
 
@@ -187,18 +186,19 @@ def crypto_test(tdesc, tpm):
       session.
     tpm: a TPM object to send extended commands to an initialized TPM
   Raises:
-    CryptoError: on various execution errors, the details are included in the
-      error message.
+    subcmd.TpmTestError: on various execution errors, the details are included
+      in the error message.
+
   """
   node_name = tdesc.get('name')
   key = get_attribute(tdesc, 'key')
   if len(key) not in (16, 24, 32):
-    raise CryptoError('wrong key size "%s:%s"' % (
+    raise subcmd.TpmTestError('wrong key size "%s:%s"' % (
         node_name,
         ''.join('%2.2x' % ord(x) for x in key)))
   iv = get_attribute(tdesc, 'iv', required=False)
   if iv and len(iv) != 16:
-    raise CryptoError('wrong iv size "%s:%s"' % (
+    raise subcmd.TpmTestError('wrong iv size "%s:%s"' % (
         node_name,
         ''.join('%2.2x' % ord(x) for x in iv)))
   clear_text = get_attribute(tdesc, 'clear_text')
diff --git a/test/tpm_test/ecc_test.py b/test/tpm_test/ecc_test.py
index cd4bb0a261..de7f7e2d64 100644
--- a/test/tpm_test/ecc_test.py
+++ b/test/tpm_test/ecc_test.py
@@ -108,10 +108,6 @@ _KEYDERIVE_INPUTS = (
 )
 
 
-class ECError(Exception):
-  pass
-
-
 def _sign_test(tpm):
   msg = 'Hello CR50'
 
@@ -129,7 +125,7 @@ def _sign_test(tpm):
     verified = tpm.unwrap_ext_response(subcmd.EC, wrapped_response)
     expected = '\x01'
     if verified != expected:
-      raise ECError('%s error:%s:%s' % (
+      raise subcmd.TpmTestError('%s error:%s:%s' % (
         test_name, utils.hex_dump(verified), utils.hex_dump(expected)))
     print('%sSUCCESS: %s' % (utils.cursor_back(), test_name))
 
@@ -143,7 +139,7 @@ def _keygen_test(tpm):
     valid = tpm.unwrap_ext_response(subcmd.EC, wrapped_response)
     expected = '\x01'
     if valid != expected:
-      raise ECError('%s error:%s:%s' % (
+      raise subcmd.TpmTestError('%s error:%s:%s' % (
         test_name, utils.hex_dump(valid), utils.hex_dump(expected)))
     print('%sSUCCESS: %s' % (utils.cursor_back(), test_name))
 
@@ -158,7 +154,7 @@ def _keyderive_test(tpm):
     valid = tpm.unwrap_ext_response(subcmd.EC, wrapped_response)
     expected = '\x01'
     if valid != expected:
-      raise ECError('%s error:%s:%s' % (
+      raise subcmd.TpmTestError('%s error:%s:%s' % (
         test_name, utils.hex_dump(valid), utils.hex_dump(expected)))
     print('%sSUCCESS: %s' % (utils.cursor_back(), test_name))
 
diff --git a/test/tpm_test/hash_test.py b/test/tpm_test/hash_test.py
index a5a1e522c7..66c5fbf00b 100644
--- a/test/tpm_test/hash_test.py
+++ b/test/tpm_test/hash_test.py
@@ -23,10 +23,6 @@ CMD_SINGLE = 3
 MODE_SHA1 = 0
 MODE_SHA256 = 1
 
-
-class HashError(Exception):
-  pass
-
 # A standard empty response to HASH extended commands.
 EMPTY_RESPONSE = ''.join('%c' % x for x in (0x80, 0x01, 0x00, 0x00, 0x00, 0x0c,
                                             0xba, 0xcc, 0xd0, 0x0a, 0x00, 0x01))
@@ -64,7 +60,7 @@ def hash_test(tpm):
     tpm: a tpm object used to communicate with the device
 
   Raises:
-    HashError: on unexpected target responses
+    subcmd.TpmTestError: on unexpected target responses
   """
 
   contexts = {}
@@ -100,7 +96,7 @@ def hash_test(tpm):
       h = contexts[handle]
       h.update(text)
       if wrapped_response != EMPTY_RESPONSE:
-        raise HashError("Unexpected response to '%s': %s" %
+        raise subcmd.TpmTestError("Unexpected response to '%s': %s" %
                         (test_name, utils.hex_dump(wrapped_response)))
       continue
     if hash_cmd == CMD_FINISH:
@@ -108,12 +104,12 @@ def hash_test(tpm):
     elif hash_cmd == CMD_SINGLE:
       h = hash_func()
     else:
-      raise HashError('Unknown command %d' % hash_cmd)
+      raise subcmd.TpmTestError('Unknown command %d' % hash_cmd)
     h.update(text)
     digest = h.digest()
     result = wrapped_response[12:]
     if result != h.digest():
-      raise HashError('%s error:%s%s' % (test_name,
+      raise subcmd.TpmTestError('%s error:%s%s' % (test_name,
                                          utils.hex_dump(digest),
                                          utils.hex_dump(result)))
     print('%sSUCCESS: %s' % (utils.cursor_back(), test_name))
diff --git a/test/tpm_test/rsa_test.py b/test/tpm_test/rsa_test.py
index c337c524cc..0783f957fd 100644
--- a/test/tpm_test/rsa_test.py
+++ b/test/tpm_test/rsa_test.py
@@ -119,10 +119,6 @@ _SIGN_INPUTS = (
 )
 
 
-class RSAError(Exception):
-  pass
-
-
 def _encrypt_tests(tpm):
   msg = 'Hello CR50!'
 
@@ -138,7 +134,7 @@ def _encrypt_tests(tpm):
     wrapped_response = tpm.command(tpm.wrap_ext_command(subcmd.RSA, cmd))
     plaintext = tpm.unwrap_ext_response(subcmd.RSA, wrapped_response)
     if msg != plaintext:
-      raise RSAError('%s error:%s%s' % (
+      raise subcmd.TpmTestError('%s error:%s%s' % (
           test_name, utils.hex_dump(msg), utils.hex_dump(plaintext)))
     print('%sSUCCESS: %s' % (utils.cursor_back(), test_name))
 
@@ -159,7 +155,7 @@ def _sign_tests(tpm):
     verified = tpm.unwrap_ext_response(subcmd.RSA, wrapped_response)
     expected = '\x01'
     if verified != expected:
-      raise RSAError('%s error:%s%s' % (
+      raise subcmd.TpmTestError('%s error:%s%s' % (
           test_name, utils.hex_dump(verified), utils.hex_dump(expected)))
     print('%sSUCCESS: %s' % (utils.cursor_back(), test_name))
 
diff --git a/test/tpm_test/subcmd.py b/test/tpm_test/subcmd.py
index d14c2e9b9c..916767222d 100644
--- a/test/tpm_test/subcmd.py
+++ b/test/tpm_test/subcmd.py
@@ -10,3 +10,8 @@ AES = 0
 HASH = 1
 RSA = 2
 EC = 3
+
+
+# The same exception class used by all tpmtest modules.
+class TpmTestError(Exception):
+  pass
diff --git a/test/tpm_test/tpmtest.py b/test/tpm_test/tpmtest.py
index 173f0ec693..46c03ef23b 100755
--- a/test/tpm_test/tpmtest.py
+++ b/test/tpm_test/tpmtest.py
@@ -29,9 +29,6 @@ import subcmd
 # Extension command for dcypto testing
 EXT_CMD = 0xbaccd00a
 
-class TpmError(Exception):
-  pass
-
 
 class TPM(object):
   """TPM accessor class.
@@ -55,11 +52,11 @@ class TPM(object):
     self._debug_enabled = debug_mode
     self._handle = ftdi_spi_tpm
     if not self._handle.FtdiSpiInit(freq, debug_mode):
-      raise TpmError()
+      raise subcmd.TpmTestError('Failed to connect')
     response = self.command(''.join('%c' % int('0x%s' % x, 16)
                                     for x in self.STARTUP_CMD.split()))
     if ' '.join('%2.2x' % ord(x) for x in response) not in self.STARTUP_RSP:
-      raise TpmError('init failed')
+      raise subcmd.TpmTestError('init failed')
 
   def validate(self, data_blob, response_mode=False):
     """Check if a data blob complies with TPM command/response header format."""
@@ -67,12 +64,12 @@ class TPM(object):
         self.HEADER_FMT, data_blob + '  ')
     prefix = 'Misformatted blob: '
     if tag not in (0x8001, 0x8002):
-      raise TpmError(prefix + 'bad tag value 0x%4.4x' % tag)
+      raise subcmd.TpmTestError(prefix + 'bad tag value 0x%4.4x' % tag)
     if size != len(data_blob):
-      raise TpmError(prefix + 'size mismatch: header %d, actual %d' %
-                     (size, len(data_blob)))
+      raise subcmd.TpmTestError(prefix + 'size mismatch: header %d, actual %d'
+                                % (size, len(data_blob)))
     if size > 4096:
-      raise TpmError(prefix + 'invalid size %d' % size)
+      raise subcmd.TpmTestError(prefix + 'invalid size %d' % size)
     if response_mode:
       return
     if cmd_code >= 0x11f and cmd_code <= 0x18f:
@@ -80,7 +77,7 @@ class TPM(object):
     if cmd_code == EXT_CMD:
       return  # This is an extension command
 
-    raise TpmError(prefix + 'invalid command code 0x%x' % cmd_code)
+    raise subcmd.TpmTestError(prefix + 'invalid command code 0x%x' % cmd_code)
 
   def command(self, cmd_data):
     # Verify command header
@@ -108,21 +105,22 @@ class TPM(object):
     Returns:
       the binary string of the response payload, if validation succeeded.
     Raises:
-      TpmError: in case there are any validation problems, the error message
-        describes the problem.
+      subcmd.TpmTestError: in case there are any validation problems, the
+        error message describes the problem.
     """
     header_size = struct.calcsize(self.HEADER_FMT)
     tag, size, cmd, subcmd = struct.unpack(self.HEADER_FMT,
                                            response[:header_size])
     if tag != 0x8001:
-      raise TpmError('Wrong response tag: %4.4x' % tag)
+      raise subcmd.TpmTestError('Wrong response tag: %4.4x' % tag)
     if cmd != EXT_CMD:
-      raise TpmError('Unexpected response command field: %8.8x' % cmd)
+      raise subcmd.TpmTestError('Unexpected response command field: %8.8x' %
+                                cmd)
     if subcmd != expected_subcmd:
-      raise TpmError('Unexpected response subcommand field: %2.2x' %
+      raise subcmd.TpmTestError('Unexpected response subcommand field: %2.2x' %
                      subcmd)
     if size != len(response):
-      raise TpmError('Size mismatch: header %d, actual %d' % (
+      raise subcmd.TpmTestError('Size mismatch: header %d, actual %d' % (
           size, len(response)))
     return response[header_size:]
 
@@ -139,10 +137,9 @@ if __name__ == '__main__':
     ecc_test.ecc_test(t)
     hash_test.hash_test(t)
     rsa_test.rsa_test(t)
-  except (TpmError, crypto_test.CryptoError, hash_test.HashError,
-          rsa_test.RSAError) as e:
-    print()
-    print('Error:', e)
+  except subcmd.TpmTestError as e:
+    exc_file, exc_line = traceback.extract_tb(sys.exc_traceback)[-1][:2]
+    print('\nError in %s:%s: ' % (os.path.basename(exc_file), exc_line), e)
     if debug_needed:
       traceback.print_exc()
     sys.exit(1)