fix: Fix private key normalization error and raise a clear ValueError instead of a pyasn1 error when a private key cannot be loaded

mauth_client/lambda_helper.py

@@ -1,8 +1,10 @@
 from base64 import b64decode
 from mauth_client.config import Config
 from mauth_client.requests_mauth import MAuth
+from mauth_client.utils import to_rsa_format
 
-RSA_PRIVATE_KEY = "RSA PRIVATE KEY"
+# Present in both the PKCS#1 ("RSA PRIVATE KEY") and PKCS#8 ("PRIVATE KEY") PEM markers.
+PRIVATE_KEY_MARKER = "PRIVATE KEY"
 
 
 def generate_mauth():
@@ -11,7 +13,10 @@ def generate_mauth():
 
 def _get_private_key():
     private_key = Config.PRIVATE_KEY
-    if RSA_PRIVATE_KEY not in private_key:
+    if not private_key:
+        return private_key
+
+    if PRIVATE_KEY_MARKER not in private_key:
         try:
             import boto3
 
@@ -20,4 +25,4 @@ def _get_private_key():
         except ModuleNotFoundError:
             pass
 
-    return private_key.replace("\\n", "\n").replace(" ", "\n").replace("\nRSA\nPRIVATE\nKEY", " RSA PRIVATE KEY")
+    return to_rsa_format(private_key)

mauth_client/rsa_signer.py

@@ -6,8 +6,12 @@
 import rsa
 from .utils import make_bytes, hexdigest
 from pyasn1.codec.der import decoder
+from pyasn1.error import PyAsn1Error
 
 RSA_ALGORITHM_OID = "1.2.840.113549.1.1.1"
+# The rsa library lets pyasn1 errors bubble up when the base64 body of a PEM block does not
+# contain valid DER, so both exception types have to be handled while loading a key.
+KEY_LOAD_ERRORS = (ValueError, PyAsn1Error)
 
 
 class RSASigner:
@@ -26,10 +30,10 @@ def load_private_key(private_key_data):
         private_key_data = make_bytes(private_key_data)
         try:
             return rsa.PrivateKey.load_pkcs1(private_key_data, "PEM")
-        except ValueError:
+        except KEY_LOAD_ERRORS:
             try:
                 return RSASigner.load_pkcs8_private_key(private_key_data)
-            except ValueError as pkcs8_error:
+            except KEY_LOAD_ERRORS as pkcs8_error:
                 raise ValueError("Unable to load private key as PKCS#1 or PKCS#8 PEM") from pkcs8_error
 
     @staticmethod

mauth_client/utils.py

@@ -11,6 +11,9 @@
     (HEADER, FOOTER),
     (PKCS8_HEADER, PKCS8_FOOTER),
 )
+# Keys stored in environment variables or secret stores are often one-liners in which the
+# newlines have been escaped, e.g. "-----BEGIN RSA PRIVATE KEY-----\\nMIIE...".
+ESCAPED_NEWLINES = re.compile(r"\\r\\n|\\n|\\r")
 
 
 def make_bytes(val):
@@ -50,14 +53,26 @@ def to_rsa_format(key: str) -> str:
     Supports both PKCS#1 (``-----BEGIN RSA PRIVATE KEY-----``) and
     PKCS#8 (``-----BEGIN PRIVATE KEY-----``) PEM formats, preserving
     the original header and footer markers.
+
+    Literal ``\\n`` escape sequences are converted to real newlines first, otherwise they
+    would be treated as part of the base64 body and corrupt the decoded key.
+
+    Values that do not carry a supported PEM header/footer pair are returned unchanged:
+    they are not PEM keys (for example a KMS-encrypted, base64-encoded key), and wrapping
+    them in PEM markers would make them look like a valid key while producing garbage.
     """
 
-    header, footer = next(
+    markers = next(
         ((hdr, ftr) for hdr, ftr in SUPPORTED_PRIVATE_KEY_FORMATS if hdr in key and ftr in key),
-        (HEADER, FOOTER),
+        None,
     )
+    if markers is None:
+        return key
+
+    header, footer = markers
+    key = ESCAPED_NEWLINES.sub("\n", key)
 
-    if "\n" in key and header in key and footer in key:
+    if "\n" in key:
         return key
 
     body = key.strip()

tests/lambda_helper_test.py

@@ -0,0 +1,51 @@
+import base64
+import sys
+import unittest
+from unittest.mock import MagicMock, patch
+
+from .common import load_key
+from mauth_client.lambda_helper import _get_private_key
+
+PRIVATE_KEY = load_key("priv").strip()
+PRIVATE_KEY_PKCS8 = load_key("pkcs8").strip()
+
+
+class TestGetPrivateKey(unittest.TestCase):
+    def _get_key(self, configured_key):
+        with patch("mauth_client.lambda_helper.Config") as config:
+            config.PRIVATE_KEY = configured_key
+            return _get_private_key()
+
+    def test_returns_key_unchanged(self):
+        self.assertEqual(self._get_key(PRIVATE_KEY), PRIVATE_KEY)
+
+    def test_normalizes_one_liner_with_spaces(self):
+        self.assertEqual(self._get_key(PRIVATE_KEY.replace("\n", " ")), PRIVATE_KEY)
+
+    def test_normalizes_one_liner_with_escaped_newlines(self):
+        self.assertEqual(self._get_key(PRIVATE_KEY.replace("\n", "\\n")), PRIVATE_KEY)
+
+    def test_returns_pkcs8_key_unchanged(self):
+        self.assertEqual(self._get_key(PRIVATE_KEY_PKCS8), PRIVATE_KEY_PKCS8)
+
+    def test_normalizes_pkcs8_one_liner_with_spaces(self):
+        self.assertEqual(self._get_key(PRIVATE_KEY_PKCS8.replace("\n", " ")), PRIVATE_KEY_PKCS8)
+
+    def test_normalizes_pkcs8_one_liner_with_escaped_newlines(self):
+        self.assertEqual(self._get_key(PRIVATE_KEY_PKCS8.replace("\n", "\\n")), PRIVATE_KEY_PKCS8)
+
+    def test_missing_key(self):
+        self.assertIsNone(self._get_key(None))
+
+    def test_encrypted_key_is_decrypted_with_kms(self):
+        ciphertext = base64.b64encode(b"encrypted-key-blob").decode("ascii")
+        kms_client = MagicMock()
+        kms_client.decrypt.return_value = {"Plaintext": PRIVATE_KEY.replace("\n", "\\n").encode("ascii")}
+        boto3 = MagicMock()
+        boto3.client.return_value = kms_client
+
+        with patch.dict(sys.modules, {"boto3": boto3}):
+            key = self._get_key(ciphertext)
+
+        kms_client.decrypt.assert_called_once_with(CiphertextBlob=b"encrypted-key-blob")
+        self.assertEqual(key, PRIVATE_KEY)

tests/rsa_signer_test.py

@@ -0,0 +1,70 @@
+import base64
+import unittest
+
+import rsa
+from pyasn1.error import PyAsn1Error
+
+from .common import load_key
+from mauth_client.rsa_signer import RSASigner
+
+LOAD_ERROR_MESSAGE = "Unable to load private key as PKCS#1 or PKCS#8 PEM"
+
+# Valid base64, but the decoded bytes are not a DER structure. This is what a PEM block looks
+# like when the body has been corrupted, e.g. by escaped "\n" sequences being left in the body
+# or by a still-encrypted (KMS ciphertext) value being wrapped in PEM markers.
+GARBAGE_BODY = base64.b64encode(b"this is not DER at all" * 20).decode("ascii")
+GARBAGE_BODY = "\n".join(GARBAGE_BODY[i:i + 64] for i in range(0, len(GARBAGE_BODY), 64))
+
+PKCS1_WITH_GARBAGE_BODY = f"-----BEGIN RSA PRIVATE KEY-----\n{GARBAGE_BODY}\n-----END RSA PRIVATE KEY-----"
+PKCS8_WITH_GARBAGE_BODY = f"-----BEGIN PRIVATE KEY-----\n{GARBAGE_BODY}\n-----END PRIVATE KEY-----"
+
+
+class LoadPrivateKeyTest(unittest.TestCase):
+    def test_loads_pkcs1_key(self):
+        self.assertIsInstance(RSASigner(load_key("priv")).private_key, rsa.PrivateKey)
+
+    def test_loads_pkcs8_key(self):
+        self.assertIsInstance(RSASigner(load_key("pkcs8")).private_key, rsa.PrivateKey)
+
+    def test_pyasn1_error_is_not_a_value_error(self):
+        # Guards the assumption behind KEY_LOAD_ERRORS: catching only ValueError is not enough,
+        # because pyasn1 decode failures would escape as PyAsn1Error.
+        self.assertFalse(issubclass(PyAsn1Error, ValueError))
+
+    def test_rsa_raises_pyasn1_error_for_malformed_pkcs1_body(self):
+        # Documents the underlying failure the wrapper has to translate: the PEM markers parse,
+        # so rsa base64-decodes the body and a PyAsn1Error is triggered in decoder.decode.
+        with self.assertRaises(PyAsn1Error):
+            rsa.PrivateKey.load_pkcs1(PKCS1_WITH_GARBAGE_BODY.encode("utf-8"), "PEM")
+
+    def test_malformed_pkcs1_body_raises_value_error(self):
+        # Regression: the PKCS#1 attempt raises PyAsn1Error, which must be caught so the
+        # PKCS#8 fallback runs and a clear ValueError is raised instead of a pyasn1 error.
+        with self.assertRaises(ValueError) as ctx:
+            RSASigner(PKCS1_WITH_GARBAGE_BODY)
+        self.assertEqual(str(ctx.exception), LOAD_ERROR_MESSAGE)
+
+    def test_malformed_pkcs8_body_raises_value_error(self):
+        # Regression: here it is the PKCS#8 fallback itself that raises PyAsn1Error.
+        with self.assertRaises(ValueError) as ctx:
+            RSASigner(PKCS8_WITH_GARBAGE_BODY)
+        self.assertEqual(str(ctx.exception), LOAD_ERROR_MESSAGE)
+        self.assertIsInstance(ctx.exception.__cause__, PyAsn1Error)
+
+    def test_encrypted_key_wrapped_in_pem_markers_raises_value_error(self):
+        # The reported failure mode: a KMS-encrypted key that reached the loader still encrypted.
+        ciphertext = base64.b64encode(b"encrypted-key-blob" * 30).decode("ascii")
+        pem = f"-----BEGIN RSA PRIVATE KEY-----\n{ciphertext}\n-----END RSA PRIVATE KEY-----"
+        with self.assertRaises(ValueError) as ctx:
+            RSASigner(pem)
+        self.assertEqual(str(ctx.exception), LOAD_ERROR_MESSAGE)
+
+    def test_value_is_not_pem_at_all(self):
+        with self.assertRaises(ValueError) as ctx:
+            RSASigner("not a key")
+        self.assertEqual(str(ctx.exception), LOAD_ERROR_MESSAGE)
+
+    def test_original_error_is_chained(self):
+        with self.assertRaises(ValueError) as ctx:
+            RSASigner(PKCS1_WITH_GARBAGE_BODY)
+        self.assertIsNotNone(ctx.exception.__cause__)

tests/utils_test.py

@@ -1,7 +1,8 @@
+import base64
 import unittest
 
 from .common import load_key
-from mauth_client.utils import to_rsa_format
+from mauth_client.utils import FOOTER, to_rsa_format
 
 PRIVATE_KEY = load_key("priv").strip()
 PRIVATE_KEY_PKCS8 = load_key("pkcs8").strip()
@@ -22,6 +23,16 @@ def test_newlines_removed(self):
         key = to_rsa_format(key_no_newlines)
         self.assertEqual(key, PRIVATE_KEY)
 
+    def test_escaped_newlines(self):
+        key_escaped_newlines = PRIVATE_KEY.replace("\n", "\\n")
+        key = to_rsa_format(key_escaped_newlines)
+        self.assertEqual(key, PRIVATE_KEY)
+
+    def test_escaped_carriage_returns(self):
+        key_escaped_newlines = PRIVATE_KEY.replace("\n", "\\r\\n")
+        key = to_rsa_format(key_escaped_newlines)
+        self.assertEqual(key, PRIVATE_KEY)
+
     def test_proper_format_pkcs8(self):
         key = to_rsa_format(PRIVATE_KEY_PKCS8)
         self.assertEqual(key, PRIVATE_KEY_PKCS8)
@@ -30,3 +41,17 @@ def test_newlines_replaced_with_spaces_pkcs8(self):
         key_no_newlines = PRIVATE_KEY_PKCS8.replace("\n", " ")
         key = to_rsa_format(key_no_newlines)
         self.assertEqual(key, PRIVATE_KEY_PKCS8)
+
+    def test_escaped_newlines_pkcs8(self):
+        key_escaped_newlines = PRIVATE_KEY_PKCS8.replace("\n", "\\n")
+        key = to_rsa_format(key_escaped_newlines)
+        self.assertEqual(key, PRIVATE_KEY_PKCS8)
+
+    def test_non_pem_value_is_untouched(self):
+        # e.g. a KMS-encrypted, base64-encoded key: it must not be wrapped in PEM markers
+        ciphertext = base64.b64encode(b"not a pem key" * 40).decode("ascii")
+        self.assertEqual(to_rsa_format(ciphertext), ciphertext)
+
+    def test_key_without_footer_is_untouched(self):
+        truncated = PRIVATE_KEY.replace(FOOTER, "")
+        self.assertEqual(to_rsa_format(truncated), truncated)