Skip to content

Commit

Permalink
Rewrite AES code with cryptography
Browse files Browse the repository at this point in the history
  • Loading branch information
roshii committed Oct 6, 2023
1 parent b27c86e commit 6eded6e
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 15 deletions.
7 changes: 3 additions & 4 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,11 @@ requires-python = ">=3.7"
license = {file = "LICENSE"}
dependencies = [
"chromalog==1.0.5",
"pyaes==1.6.1",
"service-identity==21.1.0",
"txtorcon==22.0.0",
"twisted==22.4.0",
"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"',
]

[project.optional-dependencies]
Expand All @@ -33,8 +34,6 @@ jmclient = [
jmdaemon = [
"libnacl==1.8.0",
"pyopenssl==23.2.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"',
]
jmqtui = [
"PyQt5!=5.15.0,!=5.15.1,!=5.15.2,!=6.0",
Expand Down
34 changes: 23 additions & 11 deletions src/jmbase/crypto.py
Original file line number Diff line number Diff line change
@@ -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())
47 changes: 47 additions & 0 deletions test/jmbase/test_crypto.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
#! /usr/bin/env python
import os
from binascii import unhexlify

import pytest

from jmbase import crypto


@pytest.mark.parametrize("data", [b"surely a secret message", b"joinmarket"])
def test_aes_cbc_padding(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",
"7649abac8119b246cee98e9b12e9197d5086cb9b507219ee95db113a917678b273bed6b8e3c1743b7116e69e222295163ff1caa1681fac09120eca307586e1a7",
"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710",
),
(
"8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
"000102030405060708090a0b0c0d0e0f",
"4f021db243bc633d7178183a9fa071e8b4d9ada9ad7dedf4e5e738763f69145a571b242012fb7ae07fa9baac3df102e008b0e27988598881d920a9e64f5615cd",
"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710",
),
(
"603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
"000102030405060708090a0b0c0d0e0f",
"f58c4c04d6e5f1ba779eabfb5f7bfbd69cfc4e967edb808d679f777bc6702c7d39f23369a9d9bacfa530e26304231461b2eb05e2c39be9fcda6c19078c6a9d1b",
"6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e5130c81c46a35ce411e5fbc1191a0a52eff69f2445df4f9b17ad2b417be66c3710",
),
],
)
def test_aes_cbc_nist_vectors(key, iv, ciphertext, plaintext):
_key = unhexlify(key)
_iv = unhexlify(iv)
ct = unhexlify(ciphertext)
pt = unhexlify(plaintext)

assert crypto.aes_cbc_encrypt(_key, pt, _iv) == ct
assert crypto.aes_cbc_decrypt(_key, ct, _iv) == pt

0 comments on commit 6eded6e

Please sign in to comment.