uncloud/imagescanner/main.py

import os
import json
import subprocess
import sys

from config import (logging, client, IMAGE_DIR,
                    BASE_PATH, WITHOUT_CEPH,
                    IMAGE_PREFIX, IMAGE_STORE_PREFIX)


def qemu_img_type(path):
    qemu_img_info_command = ["qemu-img", "info", "--output", "json", path]
    try:
        qemu_img_info = subprocess.check_output(qemu_img_info_command)
    except Exception as e:
        logging.exception(e)
        return None
    else:
        qemu_img_info = json.loads(qemu_img_info.decode("utf-8"))
        return qemu_img_info["format"]


def main():
    # If you are using WITHOUT_CEPH FLAG in .env
    # then please make sure that IMAGE_DIR directory
    # exists otherwise this script would fail
    if WITHOUT_CEPH and not os.path.isdir(IMAGE_DIR):
        print("You have set WITHOUT_CEPH to True. So,"
            "the {} must exists. But, it don't".format(IMAGE_DIR))
        sys.exit(1)

    try:
        subprocess.check_output(['which', 'qemu-img'])
    except Exception:
        print("qemu-img missing")
        sys.exit(1)

    # We want to get images entries that requests images to be created
    images = client.get_prefix(IMAGE_PREFIX, value_in_json=True)
    images_to_be_created = list(filter(lambda im: im.value['status'] == 'TO_BE_CREATED', images))

    for image in images_to_be_created:
        try:
            image_uuid = image.key.split('/')[-1]
            image_owner = image.value['owner']
            image_filename = image.value['filename']
            image_store_name = image.value['store_name']
            image_full_path = os.path.join(BASE_PATH, image_owner, image_filename)

            image_stores = client.get_prefix(IMAGE_STORE_PREFIX, value_in_json=True)
            user_image_store = next(filter(
                lambda s, store_name=image_store_name: s.value["name"] == store_name,
                image_stores
            ))

            image_store_pool = user_image_store.value['attributes']['pool']

        except Exception as e:
            logging.exception(e)
        else:
            # At least our basic data is available

            qemu_img_convert_command = ["qemu-img", "convert", "-f", "qcow2",
                                        "-O", "raw", image_full_path, "image.raw"]


            if WITHOUT_CEPH:
                image_import_command = ["mv", "image.raw", os.path.join(IMAGE_DIR, image_uuid)]
                snapshot_creation_command = ["true"]
                snapshot_protect_command = ["true"]
            else:
                image_import_command = ["rbd", "import", "image.raw",
                                        "{}/{}".format(image_store_pool, image_uuid)]
                snapshot_creation_command = ["rbd", "snap", "create",
                                            "{}/{}@protected".format(image_store_pool, image_uuid)]
                snapshot_protect_command = ["rbd", "snap", "protect",
                                            "{}/{}@protected".format(image_store_pool, image_uuid)]

            
            # First check whether the image is qcow2

            if qemu_img_type(image_full_path) == "qcow2":
                try:
                    # Convert .qcow2 to .raw
                    subprocess.check_output(qemu_img_convert_command)

                    # Import image either to ceph/filesystem
                    subprocess.check_output(image_import_command)

                    # Create and Protect Snapshot
                    subprocess.check_output(snapshot_creation_command)
                    subprocess.check_output(snapshot_protect_command)
                
                except Exception as e:
                    logging.exception(e)
                
                else:
                    # Everything is successfully done
                    image.value["status"] = "CREATED"
                    client.put(image.key, json.dumps(image.value))
            else:
                # The user provided image is either not found or of invalid format
                image.value["status"] = "INVALID_IMAGE"
                client.put(image.key, json.dumps(image.value))


        try:
            os.remove("image.raw")
        except Exception:
            pass


if __name__ == "__main__":
    main()