Merge pull request #2 from omaha-consulting/main

feat: add arg required_key_hashes to verify(...)

Author: liying2008

src/crx3/verifier.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import base64
+import hashlib
 import os.path
 import os.path
 import struct
@@ -53,10 +54,11 @@ def _verity_signature(crx_fp, file_size, public_key_bytes, signed_header_data, s
     return public_key.verify(signature, digest, padding.PKCS1v15(), Prehashed(sha256_algorithm))
 
 
-def verify(crx_file):
+def verify(crx_file, required_key_hashes=None):
     """
     Verify the file as a valid Crx of CRX3 format.
     :param crx_file: crx file path.
+    :param required_key_hashes: a list of hex strings of public key hashes.
     :return: (VerifierResult, CrxHeaderInfo or None)
     """
     if not os.access(crx_file, os.F_OK | os.R_OK):
@@ -109,6 +111,7 @@ def return_error(verifier_error):
         return return_error(VerifierResult.ERROR_REQUIRED_PROOF_MISSING)
 
     public_key_bytes = None
+    required_key_set = set(required_key_hashes or [])
     for rsa_proof in rsa_proofs:
         # print(rsa_proof)
         if id_util.calc_crx_id_by_public_key(rsa_proof.public_key) == crx_id:
@@ -121,8 +124,10 @@ def return_error(verifier_error):
         except Exception as e:
             # print(e)
             return return_error(VerifierResult.ERROR_SIGNATURE_VERIFICATION_FAILED)
+        public_key_hash = hashlib.sha256(rsa_proof.public_key).hexdigest()
+        required_key_set.discard(public_key_hash)
 
-    if public_key_bytes is None:
+    if public_key_bytes is None or required_key_set:
         return return_error(VerifierResult.ERROR_REQUIRED_PROOF_MISSING)
     # Convert public key bytes to base64 string.
     public_key_str = base64.b64encode(public_key_bytes).decode('utf-8')

tests/test_verifier.py

@@ -11,13 +11,15 @@
 
 class VerifierTestCase(MyTest):
 
+    CRX_FILE = os.path.join(EXAMPLE_EXTENSION_DIR, EXAMPLE_EXTENSION_NAME + '.crx')
+    PUBLIC_KEY_FILE = os.path.join(EXAMPLE_EXTENSION_DIR, EXAMPLE_EXTENSION_NAME + '-public-key.pem')
+    PUBLIC_KEY_HASH = '99ec6d34093cd254def8c05e5fd5b25ebfef20f97d976eceba28b970da7d81b9'
+
     def test_01_verify_ok_full(self):
-        crx_file = os.path.join(EXAMPLE_EXTENSION_DIR, EXAMPLE_EXTENSION_NAME + '.crx')
-        public_key_file = os.path.join(EXAMPLE_EXTENSION_DIR, EXAMPLE_EXTENSION_NAME + '-public-key.pem')
-        verifier_result, header_info = verifier.verify(crx_file)
+        verifier_result, header_info = verifier.verify(self.CRX_FILE)
         # print(verifier_result, header_info)
         self.assertEqual(verifier_result, verifier.VerifierResult.OK_FULL)
-        public_key = read_public_key_str(public_key_file)
+        public_key = read_public_key_str(self.PUBLIC_KEY_FILE)
         self.assertEqual(header_info, verifier.CrxHeaderInfo(EXAMPLE_EXTENSION_CRX_ID, public_key))
 
     def test_02_verify_error_file_not_readable(self):
@@ -46,6 +48,19 @@ def test_04_verify_error_signature_verification_failed(self):
         self.assertEqual(verifier_result, verifier.VerifierResult.ERROR_SIGNATURE_VERIFICATION_FAILED)
         self.assertEqual(header_info, None)
 
+    def test_05_verify_ok_full_with_required_key_hashes(self):
+        verifier_result, header_info = verifier.verify(self.CRX_FILE, [self.PUBLIC_KEY_HASH])
+        # print(verifier_result, header_info)
+        self.assertEqual(verifier_result, verifier.VerifierResult.OK_FULL)
+        public_key = read_public_key_str(self.PUBLIC_KEY_FILE)
+        self.assertEqual(header_info, verifier.CrxHeaderInfo(EXAMPLE_EXTENSION_CRX_ID, public_key))
+
+    def test_06_verify_fail_for_required_key_hashes(self):
+        verifier_result, header_info = verifier.verify(self.CRX_FILE, ['0' * 64])
+        # print(verifier_result, header_info)
+        self.assertEqual(verifier_result, verifier.VerifierResult.ERROR_REQUIRED_PROOF_MISSING)
+        self.assertEqual(header_info, None)
+
 
 if __name__ == '__main__':
     unittest.main()