|
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) |
|
encryptor = AES.new(key, AES.MODE_CBC, IV) |
|
padding = AES.block_size - len(source) % AES.block_size |
|
source += bytes([padding]) * padding |
|
data = IV + encryptor.encrypt(source) |
|
return data |
|
|
|
|
|
def decrypt(key: bytes, source: bytes) -> bytes: |
|
IV = source[: AES.block_size] |
|
decryptor = AES.new(key, AES.MODE_CBC, IV) |
|
data = decryptor.decrypt(source[AES.block_size :]) |
|
padding = data[-1] |
|
if ( |
|
data[-padding:] != bytes([padding]) * padding |
|
): |
|
raise ValueError("Invalid padding...") |
|
return data[:-padding] |
|
|