diff --git a/Pipfile b/Pipfile index a5b5d2e..ef8427e 100644 --- a/Pipfile +++ b/Pipfile @@ -11,7 +11,7 @@ python-decouple = "*" requests = "*" flask = "*" flask-restful = "*" -python-etcd = "*" +etcd3 = "*" [requires] python_version = "3.7" diff --git a/helper.py b/helper.py index 41ea877..cbb7e71 100644 --- a/helper.py +++ b/helper.py @@ -3,7 +3,6 @@ import requests from decouple import config from pyotp import TOTP -from etcd import EtcdKeyNotFound def check_otp(name, realm, seed): @@ -27,23 +26,3 @@ def check_otp(name, realm, seed): data=data ) return response.status_code - - -def get_next_id(client, path): - try: - r = client.read(path) - except EtcdKeyNotFound: - return 0 - - max_key_result = max(r.children, key=lambda x: int(strip_nondigit(x.key))) - if max_key_result is None: - # No key found - return 0 - - max_key = strip_nondigit(max_key_result.key.split("/")[-1]) # Get the last portion of key - - return int(max_key) + 1 - - -def strip_nondigit(s): - return "".join([char for char in s if char.isdigit()]) diff --git a/main.py b/main.py index 2298351..8a619b9 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,17 @@ -import etcd +import etcd3 +import json -from helper import check_otp, get_next_id +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 = etcd.Client(host=config("ETCD_HOST"), port=int(config("ETCD_PORT"))) -etcd_lock = etcd.Lock(etcd_client, 'etcd_lock') +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) @@ -28,15 +29,11 @@ class CreateVM(Resource): if check_otp(name, realm, seed) == 200: # User is good - # Block until lock is acquired. There is no timeout associated with this lock. - with etcd_lock: - next_vm_id = get_next_id(etcd_client, "/v1/vm/") + vm_entry = {"owner": name, + "specs": specs, + "status": VmStatus.REQUESTED_NEW.value} - vm_entry = {"owner": name, - "specs": specs, - "status": VmStatus.REQUESTED_NEW} - - etcd_client.write(f"/v1/vm/{next_vm_id}", vm_entry) + etcd_client.put(f"/v1/vm/{uuid4().hex}", json.dumps(vm_entry)) return {'message': "VM Creation Queued"}, 200 else: