mirror of
https://github.com/azaion/loader.git
synced 2026-04-22 10:06:32 +00:00
Add API client and CDN manager implementation
- Introduced `ApiClient` class for handling API interactions, including authentication and resource management. - Added `CDNManager` class for managing file uploads and downloads to/from a CDN. - Implemented security features for encryption and decryption of sensitive data. - Created supporting classes for credentials, user roles, and hardware information retrieval. - Established constants for configuration and logging. This commit lays the foundation for resource management and secure communication with the API and CDN services.
This commit is contained in:
Oleksandr Bezdieniezhnykh
@@ -0,0 +1,23 @@
|
||||
from user cimport User
|
||||
from credentials cimport Credentials
|
||||
from cdn_manager cimport CDNManager
|
||||
|
||||
|
||||
cdef class ApiClient:
|
||||
cdef Credentials credentials
|
||||
cdef CDNManager cdn_manager
|
||||
cdef public str token
|
||||
cdef str folder, api_url
|
||||
cdef User user
|
||||
|
||||
cpdef set_credentials_from_dict(self, str email, str password)
|
||||
cdef set_credentials(self, Credentials credentials)
|
||||
cdef login(self)
|
||||
cdef set_token(self, str token)
|
||||
|
||||
cdef request(self, str method, str url, object payload, bint is_stream)
|
||||
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)
|
||||
cpdef load_big_small_resource(self, str resource_name, str folder)
|
||||
cpdef upload_big_small_resource(self, bytes resource, str resource_name, str folder)
|
||||
@@ -0,0 +1,187 @@
|
||||
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 credentials cimport Credentials
|
||||
from cdn_manager cimport CDNManager, CDNCredentials
|
||||
from hardware_service cimport HardwareService
|
||||
from security cimport Security
|
||||
from user cimport User, RoleEnum
|
||||
|
||||
cdef class ApiClient:
|
||||
def __init__(self, str api_url):
|
||||
self.credentials = None
|
||||
self.user = None
|
||||
self.token = None
|
||||
self.cdn_manager = None
|
||||
self.api_url = api_url
|
||||
|
||||
cpdef set_credentials_from_dict(self, str email, str password):
|
||||
self.set_credentials(Credentials(email, password))
|
||||
|
||||
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):
|
||||
if self.credentials is None:
|
||||
raise Exception("No credentials set")
|
||||
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:
|
||||
res = response.json()
|
||||
constants.logerror(str(res))
|
||||
if response.status_code == HTTPStatus.CONFLICT:
|
||||
raise Exception(f"Error {res['ErrorCode']}: {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 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)}
|
||||
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}.")
|
||||
|
||||
cdef load_bytes(self, str filename, str folder):
|
||||
if self.credentials is None:
|
||||
raise Exception("No credentials set")
|
||||
cdef str hardware = HardwareService.get_hardware_info()
|
||||
hw_hash = Security.get_hw_hash(hardware)
|
||||
key = Security.get_api_encryption_key(self.credentials, hw_hash)
|
||||
payload = json.dumps(
|
||||
{
|
||||
"password": self.credentials.password,
|
||||
"hardware": hardware,
|
||||
"fileName": filename
|
||||
}, indent=4)
|
||||
response = self.request('post', f'{self.api_url}/resources/get/{folder}', payload, is_stream=True)
|
||||
|
||||
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 request(self, str method, str url, object payload, bint is_stream):
|
||||
if self.token is None:
|
||||
self.login()
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.token}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
response = requests.request(method, url, data=payload, headers=headers, stream=is_stream)
|
||||
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.request(method, url, data=payload, headers=headers, stream=is_stream)
|
||||
|
||||
if response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR:
|
||||
raise Exception(f'Internal API error! {response.text}')
|
||||
if response.status_code == HTTPStatus.CONFLICT:
|
||||
res = response.json()
|
||||
err_code = res['ErrorCode']
|
||||
err_msg = res['Message']
|
||||
raise Exception(f"Error {err_code}: {err_msg}")
|
||||
return response
|
||||
|
||||
cdef load_big_file_cdn(self, str folder, str big_part):
|
||||
constants.log(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()
|
||||
return encrypted_bytes_big
|
||||
else:
|
||||
raise Exception(f'Cannot download file {folder}/{big_part} from CDN!')
|
||||
|
||||
cpdef load_big_small_resource(self, str resource_name, str folder):
|
||||
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)
|
||||
|
||||
key = Security.get_resource_encryption_key()
|
||||
|
||||
constants.log(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()
|
||||
constants.log(f'local file {folder}/{big_part} is found!')
|
||||
try:
|
||||
resource = Security.decrypt_to(encrypted_bytes_small + local_bytes_big, key)
|
||||
return resource
|
||||
except Exception as ex:
|
||||
constants.logerror(f'Local file {folder}/{big_part} 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)
|
||||
|
||||
cpdef upload_big_small_resource(self, bytes resource, str resource_name, str folder):
|
||||
cdef str big_part_name = f'{resource_name}.big'
|
||||
cdef str small_part_name = f'{resource_name}.small'
|
||||
key = Security.get_resource_encryption_key()
|
||||
|
||||
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]
|
||||
|
||||
part_big = resource_encrypted[part_small_size:]
|
||||
|
||||
if not self.cdn_manager.upload(folder, <str>big_part_name, part_big):
|
||||
raise Exception(f'Failed to upload {big_part_name} to CDN bucket {folder}')
|
||||
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,63 @@
|
||||
import hashlib
|
||||
import subprocess
|
||||
|
||||
import requests
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives import padding
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
|
||||
API_SERVICES = [
|
||||
"azaion/annotations",
|
||||
"azaion/flights",
|
||||
"azaion/detections",
|
||||
"azaion/gps-denied-onboard",
|
||||
"azaion/gps-denied-desktop",
|
||||
"azaion/autopilot",
|
||||
"azaion/ai-training",
|
||||
]
|
||||
|
||||
|
||||
def download_key_fragment(resource_api_url: str, token: str) -> bytes:
|
||||
resp = requests.get(
|
||||
f"{resource_api_url}/binary-split/key-fragment",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.content
|
||||
|
||||
|
||||
def decrypt_archive(encrypted_path: str, key_fragment: bytes, output_path: str):
|
||||
aes_key = hashlib.sha256(key_fragment).digest()
|
||||
|
||||
with open(encrypted_path, "rb") as f_in:
|
||||
iv = f_in.read(16)
|
||||
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
|
||||
decryptor = cipher.decryptor()
|
||||
unpadder = padding.PKCS7(128).unpadder()
|
||||
|
||||
with open(output_path, "wb") as f_out:
|
||||
while True:
|
||||
chunk = f_in.read(64 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
decrypted = decryptor.update(chunk)
|
||||
if decrypted:
|
||||
f_out.write(unpadder.update(decrypted))
|
||||
final_decrypted = decryptor.finalize()
|
||||
f_out.write(unpadder.update(final_decrypted) + unpadder.finalize())
|
||||
|
||||
|
||||
def docker_load(tar_path: str):
|
||||
subprocess.run(["docker", "load", "-i", tar_path], check=True)
|
||||
|
||||
|
||||
def check_images_loaded(version: str) -> bool:
|
||||
for svc in API_SERVICES:
|
||||
tag = f"{svc}:{version}"
|
||||
result = subprocess.run(
|
||||
["docker", "image", "inspect", tag],
|
||||
capture_output=True,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return False
|
||||
return True
|
||||
@@ -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,44 @@
|
||||
import io
|
||||
import os
|
||||
cimport constants
|
||||
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)
|
||||
constants.log(f'uploaded {filename} ({len(file_bytes)} bytes) to the {bucket}')
|
||||
return True
|
||||
except Exception as e:
|
||||
constants.logerror(e)
|
||||
return False
|
||||
|
||||
cdef download(self, str folder, str filename):
|
||||
try:
|
||||
os.makedirs(folder, exist_ok=True)
|
||||
self.download_client.download_file(folder, filename, os.path.join(folder, filename))
|
||||
constants.log(f'downloaded {filename} from the {folder}')
|
||||
return True
|
||||
except Exception as e:
|
||||
constants.logerror(e)
|
||||
return False
|
||||
@@ -0,0 +1,5 @@
|
||||
cdef str CDN_CONFIG
|
||||
cdef int SMALL_SIZE_KB
|
||||
|
||||
cdef log(str log_message)
|
||||
cdef logerror(str error)
|
||||
@@ -0,0 +1,37 @@
|
||||
import os
|
||||
import sys
|
||||
from loguru import logger
|
||||
|
||||
cdef str CDN_CONFIG = "cdn.yaml"
|
||||
cdef int SMALL_SIZE_KB = 3
|
||||
|
||||
_log_dir = os.environ.get("LOG_DIR", "Logs")
|
||||
logger.remove()
|
||||
log_format = "[{time:HH:mm:ss} {level}] {message}"
|
||||
logger.add(
|
||||
sink=f"{_log_dir}/log_loader_{{time:YYYYMMDD}}.txt",
|
||||
level="INFO",
|
||||
format=log_format,
|
||||
enqueue=True,
|
||||
rotation="1 day",
|
||||
retention="30 days",
|
||||
)
|
||||
logger.add(
|
||||
sys.stdout,
|
||||
level="DEBUG",
|
||||
format=log_format,
|
||||
filter=lambda record: record["level"].name in ("INFO", "DEBUG", "SUCCESS"),
|
||||
colorize=True
|
||||
)
|
||||
logger.add(
|
||||
sys.stderr,
|
||||
level="WARNING",
|
||||
format=log_format,
|
||||
colorize=True
|
||||
)
|
||||
|
||||
cdef log(str log_message):
|
||||
logger.info(log_message)
|
||||
|
||||
cdef logerror(str error):
|
||||
logger.error(error)
|
||||
@@ -0,0 +1,3 @@
|
||||
cdef class Credentials:
|
||||
cdef public str email
|
||||
cdef public str password
|
||||
@@ -0,0 +1,9 @@
|
||||
cdef class Credentials:
|
||||
|
||||
def __init__(self, str email, str password):
|
||||
self.email = email
|
||||
self.password = password
|
||||
|
||||
def __str__(self):
|
||||
return f'{self.email}: {self.password}'
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
cdef str _CACHED_HW_INFO
|
||||
|
||||
cdef class HardwareService:
|
||||
|
||||
@staticmethod
|
||||
cdef str get_hardware_info()
|
||||
@@ -0,0 +1,99 @@
|
||||
import os
|
||||
import platform
|
||||
import subprocess
|
||||
|
||||
import psutil
|
||||
cimport constants
|
||||
|
||||
cdef str _CACHED_HW_INFO = None
|
||||
|
||||
|
||||
def _get_cpu():
|
||||
try:
|
||||
with open("/proc/cpuinfo") as f:
|
||||
for line in f:
|
||||
if "model name" in line.lower():
|
||||
return line.split(":")[1].strip()
|
||||
except OSError:
|
||||
pass
|
||||
cdef str p = platform.processor()
|
||||
if p:
|
||||
return p
|
||||
return platform.machine()
|
||||
|
||||
|
||||
def _get_gpu():
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["lspci"], capture_output=True, text=True, timeout=5,
|
||||
)
|
||||
for line in result.stdout.splitlines():
|
||||
if "VGA" in line:
|
||||
parts = line.split(":")
|
||||
if len(parts) > 2:
|
||||
return parts[2].strip()
|
||||
return parts[-1].strip()
|
||||
except (OSError, subprocess.TimeoutExpired, FileNotFoundError):
|
||||
pass
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["system_profiler", "SPDisplaysDataType"],
|
||||
capture_output=True, text=True, timeout=5,
|
||||
)
|
||||
for line in result.stdout.splitlines():
|
||||
if "Chipset Model" in line:
|
||||
return line.split(":")[1].strip()
|
||||
except (OSError, subprocess.TimeoutExpired, FileNotFoundError):
|
||||
pass
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _get_drive_serial():
|
||||
try:
|
||||
for block in sorted(os.listdir("/sys/block")):
|
||||
for candidate in [
|
||||
f"/sys/block/{block}/device/vpd_pg80",
|
||||
f"/sys/block/{block}/device/serial",
|
||||
f"/sys/block/{block}/serial",
|
||||
]:
|
||||
try:
|
||||
with open(candidate, "rb") as f:
|
||||
serial = f.read().strip(b"\x00\x14 \t\n\r\v\f").decode("utf-8", errors="ignore")
|
||||
if serial:
|
||||
return serial
|
||||
except OSError:
|
||||
continue
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["ioreg", "-rd1", "-c", "IOPlatformExpertDevice"],
|
||||
capture_output=True, text=True, timeout=5,
|
||||
)
|
||||
for line in result.stdout.splitlines():
|
||||
if "IOPlatformSerialNumber" in line:
|
||||
return line.split('"')[-2]
|
||||
except (OSError, subprocess.TimeoutExpired, FileNotFoundError):
|
||||
pass
|
||||
return "unknown"
|
||||
|
||||
|
||||
cdef class HardwareService:
|
||||
|
||||
@staticmethod
|
||||
cdef str get_hardware_info():
|
||||
global _CACHED_HW_INFO
|
||||
|
||||
if _CACHED_HW_INFO is not None:
|
||||
constants.log(<str>"Using cached hardware info")
|
||||
return <str> _CACHED_HW_INFO
|
||||
|
||||
cdef str cpu = _get_cpu()
|
||||
cdef str gpu = _get_gpu()
|
||||
cdef str memory = str(psutil.virtual_memory().total // 1024)
|
||||
cdef str drive_serial = _get_drive_serial()
|
||||
|
||||
cdef str res = f'CPU: {cpu}. GPU: {gpu}. Memory: {memory}. DriveSerial: {drive_serial}'
|
||||
constants.log(<str>f'Gathered hardware: {res}')
|
||||
_CACHED_HW_INFO = res
|
||||
return res
|
||||
+198
@@ -0,0 +1,198 @@
|
||||
import os
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, HTTPException, UploadFile, File, Form, BackgroundTasks
|
||||
from fastapi.responses import Response
|
||||
from pydantic import BaseModel
|
||||
|
||||
from unlock_state import UnlockState
|
||||
|
||||
app = FastAPI(title="Azaion.Loader")
|
||||
|
||||
RESOURCE_API_URL = os.environ.get("RESOURCE_API_URL", "https://api.azaion.com")
|
||||
IMAGES_PATH = os.environ.get("IMAGES_PATH", "/opt/azaion/images.enc")
|
||||
API_VERSION = os.environ.get("API_VERSION", "latest")
|
||||
|
||||
_api_client = None
|
||||
_api_client_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_api_client():
|
||||
global _api_client
|
||||
if _api_client is None:
|
||||
with _api_client_lock:
|
||||
if _api_client is None:
|
||||
from api_client import ApiClient
|
||||
_api_client = ApiClient(RESOURCE_API_URL)
|
||||
return _api_client
|
||||
|
||||
|
||||
class LoginRequest(BaseModel):
|
||||
email: str
|
||||
password: str
|
||||
|
||||
|
||||
class LoadRequest(BaseModel):
|
||||
filename: str
|
||||
folder: str
|
||||
|
||||
|
||||
class HealthResponse(BaseModel):
|
||||
status: str
|
||||
|
||||
|
||||
class StatusResponse(BaseModel):
|
||||
status: str
|
||||
authenticated: bool
|
||||
modelCacheDir: str
|
||||
|
||||
|
||||
class _UnlockStateHolder:
|
||||
def __init__(self):
|
||||
self._state = UnlockState.idle
|
||||
self._error: Optional[str] = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def get(self):
|
||||
with self._lock:
|
||||
return self._state, self._error
|
||||
|
||||
def set(self, state: UnlockState, error: Optional[str] = None):
|
||||
with self._lock:
|
||||
self._state = state
|
||||
self._error = error
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
with self._lock:
|
||||
return self._state
|
||||
|
||||
|
||||
_unlock = _UnlockStateHolder()
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> HealthResponse:
|
||||
return HealthResponse(status="healthy")
|
||||
|
||||
|
||||
@app.get("/status")
|
||||
def status() -> StatusResponse:
|
||||
client = get_api_client()
|
||||
return StatusResponse(
|
||||
status="healthy",
|
||||
authenticated=client.token is not None,
|
||||
modelCacheDir="models",
|
||||
)
|
||||
|
||||
|
||||
@app.post("/login")
|
||||
def login(req: LoginRequest):
|
||||
try:
|
||||
client = get_api_client()
|
||||
client.set_credentials_from_dict(req.email, req.password)
|
||||
return {"status": "ok"}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=401, detail=str(e))
|
||||
|
||||
|
||||
@app.post("/load/{filename}")
|
||||
def load_resource(filename: str, req: LoadRequest):
|
||||
try:
|
||||
client = get_api_client()
|
||||
data = client.load_big_small_resource(req.filename, req.folder)
|
||||
return Response(content=data, media_type="application/octet-stream")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post("/upload/{filename}")
|
||||
def upload_resource(
|
||||
filename: str,
|
||||
data: UploadFile = File(...),
|
||||
folder: str = Form("models"),
|
||||
):
|
||||
try:
|
||||
client = get_api_client()
|
||||
content = data.file.read()
|
||||
client.upload_big_small_resource(content, filename, folder)
|
||||
return {"status": "ok"}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
def _run_unlock(email: str, password: str):
|
||||
from binary_split import (
|
||||
download_key_fragment,
|
||||
decrypt_archive,
|
||||
docker_load,
|
||||
check_images_loaded,
|
||||
)
|
||||
|
||||
try:
|
||||
if check_images_loaded(API_VERSION):
|
||||
_, prev_err = _unlock.get()
|
||||
_unlock.set(UnlockState.ready, prev_err)
|
||||
return
|
||||
|
||||
_unlock.set(UnlockState.authenticating)
|
||||
|
||||
client = get_api_client()
|
||||
client.set_credentials_from_dict(email, password)
|
||||
client.login()
|
||||
token = client.token
|
||||
|
||||
_unlock.set(UnlockState.downloading_key)
|
||||
|
||||
key_fragment = download_key_fragment(RESOURCE_API_URL, token)
|
||||
|
||||
_unlock.set(UnlockState.decrypting)
|
||||
|
||||
tar_path = IMAGES_PATH.replace(".enc", ".tar")
|
||||
decrypt_archive(IMAGES_PATH, key_fragment, tar_path)
|
||||
|
||||
_unlock.set(UnlockState.loading_images)
|
||||
|
||||
docker_load(tar_path)
|
||||
|
||||
try:
|
||||
os.remove(tar_path)
|
||||
except OSError as e:
|
||||
from loguru import logger
|
||||
|
||||
logger.warning(f"Failed to remove {tar_path}: {e}")
|
||||
|
||||
_unlock.set(UnlockState.ready, None)
|
||||
|
||||
except Exception as e:
|
||||
_unlock.set(UnlockState.error, str(e))
|
||||
|
||||
|
||||
@app.post("/unlock")
|
||||
def unlock(req: LoginRequest, background_tasks: BackgroundTasks):
|
||||
state, _ = _unlock.get()
|
||||
if state == UnlockState.ready:
|
||||
return {"state": state.value}
|
||||
if state not in (UnlockState.idle, UnlockState.error):
|
||||
return {"state": state.value}
|
||||
|
||||
if not os.path.exists(IMAGES_PATH):
|
||||
from binary_split import check_images_loaded
|
||||
|
||||
if check_images_loaded(API_VERSION):
|
||||
_, prev_err = _unlock.get()
|
||||
_unlock.set(UnlockState.ready, prev_err)
|
||||
return {"state": _unlock.state.value}
|
||||
raise HTTPException(status_code=404, detail="Encrypted archive not found")
|
||||
|
||||
_unlock.set(UnlockState.authenticating, None)
|
||||
|
||||
background_tasks.add_task(_run_unlock, req.email, req.password)
|
||||
return {"state": _unlock.state.value}
|
||||
|
||||
|
||||
@app.get("/unlock/status")
|
||||
def get_unlock_status():
|
||||
state, error = _unlock.get()
|
||||
return {"state": state.value, "error": error}
|
||||
@@ -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_resource_encryption_key()
|
||||
|
||||
@staticmethod
|
||||
cdef calc_hash(str key)
|
||||
@@ -0,0 +1,63 @@
|
||||
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()
|
||||
|
||||
unpadder = padding.PKCS7(128).unpadder()
|
||||
plaintext_bytes = unpadder.update(decrypted_padded_bytes) + unpadder.finalize()
|
||||
|
||||
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_resource_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,11 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class UnlockState(str, Enum):
|
||||
idle = "idle"
|
||||
authenticating = "authenticating"
|
||||
downloading_key = "downloading_key"
|
||||
decrypting = "decrypting"
|
||||
loading_images = "loading_images"
|
||||
ready = "ready"
|
||||
error = "error"
|
||||
@@ -0,0 +1,13 @@
|
||||
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
|
||||
@@ -0,0 +1,6 @@
|
||||
cdef class User:
|
||||
|
||||
def __init__(self, str id, str email, RoleEnum role):
|
||||
self.id = id
|
||||
self.email = email
|
||||
self.role = role
|
||||
Reference in New Issue
Block a user