diff --git a/Azaion.Annotator/Annotator.xaml.cs b/Azaion.Annotator/Annotator.xaml.cs index a421c5f..19be86f 100644 --- a/Azaion.Annotator/Annotator.xaml.cs +++ b/Azaion.Annotator/Annotator.xaml.cs @@ -281,7 +281,6 @@ public partial class Annotator return; Dispatcher.Invoke(async () => { - var canvasSize = Editor.RenderSize; var videoSize = _formState.CurrentVideoSize; if (showImage) { @@ -292,13 +291,7 @@ public partial class Annotator videoSize = Editor.RenderSize; } } - foreach (var detection in annotation.Detections) - { - var annClass = _appConfig.AnnotationConfig.DetectionClasses[detection.ClassNumber]; - var canvasLabel = new CanvasLabel(detection, canvasSize, videoSize, detection.Probability); - Editor.CreateDetectionControl(annClass, annotation.Time, canvasLabel); - } - + Editor.CreateDetections(annotation.Time, annotation.Detections, _appConfig.AnnotationConfig.DetectionClasses, videoSize); }); } diff --git a/Azaion.Common/Controls/CanvasEditor.cs b/Azaion.Common/Controls/CanvasEditor.cs index 379a03e..2d7104f 100644 --- a/Azaion.Common/Controls/CanvasEditor.cs +++ b/Azaion.Common/Controls/CanvasEditor.cs @@ -154,8 +154,25 @@ public class CanvasEditor : Canvas private void CanvasMouseUp(object sender, MouseButtonEventArgs e) { if (SelectionState == SelectionState.NewAnnCreating) - CreateDetectionControl(e.GetPosition(this)); - + { + var endPos = e.GetPosition(this); + _newAnnotationRect.Width = 0; + _newAnnotationRect.Height = 0; + var width = Math.Abs(endPos.X - _newAnnotationStartPos.X); + var height = Math.Abs(endPos.Y - _newAnnotationStartPos.Y); + if (width < MIN_SIZE || height < MIN_SIZE) + return; + + var time = GetTimeFunc(); + CreateDetectionControl(CurrentAnnClass, time, new CanvasLabel + { + Width = width, + Height = height, + X = Math.Min(endPos.X, _newAnnotationStartPos.X), + Y = Math.Min(endPos.Y, _newAnnotationStartPos.Y) + }); + } + SelectionState = SelectionState.None; e.Handled = true; } @@ -290,27 +307,18 @@ public class CanvasEditor : Canvas if (diff.Y < 0) SetTop(_newAnnotationRect, currentPos.Y); } - - private void CreateDetectionControl(Point endPos) - { - _newAnnotationRect.Width = 0; - _newAnnotationRect.Height = 0; - var width = Math.Abs(endPos.X - _newAnnotationStartPos.X); - var height = Math.Abs(endPos.Y - _newAnnotationStartPos.Y); - if (width < MIN_SIZE || height < MIN_SIZE) - return; - var time = GetTimeFunc(); - CreateDetectionControl(CurrentAnnClass, time, new CanvasLabel + public void CreateDetections(TimeSpan time, IEnumerable detections, List detectionClasses, Size videoSize) + { + foreach (var detection in detections) { - Width = width, - Height = height, - X = Math.Min(endPos.X, _newAnnotationStartPos.X), - Y = Math.Min(endPos.Y, _newAnnotationStartPos.Y) - }); + var annClass = DetectionClass.FromYoloId(detection.ClassNumber, detectionClasses); + var canvasLabel = new CanvasLabel(detection, RenderSize, videoSize, detection.Probability); + CreateDetectionControl(annClass, time, canvasLabel); + } } - public DetectionControl CreateDetectionControl(DetectionClass annClass, TimeSpan time, CanvasLabel canvasLabel) + private void CreateDetectionControl(DetectionClass annClass, TimeSpan time, CanvasLabel canvasLabel) { var detectionControl = new DetectionControl(annClass, time, AnnotationResizeStart, canvasLabel.Probability) { @@ -323,7 +331,6 @@ public class CanvasEditor : Canvas Children.Add(detectionControl); CurrentDetections.Add(detectionControl); _newAnnotationRect.Fill = new SolidColorBrush(annClass.Color); - return detectionControl; } #endregion diff --git a/Azaion.Common/DTO/DetectionClass.cs b/Azaion.Common/DTO/DetectionClass.cs index ae755e4..00f7cf6 100644 --- a/Azaion.Common/DTO/DetectionClass.cs +++ b/Azaion.Common/DTO/DetectionClass.cs @@ -41,6 +41,15 @@ public class DetectionClass [JsonIgnore] public SolidColorBrush ColorBrush => new(Color); + + public static DetectionClass FromYoloId(int yoloId, List detectionClasses) + { + var cls = yoloId % 20; + var photoMode = (PhotoMode)(yoloId - cls); + var detClass = detectionClasses[cls]; + detClass.PhotoMode = photoMode; + return detClass; + } } public enum PhotoMode diff --git a/Azaion.Common/Extensions/ThrottleExtensions.cs b/Azaion.Common/Extensions/ThrottleExtensions.cs index 03323bd..72cd29e 100644 --- a/Azaion.Common/Extensions/ThrottleExtensions.cs +++ b/Azaion.Common/Extensions/ThrottleExtensions.cs @@ -41,7 +41,7 @@ public static class ThrottleExt await Task.Delay(throttleTime ?? TimeSpan.FromMilliseconds(500), cancellationToken); await func(); } - catch (Exception ex) + catch (Exception) { _taskStates[actionId] = false; } diff --git a/Azaion.Common/Services/AnnotationService.cs b/Azaion.Common/Services/AnnotationService.cs index ce1fa36..c7e9cd8 100644 --- a/Azaion.Common/Services/AnnotationService.cs +++ b/Azaion.Common/Services/AnnotationService.cs @@ -120,6 +120,8 @@ public class AnnotationService : INotificationHandler await SaveAnnotationInner(DateTime.UtcNow, annotation.OriginalMediaName, annotation.Time, annotation.ImageExtension, annotation.Detections.ToList(), SourceEnum.Manual, null, _authProvider.CurrentUser.Role, _authProvider.CurrentUser.Email, generateThumbnail: false, token); + // Manual save from Validators -> Validated -> stream: azaion-annotations-confirm + // AI, Manual save from Operators -> Created -> stream: azaion-annotations private async Task SaveAnnotationInner(DateTime createdDate, string originalMediaName, TimeSpan time, string imageExtension, List detections, SourceEnum source, Stream? stream, RoleEnum userRole, string createdEmail, @@ -132,8 +134,6 @@ public class AnnotationService : INotificationHandler var annotation = await _dbFactory.Run(async db => { var ann = await db.Annotations.FirstOrDefaultAsync(x => x.Name == fName, token: token); - // Manual Save from Validators -> Validated - // otherwise Created status = userRole.IsValidator() && source == SourceEnum.Manual ? AnnotationStatus.Validated : AnnotationStatus.Created; diff --git a/Azaion.CommonSecurity/Services/HardwareService.cs b/Azaion.CommonSecurity/Services/HardwareService.cs index c09fcd4..4820025 100644 --- a/Azaion.CommonSecurity/Services/HardwareService.cs +++ b/Azaion.CommonSecurity/Services/HardwareService.cs @@ -58,7 +58,7 @@ public class HardwareService : IHardwareService : lines[2], MacAddress = macAddress }; - hardwareInfo.Hash = ToHash($"Azaion_{macAddress}_{hardwareInfo.CPU}_{hardwareInfo.GPU}"); + hardwareInfo.Hash = ToHash($"Az|{hardwareInfo.CPU}|{hardwareInfo.GPU}|{macAddress}"); return hardwareInfo; } catch (Exception ex) diff --git a/Azaion.CommonSecurity/Services/PythonResourceLoader.cs b/Azaion.CommonSecurity/Services/PythonResourceLoader.cs index 93ed31c..cd7646c 100644 --- a/Azaion.CommonSecurity/Services/PythonResourceLoader.cs +++ b/Azaion.CommonSecurity/Services/PythonResourceLoader.cs @@ -28,7 +28,7 @@ public class PythonResourceLoader : IResourceLoader, IAuthProvider public PythonResourceLoader(PythonConfig config) { - StartPython(); + //StartPython(); _dealer.Options.Identity = Encoding.UTF8.GetBytes(_clientId.ToString("N")); _dealer.Connect($"tcp://{config.ZeroMqHost}:{config.ZeroMqPort}"); } @@ -81,7 +81,7 @@ public class PythonResourceLoader : IResourceLoader, IAuthProvider { _dealer.SendFrame(RemoteCommand.Serialize(CommandType.Load, new LoadFileData(fileName, folder))); - if (!_dealer.TryReceiveFrameBytes(TimeSpan.FromSeconds(3), out var bytes)) + if (!_dealer.TryReceiveFrameBytes(TimeSpan.FromSeconds(300), out var bytes)) throw new Exception($"Unable to receive {fileName}"); return new MemoryStream(bytes); diff --git a/Azaion.CommonSecurity/Services/Security.cs b/Azaion.CommonSecurity/Services/Security.cs deleted file mode 100644 index a1f8dfb..0000000 --- a/Azaion.CommonSecurity/Services/Security.cs +++ /dev/null @@ -1,82 +0,0 @@ -using System.Runtime.InteropServices; -using System.Security; -using System.Security.Cryptography; -using System.Text; - -namespace Azaion.CommonSecurity.Services; - -public static class Security -{ - private const int BUFFER_SIZE = 524288; // 512 KB buffer size - - public static string ToHash(this string str) => - Convert.ToBase64String(SHA384.HashData(Encoding.UTF8.GetBytes(str))); - - public static string MakeEncryptionKey(string email, string password, string? hardwareHash) => - $"{email}-{password}-{hardwareHash}-#%@AzaionKey@%#---".ToHash(); - - public static SecureString ToSecureString(this string str) - { - var secureString = new SecureString(); - foreach (var c in str.ToCharArray()) - secureString.AppendChar(c); - - return secureString; - } - - public static string? ToRealString(this SecureString value) - { - var valuePtr = IntPtr.Zero; - try - { - valuePtr = Marshal.SecureStringToGlobalAllocUnicode(value); - return Marshal.PtrToStringUni(valuePtr); - } - finally - { - Marshal.ZeroFreeGlobalAllocUnicode(valuePtr); - } - } - - - public static async Task EncryptTo(this Stream stream, Stream toStream, string key, CancellationToken cancellationToken = default) - { - if (stream is { CanRead: false }) throw new ArgumentNullException(nameof(stream)); - if (key is not { Length: > 0 }) throw new ArgumentNullException(nameof(key)); - - using var aes = Aes.Create(); - aes.Key = SHA256.HashData(Encoding.UTF8.GetBytes(key)); - aes.GenerateIV(); - - using var encryptor = aes.CreateEncryptor(aes.Key, aes.IV); - await using var cs = new CryptoStream(toStream, encryptor, CryptoStreamMode.Write, leaveOpen: true); - - // Prepend IV to the encrypted data - await toStream.WriteAsync(aes.IV.AsMemory(0, aes.IV.Length), cancellationToken); - - var buffer = new byte[BUFFER_SIZE]; - int bytesRead; - while ((bytesRead = await stream.ReadAsync(buffer, cancellationToken)) > 0) - await cs.WriteAsync(buffer.AsMemory(0, bytesRead), cancellationToken); - } - - public static async Task DecryptTo(this Stream encryptedStream, Stream toStream, string key, CancellationToken cancellationToken = default) - { - using var aes = Aes.Create(); - aes.Key = SHA256.HashData(Encoding.UTF8.GetBytes(key)); - - // Read the IV from the start of the input stream - var iv = new byte[aes.BlockSize / 8]; - _ = await encryptedStream.ReadAsync(iv, cancellationToken); - aes.IV = iv; - - using var decryptor = aes.CreateDecryptor(aes.Key, aes.IV); - await using var cryptoStream = new CryptoStream(encryptedStream, decryptor, CryptoStreamMode.Read, leaveOpen: true); - - // Read and write in chunks - var buffer = new byte[BUFFER_SIZE]; - int bytesRead; - while ((bytesRead = await cryptoStream.ReadAsync(buffer, cancellationToken)) > 0) - await toStream.WriteAsync(buffer.AsMemory(0, bytesRead), cancellationToken); - } -} diff --git a/Azaion.Dataset/DatasetExplorer.xaml.cs b/Azaion.Dataset/DatasetExplorer.xaml.cs index bd4467a..c5a502b 100644 --- a/Azaion.Dataset/DatasetExplorer.xaml.cs +++ b/Azaion.Dataset/DatasetExplorer.xaml.cs @@ -223,13 +223,7 @@ public partial class DatasetExplorer var time = ann.Time; ExplorerEditor.RemoveAllAnns(); - foreach (var deetection in ann.Detections) - { - var annClass = _annotationConfig.DetectionClassesDict[deetection.ClassNumber]; - var canvasLabel = new CanvasLabel(deetection, ExplorerEditor.RenderSize, ExplorerEditor.RenderSize); - ExplorerEditor.CreateDetectionControl(annClass, time, canvasLabel); - } - + ExplorerEditor.CreateDetections(time, ann.Detections, _annotationConfig.DetectionClasses, ExplorerEditor.RenderSize); ThumbnailLoading = false; } catch (Exception e) diff --git a/Azaion.Inference/api_client.pxd b/Azaion.Inference/api_client.pxd index 4345a27..4b55a63 100644 --- a/Azaion.Inference/api_client.pxd +++ b/Azaion.Inference/api_client.pxd @@ -3,16 +3,17 @@ from credentials cimport Credentials cdef class ApiClient: - cdef public Credentials credentials + cdef Credentials credentials cdef str token, folder, api_url cdef User user - cdef get_encryption_key(self, str hardware_hash) + cdef set_credentials(self, Credentials credentials) cdef login(self) cdef set_token(self, str token) cdef get_user(self) 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 76efcb5..e16270c 100644 --- a/Azaion.Inference/api_client.pyx +++ b/Azaion.Inference/api_client.pyx @@ -1,6 +1,4 @@ import json -import os -import time from http import HTTPStatus from uuid import UUID import jwt @@ -10,7 +8,6 @@ from hardware_service cimport HardwareService, HardwareInfo from security cimport Security from io import BytesIO from user cimport User, RoleEnum -from file_data cimport FileData cdef class ApiClient: """Handles API authentication and downloading of the AI model.""" @@ -19,9 +16,8 @@ cdef class ApiClient: self.user = None self.token = None - cdef get_encryption_key(self, str hardware_hash): - cdef str key = f'{self.credentials.email}-{self.credentials.password}-{hardware_hash}-#%@AzaionKey@%#---' - return Security.calc_hash(key) + cdef set_credentials(self, Credentials credentials): + self.credentials = credentials cdef login(self): response = requests.post(f"{constants.API_URL}/login", @@ -61,6 +57,20 @@ cdef class ApiClient: self.login() return self.user + cdef upload_file(self, str filename, str folder=None): + folder = folder or self.credentials.folder + if self.token is None: + self.login() + url = f"{constants.API_URL}/resources/{folder}" + headers = { "Authorization": f"Bearer {self.token}" } + files = dict(data=open(filename, 'rb')) + try: + r = requests.post(url, headers=headers, files=files, allow_redirects=True) + r.raise_for_status() + print(f"Upload success: {r.status_code}") + except Exception as e: + print(f"Upload fail: {e}") + cdef load_bytes(self, str filename, str folder=None): folder = folder or self.credentials.folder hardware_service = HardwareService() @@ -68,7 +78,6 @@ cdef class ApiClient: if self.token is None: self.login() - url = f"{constants.API_URL}/resources/get/{folder}" headers = { "Authorization": f"Bearer {self.token}", @@ -82,7 +91,6 @@ cdef class ApiClient: "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() headers = { @@ -94,7 +102,8 @@ cdef class ApiClient: if response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR: print('500!') - key = self.get_encryption_key(hardware.hash) + 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) @@ -102,7 +111,16 @@ cdef class ApiClient: return data cdef load_ai_model(self): - return self.load_bytes(constants.AI_MODEL_FILE) + with open(constants.AI_MODEL_FILE_BIG, 'rb') as binary_file: + encrypted_bytes_big = binary_file.read() + encrypted_bytes_small = self.load_bytes(constants.AI_MODEL_FILE_SMALL) + + 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 diff --git a/Azaion.Inference/constants.pxd b/Azaion.Inference/constants.pxd index 6ed8f80..ef9abf3 100644 --- a/Azaion.Inference/constants.pxd +++ b/Azaion.Inference/constants.pxd @@ -6,7 +6,8 @@ cdef str ANNOTATIONS_QUEUE # Name of the annotations queue in rabbit cdef str API_URL # Base URL for the external API cdef str QUEUE_CONFIG_FILENAME # queue config filename to load from api -cdef str AI_MODEL_FILE # AI Model file +cdef str AI_MODEL_FILE_BIG # AI Model file (BIG part) +cdef str AI_MODEL_FILE_SMALL # AI Model file (small part) cdef bytes DONE_SIGNAL cdef int MODEL_BATCH_SIZE diff --git a/Azaion.Inference/constants.pyx b/Azaion.Inference/constants.pyx index eff149d..952956b 100644 --- a/Azaion.Inference/constants.pyx +++ b/Azaion.Inference/constants.pyx @@ -8,7 +8,8 @@ cdef str ANNOTATIONS_QUEUE = "azaion-annotations" cdef str API_URL = "https://api.azaion.com" # Base URL for the external API cdef str QUEUE_CONFIG_FILENAME = "secured-config.json" -cdef str AI_MODEL_FILE = "azaion.onnx" +cdef str AI_MODEL_FILE_BIG = "azaion.onnx.big" +cdef str AI_MODEL_FILE_SMALL = "azaion.onnx.small" cdef bytes DONE_SIGNAL = b"DONE" cdef int MODEL_BATCH_SIZE = 4 diff --git a/Azaion.Inference/hardware_service.pxd b/Azaion.Inference/hardware_service.pxd index 2cd664a..d8e2245 100644 --- a/Azaion.Inference/hardware_service.pxd +++ b/Azaion.Inference/hardware_service.pxd @@ -1,5 +1,5 @@ cdef class HardwareInfo: - cdef str cpu, gpu, memory, mac_address, hash + cdef str cpu, gpu, memory, mac_address cdef to_json_object(self) cdef class HardwareService: diff --git a/Azaion.Inference/hardware_service.pyx b/Azaion.Inference/hardware_service.pyx index 8e73f20..3177907 100644 --- a/Azaion.Inference/hardware_service.pyx +++ b/Azaion.Inference/hardware_service.pyx @@ -1,22 +1,19 @@ import subprocess -from security cimport Security import psutil cdef class HardwareInfo: - def __init__(self, str cpu, str gpu, str memory, str mac_address, str hw_hash): + def __init__(self, str cpu, str gpu, str memory, str mac_address): self.cpu = cpu self.gpu = gpu self.memory = memory 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, + "Memory": self.memory } def __str__(self): @@ -68,7 +65,5 @@ cdef class HardwareService: cdef str gpu = lines[1].replace("Name=", "").replace(" ", " ") cdef str memory = lines[2].replace("TotalVisibleMemorySize=", "").replace(" ", " ") cdef str mac_address = self.get_mac_address() - cdef str full_hw_str = f'Azaion_{mac_address}_{cpu}_{gpu}' - hw_hash = Security.calc_hash(full_hw_str) - return HardwareInfo(cpu, gpu, memory, mac_address, hw_hash) + return HardwareInfo(cpu, gpu, memory, mac_address) diff --git a/Azaion.Inference/main.pyx b/Azaion.Inference/main.pyx index 297bbcb..592ad0c 100644 --- a/Azaion.Inference/main.pyx +++ b/Azaion.Inference/main.pyx @@ -60,7 +60,7 @@ cdef class CommandProcessor: cdef login(self, RemoteCommand command): cdef User user - self.api_client.credentials = Credentials.from_msgpack(command.data) + self.api_client.set_credentials(Credentials.from_msgpack(command.data)) user = self.api_client.get_user() self.remote_handler.send(command.client_id, user.serialize()) diff --git a/Azaion.Inference/secure_model.pxd b/Azaion.Inference/secure_model.pxd deleted file mode 100644 index fe03db6..0000000 --- a/Azaion.Inference/secure_model.pxd +++ /dev/null @@ -1,12 +0,0 @@ -cdef class SecureModelLoader: - cdef: - bytes _model_bytes - str _ramdisk_path - str _temp_file_path - int _disk_size_mb - - cpdef str load_model(self, bytes model_bytes) - cdef str _get_ramdisk_path(self) - cdef void _create_ramdisk(self) - cdef void _store_model(self) - cdef void _cleanup(self) diff --git a/Azaion.Inference/secure_model.pyx b/Azaion.Inference/secure_model.pyx deleted file mode 100644 index ae1babb..0000000 --- a/Azaion.Inference/secure_model.pyx +++ /dev/null @@ -1,104 +0,0 @@ -import os -import platform -import tempfile -from pathlib import Path -from libc.stdio cimport FILE, fopen, fclose, remove -from libc.stdlib cimport free -from libc.string cimport strdup - -cdef class SecureModelLoader: - def __cinit__(self, int disk_size_mb=512): - self._disk_size_mb = disk_size_mb - self._ramdisk_path = None - self._temp_file_path = None - - cpdef str load_model(self, bytes model_bytes): - """Public method to load YOLO model securely.""" - self._model_bytes = model_bytes - self._create_ramdisk() - self._store_model() - return self._temp_file_path - - cdef str _get_ramdisk_path(self): - """Determine the RAM disk path based on the OS.""" - if platform.system() == "Windows": - return "R:\\" - elif platform.system() == "Linux": - return "/mnt/ramdisk" - elif platform.system() == "Darwin": - return "/Volumes/RAMDisk" - else: - raise RuntimeError("Unsupported OS for RAM disk") - - cdef void _create_ramdisk(self): - """Create a RAM disk securely based on the OS.""" - system = platform.system() - - if system == "Windows": - # Create RAM disk via PowerShell - command = f'powershell -Command "subst R: {tempfile.gettempdir()}"' - if os.system(command) != 0: - raise RuntimeError("Failed to create RAM disk on Windows") - self._ramdisk_path = "R:\\" - - elif system == "Linux": - # Use tmpfs for RAM disk - self._ramdisk_path = "/mnt/ramdisk" - if not Path(self._ramdisk_path).exists(): - os.mkdir(self._ramdisk_path) - if os.system(f"mount -t tmpfs -o size={self._disk_size_mb}M tmpfs {self._ramdisk_path}") != 0: - raise RuntimeError("Failed to create RAM disk on Linux") - - elif system == "Darwin": - # Use hdiutil for macOS RAM disk - block_size = 2048 # 512-byte blocks * 2048 = 1MB - num_blocks = self._disk_size_mb * block_size - result = os.popen(f"hdiutil attach -nomount ram://{num_blocks}").read().strip() - if result: - self._ramdisk_path = "/Volumes/RAMDisk" - os.system(f"diskutil eraseVolume HFS+ RAMDisk {result}") - else: - raise RuntimeError("Failed to create RAM disk on macOS") - - cdef void _store_model(self): - """Write model securely to the RAM disk.""" - cdef char* temp_path - cdef FILE* cfile - - with tempfile.NamedTemporaryFile( - dir=self._ramdisk_path, suffix='.pt', delete=False - ) as tmp_file: - tmp_file.write(self._model_bytes) - self._temp_file_path = tmp_file.name - - encoded_path = self._temp_file_path.encode('utf-8') - temp_path = strdup(encoded_path) - with nogil: - cfile = fopen(temp_path, "rb") - if cfile == NULL: - raise IOError(f"Could not open {self._temp_file_path}") - fclose(cfile) - - cdef void _cleanup(self): - """Remove the model file and unmount RAM disk securely.""" - cdef char* c_path - if self._temp_file_path: - c_path = strdup(os.fsencode(self._temp_file_path)) - with nogil: - remove(c_path) - free(c_path) - self._temp_file_path = None - - # Unmount RAM disk based on OS - if self._ramdisk_path: - if platform.system() == "Windows": - os.system("subst R: /D") - elif platform.system() == "Linux": - os.system(f"umount {self._ramdisk_path}") - elif platform.system() == "Darwin": - os.system("hdiutil detach /Volumes/RAMDisk") - self._ramdisk_path = None - - def __dealloc__(self): - """Ensure cleanup when the object is deleted.""" - self._cleanup() diff --git a/Azaion.Inference/security.pxd b/Azaion.Inference/security.pxd index 1e11e70..4dfa64c 100644 --- a/Azaion.Inference/security.pxd +++ b/Azaion.Inference/security.pxd @@ -1,3 +1,6 @@ +from credentials cimport Credentials +from hardware_service cimport HardwareInfo + cdef class Security: @staticmethod cdef encrypt_to(input_stream, key) @@ -5,5 +8,14 @@ cdef class Security: @staticmethod cdef decrypt_to(input_stream, key) + @staticmethod + cdef get_hw_hash(HardwareInfo hardware) + + @staticmethod + cdef get_api_encryption_key(Credentials credentials, str hardware_hash) + + @staticmethod + cdef get_model_encryption_key() + @staticmethod cdef calc_hash(str key) \ No newline at end of file diff --git a/Azaion.Inference/security.pyx b/Azaion.Inference/security.pyx index 61f818c..6518e18 100644 --- a/Azaion.Inference/security.pyx +++ b/Azaion.Inference/security.pyx @@ -2,6 +2,8 @@ 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 @@ -24,14 +26,14 @@ cdef class Security: encrypted_chunk = encryptor.update(chunk) res.extend(encrypted_chunk) res.extend(encryptor.finalize()) - return res + 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(aes_key), modes.CBC(iv), backend=default_backend()) + cdef cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend()) cdef decryptor = cipher.decryptor() cdef bytearray res = bytearray() @@ -40,8 +42,22 @@ cdef class Security: res.extend(decrypted_chunk) res.extend(decryptor.finalize()) - unpadder = padding.PKCS7(128).unpadder() # AES block size is 128 bits (16 bytes) - return unpadder.update(res) + unpadder.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): diff --git a/Azaion.Inference/setup.py b/Azaion.Inference/setup.py index 33802cb..6634a7b 100644 --- a/Azaion.Inference/setup.py +++ b/Azaion.Inference/setup.py @@ -7,13 +7,12 @@ extensions = [ Extension('annotation', ['annotation.pyx']), Extension('credentials', ['credentials.pyx']), Extension('file_data', ['file_data.pyx']), - Extension('security', ['security.pyx']), Extension('hardware_service', ['hardware_service.pyx'], extra_compile_args=["-g"], extra_link_args=["-g"]), + Extension('security', ['security.pyx']), Extension('remote_command', ['remote_command.pyx']), Extension('remote_command_handler', ['remote_command_handler.pyx']), Extension('user', ['user.pyx']), Extension('api_client', ['api_client.pyx']), - Extension('secure_model', ['secure_model.pyx']), Extension('ai_config', ['ai_config.pyx']), Extension('inference', ['inference.pyx'], include_dirs=[np.get_include()]), Extension('main', ['main.pyx']), diff --git a/Azaion.Inference/start.spec b/Azaion.Inference/start.spec index 7d997c0..7ce92bd 100644 --- a/Azaion.Inference/start.spec +++ b/Azaion.Inference/start.spec @@ -4,6 +4,8 @@ from PyInstaller.utils.hooks import collect_all datas = [] binaries = [] hiddenimports = ['constants', 'annotation', 'credentials', 'file_data', 'user', 'security', 'secure_model', 'api_client', 'hardware_service', 'remote_command', 'ai_config', 'inference', 'remote_command_handler'] +tmp_ret = collect_all('pyyaml') +datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] tmp_ret = collect_all('jwt') datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2] tmp_ret = collect_all('requests') diff --git a/Azaion.Inference/test/test_download_large_file.py b/Azaion.Inference/test/test_download_large_file.py new file mode 100644 index 0000000..fc7f2ba --- /dev/null +++ b/Azaion.Inference/test/test_download_large_file.py @@ -0,0 +1,236 @@ +import json +import subprocess +import psutil +import hashlib +import base64 +import os +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes +import requests +from io import BytesIO +from http import HTTPStatus +import random + + +BUFFER_SIZE = 64 * 1024 # 64 KB +API_URL = "https://api.azaion.com" + + +class HWInfo: + def __init__(self, cpu, gpu, memory, mac_address, hw_hash): + self.cpu = cpu + self.gpu = gpu + self.memory = memory + self.mac_address = mac_address + self.hash = hw_hash + + def 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}' + + @staticmethod + def encrypt_to(input_stream, key): + 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()) + encryptor = cipher.encryptor() + + 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 res + + @staticmethod + def decrypt_to(input_stream, key): + aes_key = hashlib.sha256(key.encode('utf-8')).digest() + iv = input_stream.read(16) + + cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend()) + decryptor = cipher.decryptor() + + res = bytearray() + while chunk := input_stream.read(BUFFER_SIZE): + decrypted_chunk = decryptor.update(chunk) + res.extend(decrypted_chunk) + res.extend(decryptor.finalize()) + + unpadder = padding.PKCS7(128).unpadder() # AES block size is 128 bits (16 bytes) + return unpadder.update(res) + unpadder.finalize() + + @staticmethod + def calc_hash(key): + str_bytes = key.encode('utf-8') + hash_bytes = hashlib.sha384(str_bytes).digest() + h = base64.b64encode(hash_bytes).decode('utf-8') + return h + + +class HWService: + def __init__(self): + try: + res = subprocess.check_output("ver", shell=True).decode('utf-8') + if "Microsoft Windows" in res: + self.is_windows = True + else: + self.is_windows = False + except Exception: + print('Error during os type checking') + self.is_windows = False + + def 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 + + def get_hardware_info(self): + if self.is_windows: + os_command = ( + "wmic CPU get Name /Value && " + "wmic path Win32_VideoController get Name /Value && " + "wmic OS get TotalVisibleMemorySize /Value" + ) + else: + os_command = ( + "/bin/bash -c \" lscpu | grep 'Model name:' | cut -d':' -f2 && " + "lspci | grep VGA | cut -d':' -f3 && " + "free -g | grep Mem: | awk '{print $2}' && \"" + ) + result = subprocess.check_output(os_command, shell=True).decode('utf-8') + lines = [line.strip() for line in result.splitlines() if line.strip()] + + cpu = lines[0].replace("Name=", "").replace(" ", " ") + gpu = lines[1].replace("Name=", "").replace(" ", " ") + memory = lines[2].replace("TotalVisibleMemorySize=", "").replace(" ", " ") + mac_address = self.get_mac_address() + full_hw_str = f'Azaion_{mac_address}_{cpu}_{gpu}' + + hw_hash = HWInfo.calc_hash(full_hw_str) + return HWInfo(cpu, gpu, memory, mac_address, hw_hash) + +class Credentials: + + def __init__(self, email, password, folder): + self.email = email + self.password = password + self.folder = folder + +class Api: + + def __init__(self, credentials): + self.token = None + self.credentials = credentials + + @staticmethod + def create_file(filename, size_mb=1): # chunk_size_kb is now configurable + size_bytes = size_mb * 1024 * 1024 + chunk_size = 1024 * 1024 # 1mb chunk size + bytes_written = 0 + sha256_hash = hashlib.sha256() # init hash + + with open(filename, 'wb') as f: + while bytes_written < size_bytes: + write_size = min(chunk_size, size_bytes - bytes_written) + random_bytes = os.urandom(write_size) + f.write(random_bytes) + bytes_written += write_size + sha256_hash.update(random_bytes) + return sha256_hash.hexdigest() + + def login(self): + response = requests.post(f"{API_URL}/login", + json={"email": self.credentials.email, "password": self.credentials.password}) + response.raise_for_status() + token = response.json()["token"] + self.token = token + + @staticmethod + def get_sha256(data_bytes): + sha256 = hashlib.sha256() + sha256.update(data_bytes) + return sha256.hexdigest() + + def upload_file(self, filename, folder = None): + folder = folder or self.credentials.folder + if self.token is None: + self.login() + url = f"{API_URL}/resources/{folder}" + headers = {"Authorization": f"Bearer {self.token}"} + files = dict(data=open(filename, 'rb')) + try: + r = requests.post(url, headers=headers, files=files, allow_redirects=True) + r.raise_for_status() + print(f"Upload success: {r.status_code}") + except Exception as e: + print(f"Upload fail: {e}") + + def get_encryption_key(self, hardware_hash): + key = f'{self.credentials.email}-{self.credentials.password}-{hardware_hash}-#%@AzaionKey@%#---' + return HWInfo.calc_hash(key) + + def load_bytes(self, filename, folder = None): + folder = folder or self.credentials.folder + hardware_service = HWService() + hardware = hardware_service.get_hardware_info() + + if self.token is None: + self.login() + url = f"{API_URL}/resources/get/{folder}" + headers = { + "Authorization": f"Bearer {self.token}", + "Content-Type": "application/json" + } + + payload = json.dumps( + { + "password": self.credentials.password, + "hardware": hardware.to_json_object(), + "fileName": filename + }, indent=4) + response = requests.post(url, data=payload, headers=headers, stream=True, timeout=20) + if response.status_code == HTTPStatus.UNAUTHORIZED or response.status_code == HTTPStatus.FORBIDDEN: + self.login() + 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) + + stream = BytesIO(response.raw.read()) + data = HWInfo.decrypt_to(stream, key) + print(f'Downloaded file: {filename}, {len(data)} bytes') + return data + + +credentials = Credentials('admin@azaion.com', 'Az@1on1000Odm$n', 'stage') +api = Api(credentials) +file = 'file1' +sha256_init = Api.create_file(file, size_mb=100) +api.upload_file(file) + +file_bytes = api.load_bytes(file) +print(f'received: {len(file_bytes)/1024} kb') +sha256_downloaded = Api.get_sha256(file_bytes) +print(f'{sha256_init}: sha256 initial file') +print(f'{sha256_downloaded}: sha256 downloaded file') diff --git a/Azaion.Suite.sln b/Azaion.Suite.sln index eb541b0..5380c32 100644 --- a/Azaion.Suite.sln +++ b/Azaion.Suite.sln @@ -25,6 +25,9 @@ EndProject Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Build", "Build", "{CF141A48-8002-4006-81CF-6B85AE5B0B5F}" ProjectSection(SolutionItems) = preProject build\publish.cmd = build\publish.cmd + build\download.py = build\download.py + build\requirements.txt = build\requirements.txt + build\build_downloader.cmd = build\build_downloader.cmd EndProjectSection EndProject Global diff --git a/Azaion.Suite/App.xaml.cs b/Azaion.Suite/App.xaml.cs index cc2429f..332fa60 100644 --- a/Azaion.Suite/App.xaml.cs +++ b/Azaion.Suite/App.xaml.cs @@ -51,8 +51,8 @@ public partial class App private readonly List _encryptedResources = [ - "Azaion.Annotator", - "Azaion.Dataset" + // "Azaion.Annotator", + // "Azaion.Dataset" ]; private static PythonConfig ReadPythonConfig() @@ -83,11 +83,11 @@ public partial class App { new ConfigUpdater().CheckConfig(); var login = new Login(); - var config = ReadPythonConfig(); - _resourceLoader = new PythonResourceLoader(config); + var pythonConfig = ReadPythonConfig(); + _resourceLoader = new PythonResourceLoader(pythonConfig); login.CredentialsEntered += (_, credentials) => { - credentials.Folder = config.ResourcesFolder; + credentials.Folder = pythonConfig.ResourcesFolder; _resourceLoader.Login(credentials); _securedConfig = _resourceLoader.LoadFileFromPython("secured-config.json"); @@ -152,6 +152,7 @@ public partial class App services.AddSingleton(); services.Configure(context.Configuration); + services.ConfigureSection(context.Configuration); services.ConfigureSection(context.Configuration); services.ConfigureSection(context.Configuration); services.ConfigureSection(context.Configuration); diff --git a/Azaion.Suite/Azaion.Suite.csproj b/Azaion.Suite/Azaion.Suite.csproj index b15d7b3..e042c2f 100644 --- a/Azaion.Suite/Azaion.Suite.csproj +++ b/Azaion.Suite/Azaion.Suite.csproj @@ -31,8 +31,8 @@ - - + + @@ -55,8 +55,8 @@ - - + + diff --git a/Azaion.Suite/Login.xaml b/Azaion.Suite/Login.xaml index 0978c07..64b3bce 100644 --- a/Azaion.Suite/Login.xaml +++ b/Azaion.Suite/Login.xaml @@ -74,6 +74,7 @@ BorderBrush="DimGray" BorderThickness="0,0,0,1" HorizontalAlignment="Left" + Text="admin@azaion.com" /> + HorizontalAlignment="Left" + Password="Az@1on1000Odm$n"/>