import json from common_fields import Field, VmUUIDField, SpecsField from ucloud_common.host import HostPool, HostStatus from ucloud_common.vm import VmPool, VMStatus from helper import check_otp from config import etcd_client as client from os.path import join host_pool = HostPool(client, "/v1/host") vm_pool = VmPool(client, "/v1/vm") class BaseSchema: def __init__(self, data, fields=None): self.__errors = [] if fields is None: self.fields = [] else: self.fields = fields def validation(self): # custom validation is optional return True def is_valid(self): for field in self.fields: field.is_valid() self.add_field_errors(field) for parent in self.__class__.__bases__: try: parent.validation(self) except AttributeError: pass if not self.__errors: self.validation() if self.__errors: return False return True def get_errors(self): return {"message": self.__errors} def add_field_errors(self, field: Field): self.__errors += field.get_errors() def add_error(self, error): self.__errors.append(error) class OTPSchema(BaseSchema): def __init__(self, data: dict, fields=None): self.name = Field("name", str, data.get("name", KeyError)) self.realm = Field("realm", str, data.get("realm", KeyError)) self.token = Field("token", str, data.get("token", KeyError)) _fields = [self.name, self.realm, self.token] if fields: _fields += fields super().__init__(data=data, fields=_fields) def validation(self): rc = check_otp(self.name.value, self.realm.value, self.token.value) if rc != 200: self.add_error("Wrong Credentials") class CreateVMSchema(OTPSchema): def __init__(self, data): self.specs = SpecsField(data) self.image_uuid = Field("image_uuid", str, data.get("image_uuid", KeyError)) self.image_uuid.validation = self.image_uuid_validation fields = [self.specs, self.image_uuid] super().__init__(data=data, fields=fields) def image_uuid_validation(self): images = client.get_prefix("/v1/image/") if self.image_uuid.value not in [i.key.split("/")[-1] for i in images]: self.add_error("Image UUID not valid") class VMStatusSchema(BaseSchema): def __init__(self, data): self.uuid = VmUUIDField(data) fields = [self.uuid] super().__init__(data, fields) class CreateImageSchema(BaseSchema): def __init__(self, data): # Fields self.uuid: Field = Field("uuid", str, data.get("uuid", KeyError)) self.name = Field("name", str, data.get("name", KeyError)) self.image_store = Field("image_store", str, data.get("image_store", KeyError)) # Validations self.uuid.validation = self.file_uuid_validation self.image_store.validation = self.image_store_name_validation # All Fields fields = [self.uuid, self.name, self.image_store] super().__init__(data, fields) def file_uuid_validation(self): file_entry = client.get(f"/v1/file/{self.uuid.value}") if file_entry is None: self.add_error(f"Image File with uuid '{self.uuid.value}' Not Found") def image_store_name_validation(self): image_stores = list(client.get_prefix("/v1/image_store/")) image_store = next(filter(lambda s: json.loads(s.value)["name"] == self.image_store.value, image_stores), None) if not image_store: self.add_error(f"Store '{self.image_store.value}' does not exists") class VmActionSchema(OTPSchema): def __init__(self, data): self.uuid = VmUUIDField(data) self.action = Field("action", str, data.get("action", KeyError)) self.action.validation = self.action_validation _fields = [self.uuid, self.action] super().__init__(data=data, fields=_fields) def action_validation(self): allowed_actions = ["start", "stop", "delete"] if self.action.value not in allowed_actions: self.add_error(f"Invalid Action. Allowed Actions are {allowed_actions}") def validation(self): vm = vm_pool.get(self.uuid.value) if vm.value["owner"] != self.name.value: self.add_error("Invalid User") if self.action.value == "start" and vm.status == VMStatus.running and vm.hostname != "": self.add_error("VM Already Running") if self.action.value == "stop" and vm.status == VMStatus.stopped: self.add_error("VM Already Stopped") class VmMigrationSchema(OTPSchema): def __init__(self, data): self.uuid = VmUUIDField(data) self.destination = Field("destination", str, data.get("destination", KeyError)) self.destination.validation = self.destination_validation fields = [self.destination] super().__init__(data=data, fields=fields) def destination_validation(self): host_key = self.destination.value host = host_pool.get(host_key) if not host: self.add_error(f"No Such Host ({self.destination.value}) exists") elif host.status != HostStatus.alive: self.add_error("Destination Host is dead") def validation(self): vm = vm_pool.get(self.uuid.value) if vm.owner != self.name.value: self.add_error("Invalid User") if vm.status != VMStatus.running: self.add_error("Can't migrate non-running VM") if vm.hostname == join("/v1/host", self.destination.value): self.add_error("Destination host couldn't be same as Source Host") class CreateHostSchema(OTPSchema): def __init__(self, data): self.specs = SpecsField(data) self.hostname = Field("hostname", str, data.get("hostname", KeyError)) fields = [self.specs, self.hostname] super().__init__(data=data, fields=fields) def validation(self): if self.realm.value != "ungleich-admin": self.add_error("Invalid Credentials/Insufficient Permission")