Python Secure Implementation Of AES-256-CTR Using Cryptography.Hazmat

Question

I've implemented a wrapper for AES 256 CTR mode using the cryptography.hazmat module, I am wondering if there are any vulnerabilities in my implementation, specifically about the counter and its encoding. Here is the code:

from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.algorithms import AES
from cryptography.hazmat.primitives.ciphers.modes import CTR
from cryptography.hazmat.backends import default_backend as backend
from base58 import b58encode,b58decode
import os
#AES Cipher Class
class AES_Cipher:
 #Initialise Class, Set Countner And Key
 def __init__(self,key):
 self.counter = 0
 self.key = key
 #AES 256 Requirement
 assert len(self.key) == 32
 #Encryption Function
 def encrypt(self,plain_text):
 plain_text = plain_text.encode()
 self.counter += 1
 cipher = Cipher(AES(self.key),CTR(self.padCounter()),backend())
 encryption_engine = cipher.encryptor()
 cipher_text = self.padCounter() + encryption_engine.update(plain_text) + encryption_engine.finalize()
 return b58encode(cipher_text)
 #Decryption Function
 def decrypt(self,cipher_text):
 cipher_text = b58decode(cipher_text)
 self.counter = cipher_text[:16]
 cipher = Cipher(AES(self.key),CTR(self.counter),backend())
 decryption_engine = cipher.decryptor()
 plain_text = decryption_engine.update(cipher_text[16:]) + decryption_engine.finalize()
 return plain_text.decode()
 #Pad The Counter Into 16 Bytes
 def padCounter(self):
 return bytes(str(self.counter).zfill(16),"ascii")

Usage:

 key = os.urandom(32)
 aes_engine = AES_Cipher(key)
 aes_engine.encrypt("hello world")
 aes_engine.decrypt(b"7WkHvZEJRr8yMEasvh3TESoW8nBTkEUNVu2Li")

Answer 1

A warning

You probably already saw this coming, but: I would be remiss if I didn't say it. 'Rolling your own to learn' is fine, but perhaps the most difficult and dangerous place to do it is cryptography. Cryptographic security is notoriously difficult to ascertain. The less you do yourself and the more you leave to well-reviewed, well-established libraries, the better.

Indentation

Python code typically sees three- or four-space tabs by convention; two is a little low.

Type hints

PEP484 allows for this:

def __init__(self,key):

to be (at a guess)

def __init__(self, key: bytes):

and this:

def encrypt(self,plain_text):

to become

def encrypt(self,plain_text: str) -> str:

Helpful comments

This isn't one:

#Encryption Function

you're better off either deleting it, or writing a docstring with non-obvious documentation:

def encrypt(self,plain_text):
 """
 Encrypts a string using this object's key and the AES algorithm.
 Returns the encrypted result, as a base58-encoded string.
 """

• AlexV: Do you have any reference where 3 spaces are recommended for indentation? I would be very interested to hear that :-)
• Reinderien: @AlexV Four is the standard: python.org/dev/peps/pep-0008/#indentation
• AlexV: For sure! The Google Style Guide also lists 4 as the way to got. I was just curious about 3.
• Reinderien: It's a continuum, and unless you have a really good reason, just use four. Three might make for more easily-read code if you have deeply nested structures, but then again... just don't do that.