fix id problems with day/winter switch

This commit is contained in:
Alex Bezdieniezhnykh
2025-02-26 22:09:07 +02:00
parent d1af7958f8
commit 58839933fc
28 changed files with 379 additions and 281 deletions
+3 -2
View File
@@ -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)
+28 -10
View File
@@ -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(<str>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(<str>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')
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+1 -1
View File
@@ -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:
+3 -8
View File
@@ -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)
+1 -1
View File
@@ -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())
-12
View File
@@ -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)
-104
View File
@@ -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()
+12
View File
@@ -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)
+20 -4
View File
@@ -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(<bytes>aes_key), modes.CBC(<bytes>iv), backend=default_backend())
cdef cipher = Cipher(algorithms.AES(<bytes>aes_key), modes.CFB(<bytes>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):
+1 -2
View File
@@ -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']),
+2
View File
@@ -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')
@@ -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')