uncloud/ucloud/imagescanner/main.py

import json
import os
import subprocess as sp

from os.path import join as join_path
from ucloud.settings import settings
from ucloud.shared import shared
from ucloud.imagescanner import logger


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


def main():
    # We want to get images entries that requests images to be created
    images = shared.etcd_client.get_prefix(
        settings["etcd"]["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 = join_path(
                settings["storage"]["file_dir"],
                image_owner,
                image_filename,
            )

            image_stores = shared.etcd_client.get_prefix(
                settings["etcd"]["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:
            logger.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 qemu_img_type(image_full_path) == "qcow2":
                try:
                    # Convert .qcow2 to .raw
                    sp.check_output(qemu_img_convert_command,)

                except sp.CalledProcessError:
                    logger.exception(
                        "Image convertion from .qcow2 to .raw failed."
                    )
                else:
                    # Import and Protect
                    r_status = shared.storage_handler.import_image(
                        src="image.raw", dest=image_uuid, protect=True
                    )
                    if r_status:
                        # Everything is successfully done
                        image.value["status"] = "CREATED"
                        shared.etcd_client.put(
                            image.key, json.dumps(image.value)
                        )
                finally:
                    try:
                        os.remove("image.raw")
                    except Exception:
                        pass

            else:
                # The user provided image is either not found or of invalid format
                image.value["status"] = "INVALID_IMAGE"
                shared.etcd_client.put(
                    image.key, json.dumps(image.value)
                )


if __name__ == "__main__":
    main()