fix auth, decryption, api interaction

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)
 

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(<str>constants.TOKEN_FILE, 'w') as file:
-                file.write(self.token)
+        with open(<str>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, <bint>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')

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)

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)
+    cdef get_mac_address(self, interface=*)
+    cdef HardwareInfo get_hardware_info(self)

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
+        hw_hash = Security.calc_hash(full_hw_str)
+        return HardwareInfo(cpu, gpu, memory, mac_address, hw_hash)

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)
 

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)

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(<str>json_string)
+        cdef dict config_dict = json.loads(<str>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=<str>queue_config.host,
@@ -91,8 +90,14 @@ cdef class RabbitHandler:
         )
 
     cdef start(self):
-        self.command_consumer.start()
-        self.command_consumer.subscribe(stream=<str>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=<str>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):

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)

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(<bytes>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(<bytes>aes_key), modes.CBC(<bytes>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)
+    @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

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']),
 ]
 

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)