From a5ae77d568a38d7210850c76bd2e951ad2b2fdc1 Mon Sep 17 00:00:00 2001 From: The MMGen Project Date: Sat, 15 Aug 2026 14:13:09 +0000 Subject: [PATCH] crypto: improve AES backend logic --- mmgen/crypto.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/mmgen/crypto.py b/mmgen/crypto.py index 988b42d3..2723fbfb 100755 --- a/mmgen/crypto.py +++ b/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: