ListPublicImages, StartVM, SuspendVM, ResumeVM, ListUserVM Views added.

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

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()
+# CreateImage argparser
-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 = 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,\
+        name, realm, token, specs, image_uuid = createvm_args.name, createvm_args.realm,\
-                                   createvm_args.seed, createvm_args.specs
+                                          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"}
+                            "status": "REQUESTED_NEW",
+                            "image_uuid": image_uuid}
 
-            etcd_client.put(f"/v1/vm/{uuid4().hex}", json.dumps(vm_entry))
+                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
 
@@ -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,\
+        name, realm, token, vmid = deletevm_args.name, deletevm_args.realm,\
-                                  deletevm_args.seed, deletevm_args.vmid
+                                  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)