MAINT: Prepare for adding AES enryption support (#1818)

Full credit goes to exiledkingcc

Co-authored-by: exiledkingcc <exiledkingcc@gmail.com>

pypdf/_encryption.py

@@ -24,7 +24,6 @@
 # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 # POSSIBILITY OF SUCH DAMAGE.
-
 import hashlib
 import secrets
 import struct
@@ -56,14 +55,6 @@ class CryptIdentity(CryptBase):
     pass
 
 
-def _randrange(lower_inclusive: int, upper_exclusive: int) -> int:
-    return secrets.choice(range(lower_inclusive, upper_exclusive))
-
-
-def _randint(lower_inclusive: int, upper_inclusive: int) -> int:
-    return secrets.choice(range(lower_inclusive, upper_inclusive + 1))
-
-
 try:
     from Crypto.Cipher import AES, ARC4  # type: ignore[import]
     from Crypto.Util.Padding import pad  # type: ignore[import]
@@ -83,7 +74,7 @@ def __init__(self, key: bytes) -> None:
             self.key = key
 
         def encrypt(self, data: bytes) -> bytes:
-            iv = bytes(bytearray(_randint(0, 255) for _ in range(16)))
+            iv = secrets.token_bytes(16)
             p = 16 - len(data) % 16
             data += bytes(bytearray(p for _ in range(p)))
             aes = AES.new(self.key, AES.MODE_CBC, iv)
@@ -320,7 +311,7 @@ def compute_key(
         u_hash.update(o_entry)
         u_hash.update(struct.pack("<I", P))
         u_hash.update(id1_entry)
-        if rev >= 4 and metadata_encrypted is False:
+        if rev >= 4 and not metadata_encrypted:
             u_hash.update(b"\xff\xff\xff\xff")
         u_hash_digest = u_hash.digest()
         length = key_size // 8
@@ -781,12 +772,13 @@ def compute_U_value(password: bytes, key: bytes) -> Tuple[bytes, bytes]:
         Returns:
             A tuple (u-value, ue value)
         """
-        random_bytes = bytes(_randrange(0, 256) for _ in range(16))
+        random_bytes = secrets.token_bytes(16)
         val_salt = random_bytes[:8]
         key_salt = random_bytes[8:]
-        u_value = hashlib.sha256(password + val_salt).digest() + val_salt + key_salt
+        R = 4  # only one supported so far
+        u_value = AlgV5.calculate_hash(R, password, val_salt, b"") + val_salt + key_salt
 
-        tmp_key = hashlib.sha256(password + key_salt).digest()
+        tmp_key = AlgV5.calculate_hash(R, password, key_salt, b"")
         iv = bytes(0 for _ in range(16))
         ue_value = AES_CBC_encrypt(tmp_key, iv, key)
         return u_value, ue_value
@@ -824,14 +816,14 @@ def compute_O_value(
         Returns:
             A tuple (O value, OE value)
         """
-        random_bytes = bytes(_randrange(0, 256) for _ in range(16))
+        random_bytes = secrets.token_bytes(4)
         val_salt = random_bytes[:8]
         key_salt = random_bytes[8:]
+        R = 4  # only one supported so far
         o_value = (
-            hashlib.sha256(password + val_salt + u_value).digest() + val_salt + key_salt
+            AlgV5.calculate_hash(R, password, val_salt, u_value) + val_salt + key_salt
         )
-
-        tmp_key = hashlib.sha256(password + key_salt + u_value).digest()
+        tmp_key = AlgV5.calculate_hash(R, password, key_salt, u_value[:48])
         iv = bytes(0 for _ in range(16))
         oe_value = AES_CBC_encrypt(tmp_key, iv, key)
         return o_value, oe_value
@@ -869,7 +861,7 @@ def compute_Perms_value(key: bytes, p: int, metadata_encrypted: bool) -> bytes:
             The perms value
         """
         b8 = b"T" if metadata_encrypted else b"F"
-        rr = bytes(_randrange(0, 256) for _ in range(4))
+        rr = secrets.token_bytes(4)
         data = struct.pack("<I", p) + b"\xff\xff\xff\xff" + b8 + b"adb" + rr
         perms = AES_ECB_encrypt(key, data)
         return perms

tests/test_encryption.py

@@ -5,7 +5,7 @@
 
 import pypdf
 from pypdf import PasswordType, PdfReader
-from pypdf._encryption import AlgV5, CryptRC4, _randint, _randrange
+from pypdf._encryption import AlgV5, CryptRC4
 from pypdf.errors import DependencyError, PdfReadError
 
 try:
@@ -227,23 +227,3 @@ def test_alg_v5_generate_values():
         "/OE": values["/OE"],
         "/Perms": values["/Perms"],
     }
-
-
-def test_randrange_function():
-    """
-    _randrange() function generates a range of unique random numbers.
-
-    This test might randomly fail in very rare cases.
-    """
-    random_set = {_randrange(0, 10) for _ in range(1000)}
-    assert random_set == {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
-
-
-def test_randint_function():
-    """
-    _randint() function generates a range of unique random numbers, including the upper bound.
-
-    This test might randomly fail in very rare cases.
-    """
-    random_set = {_randint(0, 10) for _ in range(1000)}
-    assert random_set == {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10}