From 2ae5cca8c4689d15216574fc6cff142e967631a2 Mon Sep 17 00:00:00 2001
From: roshii <roshii@riseup.net>
Date: Tue, 15 Aug 2023 18:52:22 +0200
Subject: [PATCH] Rewrite AES code with cryptography

---
 jmbase/jmbase/crypto.py    | 34 +++++++++++------
 jmbase/setup.py            |  9 ++++-
 jmbase/test/test_crypto.py | 76 ++++++++++++++++++++++++++++++++++++++
 jmdaemon/setup.py          |  2 -
 4 files changed, 106 insertions(+), 15 deletions(-)
 create mode 100644 jmbase/test/test_crypto.py

diff --git a/jmbase/jmbase/crypto.py b/jmbase/jmbase/crypto.py
index e1e8fba2e..6903d62d6 100644
--- a/jmbase/jmbase/crypto.py
+++ b/jmbase/jmbase/crypto.py
@@ -1,15 +1,27 @@
-import pyaes
+from cryptography.hazmat.primitives import padding
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+
+def _pad(data: bytes) -> bytes:
+    if len(data) % 16 == 0:
+        return data
+    padder = padding.PKCS7(128).padder()
+    return padder.update(data) + padder.finalize()
+
+
+def _unpad(data: bytes) -> bytes:
+    try:
+        unpadder = padding.PKCS7(128).unpadder()
+        return unpadder.update(data) + unpadder.finalize()
+    except ValueError:
+        return data
+
 
 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()
+    return encrypter.update(_pad(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()
+    return _unpad(decrypter.update(data) + decrypter.finalize())
\ No newline at end of file
diff --git a/jmbase/setup.py b/jmbase/setup.py
index 031d4305c..3bb7a3b62 100644
--- a/jmbase/setup.py
+++ b/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)
diff --git a/jmbase/test/test_crypto.py b/jmbase/test/test_crypto.py
new file mode 100644
index 000000000..a1fd7bb1b
--- /dev/null
+++ b/jmbase/test/test_crypto.py
@@ -0,0 +1,76 @@
+#! /usr/bin/env python
+import os
+from binascii import unhexlify
+
+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
+
+
+@pytest.mark.parametrize(
+    "key, iv, ciphertext, plaintext",
+    [
+        (
+            "2b7e151628aed2a6abf7158809cf4f3c",
+            "000102030405060708090a0b0c0d0e0f",
+            "7649abac8119b246cee98e9b12e9197d",
+            "6bc1bee22e409f96e93d7e117393172a",
+        ),
+        (
+            "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
+            "000102030405060708090a0b0c0d0e0f",
+            "4f021db243bc633d7178183a9fa071e8",
+            "6bc1bee22e409f96e93d7e117393172a",
+        ),
+        (
+            "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
+            "000102030405060708090a0b0c0d0e0f",
+            "f58c4c04d6e5f1ba779eabfb5f7bfbd6",
+            "6bc1bee22e409f96e93d7e117393172a",
+        ),
+    ],
+)
+def test_aes_cbc_encrypt(key, iv, ciphertext, plaintext):
+    _key = unhexlify(key)
+    _iv = unhexlify(iv)
+    data = unhexlify(plaintext)
+
+    assert crypto.aes_cbc_encrypt(_key, data, _iv) == unhexlify(ciphertext)
+
+
+@pytest.mark.parametrize(
+    "key, iv, ciphertext, plaintext",
+    [
+        (
+            "2b7e151628aed2a6abf7158809cf4f3c",
+            "000102030405060708090a0b0c0d0e0f",
+            "7649abac8119b246cee98e9b12e9197d",
+            "6bc1bee22e409f96e93d7e117393172a",
+        ),
+        (
+            "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
+            "000102030405060708090a0b0c0d0e0f",
+            "4f021db243bc633d7178183a9fa071e8",
+            "6bc1bee22e409f96e93d7e117393172a",
+        ),
+        (
+            "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
+            "000102030405060708090a0b0c0d0e0f",
+            "f58c4c04d6e5f1ba779eabfb5f7bfbd6",
+            "6bc1bee22e409f96e93d7e117393172a",
+        ),
+    ],
+)
+def test_aes_cbc_decrypt(key, iv, ciphertext, plaintext):
+    _key = unhexlify(key)
+    _iv = unhexlify(iv)
+    ct = unhexlify(ciphertext)
+
+    assert crypto.aes_cbc_decrypt(_key, ct, _iv) == unhexlify(plaintext)
diff --git a/jmdaemon/setup.py b/jmdaemon/setup.py
index 923c51f72..bf9f95a4f 100644
--- a/jmdaemon/setup.py
+++ b/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',