From 19677540e74eba4392be53ae434b561cf74ad9a6 Mon Sep 17 00:00:00 2001
From: Todd Wildey <twwildey@amazon.com>
Date: Thu, 30 May 2024 16:02:09 -0700
Subject: [PATCH] Adding `get_pem_for_key` and `normalize_pem` methods to
 normalize PEM formatting of keys in `tests/algorithms/test_EC.py` and
 updating `tests/algorithms/test_EC_compat.py` to use these methods

Test failures were occurring due to differences of line lengths generated by the `cryptography` vs `ecdsa` PIP libraries for PEM formatting of cryptographic keys.  This method removes newlines from the bodies of PEM-formated keys so that test comparisons will not fail on differentiated line lengths between PEM formattings.
---
 tests/algorithms/test_EC.py        | 28 +++++++++++++++++++++++++++-
 tests/algorithms/test_EC_compat.py | 14 +++++++-------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/tests/algorithms/test_EC.py b/tests/algorithms/test_EC.py
index 6c167d29..b9028a77 100644
--- a/tests/algorithms/test_EC.py
+++ b/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()
diff --git a/tests/algorithms/test_EC_compat.py b/tests/algorithms/test_EC_compat.py
index 05d033cc..1bb7373a 100644
--- a/tests/algorithms/test_EC_compat.py
+++ b/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,7 +37,7 @@ 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])
@@ -45,7 +45,7 @@ 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])
@@ -53,19 +53,19 @@ 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))
