diff --git a/helper.py b/helper.py index cbb7e71..4b580f1 100644 --- a/helper.py +++ b/helper.py @@ -5,7 +5,7 @@ from decouple import config from pyotp import TOTP -def check_otp(name, realm, seed): +def check_otp(name, realm, token): try: data = { "auth_name": config('AUTH_NAME', ''), @@ -13,7 +13,7 @@ def check_otp(name, realm, seed): "auth_realm": config('AUTH_REALM', ''), "name": name, "realm": realm, - "token": TOTP(seed).now() + "token": token } except binascii.Error: return 400 @@ -26,3 +26,15 @@ def check_otp(name, realm, seed): data=data ) return response.status_code + + +def add_otp_args(parser): + parser.add_argument("name", required=True) + parser.add_argument("realm", required=True) + parser.add_argument("token", required=True) + return parser + + +def add_vmid_args(parser): + parser.add_argument("vmid", required=True) + return parser diff --git a/main.py b/main.py index 624eaa5..8aa1470 100644 --- a/main.py +++ b/main.py @@ -4,7 +4,7 @@ import etcd3 import json -from helper import check_otp +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 @@ -17,44 +17,62 @@ 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("name", type=str, required=True) -createvm_argparser.add_argument("realm", type=str, required=True) -createvm_argparser.add_argument("seed", type=str, required=True) createvm_argparser.add_argument("specs", type=dict, required=True) +createvm_argparser.add_argument("image_uuid", type=str, required=True) +add_otp_args(createvm_argparser) -deletevm_argparser = reqparse.RequestParser() -deletevm_argparser.add_argument("name", type=str, required=True) -deletevm_argparser.add_argument("realm", type=str, required=True) -deletevm_argparser.add_argument("seed", type=str, required=True) -deletevm_argparser.add_argument("vmid", type=str, required=True) - -vmstatus_argparser = reqparse.RequestParser() -vmstatus_argparser.add_argument("id", type=str, required=True) - +# 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) + + +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, seed, specs = createvm_args.name, createvm_args.realm,\ - createvm_args.seed, createvm_args.specs + 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, seed) == 200: + if check_otp(name, realm, token) == 200: # User is good + if is_image_valid(image_uuid): + vm_key = f"/v1/vm/{uuid4().hex}" + vm_entry = {"owner": name, + "specs": specs, + "hostname": "", + "status": "REQUESTED_NEW", + "image_uuid": image_uuid} - vm_entry = {"owner": name, - "specs": specs, - "hostname": "", - "status": "REQUESTED_NEW"} + etcd_client.put(vm_key, json.dumps(vm_entry)) - etcd_client.put(f"/v1/vm/{uuid4().hex}", json.dumps(vm_entry)) - - return {"message": "VM Creation Queued"}, 200 + return {"message": "VM Creation Queued"}, 200 + else: + return {"message": "Image uuid not valid"} else: return {"message": "Invalid Credentials"}, 400 @@ -62,16 +80,16 @@ class CreateVM(Resource): class DeleteVM(Resource): def post(self): deletevm_args = deletevm_argparser.parse_args() - name, realm, seed, vmid = deletevm_args.name, deletevm_args.realm,\ - deletevm_args.seed, deletevm_args.vmid + name, realm, token, vmid = deletevm_args.name, deletevm_args.realm,\ + deletevm_args.token, deletevm_args.vmid - if check_otp(name, realm, seed) == 200: + 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"] = "REQUEST_DELETE" + vmentry_etcd["status"] = "REQUESTED_DELETE" etcd_client.put(f"/v1/vm/{vmid}", json.dumps(vmentry_etcd)) @@ -85,7 +103,8 @@ class DeleteVM(Resource): class VmStatus(Resource): def get(self): args = vmstatus_argparser.parse_args() - r = etcd_client.get(f"/v1/vm/{args.id}")[0] + r = etcd_client.get(f"/v1/vm/{args.vmid}")[0] + print(r) if r: r = dict(json.loads(r.decode("utf-8"))) return r @@ -117,18 +136,109 @@ class CreateImage(Resource): "owner": file_entry_value["owner"], "filename": file_entry_value["filename"], "name": args.name, - "store_name": image_store_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 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 + + api.add_resource(CreateVM, "/vm/create") api.add_resource(DeleteVM, "/vm/delete") api.add_resource(VmStatus, "/vm/status") -api.add_resource(CreateImage, "/image/create") +api.add_resource(StartVM, "/vm/start") +api.add_resource(SuspendVM, "/vm/suspend") +api.add_resource(ResumeVM, "/vm/resume") +api.add_resource(CreateImage, "/image/create") +api.add_resource(ListPublicImages, "/image/list-public") + +api.add_resource(ListUserVM, "/user/vms") if __name__ == "__main__": app.run(host="::", debug=True)