aes_crypt.py

from Crypto import Random
from Crypto.Cipher import AES
from Crypto.Cipher import Salsa20
from Crypto.Cipher import ChaCha20
from Crypto.Cipher import ARC4
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Util import Padding
from Crypto.PublicKey import RSA



class AESCipher(object):
    def __init__(self, key_size=32):
        self.key_size = key_size
        key = Random.new().read(self.key_size)
        self.key = key
        
    def encrypt(self, raw):
        iv = Random.new().read(AES.block_size)
        p = Padding.pad(raw, AES.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return iv + cipher.encrypt(p)

    def decrypt(self, enc):
        iv, enc = enc[:AES.block_size], enc[AES.block_size:]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        return Padding.unpad(cipher.decrypt(enc), AES.block_size)


class Salsa20Cipher(object):
    def __init__(self, key_size=32):
        self.key_size = key_size
        key = Random.new().read(self.key_size)
        self.key = key

    def encrypt(self, raw):
        cipher = Salsa20.new(self.key)
        return cipher.nonce + cipher.encrypt(raw)
    
    def decrypt(self, enc):
        nonce, msg = enc[:8], enc[8:]
        cipher = Salsa20.new(key=self.key, nonce=nonce)
        return cipher.decrypt(msg)


class  ChaCha20Cipher(object):
    def __init__(self, key_size=32):
        self.key_size = key_size
        key = Random.new().read(self.key_size)
        self.key = key

    def encrypt(self, raw):
        cipher = ChaCha20.new(key=self.key)
        return cipher.nonce + cipher.encrypt(raw)
    
    def decrypt(self, enc):
        nonce, msg = enc[:8], enc[8:]
        cipher = ChaCha20.new(key=self.key, nonce=nonce)
        return cipher.decrypt(msg)


class  ARC4Cipher(object):
    def __init__(self, key_size=32):
        self.key_size = key_size
        key = Random.new().read(self.key_size)
        self.key = key

    def encrypt(self, raw):
        cipher = ARC4.new(key=self.key)
        nonce = Random.new().read(self.key_size)
        return nonce + cipher.encrypt(raw)
    
    def decrypt(self, enc):
        enc = enc[self.key_size:]
        cipher = ARC4.new(key=self.key)
        return cipher.decrypt(enc)

def generate_key_pairs(size=2048):
    key = RSA.generate(size)
    return key.exportKey(), key.publickey().exportKey()
    
class RSACipher(object):
    def __init__(self, key):
        self.key = PKCS1_OAEP.new(RSA.import_key(key))
    def encrypt(self, raw):
        return self.key.encrypt(raw)
    def decrypt(self, raw):
        return self.key.decrypt(raw)