Use the correct max size for RSA

Author: awalker4

_test_unstructured_client/unit/test_encryption.py

@@ -121,4 +121,36 @@ def test_encrypt_rsa_aes(rsa_key_pair):
         secret_obj["encrypted_aes_key"],
         secret_obj["aes_iv"],
     )
-    assert decrypted_text == plaintext
+    assert decrypted_text == plaintext
+
+
+rsa_key_size_bytes = 2048 // 8
+max_payload_size = rsa_key_size_bytes - 66  # OAEP SHA256 overhead
+
+@pytest.mark.parametrize(("plaintext", "secret_type"), [
+    ("Short message", "rsa"),
+    ("A" * (max_payload_size), "rsa"),  # Just at the RSA limit
+    ("A" * (max_payload_size + 1), "rsa_aes"),  # Just over the RSA limit
+    ("A" * 500, "rsa_aes"),  # Well over the RSA limit
+])
+def test_encrypt_around_rsa_size_limit(rsa_key_pair, plaintext, secret_type):
+    """
+    Test that payloads around the RSA size limit choose the correct algorithm.
+    """
+    _, public_key_pem = rsa_key_pair
+
+    print(f"Testing plaintext of length {len(plaintext)} with expected type {secret_type}")
+
+    # Load the public key
+    public_key = serialization.load_pem_public_key(
+        public_key_pem.encode('utf-8'),
+        backend=default_backend()
+    )
+
+    client = UnstructuredClient()
+
+    secret_obj = client.users.encrypt_secret(public_key_pem, plaintext)
+
+    # Should still use direct RSA encryption
+    assert secret_obj["type"] == secret_type
+    assert secret_obj["encrypted_value"] is not None

src/unstructured_client/users.py

@@ -471,18 +471,9 @@ async def store_secret_async(
     # region sdk-class-body
     def _encrypt_rsa_aes(
             self,
-            encryption_key_pem: str,
+            public_key: rsa.RSAPublicKey,
             plaintext: str,
     ) -> dict:
-        # Load public RSA key
-        public_key = serialization.load_pem_public_key(
-            encryption_key_pem.encode('utf-8'),
-            backend=default_backend()
-        )
-
-        if not isinstance(public_key, rsa.RSAPublicKey):
-            raise TypeError("Public key must be an RSA public key for envelope encryption.")
-
         # Generate a random AES key
         aes_key = os.urandom(32)  # 256-bit AES key
 
@@ -516,18 +507,10 @@ def _encrypt_rsa_aes(
 
     def _encrypt_rsa(
             self,
-            encryption_key_pem: str,
+            public_key: rsa.RSAPublicKey,
             plaintext: str,
     ) -> dict:
         # Load public RSA key
-        public_key = serialization.load_pem_public_key(
-            encryption_key_pem.encode('utf-8'),
-            backend=default_backend()
-        )
-
-        if not isinstance(public_key, rsa.RSAPublicKey):
-            raise TypeError("Public key must be an RSA public key for encryption.")
-
         ciphertext = public_key.encrypt(
             plaintext.encode(),
             padding.OAEP(
@@ -567,25 +550,28 @@ def encrypt_secret(
                 encryption_cert_or_key_pem.encode('utf-8'),
             )
 
-            loaded_key = cert.public_key()
-
-            # Serialize back to PEM format for consistency
-            public_key_pem = loaded_key.public_bytes(
-                encoding=serialization.Encoding.PEM,
-                format=serialization.PublicFormat.SubjectPublicKeyInfo
-            ).decode('utf-8')
-
+            public_key = cert.public_key()
         else:
-            public_key_pem = encryption_cert_or_key_pem
+            public_key = serialization.load_pem_public_key(
+                encryption_cert_or_key_pem.encode('utf-8'),
+                backend=default_backend()
+            ) 
+
+        if not isinstance(public_key, rsa.RSAPublicKey):
+            raise TypeError("Public key must be a RSA public key for encryption.")
 
         # If the plaintext is short, use RSA directly
         # Otherwise, use a RSA_AES envelope hybrid
-        # The length of the public key is a good hueristic
+        # Use the length of the public key to determine the encryption type
+        key_size_bytes = public_key.key_size // 8
+        max_rsa_length = key_size_bytes - 66  # OAEP SHA256 overhead
+        print(max_rsa_length)
+
         if not encryption_type:
-            encryption_type = "rsa" if len(plaintext) <= len(public_key_pem) else "rsa_aes"
+            encryption_type = "rsa" if len(plaintext) <= max_rsa_length else "rsa_aes"
 
         if encryption_type == "rsa":
-            return self._encrypt_rsa(public_key_pem, plaintext)
+            return self._encrypt_rsa(public_key, plaintext)
 
-        return self._encrypt_rsa_aes(public_key_pem, plaintext)
+        return self._encrypt_rsa_aes(public_key, plaintext)
         # endregion sdk-class-body