4
\$\begingroup\$

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")
\$\endgroup\$

1 Answer 1

Reset to default
5
\$\begingroup\$

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.
    """
\$\endgroup\$
4
  • \$\begingroup\$ Do you have any reference where 3 spaces are recommended for indentation? I would be very interested to hear that :-) \$\endgroup\$ Commented Oct 22, 2019 at 18:31
  • \$\begingroup\$ @AlexV Four is the standard: python.org/dev/peps/pep-0008/#indentation \$\endgroup\$ Commented Oct 22, 2019 at 18:33
  • \$\begingroup\$ For sure! The Google Style Guide also lists 4 as the way to got. I was just curious about 3. \$\endgroup\$ Commented Oct 22, 2019 at 18:34
  • 1
    \$\begingroup\$ 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. \$\endgroup\$ Commented Oct 22, 2019 at 18:35

You must log in to answer this question.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.