diff --git a/Azaion.AI/api_client.pxd b/Azaion.AI/api_client.pxd index af52823..f907d8f 100644 --- a/Azaion.AI/api_client.pxd +++ b/Azaion.AI/api_client.pxd @@ -1,10 +1,9 @@ -from processor_command cimport FileCommand - cdef class ApiClient: cdef str email, password, token, folder, token_file, api_url cdef get_encryption_key(self, str hardware_hash) - cdef login(self, str email, str password, bint persist_token=*) - cdef bytes load_file(self, str filename, bint persist_token=*) - cdef bytes load_ai_model(self) + cdef login(self, str email, str password) + cdef load_bytes(self, str filename) + cdef load_ai_model(self) + cdef load_queue_config(self) diff --git a/Azaion.AI/api_client.pyx b/Azaion.AI/api_client.pyx index 5412cac..61975e7 100644 --- a/Azaion.AI/api_client.pyx +++ b/Azaion.AI/api_client.pyx @@ -1,11 +1,13 @@ import io +import json import os from http import HTTPStatus import requests cimport constants from hardware_service cimport HardwareService, HardwareInfo -from security import Security +from security cimport Security +from io import BytesIO cdef class ApiClient: """Handles API authentication and downloading of the AI model.""" @@ -20,45 +22,57 @@ cdef class ApiClient: self.token = None cdef get_encryption_key(self, str hardware_hash): - return f'{self.email}-{self.password}-{hardware_hash}-#%@AzaionKey@%#---' + cdef str key = f'{self.email}-{self.password}-{hardware_hash}-#%@AzaionKey@%#---' + return Security.calc_hash(key) - cdef login(self, str email, str password, bint persist_token=False): + cdef login(self, str email, str password): response = requests.post(f"{constants.API_URL}/login", json={"email": email, "password": password}) response.raise_for_status() self.token = response.json()["token"] - print(f'') - if persist_token: - with open(constants.TOKEN_FILE, 'w') as file: - file.write(self.token) + with open(constants.TOKEN_FILE, 'w') as file: + file.write(self.token) - cdef bytes load_file(self, str filename, bint persist_token=False): + cdef load_bytes(self, str filename): hardware_service = HardwareService() cdef HardwareInfo hardware = hardware_service.get_hardware_info() if self.token is None: - self.login(self.email, self.password, persist_token) + self.login(self.email, self.password) url = f"{constants.API_URL}/resources/get/{self.folder}" - headers = {"Authorization": f"Bearer {self.token}"} - payload = { - "password": self.password, - "hardware": hardware, - "fileName": filename + headers = { + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json" } - response = requests.post(url, json=payload, headers=headers, stream=True) + payload = json.dumps( + { + "password": self.password, + "hardware": hardware.to_json_object(), + "fileName": filename + }, indent=4) + response = requests.post(url, data=payload, headers=headers, stream=True) + if response.status_code == HTTPStatus.UNAUTHORIZED or response.status_code == HTTPStatus.FORBIDDEN: - self.login(self.email, self.password, persist_token) - response = requests.post(url, json=payload, headers=headers, stream=True) + self.login(self.email, self.password) + headers = { + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json" + } + response = requests.post(url, data=payload, headers=headers, stream=True) + + if response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR: + print('500!') key = self.get_encryption_key(hardware.hash) - encrypted_stream = io.BytesIO(response.content) - decrypted_stream = io.BytesIO() - Security.decrypt_to(encrypted_stream, decrypted_stream, key) - return decrypted_stream - cdef bytes load_ai_model(self): - return self.load_file(constants.AI_MODEL_FILE, True) + stream = BytesIO(response.raw.read()) + return Security.decrypt_to(stream, key) + cdef load_ai_model(self): + return self.load_bytes(constants.AI_MODEL_FILE) + + cdef load_queue_config(self): + return self.load_bytes(constants.QUEUE_CONFIG_FILENAME).decode(encoding='utf-8') \ No newline at end of file diff --git a/Azaion.AI/debug.py b/Azaion.AI/debug.py index e4e7f6b..b21adcc 100644 --- a/Azaion.AI/debug.py +++ b/Azaion.AI/debug.py @@ -1,6 +1,9 @@ +from io import BytesIO + import main from main import ParsedArguments + def start_server(): args = ParsedArguments('admin@azaion.com', 'Az@1on1000Odm$n', 'stage', True) processor = main.CommandProcessor(args) diff --git a/Azaion.AI/hardware_service.pxd b/Azaion.AI/hardware_service.pxd index eee5d08..2cd664a 100644 --- a/Azaion.AI/hardware_service.pxd +++ b/Azaion.AI/hardware_service.pxd @@ -1,9 +1,8 @@ cdef class HardwareInfo: cdef str cpu, gpu, memory, mac_address, hash - + cdef to_json_object(self) cdef class HardwareService: cdef bint is_windows - - cdef HardwareInfo get_hardware_info(self) - cdef str calc_hash(self, str s) \ No newline at end of file + cdef get_mac_address(self, interface=*) + cdef HardwareInfo get_hardware_info(self) \ No newline at end of file diff --git a/Azaion.AI/hardware_service.pyx b/Azaion.AI/hardware_service.pyx index 783b68b..8e73f20 100644 --- a/Azaion.AI/hardware_service.pyx +++ b/Azaion.AI/hardware_service.pyx @@ -1,6 +1,6 @@ -import base64 import subprocess -from hashlib import sha384 +from security cimport Security +import psutil cdef class HardwareInfo: def __init__(self, str cpu, str gpu, str memory, str mac_address, str hw_hash): @@ -10,6 +10,15 @@ cdef class HardwareInfo: self.mac_address = mac_address self.hash = hw_hash + cdef to_json_object(self): + return { + "CPU": self.cpu, + "GPU": self.gpu, + "MacAddress": self.mac_address, + "Memory": self.memory, + "Hash": self.hash, + } + def __str__(self): return f'CPU: {self.cpu}. GPU: {self.gpu}. Memory: {self.memory}. MAC Address: {self.mac_address}' @@ -27,6 +36,15 @@ cdef class HardwareService: print('Error during os type checking') self.is_windows = False + cdef get_mac_address(self, interface="Ethernet"): + addresses = psutil.net_if_addrs() + for interface_name, interface_info in addresses.items(): + if interface_name == interface: + for addr in interface_info: + if addr.family == psutil.AF_LINK: + return addr.address.replace('-', '') + return None + cdef HardwareInfo get_hardware_info(self): if self.is_windows: os_command = ( @@ -49,20 +67,8 @@ cdef class HardwareService: cdef str cpu = lines[0].replace("Name=", "").replace(" ", " ") cdef str gpu = lines[1].replace("Name=", "").replace(" ", " ") cdef str memory = lines[2].replace("TotalVisibleMemorySize=", "").replace(" ", " ") - - # Get MAC address - if self.is_windows: - mac_cmd = "getmac" - else: - mac_cmd = "cat /sys/class/net/*/address" - cdef str mac_address = subprocess.check_output(mac_cmd, shell=True, text=True).splitlines()[0].strip() - + cdef str mac_address = self.get_mac_address() cdef str full_hw_str = f'Azaion_{mac_address}_{cpu}_{gpu}' - hw_hash = self.calc_hash(full_hw_str) - return HardwareInfo(cpu, gpu, memory, mac_address, hw_hash) - cdef str calc_hash(self, str s): - str_bytes = s.encode('utf-8') - hash_bytes = sha384(str_bytes).digest() - cdef str h = base64.b64encode(hash_bytes).decode('utf-8') - return h \ No newline at end of file + hw_hash = Security.calc_hash(full_hw_str) + return HardwareInfo(cpu, gpu, memory, mac_address, hw_hash) diff --git a/Azaion.AI/main.pyx b/Azaion.AI/main.pyx index 4698190..35ad438 100644 --- a/Azaion.AI/main.pyx +++ b/Azaion.AI/main.pyx @@ -35,7 +35,14 @@ cdef class CommandProcessor: self.running = True def start(self): - threading.Thread(target=self.process_queue, daemon=True).start() + while self.running: + try: + command = self.command_queue.get() + print(f'command is : {command}') + model = self.api_client.load_ai_model() + Inference(model, self.on_annotations).run_inference(command) + except Exception as e: + print(f"Error processing queue: {e}") cdef on_message(self, FileCommand cmd): try: @@ -51,17 +58,8 @@ cdef class CommandProcessor: handler.send(annotations) - cdef process_queue(self): - while self.running: - try: - command = self.command_queue.get() - model = self.api_client.load_ai_model() - Inference(model, self.on_annotations).run_inference(command) - except Exception as e: - print(f"Error processing queue: {e}") - cdef process_load(self, FileCommand command): - response = self.api_client.load_file(command) + response = self.api_client.load_bytes(command.filename) handler = self.socket_handler if command.processor_type == ProcessorType.SOCKET else self.rabbit_handler handler.send(response) diff --git a/Azaion.AI/remote_handlers.pxd b/Azaion.AI/remote_handlers.pxd index c11716a..0b15c8f 100644 --- a/Azaion.AI/remote_handlers.pxd +++ b/Azaion.AI/remote_handlers.pxd @@ -15,6 +15,6 @@ cdef class RabbitHandler: cdef object annotation_producer cdef object command_consumer - cdef start(self) cdef send(self, object message) + cdef start(self) cdef close(self) \ No newline at end of file diff --git a/Azaion.AI/remote_handlers.pyx b/Azaion.AI/remote_handlers.pyx index 57167fa..c847c9e 100644 --- a/Azaion.AI/remote_handlers.pyx +++ b/Azaion.AI/remote_handlers.pyx @@ -1,3 +1,4 @@ +import asyncio import json import socket import struct @@ -13,23 +14,21 @@ from processor_command cimport FileCommand, ProcessorType from annotation cimport Annotation cdef class QueueConfig: - cdef str host + cdef str host, cdef int port - cdef str producer_user - cdef str producer_pw - cdef str consumer_user - cdef str consumer_pw + cdef str producer_user, producer_pw, consumer_user, consumer_pw @staticmethod cdef QueueConfig from_json(str json_string): - cdef dict config_dict = json.loads(json_string) + cdef dict config_dict = json.loads(json_string)["QueueConfig"] cdef QueueConfig config = QueueConfig() - config.Host = config_dict["Host"] - config.Port = config_dict["Port"] - config.ProducerUsername = config_dict["ProducerUsername"] - config.ProducerPassword = config_dict["ProducerPassword"] - config.ConsumerUsername = config_dict["ConsumerUsername"] - config.ConsumerPassword = config_dict["ConsumerPassword"] + + config.host = config_dict["Host"] + config.port = config_dict["Port"] + config.producer_user = config_dict["ProducerUsername"] + config.producer_pw = config_dict["ProducerPassword"] + config.consumer_user = config_dict["ConsumerUsername"] + config.consumer_pw = config_dict["ConsumerPassword"] return config cdef class SocketHandler: @@ -75,7 +74,7 @@ cdef class SocketHandler: cdef class RabbitHandler: def __init__(self, ApiClient api_client, object on_message): self.on_message = on_message - cdef str config_str = api_client.load_file(constants.QUEUE_CONFIG_FILENAME).decode(encoding='utf-8') + cdef str config_str = api_client.load_queue_config() queue_config = QueueConfig.from_json(config_str) self.annotation_producer = Producer( host=queue_config.host, @@ -91,8 +90,14 @@ cdef class RabbitHandler: ) cdef start(self): - self.command_consumer.start() - self.command_consumer.subscribe(stream=constants.COMMANDS_QUEUE, callback=self.on_message_inner, + threading.Thread(target=self._run_async, daemon=True).start() + + def _run_async(self): + asyncio.run(self.start_inner()) + + async def start_inner(self): + await self.command_consumer.start() + await self.command_consumer.subscribe(stream=constants.COMMANDS_QUEUE, callback=self.on_message_inner, offset_specification=ConsumerOffsetSpecification(OffsetType.FIRST, None)) # put real offset def on_message_inner(self, message: AMQPMessage, message_context: MessageContext): diff --git a/Azaion.AI/security.pxd b/Azaion.AI/security.pxd new file mode 100644 index 0000000..1e11e70 --- /dev/null +++ b/Azaion.AI/security.pxd @@ -0,0 +1,9 @@ +cdef class Security: + @staticmethod + cdef encrypt_to(input_stream, key) + + @staticmethod + cdef decrypt_to(input_stream, key) + + @staticmethod + cdef calc_hash(str key) \ No newline at end of file diff --git a/Azaion.AI/security.pyx b/Azaion.AI/security.pyx index 1b86764..578bf4a 100644 --- a/Azaion.AI/security.pyx +++ b/Azaion.AI/security.pyx @@ -1,40 +1,48 @@ -# cython: language_level=3 import hashlib import os from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives.hashes import Hash, SHA256 +from hashlib import sha384 +import base64 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() + @staticmethod + cdef encrypt_to(input_stream, key): + cdef bytes 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()) + cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend()) encryptor = cipher.encryptor() - # Read and encrypt in chunks + cdef bytearray res = bytearray() + res.extend(iv) while chunk := input_stream.read(BUFFER_SIZE): - encrypted_data = encryptor.update(chunk) - output_stream.write(encrypted_data) + encrypted_chunk = encryptor.update(chunk) + res.extend(encrypted_chunk) + res.extend(encryptor.finalize()) + return res - final_data = encryptor.finalize() - output_stream.write(final_data) + @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 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() + cdef cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) + cdef decryptor = cipher.decryptor() + cdef bytearray res = bytearray() while chunk := input_stream.read(BUFFER_SIZE): - decrypted_data = decryptor.update(chunk) - output_stream.write(decrypted_data) + decrypted_chunk = decryptor.update(chunk) + res.extend(decrypted_chunk) + res.extend(decryptor.finalize()) + return res - final_data = decryptor.finalize() - output_stream.write(final_data) \ No newline at end of file + @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 diff --git a/Azaion.AI/setup.py b/Azaion.AI/setup.py index c578d2f..139a143 100644 --- a/Azaion.AI/setup.py +++ b/Azaion.AI/setup.py @@ -4,12 +4,12 @@ from Cython.Build import cythonize extensions = [ Extension('constants', ['constants.pyx']), Extension('annotation', ['annotation.pyx']), + Extension('security', ['security.pyx']), Extension('hardware_service', ['hardware_service.pyx'], extra_compile_args=["-g"], extra_link_args=["-g"]), Extension('processor_command', ['processor_command.pyx']), Extension('api_client', ['api_client.pyx']), Extension('inference', ['inference.pyx']), Extension('remote_handlers', ['remote_handlers.pyx']), - Extension('security', ['security.pyx']), Extension('main', ['main.pyx']), ] diff --git a/Azaion.CommonSecurity/Services/HardwareService.cs b/Azaion.CommonSecurity/Services/HardwareService.cs index f2b369d..c09fcd4 100644 --- a/Azaion.CommonSecurity/Services/HardwareService.cs +++ b/Azaion.CommonSecurity/Services/HardwareService.cs @@ -46,17 +46,19 @@ public class HardwareService : IHardwareService memoryStr = $"{Math.Round(memKb / 1024.0 / 1024.0)} Gb"; } + var macAddress = MacAddress(); var hardwareInfo = new HardwareInfo { Memory = memoryStr, CPU = lines.Length > 1 && string.IsNullOrEmpty(lines[1]) - ? "Unknown RAM" - : lines[1], + ? "Unknown CPU" + : lines[1].Trim(), GPU = lines.Length > 2 && string.IsNullOrEmpty(lines[2]) ? "Unknown GPU" - : lines[2] + : lines[2], + MacAddress = macAddress }; - hardwareInfo.Hash = ToHash($"Azaion_{MacAddress()}_{hardwareInfo.CPU}_{hardwareInfo.GPU}"); + hardwareInfo.Hash = ToHash($"Azaion_{macAddress}_{hardwareInfo.CPU}_{hardwareInfo.GPU}"); return hardwareInfo; } catch (Exception ex)