uncloud/ucloud/api/main.py

import json
import pynetbox
import logging

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 ucloud.common import counters
from ucloud.common.vm import VMStatus
from ucloud.common.request import RequestEntry, RequestType
from ucloud.settings import settings
from ucloud.shared import shared

from . import schemas
from .helper import generate_mac, mac2ipv6


logger = logging.getLogger(__name__)

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


@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):
    """API Request to Handle Creation of VM"""

    @staticmethod
    def post():
        data = request.json
        validator = schemas.CreateVMSchema(data)
        if validator.is_valid():
            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(data["network"]))]
            tap_ids = [
                counters.increment_etcd_counter(
                    shared.etcd_client, "/v1/counter/tap"
                )
                for _ in range(len(data["network"]))
            ]
            vm_entry = {
                "name": data["vm_name"],
                "owner": data["name"],
                "owner_realm": data["realm"],
                "specs": specs,
                "hostname": "",
                "status": VMStatus.stopped,
                "image_uuid": validator.image_uuid,
                "log": [],
                "vnc_socket": "",
                "network": list(zip(data["network"], 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 {"message": "VM Creation Queued"}, 200
        return validator.get_errors(), 400


class VmStatus(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = schemas.VMStatusSchema(data)
        if validator.is_valid():
            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
        else:
            return validator.get_errors(), 400


class CreateImage(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = schemas.CreateImageSchema(data)
        if validator.is_valid():
            file_entry = shared.etcd_client.get(
                join_path(settings["etcd"]["file_prefix"], data["uuid"])
            )
            file_entry_value = json.loads(file_entry.value)

            image_entry_json = {
                "status": "TO_BE_CREATED",
                "owner": file_entry_value["owner"],
                "filename": file_entry_value["filename"],
                "name": data["name"],
                "store_name": data["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."}
        return validator.get_errors(), 400


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
        validator = schemas.VmActionSchema(data)

        if validator.is_valid():
            vm_entry = shared.vm_pool.get(
                join_path(settings["etcd"]["vm_prefix"], data["uuid"])
            )
            action = data["action"]

            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 {"message": "VM successfully deleted"}
                    else:
                        logger.error(
                            "Some Error Occurred while deleting VM"
                        )
                        return {"message": "VM deletion unsuccessfull"}
                else:
                    shared.etcd_client.client.delete(vm_entry.key)
                    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 (
                {"message": "VM {} Queued".format(action.title())},
                200,
            )
        else:
            return validator.get_errors(), 400


class VMMigration(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = schemas.VmMigrationSchema(data)

        if validator.is_valid():
            vm = shared.vm_pool.get(data["uuid"])
            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 (
                {"message": "VM Migration Initialization Queued"},
                200,
            )
        else:
            return validator.get_errors(), 400


class ListUserVM(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = schemas.OTPSchema(data)

        if validator.is_valid():
            vms = shared.etcd_client.get_prefix(
                settings["etcd"]["vm_prefix"], value_in_json=True
            )
            return_vms = []
            user_vms = filter(
                lambda v: v.value["owner"] == data["name"], 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),
                    }
                )
            if return_vms:
                return {"message": return_vms}, 200
            return {"message": "No VM found"}, 404

        else:
            return validator.get_errors(), 400


class ListUserFiles(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = schemas.OTPSchema(data)

        if validator.is_valid():
            files = shared.etcd_client.get_prefix(
                settings["etcd"]["file_prefix"], value_in_json=True
            )
            return_files = []
            user_files = list(
                filter(
                    lambda f: f.value["owner"] == data["name"], files
                )
            )
            for file in user_files:
                return_files.append(
                    {
                        "filename": file.value["filename"],
                        "uuid": file.key.split("/")[-1],
                    }
                )
            return {"message": return_files}, 200
        else:
            return validator.get_errors(), 400


class CreateHost(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = schemas.CreateHostSchema(data)
        if validator.is_valid():
            host_key = join_path(
                settings["etcd"]["host_prefix"], uuid4().hex
            )
            host_entry = {
                "specs": data["specs"],
                "hostname": data["hostname"],
                "status": "DEAD",
                "last_heartbeat": "",
            }
            shared.etcd_client.put(
                host_key, host_entry, value_in_json=True
            )

            return {"message": "Host Created"}, 200

        return validator.get_errors(), 400


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 get():
        data = request.json
        validator = schemas.GetSSHSchema(data)
        if validator.is_valid():
            if not validator.key_name.value:

                # {user_prefix}/{realm}/{name}/key/
                etcd_key = join_path(
                    settings["etcd"]["user_prefix"],
                    data["realm"],
                    data["name"],
                    "key",
                )
                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:

                # {user_prefix}/{realm}/{name}/key/{key_name}
                etcd_key = join_path(
                    settings["etcd"]["user_prefix"],
                    data["realm"],
                    data["name"],
                    "key",
                    data["key_name"],
                )
                etcd_entry = shared.etcd_client.get(
                    etcd_key, value_in_json=True
                )

                if etcd_entry:
                    return {
                        "keys": {
                            etcd_entry.key.split("/")[
                                -1
                            ]: etcd_entry.value
                        }
                    }
                else:
                    return {"keys": {}}
        else:
            return validator.get_errors(), 400


class AddSSHKey(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = schemas.AddSSHSchema(data)
        if validator.is_valid():

            # {user_prefix}/{realm}/{name}/key/{key_name}
            etcd_key = join_path(
                settings["etcd"]["user_prefix"],
                data["realm"],
                data["name"],
                "key",
                data["key_name"],
            )
            etcd_entry = shared.etcd_client.get(
                etcd_key, value_in_json=True
            )
            if etcd_entry:
                return {
                    "message": "Key with name '{}' already exists".format(
                        data["key_name"]
                    )
                }
            else:
                # Key Not Found. It implies user' haven't added any key yet.
                shared.etcd_client.put(
                    etcd_key, data["key"], value_in_json=True
                )
                return {"message": "Key added successfully"}
        else:
            return validator.get_errors(), 400


class RemoveSSHKey(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = schemas.RemoveSSHSchema(data)
        if validator.is_valid():

            # {user_prefix}/{realm}/{name}/key/{key_name}
            etcd_key = join_path(
                settings["etcd"]["user_prefix"],
                data["realm"],
                data["name"],
                "key",
                data["key_name"],
            )
            etcd_entry = shared.etcd_client.get(
                etcd_key, value_in_json=True
            )
            if etcd_entry:
                shared.etcd_client.client.delete(etcd_key)
                return {"message": "Key successfully removed."}
            else:
                return {
                    "message": "No Key with name '{}' Exists at all.".format(
                        data["key_name"]
                    )
                }
        else:
            return validator.get_errors(), 400


class CreateNetwork(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = schemas.CreateNetwork(data)

        if validator.is_valid():

            network_entry = {
                "id": counters.increment_etcd_counter(
                    shared.etcd_client, "/v1/counter/vxlan"
                ),
                "type": data["type"],
            }
            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(
                                data["name"], data["network_name"]
                            ),
                            "is_pool": True,
                        }
                    )
                except Exception as err:
                    app.logger.error(err)
                    return {
                        "message": "Error occured while creating network."
                    }
                else:
                    network_entry["ipv6"] = prefix["prefix"]
            else:
                network_entry["ipv6"] = "fd00::/64"

            network_key = join_path(
                settings["etcd"]["network_prefix"],
                data["name"],
                data["network_name"],
            )
            shared.etcd_client.put(
                network_key, network_entry, value_in_json=True
            )
            return {"message": "Network successfully added."}
        else:
            return validator.get_errors(), 400


class ListUserNetwork(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = schemas.OTPSchema(data)

        if validator.is_valid():
            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
        else:
            return validator.get_errors(), 400


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():
    image_stores = list(
        shared.etcd_client.get_prefix(
            settings["etcd"]["image_store_prefix"], value_in_json=True
        )
    )
    if not image_stores:
        data = {
            "is_public": True,
            "type": "ceph",
            "name": "images",
            "description": "first ever public image-store",
            "attributes": {"list": [], "key": [], "pool": "images"},
        }

        shared.etcd_client.put(
            join_path(
                settings["etcd"]["image_store_prefix"], uuid4().hex
            ),
            json.dumps(data),
        )

    app.run(host="::", debug=False)


if __name__ == "__main__":
    main()