Coverage for confy_addons/encryption/aes.py: 100%

1"""AES encryption implementation using the cryptography library. 

2 

3This module provides an AES encryption handler that implements the 

4AESEncryptionABC abstract base class. It supports AES encryption and 

5decryption in CFB mode with 256-bit keys. 

6""" 

7 

8import base64 

9import binascii 

10import logging 

11import secrets 

12from typing import Optional 

13 

14from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes 

15 

16from confy_addons.core.abstract import AESEncryptionABC 

17from confy_addons.core.constants import AES_IV_SIZE, AES_KEY_SIZE, LOGGER_LEVEL 

18from confy_addons.core.exceptions import DecryptionError, EncryptionError 

19from confy_addons.core.mixins import EncryptionMixin 

20 

21logging.basicConfig(level=LOGGER_LEVEL) 

22logger = logging.getLogger(__name__) 

23 

24 

25class AESEncryption(EncryptionMixin, AESEncryptionABC): 

26 """AES symmetric encryption handler. 

27 

28 This class provides AES encryption and decryption operations in CFB mode 

29 with 256-bit keys. It can generate a random key or use a provided key. 

30 

31 Attributes: 

32 key: The AES encryption key (32 bytes). 

33 key_size: The size of the AES key in bytes (always 32 for AES-256). 

34 

35 """ 

36 

37 def __init__(self, key: Optional[bytes] = None): 

38 """Initialize AESEncryption with a key. 

39 

40 Creates an AES encryption handler with either a provided key or a newly 

41 generated random key. 

42 

43 Args: 

44 key: An optional 32-byte AES key. If None, a random key is generated. 

45 

46 Raises: 

47 ValueError: If the provided key is not 32 bytes long. 

48 TypeError: If the provided key is not bytes or bytearray. 

49 

50 """ 

51 self._key_size = AES_KEY_SIZE 

52 

53 if key is None: 

54 self._key = secrets.token_bytes(self._key_size) 

55 else: 

56 if not isinstance(key, (bytes, bytearray)): 

57 logger.error(f'Invalid key type: {type(key)}') 

58 raise TypeError('AES key must be bytes or bytearray') 

59 

60 key_bytes = bytes(key) 

61 

62 if len(key_bytes) != self._key_size: 

63 logger.error(f'Invalid key length: {len(key_bytes)}') 

64 raise ValueError( 

65 f'AES key must be {self._key_size} bytes long ({self._key_size * 8} bits)' 

66 ) 

67 self._key = key_bytes 

68 logger.debug('AES encryption initialized with provided key') 

69 

70 def __repr__(self): 

71 """Return a string representation of the AESEncryption instance. 

72 

73 Returns: 

74 str: A detailed string representation including module, class name, 

75 key, and memory address. 

76 

77 """ 

78 class_name = type(self).__name__ 

79 return f"""{self.__module__}.{class_name}(key=<hidden>) object at {hex(id(self))}""" 

80 

81 def encrypt(self, plaintext: str) -> str: 

82 """Encrypts text using AES in CFB mode. 

83 

84 Encrypts the provided plaintext using AES-256 in CFB mode with a 

85 randomly generated initialization vector. Returns the result as a 

86 base64-encoded string. 

87 

88 Args: 

89 plaintext: The text string to encrypt. 

90 

91 Returns: 

92 str: The base64-encoded encrypted data (IV + ciphertext). 

93 

94 Raises: 

95 EncryptionError: If an error occurs during encryption. 

96 TypeError: If the plaintext is not a string. 

97 

98 """ 

99 if not isinstance(plaintext, str): 

100 logger.error(f'Invalid plaintext type: {type(plaintext)}') 

101 raise TypeError('plaintext must be a str') 

102 

103 try: 

104 iv = secrets.token_bytes(AES_IV_SIZE) 

105 cipher = Cipher(algorithms.AES(self._key), modes.CFB(iv)) 

106 encryptor = cipher.encryptor() 

107 ciphertext = ( 

108 encryptor.update(plaintext.encode(encoding='utf-8')) + encryptor.finalize() 

109 ) 

110 return base64.b64encode(iv + ciphertext).decode(encoding='ascii') 

111 except Exception as e: 

112 logger.error(f'Error occurred during encryption: {e}') 

113 raise EncryptionError('Error occurred during encryption') from e 

114 

115 def decrypt(self, b64_ciphertext: str) -> str: 

116 """Decrypts base64-encoded AES encrypted data. 

117 

118 Decrypts the provided base64-encoded encrypted data using AES-256 

119 in CFB mode. The encrypted data must be in the format produced by 

120 the encrypt method (IV + ciphertext). 

121 

122 Args: 

123 b64_ciphertext: The base64-encoded encrypted data. 

124 

125 Returns: 

126 str: The decrypted plaintext as a string. 

127 

128 Raises: 

129 TypeError: If the b64_ciphertext is not a string. 

130 ValueError: If the base64 data is invalid or too short. 

131 DecryptionError: If an error occurs during decryption. 

132 

133 """ 

134 if not isinstance(b64_ciphertext, str): 

135 logger.error(f'Invalid b64_ciphertext type: {type(b64_ciphertext)}') 

136 raise TypeError('b64_ciphertext must be a base64-encoded str') 

137 

138 try: 

139 data = base64.b64decode(b64_ciphertext) 

140 except (binascii.Error, ValueError, TypeError) as e: 

141 logger.error(f'Error occurred during base64 decoding: {e}') 

142 raise ValueError('Invalid base64 encrypted data') from e 

143 

144 if len(data) < AES_IV_SIZE: 

145 logger.error(f'Invalid encrypted data length: {len(data)}') 

146 raise ValueError('Encrypted data is too short to contain an IV and ciphertext') 

147 

148 iv, ciphertext = data[:AES_IV_SIZE], data[AES_IV_SIZE:] 

149 

150 try: 

151 cipher = Cipher(algorithms.AES(self._key), modes.CFB(iv)) 

152 decryptor = cipher.decryptor() 

153 plaintext_bytes = decryptor.update(ciphertext) + decryptor.finalize() 

154 return plaintext_bytes.decode('utf-8') 

155 except Exception as e: 

156 logger.error(f'Error occurred during decryption: {e}') 

157 raise DecryptionError('Decryption failed') from e 

158 

159 @property 

160 def key(self) -> bytes: 

161 """Returns the AES encryption key. 

162 

163 Returns: 

164 bytes: The 32-byte AES key. 

165 

166 """ 

167 return self._key 

168 

169 @property 

170 def key_size(self) -> int: 

171 """Returns the size of the AES key in bytes. 

172 

173 Returns: 

174 int: The key size in bytes (always 32 for AES-256). 

175 

176 """ 

177 return self._key_size