import json import os from ucloud_common.host import HostPool, HostStatus from ucloud_common.vm import VmPool, VMStatus from common_fields import Field, VmUUIDField, SpecsField from helper import check_otp from config import etcd_client as client from config import (HOST_PREFIX, VM_PREFIX, IMAGE_PREFIX, FILE_PREFIX, IMAGE_STORE_PREFIX) HOST_POOL = HostPool(client, HOST_PREFIX) VM_POOL = VmPool(client, VM_PREFIX) class BaseSchema: def __init__(self, data, fields=None): _ = data # suppress linter warning 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): if check_otp(self.name.value, self.realm.value, self.token.value) != 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(IMAGE_PREFIX) 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("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(os.path.join(FILE_PREFIX, self.uuid.value)) if file_entry is None: self.add_error("Image File with uuid '{}' Not Found".format(self.uuid.value)) def image_store_name_validation(self): image_stores = list(client.get_prefix(IMAGE_STORE_PREFIX)) 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("Store '{}' does not exists".format(self.image_store.value)) 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("Invalid Action. Allowed Actions are {}".format(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("No Such Host ({}) exists".format(self.destination.value)) 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 == os.path.join(HOST_PREFIX, 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")