Files
annotations/Azaion.AI/api_client.pyx
T
2025-01-20 10:17:35 +02:00

78 lines
2.7 KiB
Cython

import io
import json
import os
from http import HTTPStatus
import requests
cimport constants
from hardware_service cimport HardwareService, HardwareInfo
from security cimport Security
from io import BytesIO
cdef class ApiClient:
"""Handles API authentication and downloading of the AI model."""
def __init__(self, str email, str password, str folder):
self.email = email
self.password = password
self.folder = folder
if os.path.exists(<str>constants.TOKEN_FILE):
with open(<str>constants.TOKEN_FILE, "r") as file:
self.token = file.read().strip()
else:
self.token = None
cdef get_encryption_key(self, str hardware_hash):
cdef str key = f'{self.email}-{self.password}-{hardware_hash}-#%@AzaionKey@%#---'
return Security.calc_hash(key)
cdef login(self, str email, str password):
response = requests.post(f"{constants.API_URL}/login", json={"email": email, "password": password})
response.raise_for_status()
self.token = response.json()["token"]
with open(<str>constants.TOKEN_FILE, 'w') as file:
file.write(self.token)
cdef load_bytes(self, str filename):
hardware_service = HardwareService()
cdef HardwareInfo hardware = hardware_service.get_hardware_info()
if self.token is None:
self.login(self.email, self.password)
url = f"{constants.API_URL}/resources/get/{self.folder}"
headers = {
"Authorization": f"Bearer {self.token}",
"Content-Type": "application/json"
}
payload = json.dumps(
{
"password": self.password,
"hardware": hardware.to_json_object(),
"fileName": filename
}, indent=4)
response = requests.post(url, data=payload, headers=headers, stream=True)
if response.status_code == HTTPStatus.UNAUTHORIZED or response.status_code == HTTPStatus.FORBIDDEN:
self.login(self.email, self.password)
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())
return Security.decrypt_to(stream, key)
cdef load_ai_model(self):
return self.load_bytes(constants.AI_MODEL_FILE)
cdef load_queue_config(self):
return self.load_bytes(constants.QUEUE_CONFIG_FILENAME).decode(encoding='utf-8')