import etcd3 import json from helper import check_otp from flask import Flask from flask_restful import Resource, Api, reqparse from decouple import config from uuid import uuid4 app = Flask(__name__) api = Api(app) etcd_client = etcd3.client(host=config("ETCD_HOST"), port=int(config("ETCD_PORT"))) 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) 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) 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 if check_otp(name, realm, seed) == 200: # User is good vm_entry = {"owner": name, "specs": specs, "hostname": "", "status": "REQUESTED_NEW"} etcd_client.put(f"/v1/vm/{uuid4().hex}", json.dumps(vm_entry)) return {"message": "VM Creation Queued"}, 200 else: return {"message": "Invalid Credentials"}, 400 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 if check_otp(name, realm, seed) == 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" 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.id}")[0] if r: r = dict(json.loads(r.decode("utf-8"))) return r return {"Message": "Not Found"} api.add_resource(CreateVM, "/vm/create") api.add_resource(DeleteVM, "/vm/delete") api.add_resource(VmStatus, "/vm/status") if __name__ == "__main__": app.run(debug=True)