ucloud-api/main.py

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

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
from functools import wraps

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

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

# AddHost argparser
addhost_argparser = reqparse.RequestParser()
addhost_argparser.add_argument("hostname", type=str, required=True)
addhost_argparser.add_argument("specs", type=dict, required=True)
add_otp_args(addhost_argparser)


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,
                }

                client.put(vm_key, vm_entry, value_in_json=True)

                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 = client.get(f"/v1/vm/{vmid}").value
            if vmentry_etcd:
                vmentry_etcd = json.loads(vmentry_etcd)
                vmentry_etcd["status"] = "REQUESTED_DELETE"

                client.put(f"/v1/vm/{vmid}", vmentry_etcd,
                           value_in_json=True)

                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 = client.get(f"/v1/vm/{args.vmid}").value
        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.value["owner"] != name:
                return {"message": "Invalid User"}
            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.value["owner"] != name:
                return {"message": "Invalid User"}
            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.value["owner"] != name:
                return {"message": "Invalid User"}
            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.value["owner"] != name:
                return {"message": "Invalid User"}

            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": file.value["filename"],
                            "uuid": file.key.split("/")[-1],
                        }
                    )
                return {"message": return_files}, 200
            else:
                return {"message": "No File found"}, 404
        else:
            return {"message": "Invalid Credentials"}, 400


class CreateHost(Resource):
    def post(self):
        args = addhost_argparser.parse_args()
        name, realm, token, specs, hostname = (
            args.name,
            args.realm,
            args.token,
            args.specs,
            args.hostname,
        )

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

            print(specs)
            host_key = f"/v1/host/{uuid4().hex}"
            host_entry = {
                "specs": specs,
                "hostname": hostname,
                "status": "DEAD",
                "last_heartbeat": "",
            }
            client.put(host_key, host_entry, value_in_json=True)

            return {"message": "Host Created"}, 200
        else:
            return {"message": "Invalid Credentials/Insufficient Permission"}, 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")

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


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