Spaces:
Runtime error
Runtime error
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 | |