crpyter.py

#!/usr/bin/env python
# -*- encoding:utf-8 -*-
import json
import hashlib
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


class AEScoder():
    def __init__(self, encrypt_key=None, key=None, method=None):
        self.BS = AES.block_size
        if not method or method not in ['cbc', 'ctr', 'cfb', 'ofb']:
            self.method = AES.MODE_ECB
        else:
            self.method = getattr(AES, 'MODE_{}'.format(method.upper()))
        self.encrypt_key = encrypt_key
        self.key = key or hashlib.sha256(
            str(self.encrypt_key).encode('utf-8')
        ).hexdigest()[0:self.BS].encode('utf-8')

    def encrypt(self, data):
        if self.method == AES.MODE_ECB:
            return self.encrypt_ecb(data)
        else:
            return self.encrypt_common(data)

    def decrypt(self, encr_data):
        if self.method == AES.MODE_ECB:
            return self.decrypt_ecb(encr_data)
        else:
            return self.decrypt_common(encr_data)

    def encrypt_ecb(self, data):
        # ECB加密
        cipher = AES.new(self.key, self.method)
        encr_raw_data = cipher.encrypt(pad(data.encode('utf-8'), self.BS))
        encr_data = base64.b64encode(encr_raw_data).decode()
        return encr_data

    def decrypt_ecb(self, encr_data):
        # ECB解密
        try:
            encr_raw_data = base64.b64decode(encr_data)
            cipher = AES.new(self.key, self.method)
            decr_data = unpad(cipher.decrypt(encr_raw_data), self.BS)
            return decr_data.decode('utf-8')
        except ValueError or KeyError:
            raise Exception("Incorrect decryption")

    def encrypt_common(self, data):
        # 加密
        cipher = AES.new(self.key, self.method)
        ct_bytes = cipher.encrypt(pad(data.encode('utf-8'), self.BS))
        iv = base64.b64encode(cipher.iv).decode('utf-8') \
            if hasattr(cipher, 'iv') else None
        nonce = base64.b64encode(cipher.nonce).decode('utf-8') \
            if hasattr(cipher, 'nonce') else None
        ct = base64.b64encode(ct_bytes).decode('utf-8')
        encr_data = json.dumps({'iv': iv, 'nonce': nonce, 'ciphertext': ct})
        return encr_data

    def decrypt_common(self, encr_data):
        try:
            b64 = json.loads(encr_data)
            iv = base64.b64decode(b64['iv']) if b64.get('iv') else None
            nonce = base64.b64decode(b64['nonce']) \
                if b64.get('nonce') else None
            ct = base64.b64decode(b64['ciphertext']) \
                if b64.get('ciphertext') else None
            if nonce:
                cipher = AES.new(self.key, self.method, nonce=nonce)
            else:
                cipher = AES.new(self.key, self.method, iv=iv)
            decr_data = unpad(cipher.decrypt(ct), self.BS)
            return decr_data.decode('utf-8')
        except ValueError or KeyError:
            raise Exception("Incorrect decryption")


if __name__ == "__main__":
    key = 'key'
    data = 'data'

    for method in [None, 'ecb', 'cbc', 'ctr', 'cfb', 'ofb']:
        coder = AEScoder(key, method=method)
        secret = coder.encrypt(data)
        coder_data = coder.decrypt(secret)
        assert data == coder_data

requirements.txt

pycryptodome==3.9.9