From 227d01ba5ef6f8d253289bff1703bf7ef084c376 Mon Sep 17 00:00:00 2001 From: Alex Bezdieniezhnykh Date: Fri, 28 Feb 2025 00:49:40 +0200 Subject: [PATCH] use cbc encryption decryption - works nice with c# --- Azaion.Inference/api_client.pxd | 2 -- Azaion.Inference/api_client.pyx | 11 +++------ Azaion.Inference/security.pxd | 2 +- Azaion.Inference/security.pyx | 43 ++++++++++++++++++--------------- 4 files changed, 27 insertions(+), 31 deletions(-) diff --git a/Azaion.Inference/api_client.pxd b/Azaion.Inference/api_client.pxd index 4b55a63..ea2d265 100644 --- a/Azaion.Inference/api_client.pxd +++ b/Azaion.Inference/api_client.pxd @@ -15,5 +15,3 @@ cdef class ApiClient: cdef load_bytes(self, str filename, str folder=*) cdef upload_file(self, str filename, str folder=*) cdef load_ai_model(self) - cdef load_queue_config(self) - diff --git a/Azaion.Inference/api_client.pyx b/Azaion.Inference/api_client.pyx index e16270c..7063d41 100644 --- a/Azaion.Inference/api_client.pyx +++ b/Azaion.Inference/api_client.pyx @@ -105,8 +105,8 @@ cdef class ApiClient: hw_hash = Security.get_hw_hash(hardware) key = Security.get_api_encryption_key(self.credentials, hw_hash) - stream = BytesIO(response.raw.read()) - data = Security.decrypt_to(stream, key) + resp_bytes = response.raw.read() + data = Security.decrypt_to(resp_bytes, key) constants.log(f'Downloaded file: {filename}, {len(data)} bytes') return data @@ -118,9 +118,4 @@ cdef class ApiClient: encrypted_model_bytes = encrypted_bytes_small + encrypted_bytes_big key = Security.get_model_encryption_key() - model_bytes = Security.decrypt_to(BytesIO(encrypted_model_bytes), key) - - - - cdef load_queue_config(self): - return self.load_bytes(constants.QUEUE_CONFIG_FILENAME).decode(encoding='utf-8') \ No newline at end of file + model_bytes = Security.decrypt_to(encrypted_model_bytes, key) diff --git a/Azaion.Inference/security.pxd b/Azaion.Inference/security.pxd index 4dfa64c..ba6f4d8 100644 --- a/Azaion.Inference/security.pxd +++ b/Azaion.Inference/security.pxd @@ -6,7 +6,7 @@ cdef class Security: cdef encrypt_to(input_stream, key) @staticmethod - cdef decrypt_to(input_stream, key) + cdef decrypt_to(input_bytes, key) @staticmethod cdef get_hw_hash(HardwareInfo hardware) diff --git a/Azaion.Inference/security.pyx b/Azaion.Inference/security.pyx index 6518e18..aabc5e2 100644 --- a/Azaion.Inference/security.pyx +++ b/Azaion.Inference/security.pyx @@ -6,43 +6,46 @@ 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 +from cryptography.hazmat.primitives import padding BUFFER_SIZE = 64 * 1024 # 64 KB cdef class Security: @staticmethod - cdef encrypt_to(input_stream, key): + cdef encrypt_to(input_bytes, key): cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest() iv = os.urandom(16) - cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend()) + cipher = Cipher(algorithms.AES( aes_key), modes.CBC(iv), backend=default_backend()) encryptor = cipher.encryptor() + padder = padding.PKCS7(128).padder() - 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) + padded_plaintext = padder.update(input_bytes) + padder.finalize() + ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize() + + return iv + ciphertext @staticmethod - cdef decrypt_to(input_stream, key): + cdef decrypt_to(ciphertext_with_iv_bytes, key): cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest() - cdef bytes iv = input_stream.read(16) + iv = ciphertext_with_iv_bytes[:16] + ciphertext_bytes = ciphertext_with_iv_bytes[16:] - cdef cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend()) - cdef decryptor = cipher.decryptor() + cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) + 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()) + decrypted_padded_bytes = decryptor.update(ciphertext_bytes) + decryptor.finalize() - return bytes(res) + # Manual PKCS7 unpadding check and removal + padding_value = decrypted_padded_bytes[-1] # Get the last byte, which indicates padding length + if 1 <= padding_value <= 16: # Valid PKCS7 padding value range for AES-128 + padding_length = padding_value + plaintext_bytes = decrypted_padded_bytes[:-padding_length] # Remove padding bytes + else: + plaintext_bytes = decrypted_padded_bytes + + return bytes(plaintext_bytes) @staticmethod cdef get_hw_hash(HardwareInfo hardware):