client-application/network/ValhallaServer.py

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):
        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
        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))