annotations/Azaion.AI/security.pyx

# cython: language_level=3
import hashlib
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

BUFFER_SIZE = 64 * 1024  # 64 KB

cdef class Security:

    cdef encrypt_to(self, input_stream, output_stream, key):
        aes_key = hashlib.sha256(key.encode('utf-8')).digest()
        iv = os.urandom(16)
        output_stream.write(iv)  # Write IV to the output stream

        cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend())
        encryptor = cipher.encryptor()

        # Read and encrypt in chunks
        while chunk := input_stream.read(BUFFER_SIZE):
            encrypted_data = encryptor.update(chunk)
            output_stream.write(encrypted_data)

        final_data = encryptor.finalize()
        output_stream.write(final_data)

    cdef decrypt_to(self, input_stream, output_stream, key):
        aes_key = hashlib.sha256(key.encode('utf-8')).digest()
        iv = input_stream.read(16)  # AES block size is 16 bytes

        # Create cipher and decryptor
        cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend())
        decryptor = cipher.decryptor()

        while chunk := input_stream.read(BUFFER_SIZE):
            decrypted_data = decryptor.update(chunk)
            output_stream.write(decrypted_data)

        final_data = decryptor.finalize()
        output_stream.write(final_data)