support generic pem format

mauth_client/utils.py

@@ -3,8 +3,7 @@
 import re
 from hashlib import sha512
 
-HEADER = '-----BEGIN RSA PRIVATE KEY-----'
-FOOTER = '-----END RSA PRIVATE KEY-----'
+PEM_BOUNDARY_RE = re.compile(r"-----(?:BEGIN|END) ([A-Za-z0-9 -]+?)-----")
 
 
 def make_bytes(val):
@@ -39,19 +38,30 @@ def decode(byte_string: bytes) -> str:
 
 
 def to_rsa_format(key: str) -> str:
-    """Convert a private key to RSA format with proper newlines."""
-
-    if "\n" in key and HEADER in key and FOOTER in key:
-        return key
+    """Convert a private key to PEM format with proper newlines (RFC 7468)."""
+    stripped = key.strip()
+    labels = PEM_BOUNDARY_RE.findall(stripped)
+
+    # Already well-formed if we have at least a BEGIN and END boundary with newlines
+    if len(labels) >= 2 and "\n" in stripped:
+        return stripped
+
+    if labels:
+        label = labels[0]
+        header = f"-----BEGIN {label}-----"
+        footer = f"-----END {label}-----"
+    else:
+        # Fallback: treat as a bare RSA private key body
+        header = "-----BEGIN RSA PRIVATE KEY-----"
+        footer = "-----END RSA PRIVATE KEY-----"
 
-    body = key.strip()
-    body = body.replace(HEADER, "").replace(FOOTER, "").strip()
+    body = PEM_BOUNDARY_RE.sub("", stripped).strip()
 
     # Replace whitespace with newlines or chunk into 64-char lines
     if " " in body or "\t" in body:
-        body = re.sub(r'\s+', '\n', body)
+        body = re.sub(r"\s+", "\n", body)
     else:
         # PEM-encoded keys are typically split into lines of 64 characters as per RFC 7468 (section 2)
-        body = '\n'.join(body[i:i + 64] for i in range(0, len(body), 64))
+        body = "\n".join(body[i : i + 64] for i in range(0, len(body), 64))
 
-    return f"{HEADER}\n{body}\n{FOOTER}"
+    return f"{header}\n{body}\n{footer}"

tests/utils_test.py

@@ -20,3 +20,21 @@ def test_newlines_removed(self):
         key_no_newlines = PRIVATE_KEY.replace("\n", "")
         key = to_rsa_format(key_no_newlines)
         self.assertEqual(key, PRIVATE_KEY)
+
+    def test_non_rsa_label_well_formed(self):
+        """A well-formed PRIVATE KEY (non-RSA label) PEM should be returned unchanged."""
+        non_rsa_key = PRIVATE_KEY.replace("RSA PRIVATE KEY", "PRIVATE KEY")
+        self.assertEqual(to_rsa_format(non_rsa_key), non_rsa_key)
+
+    def test_boundary_less_body_with_embedded_newlines(self):
+        """A bare body with embedded newlines is wrapped in RSA PRIVATE KEY boundaries."""
+        lines = PRIVATE_KEY.splitlines()
+        body_with_newlines = "\n".join(lines[1:-1])
+        header = "-----BEGIN RSA PRIVATE KEY-----"
+        footer = "-----END RSA PRIVATE KEY-----"
+        result = to_rsa_format(body_with_newlines)
+        self.assertTrue(result.startswith(header + "\n"))
+        self.assertTrue(result.endswith("\n" + footer))
+        # Verify the base64 content is preserved
+        result_body = result[len(header) + 1 : -(len(footer) + 1)]
+        self.assertEqual(result_body.replace("\n", ""), body_with_newlines.replace("\n", ""))