ucloud-api/main.py

# TODO
#  1. Allow user of realm ungleich-admin to perform any action on
#     any user vm.

import json
import subprocess

from flask import Flask, request
from flask_restful import Resource, Api
from uuid import uuid4
from etcd3_wrapper import Etcd3Wrapper
from os.path import join

from ucloud_common.vm import VmPool, VMStatus
from ucloud_common.host import HostPool
from ucloud_common.request import RequestEntry, RequestPool, RequestType
from schemas import (CreateVMSchema, VMStatusSchema,
                     CreateImageSchema, VmActionSchema,
                     OTPSchema, CreateHostSchema,
                     VmMigrationSchema)

app = Flask(__name__)
api = Api(app)

client = Etcd3Wrapper()
vm_pool = VmPool(client, "/v1/vm")
host_pool = HostPool(client, "/v1/host")
request_pool = RequestPool(client, "/v1/request")


class CreateVM(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = CreateVMSchema(data)
        if validator.is_valid():
            # Create VM Entry under /v1/vm/
            vm_uuid = uuid4().hex
            vm_key = f"/v1/vm/{vm_uuid}"
            vm_entry = {
                "owner": data["name"],
                "specs": data["specs"],
                "hostname": "",
                "status": "",
                "image_uuid": data["image_uuid"],
                "log": []
            }
            client.put(vm_key, vm_entry, value_in_json=True)

            # Create ScheduleVM Request
            r = RequestEntry.from_scratch(type=RequestType.ScheduleVM,
                                          uuid=vm_uuid)
            request_pool.put(r)

            return {"message": "VM Creation Queued"}, 200
        else:
            return validator.get_errors(), 400


class VmStatus(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = VMStatusSchema(data)
        if validator.is_valid():
            vm = vm_pool.get(f"/v1/vm/{data['uuid']}")
            return str(vm)
        else:
            return validator.get_errors(), 400


class CreateImage(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = CreateImageSchema(data)
        if validator.is_valid():
            file_entry = client.get(f"/v1/file/{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",
            }
            client.put(f"/v1/image/{data['uuid']}", json.dumps(image_entry_json))

            return {"message": "Image successfully created"}
        else:
            return validator.get_errors(), 400


class ListPublicImages(Resource):
    @staticmethod
    def get():
        images = client.get_prefix("/v1/image/")
        r = {}
        for image in images:
            r[image.key.split("/")[-1]] = json.loads(image.value)

        return r, 200


class VMAction(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = VmActionSchema(data)

        if validator.is_valid():
            vm_entry = vm_pool.get(f"/v1/vm/{data['uuid']}")
            action = data["action"]

            if action == "start":
                vm_entry.status = VMStatus.requested_start
                vm_pool.put(vm_entry)
                action = "schedule"

            if action == "delete" and vm_entry.hostname == "":
                try:
                    path_without_protocol = vm_entry.path[vm_entry.path.find(":")+1:]
                    rc = subprocess.call(f"rbd rm {path_without_protocol}".split(" "))
                except FileNotFoundError:
                    return {"message": "VM image does not exists"}
                else:
                    if rc == 0:
                        client.client.delete(vm_entry.key)
                        return {"message": "VM successfully deleted"}
                    else:
                        return {"message": "Some error occurred while deleting VM"}

            r = RequestEntry.from_scratch(type=f"{action.title()}VM",
                                          uuid=data['uuid'],
                                          hostname=vm_entry.hostname)
            request_pool.put(r)
            return {"message": f"VM {action.title()} Queued"}, 200
        else:
            return validator.get_errors(), 400


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

        if validator.is_valid():
            vm = vm_pool.get(data['uuid'])

            r = RequestEntry.from_scratch(type=RequestType.ScheduleVM,
                                          uuid=vm.uuid,
                                          destination=join("/v1/host", data["destination"]),
                                          migration=True)
            request_pool.put(r)
            return {"message": f"VM Migration Initialization Queued"}, 200
        else:
            return validator.get_errors(), 400


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

        if validator.is_valid():
            vms = client.get_prefix(f"/v1/vm/", value_in_json=True)
            if vms:
                return_vms = []
                user_vms = list(filter(lambda v: v.value["owner"] == data["name"], vms))
                for vm in user_vms:
                    return_vms.append(
                        {
                            "vm_uuid": vm.key.split("/")[-1],
                            "specs": vm.value["specs"],
                            "status": vm.value["status"],
                            "hostname": vm.value["hostname"]
                        }
                    )
                return {"message": return_vms}, 200
            else:
                return {"message": "No VM found"}, 404
        else:
            return validator.get_errors(), 400


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

        if validator.is_valid():
            files = client.get_prefix(f"/v1/file/", value_in_json=True)
            if files:
                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 {"message": "No File found"}, 404
        else:
            return validator.get_errors(), 400


class CreateHost(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = CreateHostSchema(data)
        if validator.is_valid():
            host_key = f"/v1/host/{uuid4().hex}"
            host_entry = {
                "specs": data["specs"],
                "hostname": data["hostname"],
                "status": "DEAD",
                "last_heartbeat": "",
            }
            client.put(host_key, host_entry, value_in_json=True)

            return {"message": "Host Created"}, 200
        else:
            return validator.get_errors(), 400


class ListHost(Resource):
    @staticmethod
    def get():
        hosts = host_pool.hosts
        r = {host.key: {"status": host.status, "specs": host.specs, "hostname": host.hostname} for host in hosts}
        return r, 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(CreateHost, "/host/create")
api.add_resource(ListHost, "/host/list")

if __name__ == "__main__":
    app.run(host="::", debug=True)