mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 16:06:30 +00:00
add azaion loader
This commit is contained in:
Alex Bezdieniezhnykh
@@ -0,0 +1,21 @@
|
||||
from user cimport User
|
||||
from credentials cimport Credentials
|
||||
from cdn_manager cimport CDNManager
|
||||
|
||||
|
||||
cdef class ApiClient:
|
||||
cdef Credentials credentials
|
||||
cdef CDNManager cdn_manager
|
||||
cdef str token, folder, api_url
|
||||
cdef User user
|
||||
|
||||
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, bytes resource, str folder)
|
||||
cdef load_big_file_cdn(self, str folder, str big_part)
|
||||
cdef load_big_small_resource(self, str resource_name, str folder, str key)
|
||||
cdef upload_big_small_resource(self, bytes resource, str resource_name, str folder, str key)
|
||||
@@ -0,0 +1,177 @@
|
||||
import json
|
||||
import os
|
||||
from http import HTTPStatus
|
||||
from os import path
|
||||
from uuid import UUID
|
||||
import jwt
|
||||
import requests
|
||||
cimport constants
|
||||
import yaml
|
||||
from requests import HTTPError
|
||||
|
||||
from cdn_manager cimport CDNManager, CDNCredentials
|
||||
from hardware_service cimport HardwareService
|
||||
from security cimport Security
|
||||
from user cimport User, RoleEnum
|
||||
|
||||
cdef class ApiClient:
|
||||
"""Handles API authentication and downloading of the AI model."""
|
||||
def __init__(self, str api_url):
|
||||
self.credentials = None
|
||||
self.user = None
|
||||
self.token = None
|
||||
self.cdn_manager = None
|
||||
self.api_url = api_url
|
||||
|
||||
cdef set_credentials(self, Credentials credentials):
|
||||
self.credentials = credentials
|
||||
if self.cdn_manager is not None:
|
||||
return
|
||||
|
||||
yaml_bytes = self.load_bytes(constants.CDN_CONFIG, <str>'')
|
||||
yaml_config = yaml.safe_load(yaml_bytes)
|
||||
creds = CDNCredentials(yaml_config["host"],
|
||||
yaml_config["downloader_access_key"],
|
||||
yaml_config["downloader_access_secret"],
|
||||
yaml_config["uploader_access_key"],
|
||||
yaml_config["uploader_access_secret"])
|
||||
|
||||
self.cdn_manager = CDNManager(creds)
|
||||
|
||||
cdef login(self):
|
||||
response = None
|
||||
try:
|
||||
response = requests.post(f"{self.api_url}/login",
|
||||
json={"email": self.credentials.email, "password": self.credentials.password})
|
||||
response.raise_for_status()
|
||||
token = response.json()["token"]
|
||||
self.set_token(token)
|
||||
except HTTPError as e:
|
||||
print(response.json())
|
||||
if response.status_code == HTTPStatus.CONFLICT:
|
||||
res = response.json()
|
||||
raise Exception(res['Message'])
|
||||
|
||||
|
||||
cdef set_token(self, str token):
|
||||
self.token = token
|
||||
claims = jwt.decode(token, options={"verify_signature": False})
|
||||
|
||||
try:
|
||||
id = str(UUID(claims.get("nameid", "")))
|
||||
except ValueError:
|
||||
raise ValueError("Invalid GUID format in claims")
|
||||
|
||||
email = claims.get("unique_name", "")
|
||||
|
||||
role_str = claims.get("role", "")
|
||||
if role_str == "ApiAdmin":
|
||||
role = RoleEnum.ApiAdmin
|
||||
elif role_str == "Admin":
|
||||
role = RoleEnum.Admin
|
||||
elif role_str == "ResourceUploader":
|
||||
role = RoleEnum.ResourceUploader
|
||||
elif role_str == "Validator":
|
||||
role = RoleEnum.Validator
|
||||
elif role_str == "Operator":
|
||||
role = RoleEnum.Operator
|
||||
else:
|
||||
role = RoleEnum.NONE
|
||||
self.user = User(id, email, role)
|
||||
|
||||
cdef get_user(self):
|
||||
if self.user is None:
|
||||
self.login()
|
||||
return self.user
|
||||
|
||||
cdef upload_file(self, str filename, bytes resource, str folder):
|
||||
if self.token is None:
|
||||
self.login()
|
||||
url = f"{self.api_url}/resources/{folder}"
|
||||
headers = { "Authorization": f"Bearer {self.token}" }
|
||||
files = {'data': (filename, resource)}
|
||||
try:
|
||||
r = requests.post(url, headers=headers, files=files, allow_redirects=True)
|
||||
r.raise_for_status()
|
||||
constants.log(f"Uploaded {filename} to {self.api_url}/{folder} successfully: {r.status_code}.")
|
||||
except Exception as e:
|
||||
constants.log(f"Upload fail: {e}")
|
||||
|
||||
cdef load_bytes(self, str filename, str folder):
|
||||
hardware_service = HardwareService()
|
||||
cdef str hardware = hardware_service.get_hardware_info()
|
||||
if self.token is None:
|
||||
self.login()
|
||||
url = f"{self.api_url}/resources/get/{folder}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.token}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = json.dumps(
|
||||
{
|
||||
"password": self.credentials.password,
|
||||
"hardware": hardware,
|
||||
"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 = {
|
||||
"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!')
|
||||
|
||||
hw_hash = Security.get_hw_hash(hardware)
|
||||
key = Security.get_api_encryption_key(self.credentials, hw_hash)
|
||||
|
||||
resp_bytes = response.raw.read()
|
||||
data = Security.decrypt_to(resp_bytes, key)
|
||||
constants.log(<str>f'Downloaded file: {filename}, {len(data)} bytes')
|
||||
return data
|
||||
|
||||
cdef load_big_file_cdn(self, str folder, str big_part):
|
||||
print(f'downloading file {folder}\\{big_part} from cdn...')
|
||||
if self.cdn_manager.download(folder, big_part):
|
||||
with open(path.join(<str> folder, big_part), 'rb') as binary_file:
|
||||
encrypted_bytes_big = binary_file.read()
|
||||
else:
|
||||
return None
|
||||
|
||||
cdef load_big_small_resource(self, str resource_name, str folder, str key):
|
||||
cdef str big_part = f'{resource_name}.big'
|
||||
cdef str small_part = f'{resource_name}.small'
|
||||
|
||||
encrypted_bytes_small = self.load_bytes(small_part, folder)
|
||||
|
||||
print(f'checking on existence for {folder}\\{big_part}')
|
||||
if os.path.exists(os.path.join(<str> folder, big_part)):
|
||||
with open(path.join(<str> folder, big_part), 'rb') as binary_file:
|
||||
local_bytes_big = binary_file.read()
|
||||
print(f'local file {folder}\\{big_part} is found!')
|
||||
try:
|
||||
return Security.decrypt_to(encrypted_bytes_small + local_bytes_big, key)
|
||||
except Exception as ex:
|
||||
print('Local file doesnt match with api file, old version')
|
||||
|
||||
remote_bytes_big = self.load_big_file_cdn(folder, big_part)
|
||||
return Security.decrypt_to(encrypted_bytes_small + remote_bytes_big, key)
|
||||
|
||||
cdef upload_big_small_resource(self, bytes resource, str resource_name, str folder, str key):
|
||||
cdef str big_part_name = f'{resource_name}.big'
|
||||
cdef str small_part_name = f'{resource_name}.small'
|
||||
|
||||
resource_encrypted = Security.encrypt_to(<bytes>resource, key)
|
||||
part_small_size = min(constants.SMALL_SIZE_KB * 1024, int(0.3 * len(resource_encrypted)))
|
||||
part_small = resource_encrypted[:part_small_size] # slice bytes for part1
|
||||
|
||||
part_big = resource_encrypted[part_small_size:]
|
||||
|
||||
self.cdn_manager.upload(folder, <str>big_part_name, part_big)
|
||||
with open(path.join(<str>folder, <str>big_part_name), 'wb') as f:
|
||||
f.write(part_big)
|
||||
self.upload_file(small_part_name, part_small, folder)
|
||||
@@ -0,0 +1,24 @@
|
||||
echo Build Cython app
|
||||
set CURRENT_DIR=%cd%
|
||||
|
||||
REM Change to the parent directory of the current location
|
||||
cd /d %~dp0
|
||||
|
||||
echo remove dist folder:
|
||||
if exist dist rmdir dist /s /q
|
||||
if exist build rmdir build /s /q
|
||||
|
||||
echo install python and dependencies
|
||||
if not exist venv (
|
||||
python -m venv venv
|
||||
)
|
||||
|
||||
venv\Scripts\python -m pip install --upgrade pip
|
||||
venv\Scripts\pip install -r requirements.txt
|
||||
venv\Scripts\pip install --upgrade pyinstaller pyinstaller-hooks-contrib
|
||||
|
||||
venv\Scripts\python setup.py build_ext --inplace
|
||||
|
||||
echo install azaion-loader
|
||||
venv\Scripts\pyinstaller --name=azaion-loader ^
|
||||
--collect-all boto3 ^
|
||||
@@ -0,0 +1,14 @@
|
||||
cdef class CDNCredentials:
|
||||
cdef str host
|
||||
cdef str downloader_access_key
|
||||
cdef str downloader_access_secret
|
||||
cdef str uploader_access_key
|
||||
cdef str uploader_access_secret
|
||||
|
||||
cdef class CDNManager:
|
||||
cdef CDNCredentials creds
|
||||
cdef object download_client
|
||||
cdef object upload_client
|
||||
|
||||
cdef upload(self, str bucket, str filename, bytes file_bytes)
|
||||
cdef download(self, str bucket, str filename)
|
||||
@@ -0,0 +1,41 @@
|
||||
import io
|
||||
import boto3
|
||||
|
||||
|
||||
cdef class CDNCredentials:
|
||||
def __init__(self, host, downloader_access_key, downloader_access_secret, uploader_access_key, uploader_access_secret):
|
||||
self.host = host
|
||||
self.downloader_access_key = downloader_access_key
|
||||
self.downloader_access_secret = downloader_access_secret
|
||||
self.uploader_access_key = uploader_access_key
|
||||
self.uploader_access_secret = uploader_access_secret
|
||||
|
||||
|
||||
cdef class CDNManager:
|
||||
def __init__(self, CDNCredentials credentials):
|
||||
|
||||
self.creds = credentials
|
||||
self.download_client = boto3.client('s3', endpoint_url=self.creds.host,
|
||||
aws_access_key_id=self.creds.downloader_access_key,
|
||||
aws_secret_access_key=self.creds.downloader_access_secret)
|
||||
self.upload_client = boto3.client('s3', endpoint_url=self.creds.host,
|
||||
aws_access_key_id=self.creds.uploader_access_key,
|
||||
aws_secret_access_key=self.creds.uploader_access_secret)
|
||||
|
||||
cdef upload(self, str bucket, str filename, bytes file_bytes):
|
||||
try:
|
||||
self.upload_client.upload_fileobj(io.BytesIO(file_bytes), bucket, filename)
|
||||
print(f'uploaded {filename} ({len(file_bytes)} bytes) to the {bucket}')
|
||||
return True
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return False
|
||||
|
||||
cdef download(self, str folder, str filename):
|
||||
try:
|
||||
self.download_client.download_file(folder, filename, f'{folder}\\{filename}')
|
||||
print(f'downloaded {filename} from the {folder} to current folder')
|
||||
return True
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return False
|
||||
@@ -0,0 +1,17 @@
|
||||
cdef str CONFIG_FILE # Port for the zmq
|
||||
|
||||
cdef int QUEUE_MAXSIZE # Maximum size of the command queue
|
||||
cdef str COMMANDS_QUEUE # Name of the commands queue in rabbit
|
||||
cdef str ANNOTATIONS_QUEUE # Name of the annotations queue in rabbit
|
||||
|
||||
cdef str QUEUE_CONFIG_FILENAME # queue config filename to load from api
|
||||
|
||||
cdef str AI_ONNX_MODEL_FILE
|
||||
|
||||
cdef str CDN_CONFIG
|
||||
cdef str MODELS_FOLDER
|
||||
|
||||
cdef int SMALL_SIZE_KB
|
||||
|
||||
|
||||
cdef log(str log_message, bytes client_id=*)
|
||||
@@ -0,0 +1,16 @@
|
||||
import time
|
||||
|
||||
cdef str CONFIG_FILE = "config.yaml" # Port for the zmq
|
||||
|
||||
cdef str QUEUE_CONFIG_FILENAME = "secured-config.json"
|
||||
cdef str AI_ONNX_MODEL_FILE = "azaion.onnx"
|
||||
|
||||
cdef str CDN_CONFIG = "cdn.yaml"
|
||||
cdef str MODELS_FOLDER = "models"
|
||||
|
||||
cdef int SMALL_SIZE_KB = 3
|
||||
|
||||
cdef log(str log_message, bytes client_id=None):
|
||||
local_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
|
||||
client_str = '' if client_id is None else f' {client_id}'
|
||||
print(f'[{local_time}{client_str}]: {log_message}')
|
||||
@@ -0,0 +1,6 @@
|
||||
cdef class Credentials:
|
||||
cdef public str email
|
||||
cdef public str password
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data)
|
||||
@@ -0,0 +1,15 @@
|
||||
from msgpack import unpackb
|
||||
|
||||
cdef class Credentials:
|
||||
|
||||
def __init__(self, str email, str password):
|
||||
self.email = email
|
||||
self.password = password
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data):
|
||||
unpacked = unpackb(data, strict_map_key=False)
|
||||
return Credentials(
|
||||
unpacked.get("Email"),
|
||||
unpacked.get("Password"))
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
cdef class FileData:
|
||||
cdef public str folder
|
||||
cdef public str filename
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data)
|
||||
@@ -0,0 +1,14 @@
|
||||
from msgpack import unpackb
|
||||
|
||||
cdef class FileData:
|
||||
|
||||
def __init__(self, str folder, str filename):
|
||||
self.folder = folder
|
||||
self.filename = filename
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data):
|
||||
unpacked = unpackb(data, strict_map_key=False)
|
||||
return FileData(
|
||||
unpacked.get("Folder"),
|
||||
unpacked.get("Filename"))
|
||||
@@ -0,0 +1,5 @@
|
||||
cdef class HardwareService:
|
||||
|
||||
@staticmethod
|
||||
cdef has_nvidia_gpu()
|
||||
cdef str get_hardware_info(self)
|
||||
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
import subprocess
|
||||
import pynvml
|
||||
|
||||
|
||||
cdef class HardwareService:
|
||||
|
||||
@staticmethod
|
||||
cdef has_nvidia_gpu():
|
||||
try:
|
||||
pynvml.nvmlInit()
|
||||
device_count = pynvml.nvmlDeviceGetCount()
|
||||
|
||||
if device_count > 0:
|
||||
print(f"Found NVIDIA GPU(s).")
|
||||
return True
|
||||
else:
|
||||
print("No NVIDIA GPUs found by NVML.")
|
||||
return False
|
||||
|
||||
except pynvml.NVMLError as error:
|
||||
print(f"Failed to find NVIDIA GPU")
|
||||
return False
|
||||
finally:
|
||||
try:
|
||||
pynvml.nvmlShutdown()
|
||||
except:
|
||||
print('Failed to shutdown pynvml cause probably no NVidia GPU')
|
||||
pass
|
||||
|
||||
cdef str get_hardware_info(self):
|
||||
if os.name == 'nt': # windows
|
||||
os_command = (
|
||||
"powershell -Command \""
|
||||
"Get-CimInstance -ClassName Win32_Processor | Select-Object -ExpandProperty Name | Write-Output; "
|
||||
"Get-CimInstance -ClassName Win32_VideoController | Select-Object -ExpandProperty Name | Write-Output; "
|
||||
"Get-CimInstance -ClassName Win32_OperatingSystem | Select-Object -ExpandProperty TotalVisibleMemorySize | Write-Output; "
|
||||
"(Get-Disk | Where-Object {$_.IsSystem -eq $true}).SerialNumber"
|
||||
"\""
|
||||
)
|
||||
else:
|
||||
os_command = (
|
||||
"lscpu | grep 'Model name:' | cut -d':' -f2 && "
|
||||
"lspci | grep VGA | cut -d':' -f3 && "
|
||||
"free -k | awk '/^Mem:/ {print $2}' && "
|
||||
"cat /sys/block/sda/device/vpd_pg80 2>/dev/null || cat /sys/block/sda/device/serial 2>/dev/null"
|
||||
)
|
||||
|
||||
result = subprocess.check_output(os_command, shell=True).decode('utf-8', errors='ignore')
|
||||
lines = [line.replace(" ", " ").replace("Name=", "").strip('\x00\x14 \t\n\r\v\f') for line in result.splitlines() if line.strip()]
|
||||
|
||||
cdef str cpu = lines[0]
|
||||
cdef str gpu = lines[1]
|
||||
# could be multiple gpus
|
||||
|
||||
len_lines = len(lines)
|
||||
cdef str memory = lines[len_lines-2].replace("TotalVisibleMemorySize=", "").replace(" ", " ")
|
||||
cdef str drive_serial = lines[len_lines-1]
|
||||
|
||||
cdef str res = f'CPU: {cpu}. GPU: {gpu}. Memory: {memory}. DriveSerial: {drive_serial}'
|
||||
return res
|
||||
@@ -0,0 +1,61 @@
|
||||
import threading
|
||||
from threading import Thread
|
||||
import traceback
|
||||
from credentials cimport Credentials
|
||||
from remote_command cimport RemoteCommand, CommandType
|
||||
from remote_command_handler cimport RemoteCommandHandler
|
||||
from file_data cimport FileData
|
||||
from api_client cimport ApiClient
|
||||
|
||||
cdef class CommandProcessor:
|
||||
cdef RemoteCommandHandler remote_handler
|
||||
cdef ApiClient api_client
|
||||
cdef bint running
|
||||
cdef object shutdown_event
|
||||
|
||||
def __init__(self, int zmq_port, str api_url):
|
||||
self.api_client = ApiClient(api_url)
|
||||
self.shutdown_event = threading.Event()
|
||||
self.remote_handler = RemoteCommandHandler(zmq_port, self.on_command)
|
||||
self.remote_handler.start()
|
||||
self.running = True
|
||||
|
||||
def start(self):
|
||||
while self.running:
|
||||
try:
|
||||
while not self.shutdown_event.is_set():
|
||||
self.shutdown_event.wait(timeout=1.0)
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
print('EXIT!')
|
||||
|
||||
cdef on_command(self, RemoteCommand command):
|
||||
try:
|
||||
if command.command_type == CommandType.LOGIN:
|
||||
self.api_client.set_credentials(Credentials.from_msgpack(command.data))
|
||||
elif command.command_type == CommandType.LOAD:
|
||||
self.load_file(command)
|
||||
elif command.command_type == CommandType.EXIT:
|
||||
t = Thread(target=self.stop) # non-block worker:
|
||||
t.start()
|
||||
else:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Error handling client: {e}")
|
||||
|
||||
cdef load_file(self, RemoteCommand command):
|
||||
cdef RemoteCommand response
|
||||
cdef FileData file_data
|
||||
cdef bytes file_bytes
|
||||
try:
|
||||
file_data = FileData.from_msgpack(command.data)
|
||||
file_bytes = self.api_client.load_bytes(file_data.filename, file_data.folder)
|
||||
response = RemoteCommand(CommandType.DATA_BYTES, file_bytes)
|
||||
except Exception as e:
|
||||
response = RemoteCommand(CommandType.DATA_BYTES, None, str(e))
|
||||
self.remote_handler.send(command.client_id, response.serialize())
|
||||
|
||||
def stop(self):
|
||||
self.shutdown_event.set()
|
||||
self.remote_handler.stop()
|
||||
self.running = False
|
||||
@@ -0,0 +1,22 @@
|
||||
cdef enum CommandType:
|
||||
LOGIN = 10
|
||||
LOAD = 20
|
||||
DATA_BYTES = 25
|
||||
INFERENCE = 30
|
||||
INFERENCE_DATA = 35
|
||||
STOP_INFERENCE = 40
|
||||
AI_AVAILABILITY_CHECK = 80
|
||||
AI_AVAILABILITY_RESULT = 85
|
||||
ERROR = 90
|
||||
EXIT = 100
|
||||
|
||||
cdef class RemoteCommand:
|
||||
cdef public bytes client_id
|
||||
cdef CommandType command_type
|
||||
cdef str message
|
||||
cdef bytes data
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data)
|
||||
|
||||
cdef bytes serialize(self)
|
||||
@@ -0,0 +1,35 @@
|
||||
import msgpack
|
||||
|
||||
cdef class RemoteCommand:
|
||||
def __init__(self, CommandType command_type, bytes data, str message=None):
|
||||
self.command_type = command_type
|
||||
self.data = data
|
||||
self.message = message
|
||||
|
||||
def __str__(self):
|
||||
command_type_names = {
|
||||
10: "LOGIN",
|
||||
20: "LOAD",
|
||||
25: "DATA_BYTES",
|
||||
30: "INFERENCE",
|
||||
35: "INFERENCE_DATA",
|
||||
40: "STOP_INFERENCE",
|
||||
80: "AI_AVAILABILITY_CHECK",
|
||||
85: "AI_AVAILABILITY_RESULT",
|
||||
90: "ERROR",
|
||||
100: "EXIT"
|
||||
}
|
||||
data_str = f'{len(self.data)} bytes' if self.data else ''
|
||||
return f'{command_type_names[self.command_type]} ({data_str})'
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data):
|
||||
unpacked = msgpack.unpackb(data, strict_map_key=False)
|
||||
return RemoteCommand(unpacked.get("CommandType"), unpacked.get("Data"), unpacked.get("Message"))
|
||||
|
||||
cdef bytes serialize(self):
|
||||
return msgpack.packb({
|
||||
"CommandType": self.command_type,
|
||||
"Data": self.data,
|
||||
"Message": self.message
|
||||
})
|
||||
@@ -0,0 +1,16 @@
|
||||
cdef class RemoteCommandHandler:
|
||||
cdef object _context
|
||||
cdef object _router
|
||||
cdef object _dealer
|
||||
cdef object _control
|
||||
cdef object _shutdown_event
|
||||
cdef object _on_command
|
||||
|
||||
cdef object _proxy_thread
|
||||
cdef object _workers
|
||||
|
||||
cdef start(self)
|
||||
cdef _proxy_loop(self)
|
||||
cdef _worker_loop(self)
|
||||
cdef send(self, bytes client_id, bytes data)
|
||||
cdef stop(self)
|
||||
@@ -0,0 +1,94 @@
|
||||
import time
|
||||
import zmq
|
||||
from threading import Thread, Event
|
||||
from remote_command cimport RemoteCommand
|
||||
cimport constants
|
||||
import yaml
|
||||
|
||||
cdef class RemoteCommandHandler:
|
||||
def __init__(self, int zmq_port, object on_command):
|
||||
self._on_command = on_command
|
||||
self._context = zmq.Context.instance()
|
||||
|
||||
self._router = self._context.socket(zmq.ROUTER)
|
||||
self._router.setsockopt(zmq.LINGER, 0)
|
||||
self._router.bind(f'tcp://*:{zmq_port}')
|
||||
|
||||
self._dealer = self._context.socket(zmq.DEALER)
|
||||
self._dealer.setsockopt(zmq.LINGER, 0)
|
||||
self._dealer.bind("inproc://backend")
|
||||
|
||||
self._control = self._context.socket(zmq.PAIR)
|
||||
self._control.bind("inproc://control")
|
||||
self._shutdown_event = Event()
|
||||
|
||||
self._proxy_thread = Thread(target=self._proxy_loop, daemon=True)
|
||||
|
||||
self._workers = []
|
||||
for _ in range(4): # 4 worker threads
|
||||
worker = Thread(target=self._worker_loop, daemon=True)
|
||||
self._workers.append(worker)
|
||||
print(f'Listening to commands on port {zmq_port}...')
|
||||
|
||||
cdef start(self):
|
||||
self._proxy_thread.start()
|
||||
for worker in self._workers:
|
||||
worker.start()
|
||||
|
||||
cdef _proxy_loop(self):
|
||||
try:
|
||||
zmq.proxy_steerable(self._router, self._dealer, control=self._control)
|
||||
except zmq.error.ZMQError as e:
|
||||
if self._shutdown_event.is_set():
|
||||
print("Shutdown, exit proxy loop.")
|
||||
else:
|
||||
raise
|
||||
|
||||
cdef _worker_loop(self):
|
||||
worker_socket = self._context.socket(zmq.DEALER)
|
||||
worker_socket.setsockopt(zmq.LINGER, 0)
|
||||
worker_socket.connect("inproc://backend")
|
||||
poller = zmq.Poller()
|
||||
poller.register(worker_socket, zmq.POLLIN)
|
||||
try:
|
||||
|
||||
while not self._shutdown_event.is_set():
|
||||
try:
|
||||
socks = dict(poller.poll(500))
|
||||
if worker_socket in socks:
|
||||
client_id, message = worker_socket.recv_multipart()
|
||||
cmd = RemoteCommand.from_msgpack(<bytes> message)
|
||||
cmd.client_id = client_id
|
||||
constants.log(<str>f'{cmd}', client_id)
|
||||
self._on_command(cmd)
|
||||
except Exception as e:
|
||||
if not self._shutdown_event.is_set():
|
||||
constants.log(f"Worker error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
worker_socket.close()
|
||||
|
||||
cdef send(self, bytes client_id, bytes data):
|
||||
self._router.send_multipart([client_id, data])
|
||||
|
||||
# with self._context.socket(zmq.DEALER) as socket:
|
||||
# socket.connect("inproc://backend")
|
||||
# socket.send_multipart([client_id, data])
|
||||
# # constants.log(<str>f'Sent {len(data)} bytes.', client_id)
|
||||
|
||||
cdef stop(self):
|
||||
self._shutdown_event.set()
|
||||
try:
|
||||
self._control.send(b"TERMINATE", flags=zmq.DONTWAIT)
|
||||
except zmq.error.ZMQError:
|
||||
pass
|
||||
self._router.close(linger=0)
|
||||
self._dealer.close(linger=0)
|
||||
self._control.close(linger=0)
|
||||
|
||||
self._proxy_thread.join(timeout=2)
|
||||
while any(w.is_alive() for w in self._workers):
|
||||
time.sleep(0.1)
|
||||
|
||||
self._context.term()
|
||||
@@ -0,0 +1,9 @@
|
||||
pyinstaller
|
||||
Cython
|
||||
psutil
|
||||
msgpack
|
||||
pyjwt
|
||||
zmq
|
||||
requests
|
||||
pyyaml
|
||||
boto3
|
||||
@@ -0,0 +1,20 @@
|
||||
from credentials cimport Credentials
|
||||
|
||||
cdef class Security:
|
||||
@staticmethod
|
||||
cdef encrypt_to(input_stream, key)
|
||||
|
||||
@staticmethod
|
||||
cdef decrypt_to(input_bytes, key)
|
||||
|
||||
@staticmethod
|
||||
cdef get_hw_hash(str hardware)
|
||||
|
||||
@staticmethod
|
||||
cdef get_api_encryption_key(Credentials credentials, str hardware_hash)
|
||||
|
||||
@staticmethod
|
||||
cdef get_model_encryption_key()
|
||||
|
||||
@staticmethod
|
||||
cdef calc_hash(str key)
|
||||
@@ -0,0 +1,68 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
from hashlib import sha384
|
||||
from credentials cimport Credentials
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
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_bytes, key):
|
||||
cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
|
||||
iv = os.urandom(16)
|
||||
|
||||
cipher = Cipher(algorithms.AES(<bytes> aes_key), modes.CBC(iv), backend=default_backend())
|
||||
encryptor = cipher.encryptor()
|
||||
padder = padding.PKCS7(128).padder()
|
||||
|
||||
padded_plaintext = padder.update(input_bytes) + padder.finalize()
|
||||
ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize()
|
||||
|
||||
return iv + ciphertext
|
||||
|
||||
@staticmethod
|
||||
cdef decrypt_to(ciphertext_with_iv_bytes, key):
|
||||
cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
|
||||
iv = ciphertext_with_iv_bytes[:16]
|
||||
ciphertext_bytes = ciphertext_with_iv_bytes[16:]
|
||||
|
||||
cipher = Cipher(algorithms.AES(<bytes>aes_key), modes.CBC(<bytes>iv), backend=default_backend())
|
||||
decryptor = cipher.decryptor()
|
||||
|
||||
decrypted_padded_bytes = decryptor.update(ciphertext_bytes) + decryptor.finalize()
|
||||
|
||||
# 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(str hardware):
|
||||
cdef str key = f'Azaion_{hardware}_%$$$)0_'
|
||||
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):
|
||||
str_bytes = key.encode('utf-8')
|
||||
hash_bytes = sha384(str_bytes).digest()
|
||||
cdef str h = base64.b64encode(hash_bytes).decode('utf-8')
|
||||
return h
|
||||
@@ -0,0 +1,31 @@
|
||||
from setuptools import setup, Extension
|
||||
from Cython.Build import cythonize
|
||||
|
||||
extensions = [
|
||||
Extension('constants', ['constants.pyx']),
|
||||
Extension('credentials', ['credentials.pyx']),
|
||||
Extension('file_data', ['file_data.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('cdn_manager', ['cdn_manager.pyx']),
|
||||
Extension('api_client', ['api_client.pyx']),
|
||||
Extension('main', ['main.pyx']),
|
||||
]
|
||||
|
||||
setup(
|
||||
name="azaion.loader",
|
||||
ext_modules=cythonize(
|
||||
extensions,
|
||||
compiler_directives={
|
||||
"language_level": 3,
|
||||
"emit_code_comments" : False,
|
||||
"binding": True,
|
||||
'boundscheck': False,
|
||||
'wraparound': False
|
||||
}
|
||||
),
|
||||
zip_safe=False
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
from main import CommandProcessor
|
||||
import argparse
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-p", "--port", type=str, default="5025", help="zero mq port")
|
||||
parser.add_argument("-a", "--api", type=str, default="https://api.azaion.com", help="api url")
|
||||
args = parser.parse_args()
|
||||
|
||||
processor = CommandProcessor(int(args.port), args.api)
|
||||
try:
|
||||
processor.start()
|
||||
except KeyboardInterrupt:
|
||||
processor.stop()
|
||||
@@ -0,0 +1,15 @@
|
||||
cdef enum RoleEnum:
|
||||
NONE = 0
|
||||
Operator = 10
|
||||
Validator = 20
|
||||
CompanionPC = 30
|
||||
Admin = 40
|
||||
ResourceUploader = 50
|
||||
ApiAdmin = 1000
|
||||
|
||||
cdef class User:
|
||||
cdef public str id
|
||||
cdef public str email
|
||||
cdef public RoleEnum role
|
||||
|
||||
cdef bytes serialize(self)
|
||||
@@ -0,0 +1,15 @@
|
||||
import msgpack
|
||||
|
||||
cdef class User:
|
||||
|
||||
def __init__(self, str id, str email, RoleEnum role):
|
||||
self.id = id
|
||||
self.email = email
|
||||
self.role = role
|
||||
|
||||
cdef bytes serialize(self):
|
||||
return msgpack.packb({
|
||||
"i": self.id,
|
||||
"e": self.email,
|
||||
"r": self.role
|
||||
})
|
||||
Reference in New Issue
Block a user