pyinfra/pyinfra/utils/cipher.py

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")