import requests as req import logging import os import time import urllib3 import shutil from utils.MachineData import MachineData class ValhallaServer: def __init__( self, server_address: str, server_port: str, server_user: str, server_access_password: str, logging_level: str, ): self.server_address = server_address self.server_port = server_port self.server_user = server_user self.server_name = "not connected" self.server_version = "not_connected" self.server_access_password = server_access_password self.logging_level = logging_level self.logger = logging.getLogger(__name__) self.bearer_token = "" self.auth_headers = {} self.user_id = 0 self.vm_ids_list_from_server = [] self.vms_metadata = {} log_level_mapping_dict = { "NOTSET": 0, "DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40, "CRITICAL": 50, } self.logger.setLevel(log_level_mapping_dict[logging_level]) def authenticate(self): try: login_data = { "username": self.server_user, "password": self.server_access_password, } headers = {"Content-Type": "application/json"} response = req.post( url=f"{self.server_address}:{self.server_port}/login", json=login_data, headers=headers, ) # self.logger.info(str(response.json())) if response.status_code == 202: user_data = response.json() self.bearer_token = user_data["token"] self.user_id = user_data["user_id"] self.auth_headers = {"Authorization": f"Bearer {self.bearer_token}"} else: raise Exception("Bad input or server not responding") except Exception as ex: self.logger.error(f"{str(ex)}") def get_server_data(self): try: response = req.get( url=f"{self.server_address}:{self.server_port}/", headers=self.auth_headers, ) print(str(response.json())) if response.status_code == 200: server_data = response.json() self.server_version = server_data["server_version"] self.server_name = server_data["server_name"] except Exception as ex: self.logger.error(f"{str(ex)}") def update_machine_data_on_server(self, machine_data: MachineData): try: response = req.get( url=f"{self.server_address}:{self.server_port}/clients/{machine_data.mac_address}", headers=self.auth_headers, ) if response.status_code == 401: self.authenticate() print(self.bearer_token) response = req.get( url=f"{self.server_address}:{self.server_port}/clients/{machine_data.mac_address}", headers=self.auth_headers, ) print(f"{str(response.json())}") client_data = { "mac_address": machine_data.mac_address, "ip_address": machine_data.ip_address, "hostname": machine_data.hostname, "client_version": machine_data.client_version, "vm_list_on_machine": [], } if response.status_code == 404: print("Client not registered, registering now...") for i in range(1, 5): register_response = req.post( url=f"{self.server_address}:{self.server_port}/clients", headers=self.auth_headers, json=client_data, ) if register_response == 201: print("Registered successfully") break else: print("Failed to register, waiting 5s before trying again...") time.sleep(5) continue if response.status_code == 200: client_data_from_server = response.json() try: client_data.pop("vm_list_on_machine") except KeyError: pass if client_data_from_server != client_data: print("Client data needs updating, updating now...") for i in range(1, 5): register_response = req.put( url=f"{self.server_address}:{self.server_port}/clients", headers=self.auth_headers, json=client_data, ) if register_response.status_code == 202: print("Updated successfully") break elif response.status_code == 401: print("Reauthentication needed...") self.authenticate() else: print("Failed to update, waiting 5s before trying again...") time.sleep(5) continue except Exception as ex: self.logger.error(f"{str(ex)}") def get_vm_list_from_server(self, machine_data: MachineData): try: response = req.get( url=f"{self.server_address}:{self.server_port}/clients/{machine_data.mac_address}/vms", headers=self.auth_headers, ) if response.status_code == 401: self.authenticate() response = req.get( url=f"{self.server_address}:{self.server_port}/clients/{machine_data.mac_address}/vms", headers=self.auth_headers, ) if response.status_code == 200: vm_list = response.json() self.vm_ids_list_from_server = vm_list return vm_list else: print(f"Error getting data from remote server: {response.json()}") except Exception as ex: pass def get_vms_data_from_server(self) ->dict(): try: vms_data = {} for image_id in self.vm_ids_list_from_server: response = req.get( url=f"{self.server_address}:{self.server_port}/images/{image_id}", headers=self.auth_headers, ) if response.status_code == 200: data = response.json() vms_data[f"{image_id}"] = data self.vms_metadata = vms_data return vms_data except Exception as ex: print(str(ex)) def update_vms_images_from_server(self): try: for image_id in self.vm_ids_list_from_server: image_data = self.vms_metadata[f"{image_id}"] print(f"{image_data['image_name']} download") http = urllib3.PoolManager() self.authenticate() filename = f"images/{self.server_name}/{image_data['image_name']}/{image_data['image_version']}/image.qcow2" os.makedirs(os.path.dirname(filename), exist_ok=True) with http.request( method="GET", url=f"{self.server_address}:{self.server_port}/images/{image_id}/download", preload_content=False, headers=self.auth_headers, ) as r, open( filename, "wb", ) as out_file: shutil.copyfileobj(r, out_file) except Exception as ex: print(str(ex))