diff --git a/pypdf/_crypt_providers/_cryptography.py b/pypdf/_crypt_providers/_cryptography.py index cf04ad8ccc..b7e126da02 100644 --- a/pypdf/_crypt_providers/_cryptography.py +++ b/pypdf/_crypt_providers/_cryptography.py @@ -62,21 +62,21 @@ def encrypt(self, data: bytes) -> bytes: return iv + encryptor.update(data) + encryptor.finalize() def decrypt(self, data: bytes) -> bytes: - if len(data) == 0: - return data iv = data[:16] data = data[16:] - if len(data) % 16: + # for empty encrypted data + if not data: + return data + + # just for robustness, it does not happen under normal circumstances + if len(data) % 16 != 0: pad = padding.PKCS7(128).padder() data = pad.update(data) + pad.finalize() cipher = Cipher(self.alg, CBC(iv)) decryptor = cipher.decryptor() d = decryptor.update(data) + decryptor.finalize() - if len(d) == 0: - return d - else: - return d[: -d[-1]] + return d[: -d[-1]] def rc4_encrypt(key: bytes, data: bytes) -> bytes: diff --git a/pypdf/_crypt_providers/_pycryptodome.py b/pypdf/_crypt_providers/_pycryptodome.py index b2425860e8..05d375119d 100644 --- a/pypdf/_crypt_providers/_pycryptodome.py +++ b/pypdf/_crypt_providers/_pycryptodome.py @@ -55,12 +55,16 @@ def encrypt(self, data: bytes) -> bytes: return iv + aes.encrypt(data) def decrypt(self, data: bytes) -> bytes: - if len(data) == 0: - return data iv = data[:16] data = data[16:] - if len(data) % 16: + # for empty encrypted data + if not data: + return data + + # just for robustness, it does not happen under normal circumstances + if len(data) % 16 != 0: data = pad(data, 16) + aes = AES.new(self.key, AES.MODE_CBC, iv) d = aes.decrypt(data) if len(d) == 0: diff --git a/tests/test_encryption.py b/tests/test_encryption.py index 74d0c99f93..ff33d21211 100644 --- a/tests/test_encryption.py +++ b/tests/test_encryption.py @@ -1,4 +1,4 @@ -"""Test the pypdf.encryption module.""" +"""Test the pypdf._encryption module.""" import secrets from pathlib import Path @@ -6,7 +6,7 @@ import pypdf from pypdf import PasswordType, PdfReader, PdfWriter -from pypdf._encryption import AlgV5, CryptRC4 +from pypdf._encryption import AlgV5, CryptAES, CryptRC4 from pypdf.errors import DependencyError, PdfReadError try: @@ -344,3 +344,11 @@ def test_pdf_encrypt_multiple(pdf_file_path, count): page = reader.pages[0] text1 = page.extract_text() assert text0 == text1 + + +@pytest.mark.skipif(not HAS_PYCRYPTODOME and not HAS_CRYPTOGRAPHY, reason="No pycryptodome / cryptography") +def test_aes_decrypt_corrupted_data(): + """Just for robustness""" + aes = CryptAES(secrets.token_bytes(16)) + for num in [0, 17, 32]: + aes.decrypt(secrets.token_bytes(num))