File size: 1,289 Bytes
443d045 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from Crypto import Random
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
def get_key(password: str) -> bytes:
"""Generates an encryption key based on the password provided."""
key = SHA256.new(password.encode()).digest()
return key
def encrypt(key: bytes, source: bytes) -> bytes:
"""Encrypts source data using the provided encryption key"""
IV = Random.new().read(AES.block_size) # generate IV
encryptor = AES.new(key, AES.MODE_CBC, IV)
padding = AES.block_size - len(source) % AES.block_size # calculate needed padding
source += bytes([padding]) * padding # Python 2.x: source += chr(padding) * padding
data = IV + encryptor.encrypt(source) # store the IV at the beginning and encrypt
return data
def decrypt(key: bytes, source: bytes) -> bytes:
IV = source[: AES.block_size] # extract the IV from the beginning
decryptor = AES.new(key, AES.MODE_CBC, IV)
data = decryptor.decrypt(source[AES.block_size :]) # decrypt
padding = data[-1] # pick the padding value from the end; Python 2.x: ord(data[-1])
if (
data[-padding:] != bytes([padding]) * padding
): # Python 2.x: chr(padding) * padding
raise ValueError("Invalid padding...")
return data[:-padding] # remove the padding
|