Skip to content

Commit 49eb79d

Browse files
committed
Throw an error if wrong key type is used
1 parent 6009134 commit 49eb79d

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

src/unstructured_client/users.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
# region imports
1111
from cryptography import x509
1212
from cryptography.hazmat.primitives import serialization, hashes
13-
from cryptography.hazmat.primitives.serialization import load_pem_public_key
1413
from cryptography.hazmat.primitives.asymmetric import padding, rsa
1514
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
1615
from cryptography.hazmat.backends import default_backend
@@ -481,6 +480,9 @@ def _encrypt_rsa_aes(
481480
backend=default_backend()
482481
)
483482

483+
if not isinstance(public_key, rsa.RSAPublicKey):
484+
raise TypeError("Public key must be an RSA public key for envelope encryption.")
485+
484486
# Generate a random AES key
485487
aes_key = os.urandom(32) # 256-bit AES key
486488

@@ -523,6 +525,9 @@ def _encrypt_rsa(
523525
backend=default_backend()
524526
)
525527

528+
if not isinstance(public_key, rsa.RSAPublicKey):
529+
raise TypeError("Public key must be an RSA public key for encryption.")
530+
526531
ciphertext = public_key.encrypt(
527532
plaintext.encode(),
528533
padding.OAEP(

0 commit comments

Comments
 (0)