pyinfra/pyinfra/utils/cipher.py
Julius Unverfehrt 52c047c47b add AES/GCM cipher functions
- decrypt x-tenant storage connection strings
2023-03-22 09:39:19 +01:00

50 lines
2.0 KiB
Python

import base64
import os
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
def build_aes_gcm_cipher(public_key, iv=None):
encoded_key = public_key.encode("utf-8")
kdf = PBKDF2HMAC(
algorithm=hashes.SHA1(),
length=16,
salt=iv,
iterations=65536,
)
private_key = kdf.derive(encoded_key)
return AESGCM(private_key)
def encrypt(public_key: str, plaintext: str, iv: int = None) -> str:
"""Encrypt a text with AES/GCS using a public key.
The byte-converted ciphertext consists of an unsigned 32-bit integer big-endian byteorder header i.e. the first 4
bytes, specifying the length of the following initialization vector (iv). The rest of the text contains the
encrypted message.
"""
iv = iv or os.urandom(12)
plaintext_bytes = plaintext.encode("utf-8")
cipher = build_aes_gcm_cipher(public_key, iv)
header = len(iv).to_bytes(length=4, byteorder="big")
encrypted = header + iv + cipher.encrypt(nonce=iv, data=plaintext_bytes, associated_data=None)
return base64.b64encode(encrypted).decode("utf-8")
def decrypt(public_key: str, ciphertext: str) -> str:
"""Decrypt an AES/GCS encrypted text with a public key.
The byte-converted ciphertext consists of an unsigned 32-bit integer big-endian byteorder header i.e. the first 4
bytes, specifying the length of the following initialization vector (iv). The rest of the text contains the
encrypted message.
"""
ciphertext_bytes = base64.b64decode(ciphertext)
header, rest = ciphertext_bytes[:4], ciphertext_bytes[4:]
iv_length = int.from_bytes(header, "big")
iv, ciphertext_bytes = rest[:iv_length], rest[iv_length:]
cipher = build_aes_gcm_cipher(public_key, iv)
decrypted_text = cipher.decrypt(nonce=iv, data=ciphertext_bytes, associated_data=None)
return decrypted_text.decode("utf-8")