uncloud-mravi/uncloud/api/main.py

import json
import pynetbox
import logging
import argparse

from uuid import uuid4
from os.path import join as join_path

from flask import Flask, request
from flask_restful import Resource, Api
from werkzeug.exceptions import HTTPException

from uncloud.common.shared import shared

from uncloud.common import counters
from uncloud.common.vm import VMStatus
from uncloud.common.host import HostStatus
from uncloud.common.request import RequestEntry, RequestType
from uncloud.common.settings import settings
from uncloud.api import schemas
from uncloud.api.schemas import ValidationException
from uncloud.common.network import generate_mac, mac2ipv6
from uncloud.api.helper import make_return_message
from uncloud import UncloudException

logger = logging.getLogger(__name__)

app = Flask(__name__)
api = Api(app)
app.logger.handlers.clear()

arg_parser = argparse.ArgumentParser('api', add_help=False)
arg_parser.add_argument('--port', '-p')


@app.errorhandler(Exception)
def handle_exception(e):
    app.logger.error(e)

    # pass through HTTP errors
    if isinstance(e, HTTPException):
        return e

    # now you're handling non-HTTP exceptions only
    return {'message': 'Server Error'}, 500


class CreateVM(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.CreateVMSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            vm_uuid = uuid4().hex
            vm_key = join_path(settings['etcd']['vm_prefix'], vm_uuid)
            specs = {
                'cpu': validator.specs['cpu'],
                'ram': validator.specs['ram'],
                'os-ssd': validator.specs['os-ssd'],
                'hdd': validator.specs['hdd'],
            }
            macs = [generate_mac() for _ in range(len(validator.network.value))]
            tap_ids = [
                counters.increment_etcd_counter(settings['etcd']['tap_counter'])
                for _ in range(len(validator.network.value))
            ]
            vm_entry = {
                'name': validator.vm_name.value,
                'owner': validator.name.value,
                'owner_realm': validator.realm.value,
                'specs': specs,
                'hostname': '',
                'status': VMStatus.stopped,
                'image_uuid': validator.image_uuid,
                'log': [],
                'vnc_socket': '',
                'network': list(zip(validator.network.value, macs, tap_ids)),
                'metadata': {'ssh-keys': []},
                'in_migration': False,
            }
            shared.etcd_client.put(vm_key, vm_entry, value_in_json=True)

            # Create ScheduleVM Request
            r = RequestEntry.from_scratch(
                type=RequestType.ScheduleVM,
                uuid=vm_uuid,
                request_prefix=settings['etcd']['request_prefix'],
            )
            shared.request_pool.put(r)
            return make_return_message('VM Creation Queued')


class VmStatus(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.VMStatusSchema(data)
            validator.is_valid()
        except (ValidationException, Exception) as err:
            return make_return_message(err, 400)
        else:
            vm = shared.vm_pool.get(join_path(settings['etcd']['vm_prefix'], data['uuid']))
            vm_value = vm.value.copy()
            vm_value['ip'] = []
            for network_mac_and_tap in vm.network:
                network_name, mac, tap = network_mac_and_tap
                network = shared.etcd_client.get(
                    join_path(settings['etcd']['network_prefix'], data['name'], network_name),
                    value_in_json=True,
                )
                ipv6_addr = (network.value.get('ipv6').split('::')[0] + '::')
                vm_value['ip'].append(mac2ipv6(mac, ipv6_addr))
            vm.value = vm_value
            return vm.value, 200


class CreateImage(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.CreateImageSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            try:
                file_entry = shared.etcd_client.get(
                    join_path(settings['etcd']['file_prefix'], data['uuid']), value_in_json=True
                )
            except KeyError:
                # TODO: Add some message
                pass
            else:
                image_entry_json = {
                    'status': 'TO_BE_CREATED',
                    'owner': file_entry.value['owner'],
                    'filename': file_entry.value['filename'],
                    'name': validator.name,
                    'store_name': validator.image_store,
                    'visibility': 'public',
                }
                shared.etcd_client.put(
                    join_path(settings['etcd']['image_prefix'], data['uuid']),
                    json.dumps(image_entry_json),
                )

                return {'message': 'Image queued for creation.'}, 200


class ListPublicImages(Resource):
    @staticmethod
    def get():
        images = shared.etcd_client.get_prefix(settings['etcd']['image_prefix'], value_in_json=True)
        r = {'images': []}
        for image in images:
            image_key = '{}:{}'.format(image.value['store_name'], image.value['name'])
            r['images'].append({'name': image_key, 'status': image.value['status']})
        return r, 200


class VMAction(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.VmActionSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            vm_entry = shared.vm_pool.get(join_path(settings['etcd']['vm_prefix'], data['uuid']))
            action = validator.action.value

            if action == 'start':
                action = 'schedule'

            if action == 'delete' and vm_entry.hostname == '':
                if shared.storage_handler.is_vm_image_exists(vm_entry.uuid):
                    r_status = shared.storage_handler.delete_vm_image(vm_entry.uuid)
                    if r_status:
                        shared.etcd_client.client.delete(vm_entry.key)
                        return make_return_message('VM successfully deleted')
                    else:
                        logger.error('Some Error Occurred while deleting VM')
                        return make_return_message('VM deletion unsuccessfull')
                else:
                    shared.etcd_client.client.delete(vm_entry.key)
                    return make_return_message('VM successfully deleted')

            r = RequestEntry.from_scratch(
                type='{}VM'.format(action.title()),
                uuid=data['uuid'],
                hostname=vm_entry.hostname,
                request_prefix=settings['etcd']['request_prefix'],
            )
            shared.request_pool.put(r)
            return make_return_message('VM {} Queued'.format(action.title()))


class VMMigration(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.VmMigrationSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err), 400
        else:
            vm = shared.vm_pool.get(validator.uuid.value)
            r = RequestEntry.from_scratch(
                type=RequestType.InitVMMigration,
                uuid=vm.uuid,
                hostname=join_path(
                    settings['etcd']['host_prefix'],
                    validator.destination.value,
                ),
                request_prefix=settings['etcd']['request_prefix'],
            )

            shared.request_pool.put(r)
            return make_return_message('VM Migration Initialization Queued')


class ListUserVM(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.OTPSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            vms = shared.etcd_client.get_prefix(settings['etcd']['vm_prefix'], value_in_json=True)
            return_vms = []
            user_vms = filter(lambda v: v.value['owner'] == validator.name.value, vms)
            for vm in user_vms:
                return_vms.append(
                    {
                        'name': vm.value['name'],
                        'vm_uuid': vm.key.split('/')[-1],
                        'specs': vm.value['specs'],
                        'status': vm.value['status'],
                        'hostname': vm.value['hostname'],
                        'vnc_socket': vm.value.get('vnc_socket', None),
                    }
                )
            return make_return_message(return_vms)


class ListUserFiles(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.OTPSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            files = shared.etcd_client.get_prefix(settings['etcd']['file_prefix'], value_in_json=True)
            return_files = []
            user_files = [f for f in files if f.value['owner'] == data['name']]
            for file in user_files:
                file_uuid = file.key.split('/')[-1]
                file = file.value
                file['uuid'] = file_uuid

                file.pop('sha512sum', None)
                file.pop('owner', None)

                return_files.append(file)
            return make_return_message(return_files)


class CreateHost(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.CreateHostSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            host_key = join_path(settings['etcd']['host_prefix'], uuid4().hex)
            host_entry = {
                'specs': validator.specs.value,
                'hostname': validator.hostname.value,
                'status': HostStatus.dead,
                'last_heartbeat': '',
            }
            shared.etcd_client.put(host_key, host_entry, value_in_json=True)
            return make_return_message('Host Created.')


class ListHost(Resource):
    @staticmethod
    def get():
        hosts = shared.host_pool.hosts
        r = {
            host.key: {
                'status': host.status,
                'specs': host.specs,
                'hostname': host.hostname,
            }
            for host in hosts
        }
        return r, 200


class GetSSHKeys(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.GetSSHSchema(data)
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            etcd_key = join_path(settings['etcd']['user_prefix'], validator.realm.value,
                                 validator.name.value, 'key')
            if not validator.key_name.value:
                etcd_entry = shared.etcd_client.get_prefix(etcd_key, value_in_json=True)
                keys = {
                    key.key.split('/')[-1]: key.value
                    for key in etcd_entry
                }
                return {'keys': keys}
            else:
                etcd_key = join_path(validator.key_name.value)
                try:
                    etcd_entry = shared.etcd_client.get(etcd_key, value_in_json=True)
                except KeyError:
                    return make_return_message('No such key found.', 400)
                else:
                    return {
                        'keys': {etcd_entry.key.split('/')[-1]: etcd_entry.value}
                    }


class AddSSHKey(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.AddSSHSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            # {user_prefix}/{realm}/{name}/key/{key_name}
            etcd_key = join_path(
                settings['etcd']['user_prefix'], validator.realm.value,
                validator.name.value, 'key', validator.key_name.value
            )
            try:
                shared.etcd_client.get(etcd_key, value_in_json=True)
            except KeyError:
                # Key Not Found. It implies user' haven't added any key yet.
                shared.etcd_client.put(etcd_key, validator.key.value, value_in_json=True)
                return make_return_message('Key added successfully')
            else:
                return make_return_message('Key "{}" already exists'.format(validator.key_name.value))


class RemoveSSHKey(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.RemoveSSHSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            # {user_prefix}/{realm}/{name}/key/{key_name}
            etcd_key = join_path(settings['etcd']['user_prefix'], validator.realm.value,
                                 validator.name.value, 'key', validator.key_name.value)
            try:
                etcd_entry = shared.etcd_client.get(etcd_key, value_in_json=True)
            except KeyError:
                return make_return_message('No Key "{}" exists.'.format(validator.key_name.value))
            if etcd_entry:
                shared.etcd_client.client.delete(etcd_key)
                return {'message': 'Key successfully removed.'}


class CreateNetwork(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.CreateNetwork(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            network_entry = {
                'id': counters.increment_etcd_counter(settings['etcd']['vxlan_counter']),
                'type': validator.type.value,
            }
            if validator.user.value:
                try:
                    nb = pynetbox.api(url=settings['netbox']['url'], token=settings['netbox']['token'])
                    nb_prefix = nb.ipam.prefixes.get(prefix=settings['network']['prefix'])
                    prefix = nb_prefix.available_prefixes.create(
                        data={
                            'prefix_length': int(settings['network']['prefix_length']),
                            'description': '{}\'s network "{}"'.format(
                                validator.name.value,
                                validator.network_name.value
                            ),
                            'is_pool': True,
                        }
                    )
                except Exception as err:
                    app.logger.error(err)
                    return make_return_message('Error occured while creating network.', 400)
                else:
                    network_entry['ipv6'] = prefix['prefix']
            else:
                network_entry['ipv6'] = 'fd00::/64'

            network_key = join_path(settings['etcd']['network_prefix'], validator.name.value,
                                    validator.network_name.value)
            shared.etcd_client.put(network_key, network_entry, value_in_json=True)
            return make_return_message('Network successfully added.')


class ListUserNetwork(Resource):
    @staticmethod
    def post():
        data = request.json
        try:
            validator = schemas.OTPSchema(data)
            validator.is_valid()
        except ValidationException as err:
            return make_return_message(err, 400)
        else:
            prefix = join_path(settings['etcd']['network_prefix'], data['name'])
            networks = shared.etcd_client.get_prefix(prefix, value_in_json=True)
            user_networks = []
            for net in networks:
                net.value['name'] = net.key.split('/')[-1]
                user_networks.append(net.value)
            return {'networks': user_networks}, 200


api.add_resource(CreateVM, '/vm/create')
api.add_resource(VmStatus, '/vm/status')

api.add_resource(VMAction, '/vm/action')
api.add_resource(VMMigration, '/vm/migrate')

api.add_resource(CreateImage, '/image/create')
api.add_resource(ListPublicImages, '/image/list-public')

api.add_resource(ListUserVM, '/user/vms')
api.add_resource(ListUserFiles, '/user/files')
api.add_resource(ListUserNetwork, '/user/networks')

api.add_resource(AddSSHKey, '/user/add-ssh')
api.add_resource(RemoveSSHKey, '/user/remove-ssh')
api.add_resource(GetSSHKeys, '/user/get-ssh')

api.add_resource(CreateHost, '/host/create')
api.add_resource(ListHost, '/host/list')

api.add_resource(CreateNetwork, '/network/create')


def main(debug=False, port=None):
    try:
        app.run(host='::', port=port, debug=debug)
    except OSError as e:
        raise UncloudException('Failed to start Flask: {}'.format(e))


if __name__ == '__main__':
    main()