ucloud-api/main.py

# TODO
#  convert etcd3 usage to etcd3_wrapper
import etcd3
import json

from helper import check_otp, add_otp_args, add_vmid_args
from flask import Flask
from flask_restful import Resource, Api, reqparse
from decouple import config
from uuid import uuid4
from etcd3_wrapper import Etcd3Wrapper
from specs_parser import SpecsParser

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

etcd_client = etcd3.client(host=config("ETCD_HOST"), port=int(config("ETCD_PORT")))
client = Etcd3Wrapper()

# CreateVM argparser
createvm_argparser = reqparse.RequestParser()
createvm_argparser.add_argument("specs", type=dict, required=True)
createvm_argparser.add_argument("image_uuid", type=str, required=True)
add_otp_args(createvm_argparser)

# CreateImage argparser
createimage_argparser = reqparse.RequestParser()
createimage_argparser.add_argument("uuid", type=str, required=True)
createimage_argparser.add_argument("name", type=str, required=True)
createimage_argparser.add_argument("image_store", type=str, required=True)

# DeleteVM argparser
deletevm_argparser = reqparse.RequestParser()
add_vmid_args(add_otp_args(deletevm_argparser))

# VMStatus argparser
vmstatus_argparser = reqparse.RequestParser()
add_vmid_args(vmstatus_argparser)

# StartVM argparser
startvm_argparser = reqparse.RequestParser()
add_vmid_args(add_otp_args(startvm_argparser))

# UserVM argparser
uservm_argparser = reqparse.RequestParser()
add_otp_args(uservm_argparser)

# Specs parser
specs_parser = SpecsParser(exceptional_devices=["cpu"])


def is_image_valid(image_uuid):
    images = client.get_prefix("/v1/image/")
    return image_uuid in [i.key.split("/")[-1] for i in images]


class CreateVM(Resource):
    def post(self):
        createvm_args = createvm_argparser.parse_args()
        name, realm, token, specs, image_uuid = createvm_args.name, createvm_args.realm,\
                                          createvm_args.token, createvm_args.specs,\
                                          createvm_args.image_uuid

        if check_otp(name, realm, token) == 200:
            # User is good
            if is_image_valid(image_uuid):
                if not specs_parser.transform_specs(specs):
                    return {"message": f"""Invalid unit - Please use following units {specs_parser.get_allowed_units()}"""}, 400

                print(specs)
                vm_key = f"/v1/vm/{uuid4().hex}"
                vm_entry = {"owner": name,
                            "specs": specs,
                            "hostname": "",
                            "status": "REQUESTED_NEW",
                            "image_uuid": image_uuid}

                etcd_client.put(vm_key, json.dumps(vm_entry))

                return {"message": "VM Creation Queued"}, 200
            else:
                return {"message": "Image uuid not valid"}
        else:
            return {"message": "Invalid Credentials"}, 400


class DeleteVM(Resource):
    def post(self):
        deletevm_args = deletevm_argparser.parse_args()
        name, realm, token, vmid = deletevm_args.name, deletevm_args.realm,\
                                  deletevm_args.token, deletevm_args.vmid

        if check_otp(name, realm, token) == 200:
            # User is good

            vmentry_etcd = etcd_client.get(f"/v1/vm/{vmid}")[0]
            if vmentry_etcd:
                vmentry_etcd = json.loads(vmentry_etcd)
                vmentry_etcd["status"] = "REQUESTED_DELETE"

                etcd_client.put(f"/v1/vm/{vmid}", json.dumps(vmentry_etcd))

                return {"message": "VM Deletion Queued"}, 200
            else:
                return {"message": "Invalid VM ID"}
        else:
            return {"message": "Invalid Credentials"}, 400


class VmStatus(Resource):
    def get(self):
        args = vmstatus_argparser.parse_args()
        r = etcd_client.get(f"/v1/vm/{args.vmid}")[0]
        print(r)
        if r:
            r = dict(json.loads(r.decode("utf-8")))
            return r
        return {"Message": "Not Found"}


class CreateImage(Resource):
    def post(self):
        image_stores = list(client.get_prefix("/v1/image_store/"))
        args = createimage_argparser.parse_args()
        image_file_uuid = args.uuid
        image_store_name = args.image_store

        file_entry = client.get(f"/v1/files/{image_file_uuid}")
        if file_entry is None:
            return {
                "Message":
                f"Image File with uuid '{image_file_uuid}' Not Found"}, 400

        file_entry_value = json.loads(file_entry.value)

        image_store = list(filter(lambda s: json.loads(s.value)["name"] == image_store_name, image_stores))
        if not image_store:
            return {"Message": f"Store '{image_store_name}' does not exists"}, 400

        image_store = image_store[0]
        image_entry_json = {
            "status": "TO_BE_CREATED",
            "owner": file_entry_value["owner"],
            "filename": file_entry_value["filename"],
            "name": args.name,
            "store_name": image_store_name,
            "visibility": "public"
        }
        client.put(f"/v1/image/{image_file_uuid}", json.dumps(image_entry_json))

        return {"Message": "Image successfully created"}


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

        return r, 200


class StartVM(Resource):
    def post(self):
        args = startvm_argparser.parse_args()
        name, realm, token, vm_uuid = args.name, args.realm, args.token, args.vmid

        if check_otp(name, realm, token) == 200:
            vm = client.get(f"/v1/vm/{vm_uuid}", value_in_json=True)
            if vm:
                vm.value["status"] = "REQUESTED_START"
                client.put(vm.key, json.dumps(vm.value))
                return {"message": f"VM Start Queued"}
            else:
                return {"message": "No such VM found"}
        else:
            return {"message": "Invalid Credentials"}, 400


class SuspendVM(Resource):
    def post(self):
        args = startvm_argparser.parse_args()
        name, realm, token, vm_uuid = args.name, args.realm, args.token, args.vmid

        if check_otp(name, realm, token) == 200:
            vm = client.get(f"/v1/vm/{vm_uuid}", value_in_json=True)
            if vm:
                vm.value["status"] = "REQUESTED_SUSPEND"
                client.put(vm.key, json.dumps(vm.value))
                return {"message": f"VM Suspension Queued"}
            else:
                return {"message": "No such VM found"}
        else:
            return {"message": "Invalid Credentials"}, 400


class ResumeVM(Resource):
    def post(self):
        args = startvm_argparser.parse_args()
        name, realm, token, vm_uuid = args.name, args.realm, args.token, args.vmid

        if check_otp(name, realm, token) == 200:
            vm = client.get(f"/v1/vm/{vm_uuid}", value_in_json=True)
            if vm:
                vm.value["status"] = "REQUESTED_RESUME"
                client.put(vm.key, json.dumps(vm.value))
                return {"message": f"VM Resume Queued"}
            else:
                return {"message": "No such VM found"}
        else:
            return {"message": "Invalid Credentials"}, 400


class ShutdownVM(Resource):
    def post(self):
        args = startvm_argparser.parse_args()
        name, realm, token, vm_uuid = args.name, args.realm, args.token, args.vmid

        if check_otp(name, realm, token) == 200:
            vm = client.get(f"/v1/vm/{vm_uuid}", value_in_json=True)
            if vm:
                vm.value["status"] = "REQUESTED_SHUTDOWN"
                client.put(vm.key, json.dumps(vm.value))
                return {"message": f"VM Shutdown Queued"}
            else:
                return {"message": "No such VM found"}
        else:
            return {"message": "Invalid Credentials"}, 400


class ListUserVM(Resource):
    def get(self):
        args = uservm_argparser.parse_args()
        name, realm, token = args.name, args.realm, args.token

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


class ListUserFiles(Resource):
    def get(self):
        args = uservm_argparser.parse_args()
        name, realm, token = args.name, args.realm, args.token

        if check_otp(name, realm, token) == 200:
            files = client.get_prefix(f"/v1/files/", value_in_json=True)
            if files:
                return_files = []
                user_files = list(filter(lambda f: f.value["owner"] == name, files))
                for file in user_files:
                    return_files.append({
                        "filename": vm.value["filename"]
                    })
                return {"message": return_files}, 200
            else:
                return {"message": "No File found"}, 404
        else:
            return {"message": "Invalid Credentials"}, 400

api.add_resource(CreateVM, "/vm/create")
api.add_resource(DeleteVM, "/vm/delete")
api.add_resource(VmStatus, "/vm/status")

api.add_resource(StartVM, "/vm/start")
api.add_resource(SuspendVM, "/vm/suspend")
api.add_resource(ResumeVM, "/vm/resume")
api.add_resource(ShutdownVM, "/vm/shutdown")

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")

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