ucloud-api/main.py

# TODO:  realm should be part of user's name. Because, there can be multiple
#       name but with different realms. For Example, nico exists in both
#       ungleich-admin and ungleich-twitter realm. So, to differentiate between
#       them, we need to make realm part of user's name
import json
import subprocess
import os

from uuid import uuid4

from flask import Flask, request
from flask_restful import Resource, Api
from ucloud_common.vm import VMStatus
from ucloud_common.request import RequestEntry, RequestType

from config import etcd_client as client
from config import (
    WITHOUT_CEPH,
    VM_PREFIX,
    HOST_PREFIX,
    FILE_PREFIX,
    IMAGE_PREFIX,
    logging,
    REQUEST_POOL,
    VM_POOL,
    HOST_POOL,
)
from schemas import (
    CreateVMSchema,
    VMStatusSchema,
    CreateImageSchema,
    VmActionSchema,
    OTPSchema,
    CreateHostSchema,
    VmMigrationSchema,
    AddSSHSchema,
    RemoveSSHSchema
)

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


class CreateVM(Resource):
    @staticmethod
    def post():
        data = request.json
        validator = CreateVMSchema(data)
        if validator.is_valid():
            vm_uuid = uuid4().hex
            vm_key = os.path.join(VM_PREFIX, vm_uuid)
            vm_entry = {
                "name": data["vm_name"],
                "owner": data["name"],
                "specs": data["specs"],
                "hostname": "",
                "status": "",
                "image_uuid": data["image_uuid"],
                "log": [],
                "storage_attachment": [],
                "vnc_socket": "",
            }
            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
        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(os.path.join(VM_PREFIX, 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(os.path.join(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",
            }
            client.put(
                os.path.join(IMAGE_PREFIX, data["uuid"]), json.dumps(image_entry_json)
            )

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


class ListPublicImages(Resource):
    @staticmethod
    def get():
        images = client.get_prefix(IMAGE_PREFIX)
        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(os.path.join(VM_PREFIX, 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 :]

                    if WITHOUT_CEPH:
                        command_to_delete = [
                            "rm",
                            os.path.join("/var/vm", vm_entry.uuid),
                        ]
                    else:
                        command_to_delete = ["rbd", "rm", path_without_protocol]

                    subprocess.check_output(command_to_delete, stderr=subprocess.PIPE)
                except subprocess.CalledProcessError as e:
                    if "No such file" in e.stderr.decode("utf-8"):
                        client.client.delete(vm_entry.key)
                        return {"message": "VM successfully deleted"}
                    else:
                        logging.exception(e)
                        return {"message": "Some error occurred while deleting VM"}
                else:
                    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_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 = VmMigrationSchema(data)

        if validator.is_valid():
            vm = VM_POOL.get(data["uuid"])

            r = RequestEntry.from_scratch(
                type=RequestType.ScheduleVM,
                uuid=vm.uuid,
                destination=os.path.join(HOST_PREFIX, data["destination"]),
                migration=True,
            )
            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 = OTPSchema(data)

        if validator.is_valid():
            vms = client.get_prefix(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": None
                        if vm.value.get("vnc_socket", None) == None
                        else vm.value["vnc_socket"],
                    }
                )
            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 = OTPSchema(data)

        if validator.is_valid():
            files = client.get_prefix(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 = CreateHostSchema(data)
        if validator.is_valid():
            host_key = os.path.join(HOST_PREFIX, 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

        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


class AddSSHKey(Resource):
    @staticmethod
    def get():
        data = request.json
        validator = AddSSHSchema(data)
        if validator.is_valid():
            etcd_key = "/v1/user/{}/keys".format(data["name"])
            etcd_entry = client.get(etcd_key, value_in_json=True)
            if etcd_entry:
                if data["key_name"] not in etcd_entry.value:
                    etcd_entry.value[data["key_name"]] = data["key"]
                    client.put(etcd_entry.key, etcd_entry.value, value_in_json=True)
                else:
                    return {"message": "Key with name '{}' already exists".format(data["key_name"])}
            else:
                # Key Not Found. It implies user' haven't added any key yet.
                key_entry = {data["key_name"]: data["key"]}
                client.put(etcd_key, key_entry, 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 = RemoveSSHSchema(data)
        if validator.is_valid():
            etcd_key = "/v1/user/{}/keys".format(data["name"])
            etcd_entry = client.get(etcd_key, value_in_json=True)
            if etcd_entry:
                try:
                    del etcd_entry.value[data["key_name"]]
                except KeyError:
                    return {"message": "No such key exists."}
                else:
                    client.put(etcd_entry.key, etcd_entry.value, value_in_json=True)
                    return {"message": "Key successfully removed."}
            else:
                return {"message": "No Key Exists at all."}
        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(AddSSHKey, "/user/add-ssh")
api.add_resource(RemoveSSHKey, "/user/remove-ssh")

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

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