Spaces:
Configuration error
Configuration error
| # Author: Trevor Perrin | |
| # See the LICENSE file for legal information regarding use of this file. | |
| """Pure-Python AES implementation.""" | |
| import sys | |
| from .aes import AES | |
| from .rijndael import Rijndael | |
| from .cryptomath import bytesToNumber, numberToByteArray | |
| __all__ = ["new", "Python_AES"] | |
| def new(key, mode, IV): | |
| # IV argument name is a part of the interface | |
| # pylint: disable=invalid-name | |
| if mode == 2: | |
| return Python_AES(key, mode, IV) | |
| elif mode == 6: | |
| return Python_AES_CTR(key, mode, IV) | |
| else: | |
| raise NotImplementedError() | |
| class Python_AES(AES): | |
| def __init__(self, key, mode, IV): | |
| # IV argument/field names are a part of the interface | |
| # pylint: disable=invalid-name | |
| key, IV = bytearray(key), bytearray(IV) | |
| super(Python_AES, self).__init__(key, mode, IV, "python") | |
| self.rijndael = Rijndael(key, 16) | |
| self.IV = IV | |
| def encrypt(self, plaintext): | |
| super(Python_AES, self).encrypt(plaintext) | |
| plaintextBytes = bytearray(plaintext) | |
| chainBytes = self.IV[:] | |
| # CBC Mode: For each block... | |
| for x in range(len(plaintextBytes) // 16): | |
| # XOR with the chaining block | |
| blockBytes = plaintextBytes[x * 16 : (x * 16) + 16] | |
| for y in range(16): | |
| blockBytes[y] ^= chainBytes[y] | |
| # Encrypt it | |
| encryptedBytes = self.rijndael.encrypt(blockBytes) | |
| # Overwrite the input with the output | |
| for y in range(16): | |
| plaintextBytes[(x * 16) + y] = encryptedBytes[y] | |
| # Set the next chaining block | |
| chainBytes = encryptedBytes | |
| self.IV = chainBytes[:] | |
| return plaintextBytes | |
| def decrypt(self, ciphertext): | |
| super(Python_AES, self).decrypt(ciphertext) | |
| ciphertextBytes = ciphertext[:] | |
| chainBytes = self.IV[:] | |
| # CBC Mode: For each block... | |
| for x in range(len(ciphertextBytes) // 16): | |
| # Decrypt it | |
| blockBytes = ciphertextBytes[x * 16 : (x * 16) + 16] | |
| decryptedBytes = self.rijndael.decrypt(blockBytes) | |
| # XOR with the chaining block and overwrite the input with output | |
| for y in range(16): | |
| decryptedBytes[y] ^= chainBytes[y] | |
| ciphertextBytes[(x * 16) + y] = decryptedBytes[y] | |
| # Set the next chaining block | |
| chainBytes = blockBytes | |
| self.IV = chainBytes[:] | |
| return ciphertextBytes | |
| class Python_AES_CTR(AES): | |
| def __init__(self, key, mode, IV): | |
| super(Python_AES_CTR, self).__init__(key, mode, IV, "python") | |
| self.rijndael = Rijndael(key, 16) | |
| self.IV = IV | |
| self._counter_bytes = 16 - len(self.IV) | |
| self._counter = self.IV + bytearray(b"\x00" * self._counter_bytes) | |
| def counter(self): | |
| return self._counter | |
| def counter(self, ctr): | |
| self._counter = ctr | |
| def _counter_update(self): | |
| counter_int = bytesToNumber(self._counter) + 1 | |
| self._counter = numberToByteArray(counter_int, 16) | |
| if self._counter_bytes > 0 and self._counter[-self._counter_bytes :] == bytearray( | |
| b"\xff" * self._counter_bytes | |
| ): | |
| raise OverflowError("CTR counter overflowed") | |
| def encrypt(self, plaintext): | |
| mask = bytearray() | |
| while len(mask) < len(plaintext): | |
| mask += self.rijndael.encrypt(self._counter) | |
| self._counter_update() | |
| if sys.version_info < (3, 0): | |
| inp_bytes = bytearray(ord(i) ^ j for i, j in zip(plaintext, mask)) | |
| else: | |
| inp_bytes = bytearray(i ^ j for i, j in zip(plaintext, mask)) | |
| return inp_bytes | |
| def decrypt(self, ciphertext): | |
| return self.encrypt(ciphertext) | |