crypto: improve AES backend logic

mmgen/crypto.py

@@ -67,29 +67,34 @@ class Crypto:
 	def __init__(self, cfg):
 		self.cfg = cfg
 		self.util = cfg._util
-		if cfg.test_suite and self.cfg.aes_backend == 'pyaes':
-			self.get_aes_ctr = self.get_aes_ctr_pyaes
-			self.encrypt_aes_ctr = self.encrypt_aes_ctr_pyaes
+		match self.cfg.aes_backend:
+			case 'cryptography':
+				self.get_aes_ctr = self.get_aes_ctr_cryptography
+			case 'pyaes':
+				assert cfg.test_suite, '`pyaes` module is insecure and suitable only for testing'
+				self.get_aes_ctr = self.get_aes_ctr_pyaes
+			case s:
+				die(3, f'{s}: unrecognized AES backend')
 
 	@staticmethod
-	def get_aes_ctr(key, iv):
+	def get_aes_ctr_cryptography(key, iv):
 		from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 		from cryptography.hazmat.backends import default_backend
 		return Cipher(algorithms.AES(key), modes.CTR(iv), backend=default_backend()).encryptor()
 
-	def encrypt_aes_ctr(self, key, iv, data):
-		encryptor = self.get_aes_ctr(key, iv)
-		return encryptor.update(data) + encryptor.finalize()
-
 	@staticmethod
 	def get_aes_ctr_pyaes(key, iv):
 		import pyaes
 		class MyAES(pyaes.AESModeOfOperationCTR):
 			update = pyaes.AESModeOfOperationCTR.encrypt
+			@staticmethod
+			def finalize():
+				return b''
 		return MyAES(key, pyaes.Counter(int.from_bytes(iv)))
 
-	def encrypt_aes_ctr_pyaes(self, key, iv, data):
-		return self.get_aes_ctr_pyaes(key, iv).encrypt(data)
+	def encrypt_aes_ctr(self, key, iv, data):
+		encryptor = self.get_aes_ctr(key, iv)
+		return encryptor.update(data) + encryptor.finalize()
 
 	def get_hash_params(self, hash_preset):
 		if hash_preset in self.hash_presets: