1# Copyright (c) 2023, exiledkingcc 
    2# All rights reserved. 
    3# 
    4# Redistribution and use in source and binary forms, with or without 
    5# modification, are permitted provided that the following conditions are 
    6# met: 
    7# 
    8# * Redistributions of source code must retain the above copyright notice, 
    9# this list of conditions and the following disclaimer. 
    10# * Redistributions in binary form must reproduce the above copyright notice, 
    11# this list of conditions and the following disclaimer in the documentation 
    12# and/or other materials provided with the distribution. 
    13# * The name of the author may not be used to endorse or promote products 
    14# derived from this software without specific prior written permission. 
    15# 
    16# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
    17# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
    18# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
    19# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 
    20# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
    21# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
    22# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
    23# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
    24# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
    25# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
    26# POSSIBILITY OF SUCH DAMAGE. 
    27 
    28import secrets 
    29 
    30from Crypto import __version__ 
    31from Crypto.Cipher import AES, ARC4 
    32from Crypto.Util.Padding import pad 
    33 
    34from pypdf._crypt_providers._base import CryptBase 
    35 
    36crypt_provider = ("pycryptodome", __version__) 
    37 
    38 
    39class CryptRC4(CryptBase): 
    40    def __init__(self, key: bytes) -> None: 
    41        self.key = key 
    42 
    43    def encrypt(self, data: bytes) -> bytes: 
    44        return ARC4.ARC4Cipher(self.key).encrypt(data) 
    45 
    46    def decrypt(self, data: bytes) -> bytes: 
    47        return ARC4.ARC4Cipher(self.key).decrypt(data) 
    48 
    49 
    50class CryptAES(CryptBase): 
    51    def __init__(self, key: bytes) -> None: 
    52        self.key = key 
    53 
    54    def encrypt(self, data: bytes) -> bytes: 
    55        iv = secrets.token_bytes(16) 
    56        data = pad(data, 16) 
    57        aes = AES.new(self.key, AES.MODE_CBC, iv) 
    58        return iv + aes.encrypt(data) 
    59 
    60    def decrypt(self, data: bytes) -> bytes: 
    61        iv = data[:16] 
    62        data = data[16:] 
    63        # for empty encrypted data 
    64        if not data: 
    65            return data 
    66 
    67        # just for robustness, it does not happen under normal circumstances 
    68        if len(data) % 16 != 0: 
    69            data = pad(data, 16) 
    70 
    71        aes = AES.new(self.key, AES.MODE_CBC, iv) 
    72        d = aes.decrypt(data) 
    73        return d[: -d[-1]] 
    74 
    75 
    76def rc4_encrypt(key: bytes, data: bytes) -> bytes: 
    77    return ARC4.ARC4Cipher(key).encrypt(data) 
    78 
    79 
    80def rc4_decrypt(key: bytes, data: bytes) -> bytes: 
    81    return ARC4.ARC4Cipher(key).decrypt(data) 
    82 
    83 
    84def aes_ecb_encrypt(key: bytes, data: bytes) -> bytes: 
    85    return AES.new(key, AES.MODE_ECB).encrypt(data) 
    86 
    87 
    88def aes_ecb_decrypt(key: bytes, data: bytes) -> bytes: 
    89    return AES.new(key, AES.MODE_ECB).decrypt(data) 
    90 
    91 
    92def aes_cbc_encrypt(key: bytes, iv: bytes, data: bytes) -> bytes: 
    93    return AES.new(key, AES.MODE_CBC, iv).encrypt(data) 
    94 
    95 
    96def aes_cbc_decrypt(key: bytes, iv: bytes, data: bytes) -> bytes: 
    97    return AES.new(key, AES.MODE_CBC, iv).decrypt(data)