annotations/Azaion.Inference/security.pyx

import base64
import hashlib
import os
from hashlib import sha384
from credentials cimport Credentials
from hardware_service cimport HardwareInfo

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

BUFFER_SIZE = 64 * 1024  # 64 KB

cdef class Security:
    @staticmethod
    cdef encrypt_to(input_stream, key):
        cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
        iv = os.urandom(16)

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

        cdef bytearray res = bytearray()
        res.extend(iv)
        while chunk := input_stream.read(BUFFER_SIZE):
            encrypted_chunk = encryptor.update(chunk)
            res.extend(encrypted_chunk)
        res.extend(encryptor.finalize())
        return bytes(res)

    @staticmethod
    cdef decrypt_to(input_stream, key):
        cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
        cdef bytes iv = input_stream.read(16)

        cdef cipher = Cipher(algorithms.AES(<bytes>aes_key), modes.CFB(<bytes>iv), backend=default_backend())
        cdef decryptor = cipher.decryptor()

        cdef bytearray res = bytearray()
        while chunk := input_stream.read(BUFFER_SIZE):
            decrypted_chunk = decryptor.update(chunk)
            res.extend(decrypted_chunk)
        res.extend(decryptor.finalize())

        return bytes(res)

    @staticmethod
    cdef get_hw_hash(HardwareInfo hardware):
        cdef str key = f'Azaion_{hardware.mac_address}_{hardware.cpu}_{hardware.gpu}'
        return Security.calc_hash(key)

    @staticmethod
    cdef get_api_encryption_key(Credentials creds, str hardware_hash):
        cdef str key = f'{creds.email}-{creds.password}-{hardware_hash}-#%@AzaionKey@%#---'
        return Security.calc_hash(key)

    @staticmethod
    cdef get_model_encryption_key():
        cdef str key = '-#%@AzaionKey@%#---234sdfklgvhjbnn'
        return Security.calc_hash(key)

    @staticmethod
    cdef calc_hash(str key):
        str_bytes = key.encode('utf-8')
        hash_bytes = sha384(str_bytes).digest()
        cdef str h = base64.b64encode(hash_bytes).decode('utf-8')
        return h