Adding get_pem_for_key and normalize_pem methods to normalize PEM formatting of keys

tests/algorithms/test_EC.py

@@ -1,4 +1,5 @@
 import json
+import re
 
 from jose.backends import ECKey
 from jose.constants import ALGORITHMS
@@ -48,6 +49,31 @@
     b"\xfeMO\x04\xb2[\x86A\xbd\xc6hu\x953X\x1e"
 )
 
+# Define the regex pattern to capture the header, body, and footer of the PEM file
+PEM_REGEX = re.compile(r"(-----BEGIN [A-Z ]+-----)(.*?)(-----END [A-Z ]+-----)", re.DOTALL)
+WHITE_SPACE_REGEX = re.compile(r"\s+")
+
+
+def get_pem_for_key(key):
+    return key.to_pem().strip().decode("utf-8")
+
+
+def normalize_pem(key_pem_str):
+    # Search for the PEM sections
+    pem_match = PEM_REGEX.search(key_pem_str)
+    if not pem_match:
+        raise ValueError("The provided string does not contain a valid PEM formatted data.")
+
+    header = pem_match.group(1)
+    body = pem_match.group(2)
+    footer = pem_match.group(3)
+
+    # Remove all newlines and spaces from the body
+    clean_body = WHITE_SPACE_REGEX.sub("", body)
+
+    # Reassemble the PEM string
+    return f"{header}\n{clean_body}\n{footer}"
+
 
 def _backend_exception_types():
     """Build the backend exception types based on available backends."""
@@ -104,7 +130,7 @@ def test_key_from_pem(self):
     def test_to_pem(self):
         key = ECKey(private_key, ALGORITHMS.ES256)
         assert not key.is_public()
-        assert key.to_pem().strip() == private_key.strip().encode("utf-8")
+        assert normalize_pem(get_pem_for_key(key)) == normalize_pem(private_key.strip())
 
         public_pem = key.public_key().to_pem()
         assert ECKey(public_pem, ALGORITHMS.ES256).is_public()

tests/algorithms/test_EC_compat.py

@@ -7,7 +7,7 @@
     ECDSAECKey = CryptographyECKey = None
 from jose.constants import ALGORITHMS
 
-from .test_EC import private_key
+from .test_EC import get_pem_for_key, normalize_pem, private_key
 
 
 @pytest.mark.backend_compatibility
@@ -37,35 +37,35 @@ def test_public_key_to_pem(self, BackendFrom, BackendTo):
         key = BackendFrom(private_key, ALGORITHMS.ES256)
         key2 = BackendTo(private_key, ALGORITHMS.ES256)
 
-        assert key.public_key().to_pem().strip() == key2.public_key().to_pem().strip()
+        assert normalize_pem(get_pem_for_key(key.public_key())) == normalize_pem(get_pem_for_key(key2.public_key()))
 
     @pytest.mark.parametrize("BackendFrom", [ECDSAECKey, CryptographyECKey])
     @pytest.mark.parametrize("BackendTo", [ECDSAECKey, CryptographyECKey])
     def test_private_key_to_pem(self, BackendFrom, BackendTo):
         key = BackendFrom(private_key, ALGORITHMS.ES256)
         key2 = BackendTo(private_key, ALGORITHMS.ES256)
 
-        assert key.to_pem().strip() == key2.to_pem().strip()
+        assert normalize_pem(get_pem_for_key(key)) == normalize_pem(get_pem_for_key(key2))
 
     @pytest.mark.parametrize("BackendFrom", [ECDSAECKey, CryptographyECKey])
     @pytest.mark.parametrize("BackendTo", [ECDSAECKey, CryptographyECKey])
     def test_public_key_load_cycle(self, BackendFrom, BackendTo):
         key = BackendFrom(private_key, ALGORITHMS.ES256)
         pubkey = key.public_key()
 
-        pub_pem_source = pubkey.to_pem().strip()
+        pub_pem_source = normalize_pem(get_pem_for_key(pubkey))
 
         pub_target = BackendTo(pub_pem_source, ALGORITHMS.ES256)
 
-        assert pub_pem_source == pub_target.to_pem().strip()
+        assert pub_pem_source == normalize_pem(get_pem_for_key(pub_target))
 
     @pytest.mark.parametrize("BackendFrom", [ECDSAECKey, CryptographyECKey])
     @pytest.mark.parametrize("BackendTo", [ECDSAECKey, CryptographyECKey])
     def test_private_key_load_cycle(self, BackendFrom, BackendTo):
         key = BackendFrom(private_key, ALGORITHMS.ES256)
 
-        pem_source = key.to_pem().strip()
+        pem_source = normalize_pem(get_pem_for_key(key))
 
         target = BackendTo(pem_source, ALGORITHMS.ES256)
 
-        assert pem_source == target.to_pem().strip()
+        assert pem_source == normalize_pem(get_pem_for_key(target))