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. """