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 enums import VmStatus 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) 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": VmStatus.REQUESTED_NEW.value} 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 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(VmStatus, '/vm/status') if __name__ == '__main__': app.run(debug=True)