Rewrite AES code with cryptography

jmbase/jmbase/crypto.py

@@ -1,15 +1,16 @@
-import pyaes
+from cryptography.hazmat.primitives import padding
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
 
 def aes_cbc_encrypt(key: bytes, data: bytes, iv: bytes) -> bytes:
-    encrypter = pyaes.Encrypter(
-        pyaes.AESModeOfOperationCBC(key, iv=iv))
-    enc_data = encrypter.feed(data)
-    enc_data += encrypter.feed()
-    return enc_data
+    encrypter = Cipher(algorithms.AES(key), modes.CBC(iv)).encryptor()
+    padder = padding.PKCS7(len(iv) * 8).padder()
+    padded_data = padder.update(data) + padder.finalize()
+    return encrypter.update(padded_data) + encrypter.finalize()
+
 
 def aes_cbc_decrypt(key: bytes, data: bytes, iv: bytes) -> bytes:
-    decrypter = pyaes.Decrypter(
-        pyaes.AESModeOfOperationCBC(key, iv=iv))
-    dec_data = decrypter.feed(data)
-    dec_data += decrypter.feed()
-    return dec_data
+    decrypter = Cipher(algorithms.AES(key), modes.CBC(iv)).decryptor()
+    unpadder = padding.PKCS7(len(iv) * 8).unpadder()
+    padded_data = decrypter.update(data) + decrypter.finalize()
+    return unpadder.update(padded_data) + unpadder.finalize()

jmbase/setup.py

@@ -9,7 +9,12 @@
       author_email='',
       license='GPL',
       packages=['jmbase'],
-      install_requires=['twisted==22.4.0', 'service-identity==21.1.0',
-                        'chromalog==1.0.5', 'pyaes==1.6.1'],
+      install_requires=[
+        "chromalog==1.0.5",
+        "service-identity==21.1.0",
+        "twisted==22.4.0",
+        'cryptography==3.3.2; platform_machine != "aarch64" and platform_machine != "amd64" and platform_machine != "x86_64"',
+        'cryptography==41.0.2; platform_machine == "aarch64" or platform_machine == "amd64" or platform_machine == "x86_64"',
+      ],
       python_requires='>=3.7',
       zip_safe=False)

jmbase/test/test_crypto.py

@@ -0,0 +1,13 @@
+#! /usr/bin/env python
+import os
+
+import pytest
+
+from jmbase import crypto
+
+
+@pytest.mark.parametrize("data", [b"a secret message", b"joinmarket"])
+def test_aes_cbc(data):
+    key, iv = os.urandom(32), os.urandom(16)
+    encrypted = crypto.aes_cbc_encrypt(key, data, iv)
+    assert crypto.aes_cbc_decrypt(key, encrypted, iv) == data

jmdaemon/setup.py

@@ -10,8 +10,6 @@
       license='GPL',
       packages=['jmdaemon'],
       install_requires=['txtorcon==22.0.0',
-                        'cryptography==3.3.2; platform_machine != "aarch64" and platform_machine != "amd64" and platform_machine != "x86_64"',
-                        'cryptography==41.0.2; platform_machine == "aarch64" or platform_machine == "amd64" or platform_machine == "x86_64"',
                         'pyopenssl==23.2.0', 'libnacl==1.8.0',
                         'joinmarketbase==0.9.10dev'],
       python_requires='>=3.7',