diff --git a/utils/database/database.py b/utils/database/database.py index 8bf58de..4458de4 100644 --- a/utils/database/database.py +++ b/utils/database/database.py @@ -1,6 +1,6 @@ from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker -from utils.models import Client, VMImage +from utils.models.models import Client, VMImage import logging @@ -27,15 +27,15 @@ class Database: print(ex) exit(-1) - def get_clients(self) -> list(Client): - result = None + def get_clients(self) -> list[Client]: + result = [] try: with self.session.begin(): result = self.session.query(Client).all() except Exception as ex: self.logger.error( f"Error getting list of clients from database: {ex}") - result = None + result = [] return result def get_client_by_mac_address(self, mac_address: str) -> Client: @@ -45,8 +45,29 @@ class Database: result = self.session.query( Client, mac_address=mac_address).first() except Exception as ex: - self.logger.error(f"Error getting client by mac address: {ex}") + self.logger.warn(f"Error getting client by mac address: {ex}") return result - def get_clients_by_client_version(self, client_version: str) -> list(Client): - result = None + def get_clients_by_client_version(self, client_version: str) -> list[Client]: + result = [] + try: + with self.session.begin(): + result = self.session.query( + Client, client_version=client_version).all() + except Exception as ex: + self.logger.warn( + f"Error getting client list by software version: {ex}") + return result + + def get_clients_by_vm_image(self, vm_image: VMImage) -> list[Client]: + result = [] + try: + clients_list = self.get_clients() + for client in clients_list: + if client.has_vm_installed(vm_image.image_hash): + result.append() + except Exception as ex: + self.logger.warn( + f"Error getting list of clients with VM installed: {ex}") + result = [] + return result diff --git a/utils/models/__init__.py b/utils/models/__init__.py index 20f2059..9028d94 100644 --- a/utils/models/__init__.py +++ b/utils/models/__init__.py @@ -1 +1 @@ -from utils.models import Client, VMImage \ No newline at end of file +from utils.models import Client, VMImage