# TODO: realm should be part of user's name. Because, there can be multiple # name but with different realms. For Example, nico exists in both # ungleich-admin and ungleich-twitter realm. So, to differentiate between # them, we need to make realm part of user's name import json import subprocess import os from uuid import uuid4 from flask import Flask, request from flask_restful import Resource, Api from ucloud_common.vm import VMStatus from ucloud_common.request import RequestEntry, RequestType from config import etcd_client as client from config import ( WITHOUT_CEPH, VM_PREFIX, HOST_PREFIX, FILE_PREFIX, IMAGE_PREFIX, logging, REQUEST_POOL, VM_POOL, HOST_POOL, ) from schemas import ( CreateVMSchema, VMStatusSchema, CreateImageSchema, VmActionSchema, OTPSchema, CreateHostSchema, VmMigrationSchema, AddSSHSchema, RemoveSSHSchema ) app = Flask(__name__) api = Api(app) class CreateVM(Resource): @staticmethod def post(): data = request.json validator = CreateVMSchema(data) if validator.is_valid(): vm_uuid = uuid4().hex vm_key = os.path.join(VM_PREFIX, vm_uuid) vm_entry = { "name": data["vm_name"], "owner": data["name"], "specs": data["specs"], "hostname": "", "status": "", "image_uuid": data["image_uuid"], "log": [], "storage_attachment": [], "vnc_socket": "", } client.put(vm_key, vm_entry, value_in_json=True) # Create ScheduleVM Request r = RequestEntry.from_scratch(type=RequestType.ScheduleVM, uuid=vm_uuid) REQUEST_POOL.put(r) return {"message": "VM Creation Queued"}, 200 return validator.get_errors(), 400 class VmStatus(Resource): @staticmethod def get(): data = request.json validator = VMStatusSchema(data) if validator.is_valid(): vm = VM_POOL.get(os.path.join(VM_PREFIX, data["uuid"])) return str(vm) else: return validator.get_errors(), 400 class CreateImage(Resource): @staticmethod def post(): data = request.json validator = CreateImageSchema(data) if validator.is_valid(): file_entry = client.get(os.path.join(FILE_PREFIX, data["uuid"])) file_entry_value = json.loads(file_entry.value) image_entry_json = { "status": "TO_BE_CREATED", "owner": file_entry_value["owner"], "filename": file_entry_value["filename"], "name": data["name"], "store_name": data["image_store"], "visibility": "public", } client.put( os.path.join(IMAGE_PREFIX, data["uuid"]), json.dumps(image_entry_json) ) return {"message": "Image successfully created"} return validator.get_errors(), 400 class ListPublicImages(Resource): @staticmethod def get(): images = client.get_prefix(IMAGE_PREFIX) r = {} for image in images: r[image.key.split("/")[-1]] = json.loads(image.value) return r, 200 class VMAction(Resource): @staticmethod def post(): data = request.json validator = VmActionSchema(data) if validator.is_valid(): vm_entry = VM_POOL.get(os.path.join(VM_PREFIX, data["uuid"])) action = data["action"] if action == "start": vm_entry.status = VMStatus.requested_start VM_POOL.put(vm_entry) action = "schedule" if action == "delete" and vm_entry.hostname == "": try: path_without_protocol = vm_entry.path[vm_entry.path.find(":") + 1 :] if WITHOUT_CEPH: command_to_delete = [ "rm", os.path.join("/var/vm", vm_entry.uuid), ] else: command_to_delete = ["rbd", "rm", path_without_protocol] subprocess.check_output(command_to_delete, stderr=subprocess.PIPE) except subprocess.CalledProcessError as e: if "No such file" in e.stderr.decode("utf-8"): client.client.delete(vm_entry.key) return {"message": "VM successfully deleted"} else: logging.exception(e) return {"message": "Some error occurred while deleting VM"} else: client.client.delete(vm_entry.key) return {"message": "VM successfully deleted"} r = RequestEntry.from_scratch( type="{}VM".format(action.title()), uuid=data["uuid"], hostname=vm_entry.hostname, ) REQUEST_POOL.put(r) return {"message": "VM {} Queued".format(action.title())}, 200 else: return validator.get_errors(), 400 class VMMigration(Resource): @staticmethod def post(): data = request.json validator = VmMigrationSchema(data) if validator.is_valid(): vm = VM_POOL.get(data["uuid"]) r = RequestEntry.from_scratch( type=RequestType.ScheduleVM, uuid=vm.uuid, destination=os.path.join(HOST_PREFIX, data["destination"]), migration=True, ) REQUEST_POOL.put(r) return {"message": "VM Migration Initialization Queued"}, 200 else: return validator.get_errors(), 400 class ListUserVM(Resource): @staticmethod def get(): data = request.json validator = OTPSchema(data) if validator.is_valid(): vms = client.get_prefix(VM_PREFIX, value_in_json=True) return_vms = [] user_vms = filter(lambda v: v.value["owner"] == data["name"], vms) for vm in user_vms: return_vms.append( { "name": vm.value["name"], "vm_uuid": vm.key.split("/")[-1], "specs": vm.value["specs"], "status": vm.value["status"], "hostname": vm.value["hostname"], "vnc_socket": None if vm.value.get("vnc_socket", None) == None else vm.value["vnc_socket"], } ) if return_vms: return {"message": return_vms}, 200 return {"message": "No VM found"}, 404 else: return validator.get_errors(), 400 class ListUserFiles(Resource): @staticmethod def get(): data = request.json validator = OTPSchema(data) if validator.is_valid(): files = client.get_prefix(FILE_PREFIX, value_in_json=True) return_files = [] user_files = list(filter(lambda f: f.value["owner"] == data["name"], files)) for file in user_files: return_files.append( { "filename": file.value["filename"], "uuid": file.key.split("/")[-1], } ) return {"message": return_files}, 200 else: return validator.get_errors(), 400 class CreateHost(Resource): @staticmethod def post(): data = request.json validator = CreateHostSchema(data) if validator.is_valid(): host_key = os.path.join(HOST_PREFIX, uuid4().hex) host_entry = { "specs": data["specs"], "hostname": data["hostname"], "status": "DEAD", "last_heartbeat": "", } client.put(host_key, host_entry, value_in_json=True) return {"message": "Host Created"}, 200 return validator.get_errors(), 400 class ListHost(Resource): @staticmethod def get(): hosts = HOST_POOL.hosts r = { host.key: { "status": host.status, "specs": host.specs, "hostname": host.hostname, } for host in hosts } return r, 200 class AddSSHKey(Resource): @staticmethod def get(): data = request.json validator = AddSSHSchema(data) if validator.is_valid(): etcd_key = "/v1/user/{}/keys".format(data["name"]) etcd_entry = client.get(etcd_key, value_in_json=True) if etcd_entry: if data["key_name"] not in etcd_entry.value: etcd_entry.value[data["key_name"]] = data["key"] client.put(etcd_entry.key, etcd_entry.value, value_in_json=True) else: return {"message": "Key with name '{}' already exists".format(data["key_name"])} else: # Key Not Found. It implies user' haven't added any key yet. key_entry = {data["key_name"]: data["key"]} client.put(etcd_key, key_entry, value_in_json=True) return {"message": "Key added successfully"} else: return validator.get_errors(), 400 class RemoveSSHKey(Resource): @staticmethod def get(): data = request.json validator = RemoveSSHSchema(data) if validator.is_valid(): etcd_key = "/v1/user/{}/keys".format(data["name"]) etcd_entry = client.get(etcd_key, value_in_json=True) if etcd_entry: try: del etcd_entry.value[data["key_name"]] except KeyError: return {"message": "No such key exists."} else: client.put(etcd_entry.key, etcd_entry.value, value_in_json=True) return {"message": "Key successfully removed."} else: return {"message": "No Key Exists at all."} else: return validator.get_errors(), 400 api.add_resource(CreateVM, "/vm/create") api.add_resource(VmStatus, "/vm/status") api.add_resource(VMAction, "/vm/action") api.add_resource(VMMigration, "/vm/migrate") api.add_resource(CreateImage, "/image/create") api.add_resource(ListPublicImages, "/image/list-public") api.add_resource(ListUserVM, "/user/vms") api.add_resource(ListUserFiles, "/user/files") api.add_resource(AddSSHKey, "/user/add-ssh") api.add_resource(RemoveSSHKey, "/user/remove-ssh") api.add_resource(CreateHost, "/host/create") api.add_resource(ListHost, "/host/list") if __name__ == "__main__": app.run(host="::", debug=True)