130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
import base64
|
|
import hashlib
|
|
import os
|
|
import json
|
|
|
|
from cryptography.fernet import Fernet
|
|
from cryptography.fernet import InvalidToken
|
|
from cryptography.hazmat.backends import default_backend
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives import padding
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.asymmetric import rsa
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
from functools import lru_cache
|
|
|
|
|
|
__all__ = [
|
|
"InvalidToken",
|
|
"fernet_encrypt",
|
|
"fernet_decrypt",
|
|
"aes_cbc_encrypt",
|
|
"aes_cbc_decrypt",
|
|
]
|
|
|
|
|
|
def generate_rsa_dkim(key_size, public_exponent=65537):
|
|
private_key = rsa.generate_private_key(
|
|
public_exponent=public_exponent,
|
|
key_size=key_size,
|
|
)
|
|
|
|
public_key_bytes = private_key.public_key().public_bytes(
|
|
encoding=serialization.Encoding.DER,
|
|
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
|
)
|
|
private_key_pem_bytes = private_key.private_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PrivateFormat.PKCS8,
|
|
encryption_algorithm=serialization.NoEncryption(),
|
|
)
|
|
|
|
public_key_string = base64.b64encode(public_key_bytes).decode("utf-8")
|
|
private_key_pem_string = private_key_pem_bytes.decode("utf-8")
|
|
|
|
return private_key_pem_string, public_key_string
|
|
|
|
|
|
def file_digest_sha256(filename: str):
|
|
with open(filename, "rb", buffering=0) as f:
|
|
return hashlib.file_digest(f, "256").hexdigest()
|
|
|
|
|
|
def dict_digest_sha1(d: str):
|
|
return hashlib.sha1(json.dumps(d, sort_keys=True).encode("utf-8")).hexdigest()
|
|
|
|
|
|
def dict_digest_sha256(d: str):
|
|
return hashlib.sha256(json.dumps(d, sort_keys=True).encode("utf-8")).hexdigest()
|
|
|
|
|
|
def aes_cbc_encrypt(
|
|
data: str, code: str, iv: bytes = os.urandom(16), salt: bytes = os.urandom(16)
|
|
) -> bytes:
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
backend=default_backend(),
|
|
)
|
|
key = kdf.derive(code.encode("utf-8"))
|
|
cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
|
|
padder = padding.PKCS7(algorithms.AES.block_size).padder()
|
|
padded_data = padder.update(data.encode("utf-8")) + padder.finalize()
|
|
encryptor = cipher.encryptor()
|
|
ciphertext = encryptor.update(padded_data) + encryptor.finalize()
|
|
return salt + iv + ciphertext
|
|
|
|
|
|
@lru_cache
|
|
def aes_cbc_decrypt(data: bytes, code: str) -> str:
|
|
salt = data[:16]
|
|
iv = data[16:32]
|
|
ciphertext = data[32:]
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
backend=default_backend(),
|
|
)
|
|
key = kdf.derive(code.encode("utf-8"))
|
|
cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
|
|
decryptor = cipher.decryptor()
|
|
padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()
|
|
unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
|
|
plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()
|
|
return plaintext.decode(encoding="utf-8")
|
|
|
|
|
|
def fernet_encrypt(data: str, code: str, salt: bytes = os.urandom(16)) -> bytes:
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
backend=default_backend(),
|
|
)
|
|
fernet = Fernet(
|
|
base64.urlsafe_b64encode(kdf.derive(code.encode("utf-8"))),
|
|
)
|
|
return base64.urlsafe_b64encode(salt + fernet.encrypt(data.encode("utf-8")))
|
|
|
|
|
|
@lru_cache
|
|
def fernet_decrypt(data: str, code: str) -> str:
|
|
data = base64.urlsafe_b64decode(data)
|
|
salt = data[:16]
|
|
encrypted_data = data[16:]
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=salt,
|
|
iterations=100000,
|
|
backend=default_backend(),
|
|
)
|
|
fernet = Fernet(base64.urlsafe_b64encode(kdf.derive(code.encode("utf-8"))))
|
|
return fernet.decrypt(encrypted_data).decode(encoding="utf-8")
|