Changes after second round of code review

- Reduced code duplication by using a pre-compiled regular expression

middleware/admin/verify_attestation.py

@@ -29,8 +29,8 @@
 from .certificate import HSMCertificate
 
 
-UI_MESSAGE_HEADER = b"HSM:UI:5.X"
-SIGNER_MESSAGE_HEADER = b"HSM:SIGNER:5.X"
+UI_MESSAGE_HEADER_REGEX = re.compile(b"^HSM:UI:(5.[0-9])")
+SIGNER_MESSAGE_HEADER_REGEX = re.compile(b"^HSM:SIGNER:(5.[0-9])")
 UI_DERIVATION_PATH = "m/44'/0'/0'/0/0"
 UD_VALUE_LENGTH = 32
 PUBKEY_COMPRESSED_LENGTH = 33
@@ -46,14 +46,12 @@
                          "dad609"
 
 
-def validate_ui_message_header(ui_message):
-    header = ui_message[:len(UI_MESSAGE_HEADER)]
-    return re.compile(b"^HSM:UI:5.[0-9]$").match(header) is not None
+def match_ui_message_header(ui_message):
+    return UI_MESSAGE_HEADER_REGEX.match(ui_message)
 
 
-def validate_signer_message_header(signer_message):
-    header = signer_message[:len(SIGNER_MESSAGE_HEADER)]
-    return re.compile(b"^HSM:SIGNER:5.[0-9]$").match(header) is not None
+def match_signer_message_header(signer_message):
+    return SIGNER_MESSAGE_HEADER_REGEX.match(signer_message)
 
 
 def do_verify_attestation(options):
@@ -132,13 +130,14 @@ def do_verify_attestation(options):
 
     ui_message = bytes.fromhex(ui_result[1])
     ui_hash = bytes.fromhex(ui_result[2])
-    mh_len = len(UI_MESSAGE_HEADER)
-    if not validate_ui_message_header(ui_message):
+    mh_match = match_ui_message_header(ui_message)
+    if mh_match is None:
         raise AdminError(
-            f"Invalid UI attestation message header: {ui_message[:mh_len].hex()}")
+            f"Invalid UI attestation message header: {ui_message.hex()}")
+    mh_len = len(mh_match.group(0))
 
     # Extract UI version, UD value, UI public key and signer version from message
-    ui_version = re.match(b"^HSM:UI:(5.[0-9])$", ui_message[:mh_len]).group(1)
+    ui_version = mh_match.group(1)
     ud_value = ui_message[mh_len:mh_len + UD_VALUE_LENGTH].hex()
     ui_public_key = ui_message[mh_len + UD_VALUE_LENGTH:mh_len + UD_VALUE_LENGTH +
                                PUBKEY_COMPRESSED_LENGTH].hex()
@@ -175,19 +174,20 @@ def do_verify_attestation(options):
 
     signer_message = bytes.fromhex(signer_result[1])
     signer_hash = bytes.fromhex(signer_result[2])
-    mh_len = len(SIGNER_MESSAGE_HEADER)
-    if not validate_signer_message_header(signer_message):
+    mh_match = match_signer_message_header(signer_message)
+    if mh_match is None:
         raise AdminError(
-            f"Invalid Signer attestation message header: {signer_message[:mh_len].hex()}")
+            f"Invalid Signer attestation message header: {signer_message.hex()}")
 
-    signer_version = re.match(b"^HSM:SIGNER:(5.[0-9])$", signer_message[:mh_len]).group(1)
+    mh_len = len(mh_match.group(0))
     if signer_message[mh_len:] != pubkeys_hash:
         reported = signer_message[mh_len:].hex()
         raise AdminError(
             f"Signer attestation public keys hash mismatch: expected {pubkeys_hash.hex()}"
             f" but attestation reports {reported}"
         )
 
+    signer_version = mh_match.group(1)
     head(
         ["Signer verified with public keys:"] + pubkeys_output + [
             "",

middleware/tests/admin/test_verify_attestation.py

@@ -27,8 +27,8 @@
 from admin.pubkeys import PATHS
 from admin.verify_attestation import (
     do_verify_attestation,
-    validate_ui_message_header,
-    validate_signer_message_header
+    match_ui_message_header,
+    match_signer_message_header
 )
 import ecdsa
 import hashlib
@@ -285,7 +285,7 @@ def test_verify_attestation_invalid_signer_att(self,
         self.assertEqual(("Invalid Signer attestation: error validating 'signer'"),
                          str(e.exception))
 
-    def test_validate_ui_message_header_valid_header(self, _):
+    def test_match_ui_message_header_valid_header(self, _):
         valid_headers = [
             UI_HEADER,
             b"HSM:UI:5.0",
@@ -294,19 +294,19 @@ def test_validate_ui_message_header_valid_header(self, _):
         ]
         for header in valid_headers:
             ui_message = header + self.ui_msg[len(UI_HEADER):]
-            self.assertTrue(validate_ui_message_header(ui_message))
+            self.assertTrue(match_ui_message_header(ui_message))
 
-    def test_validate_ui_message_header_invalid_header(self, _):
+    def test_match_ui_message_header_invalid_header(self, _):
         invalid_headers = [
             SIGNER_HEADER,
             b"HSM:UI:4.0",
             b"HSM:UI:5.X",
         ]
         for header in invalid_headers:
             ui_message = header + self.ui_msg[len(UI_HEADER):]
-            self.assertFalse(validate_ui_message_header(ui_message))
+            self.assertFalse(match_ui_message_header(ui_message))
 
-    def test_validate_signer_message_header_valid_header(self, _):
+    def test_match_signer_message_header_valid_header(self, _):
         valid_headers = [
             SIGNER_HEADER,
             b"HSM:SIGNER:5.0",
@@ -315,14 +315,14 @@ def test_validate_signer_message_header_valid_header(self, _):
         ]
         for header in valid_headers:
             signer_message = header + self.signer_msg[len(SIGNER_HEADER):]
-            self.assertTrue(validate_signer_message_header(signer_message))
+            self.assertTrue(match_signer_message_header(signer_message))
 
-    def test_validate_signer_message_header_invalid_header(self, _):
+    def test_match_signer_message_header_invalid_header(self, _):
         invalid_headers = [
             UI_HEADER,
             b"HSM:SIGNER:4.0",
             b"HSM:SIGNER:5.X",
         ]
         for header in invalid_headers:
             signer_message = header + self.signer_msg[len(SIGNER_HEADER):]
-            self.assertFalse(validate_signer_message_header(signer_message))
+            self.assertFalse(match_signer_message_header(signer_message))