uncloud/uncloud/api/schemas.py

import json
import os

import bitmath

from uncloud.common.host import HostStatus
from uncloud.common.vm import VMStatus
from uncloud.common.shared import shared
from uncloud.common.settings import settings
from uncloud.api import helper
from uncloud.api.helper import check_otp, resolve_vm_name


class ValidationException(Exception):
    """Validation Error"""


class Field:
    def __init__(self, _name, _type, _value=None, validators=None):
        if validators is None:
            validators = []

        assert isinstance(validators, list)

        self.name = _name
        self.value = _value
        self.type = _type
        self.validators = validators

    def is_valid(self):
        if not isinstance(self.value, self.type):
            raise ValidationException("Incorrect Type for '{}' field".format(self.name))

        for validator in self.validators:
            validator()

    def __repr__(self):
        return self.name


class VmUUIDField(Field):
    def __init__(self, data):
        self.uuid = data.get('uuid', KeyError)

        super().__init__('uuid', str, self.uuid, validators=[self.vm_uuid_validation])

    def vm_uuid_validation(self):
        try:
            shared.etcd_client.get(os.path.join(settings['etcd']['vm_prefix'], self.uuid))
        except KeyError:
            raise ValidationException('VM with uuid {} does not exists'.format(self.uuid))


class BaseSchema:
    def __init__(self):
        self.fields = [getattr(self, field) for field in dir(self) if isinstance(getattr(self, field), Field)]

    def validation(self):
        # custom validation is optional
        return True

    def is_valid(self):
        for field in self.fields:
            field.is_valid()

        for parent in self.__class__.__bases__:
            parent.validation(self)

        self.validation()

        for field in self.fields:
            setattr(self, field.name, field.value)


def get(dictionary: dict, key: str, return_default=False, default=None):
    if dictionary is None:
        raise ValidationException('No data provided at all.')
    try:
        value = dictionary[key]
    except KeyError:
        if return_default:
            return default
        raise ValidationException("Missing data for '{}' field.".format(key))
    else:
        return value


class OTPSchema(BaseSchema):
    def __init__(self, data: dict):
        self.name = Field('name', str, get(data, 'name'))
        self.realm = Field('realm', str, get(data, 'realm'))
        self.token = Field('token', str, get(data, 'token'))
        super().__init__()

    def validation(self):
        if check_otp(self.name.value, self.realm.value, self.token.value) != 200:
            raise ValidationException('Wrong Credentials')


class CreateImageSchema(BaseSchema):
    def __init__(self, data):
        self.uuid = Field('uuid', str, get(data, 'uuid'), validators=[self.file_uuid_validation])
        self.name = Field('name', str, get(data, 'name'))
        self.image_store = Field('image_store', str, get(data, 'image_store'),
                                 validators=[self.image_store_name_validation])
        super().__init__()

    def file_uuid_validation(self):
        try:
            shared.etcd_client.get(os.path.join(settings['etcd']['file_prefix'], self.uuid.value))
        except KeyError:
            raise ValidationException("Image File with uuid '{}' Not Found".format(self.uuid.value))

    def image_store_name_validation(self):
        image_stores = list(shared.etcd_client.get_prefix(settings['etcd']['image_store_prefix']))
        try:
            next(filter(lambda s: json.loads(s.value)['name'] == self.image_store.value, image_stores))
        except StopIteration:
            raise ValidationException("Store '{}' does not exists".format(self.image_store.value))


class CreateHostSchema(OTPSchema):
    def __init__(self, data):
        self.specs = Field('specs', dict, get(data, 'specs'), validators=[self.specs_validation])
        self.hostname = Field('hostname', str, get(data, 'hostname'))

        super().__init__(data)

    def specs_validation(self):
        allowed_base = 10

        _cpu = self.specs.value.get('cpu', KeyError)
        _ram = self.specs.value.get('ram', KeyError)
        _os_ssd = self.specs.value.get('os-ssd', KeyError)
        _hdd = self.specs.value.get('hdd', KeyError)

        if KeyError in [_cpu, _ram, _os_ssd]:
            raise ValidationException('You must specify CPU, RAM and OS-SSD in your specs')
        try:
            parsed_ram = bitmath.parse_string_unsafe(_ram)
            parsed_os_ssd = bitmath.parse_string_unsafe(_os_ssd)

            if parsed_ram.base != allowed_base:
                raise ValidationException('Your specified RAM is not in correct units')

            if parsed_os_ssd.base != allowed_base:
                raise ValidationException('Your specified OS-SSD is not in correct units')

            if _cpu < 1:
                raise ValidationException('CPU must be atleast 1')

            if parsed_ram < bitmath.GB(1):
                raise ValidationException('RAM must be atleast 1 GB')

            if parsed_os_ssd < bitmath.GB(10):
                raise ValidationException('OS-SSD must be atleast 10 GB')

            parsed_hdd = []
            for hdd in _hdd:
                _parsed_hdd = bitmath.parse_string_unsafe(hdd)
                if _parsed_hdd.base != allowed_base:
                    raise ValidationException('Your specified HDD is not in correct units')
                else:
                    parsed_hdd.append(str(_parsed_hdd))

        except ValueError:
            raise ValidationException('Specs are not correct.')
        else:
            self.specs = {
                'cpu': _cpu,
                'ram': str(parsed_ram),
                'os-ssd': str(parsed_os_ssd),
                'hdd': parsed_hdd,
            }

    def validation(self):
        if self.realm.value != 'ungleich-admin':
            raise ValidationException('Invalid Credentials/Insufficient Permission')


class CreateVMSchema(OTPSchema):
    def __init__(self, data):
        self.specs = Field('specs', dict, get(data, 'specs'), validators=[self.specs_validation])
        self.vm_name = Field('vm_name', str, get(data, 'vm_name'), validators=[self.vm_name_validation])
        self.image = Field('image', str, get(data, 'image'), validators=[self.image_validation])
        self.network = Field('network', list, get(data, 'network', return_default=True, default=[]),
                             validators=[self.network_validation])
        self.image_uuid = None

        super().__init__(data=data)

    def image_validation(self):
        try:
            image_uuid = helper.resolve_image_name(self.image.value)
        except Exception:
            raise ValidationException('No image of name \'{}\' found'.format(self.image.value))
        else:
            self.image_uuid = image_uuid

    def vm_name_validation(self):
        if resolve_vm_name(name=self.vm_name.value, owner=self.name.value):
            raise ValidationException("VM with same name '{}' already exists".format(self.vm_name.value))

    def network_validation(self):
        _network = self.network.value

        if _network:
            for net in _network:
                try:
                    shared.etcd_client.get(
                        os.path.join(settings['etcd']['network_prefix'], self.name.value, net),
                        value_in_json=True
                    )
                except KeyError:
                    raise ValidationException('Network with name {} does not exists'.format(net))

    def specs_validation(self):
        allowed_base = 10

        try:
            _cpu = get(self.specs.value, 'cpu')
            _ram = get(self.specs.value, 'ram')
            _os_ssd = get(self.specs.value, 'os-ssd')
            _hdd = get(self.specs.value, 'hdd', return_default=True, default=[])
        except (KeyError, Exception):
            raise ValidationException('You must specify CPU, RAM and OS-SSD in your specs')
        else:
            try:
                parsed_ram = bitmath.parse_string_unsafe(_ram)
                parsed_os_ssd = bitmath.parse_string_unsafe(_os_ssd)

                if parsed_ram.base != allowed_base:
                    raise ValidationException('Your specified RAM is not in correct units')

                if parsed_os_ssd.base != allowed_base:
                    raise ValidationException('Your specified OS-SSD is not in correct units')

                if int(_cpu) < 1:
                    raise ValidationException('CPU must be atleast 1')

                if parsed_ram < bitmath.GB(1):
                    raise ValidationException('RAM must be atleast 1 GB')

                if parsed_os_ssd < bitmath.GB(1):
                    raise ValidationException('OS-SSD must be atleast 1 GB')

                parsed_hdd = []
                for hdd in _hdd:
                    _parsed_hdd = bitmath.parse_string_unsafe(hdd)
                    if _parsed_hdd.base != allowed_base:
                        raise ValidationException('Your specified HDD is not in correct units')
                    else:
                        parsed_hdd.append(str(_parsed_hdd))

            except ValueError:
                raise ValidationException('Specs are not correct.')
            else:
                self.specs = {
                    'cpu': _cpu,
                    'ram': str(parsed_ram),
                    'os-ssd': str(parsed_os_ssd),
                    'hdd': parsed_hdd,
                }


class VMStatusSchema(OTPSchema):
    def __init__(self, data):
        data['uuid'] = (
            resolve_vm_name(
                name=get(data, 'vm_name', return_default=True),
                owner=(
                        get(data, 'in_support_of', return_default=True) or
                        get(data, 'name', return_default=True)
                )
            )
            or KeyError
        )
        self.uuid = VmUUIDField(data)

        super().__init__(data)

    def validation(self):
        vm = shared.vm_pool.get(self.uuid.value)
        if not (vm.value['owner'] == self.name.value or self.realm.value == 'ungleich-admin'):
            raise ValidationException('Invalid User')


class VmActionSchema(OTPSchema):
    def __init__(self, data):
        data['uuid'] = (
            resolve_vm_name(
                name=get(data, 'vm_name', return_default=True),
                owner=(
                        get(data, 'in_support_of', return_default=True) or
                        get(data, 'name', return_default=True)
                )
            )
            or KeyError
        )
        self.uuid = VmUUIDField(data)
        self.action = Field('action', str, get(data, 'action'), validators=[self.action_validation])

        super().__init__(data=data)

    def action_validation(self):
        allowed_actions = ['start', 'stop', 'delete']
        if self.action.value not in allowed_actions:
            raise ValidationException('Invalid Action. Allowed Actions are {}'.format(allowed_actions))

    def validation(self):
        vm = shared.vm_pool.get(self.uuid.value)
        if not (vm.value['owner'] == self.name.value or self.realm.value == 'ungleich-admin'):
            raise ValidationException('Invalid User.')

        if self.action.value == 'start' and vm.status == VMStatus.running and vm.hostname != '':
            raise ValidationException('VM Already Running')

        if self.action.value == 'stop':
            if vm.status == VMStatus.stopped:
                raise ValidationException('VM Already Stopped')
            elif vm.status != VMStatus.running:
                raise ValidationException('Cannot stop non-running VM')


class VmMigrationSchema(OTPSchema):
    def __init__(self, data):
        data['uuid'] = (
            resolve_vm_name(
                name=get(data, 'vm_name', return_default=True),
                owner=(
                        get(data, 'in_support_of', return_default=True) or
                        get(data, 'name', return_default=True)
                )
            ) or KeyError
        )

        self.uuid = VmUUIDField(data)
        self.destination = Field('destination', str, get(data, 'destination'),
                                 validators=[self.destination_validation])

        super().__init__(data=data)

    def destination_validation(self):
        hostname = self.destination.value
        host = next(filter(lambda h: h.hostname == hostname, shared.host_pool.hosts), None,)
        if not host:
            raise ValidationException('No Such Host ({}) exists'.format(self.destination.value))
        elif host.status != HostStatus.alive:
            raise ValidationException('Destination Host is dead')
        else:
            self.destination.value = host.key

    def validation(self):
        vm = shared.vm_pool.get(self.uuid.value)
        if not (vm.value['owner'] == self.name.value or self.realm.value == 'ungleich-admin'):
            raise ValidationException('Invalid User')

        if vm.status != VMStatus.running:
            raise ValidationException("Can't migrate non-running VM")

        if vm.hostname == os.path.join(settings['etcd']['host_prefix'], self.destination.value):
            raise ValidationException("Destination host couldn't be same as Source Host")


class AddSSHSchema(OTPSchema):
    def __init__(self, data):
        self.key_name = Field('key_name', str, get(data, 'key_name'))
        self.key = Field('key', str, get(data, 'key'))
        super().__init__(data=data)


class RemoveSSHSchema(OTPSchema):
    def __init__(self, data):
        self.key_name = Field('key_name', str, get(data, 'key_name'))
        super().__init__(data=data)


class GetSSHSchema(OTPSchema):
    def __init__(self, data):
        self.key_name = Field('key_name', str, get(data, 'key_name', return_default=True))
        super().__init__(data=data)


class CreateNetwork(OTPSchema):
    def __init__(self, data):
        self.network_name = Field('network_name', str, get(data, 'name'),
                                  validators=[self.network_name_validation])
        self.type = Field('type', str, get(data, 'type'), validators=[self.network_type_validation])
        self.user = Field('user', bool, bool(get(data, 'user', return_default=True, default=False)))
        super().__init__(data)

    def network_name_validation(self):
        key = os.path.join(settings['etcd']['network_prefix'], self.name.value, self.network_name.value)
        network = shared.etcd_client.get(key, value_in_json=True)
        if network:
            raise ValidationException('Network with name {} already exists'.format(self.network_name.value))

    def network_type_validation(self):
        supported_network_types = ['vxlan']
        if self.type.value not in supported_network_types:
            raise ValidationException('Unsupported Network Type. Supported network types are {}'.format(supported_network_types))