fetch.py

#!/usr/bin/python3
"""
fetch.py

Fetches and deprecates images defined in input.json
For more details refer to README.md


"""
import sys
import os
import json
import time
import hashlib
import logging
from logging.handlers import SysLogHandler
import yaml
import requests
import openstack
import openstack.exceptions


class ImgBuildLogger:
    """
    Set value for exitting the program and handle logging
    0 = OK
    1 = Warning
    2 = Error
    """

    def __init__(self, **kwargs):
        """Initialize log object"""
        self.config = kwargs
        self._code = 0
        self._log = logging.getLogger("imagebuilder")
        self._log.setLevel(logging.DEBUG)

        sysloghandler = SysLogHandler()
        sysloghandler.setLevel(logging.DEBUG)
        self._log.addHandler(sysloghandler)

        streamhandler = logging.StreamHandler(sys.stdout)
        if self.config["output_format"] == "PLAIN":
            stream_formatter = logging.Formatter("%(name)s %(levelname)s %(message)s")
            streamhandler.setFormatter(stream_formatter)
        streamhandler.setLevel(
            logging.getLevelName(self.config.get("debug_level", "INFO"))
        )
        self._log.addHandler(streamhandler)

        if "log_file" in self.config:
            log_file = self.config["log_file"]
        else:
            home_folder = os.environ.get("HOME", os.environ.get("USERPROFILE", ""))
            log_folder = os.path.join(home_folder, "log")
            log_file = os.path.join(log_folder, "imagebuilder.log")

        if not os.path.exists(os.path.dirname(log_file)):
            os.mkdir(os.path.dirname(log_file))

        filehandler = logging.handlers.RotatingFileHandler(log_file, maxBytes=102400000)
        if self.config["output_format"] == "PLAIN":
            plain_file_formatter = logging.Formatter(
                "%(asctime)s %(name)s %(levelname)s %(message)s"
            )
            filehandler.setFormatter(plain_file_formatter)
        filehandler.setLevel(logging.DEBUG)
        self._log.addHandler(filehandler)

    def _output(self, message, severity):
        if self.config["output_format"] == "PLAIN":
            return f"{time.time()} {message}"
        if isinstance(message, str):
            message = {"message": message}
        if "timestamp" not in message:
            message["timestamp"] = time.time()
        message["severity"] = severity
        message["log_name"] = self._log.name

        if self.config["output_format"] == "JSON":
            return json.dumps(message)
        if self.config["output_format"] == "YAML":
            return yaml.dump(message, Dumper=yaml.Dumper)

        self._log.warning(
            {"message": f"Output format '{self.config['output_format']}' not supported"}
        )
        return message

    def info(self, message):
        """Output information message"""
        return self._log.info(self._output(message, "info"))

    def warning(self, message):
        """Output warning message"""
        self._code = max(self._code, 1)
        return self._log.warning(self._output(message, "warning"))

    def error(self, message):
        """Output error message"""
        self._code = 2
        return self._log.error(self._output(message, "error"))

    def debug(self, message):
        """Output debugging message"""
        return self._log.debug(self._output(message, "debug"))

    @property
    def exit_code(self) -> int:
        """
        Return the set code
        Returns
        -------
        int
            Current exit code
        """
        return self._code


cloud_name = os.getenv(
    "IMAGEBUILDER_CLOUD"
)  # get it once to avoid potential race conditions
network_name = os.getenv("IMAGEBUILDER_NETWORK")

logger = ImgBuildLogger(
    log_file=os.getenv("IMAGEBUILDER_LOG_FILE", f"./{cloud_name}_log.json"),
    output_format=os.getenv("IMAGEBUILDER_OUTPUT_FORMAT", "JSON"),
    debug_level=os.getenv("IMAGEBUILDER_DEBUG_LEVEL", "INFO"),
)


def get_file_hash(filename: str, cur_hash: hashlib._hashlib.HASH) -> str:
    """
    Calculate the hash of a file chunked
    Returns the hexdigest in lowercase

    Returns
    -------
    str
        Hash of filename in the format of cur_hash
        Empty string if file does not exist
    """

    if not os.path.exists(filename):
        return ""

    with open(filename, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            cur_hash.update(chunk)

    return cur_hash.hexdigest().lower()


def download_image(url: str, filename: str, new_checksum: str) -> bool:
    """
    Downloads a file from a specified url using chucking
    Also converts it to a .raw file

    Returns
    -------
    bool
        True if everything is successful
        False if something goes wrong
    """

    print_progressbar = sys.stdout.isatty() or (
        os.getenv("IMAGEBUILDER_OUTPUT_FORMAT", "JSON") == "plain"
    )

    # Check if the same file already exists on disk
    cur_hash = get_file_hash("tmp/" + filename, hashlib.sha256())
    if cur_hash != "" and cur_hash in new_checksum.lower():

        logger.info(f"The newest version of {filename} already exists on disk.")

        # Local file exists but raw doesn't, create one
        if not os.path.exists("tmp/" + filename + ".raw"):
            logger.info("RAW file did not exist. Creating...")
            os.system(
                f"qemu-img convert {'-p' if print_progressbar else ''} \
                      -O raw tmp/{filename} tmp/{filename}.raw"
            )

        return True

    # Remove old files that may exist before downloading new ones
    cleanup_files(filename)

    try:
        with requests.get(
            url,
            allow_redirects=True,
            timeout=600,
            stream=True,
            headers={'Accept-Encoding': 'deflate'}) as r:
            with open("tmp/" + filename, "wb") as f:
                file_size = int(r.headers.get("content-length"))
                progress = 0

                if print_progressbar:
                    disp_filename = filename[
                        0 : (int(os.get_terminal_size().columns * 0.5) - 1)
                    ]
                    if disp_filename != filename:
                        disp_filename = disp_filename[:-3] + "..."

                    bar_width = min(
                        50, int(os.get_terminal_size().columns - len(disp_filename) - 3)
                    )

                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)
                    progress += len(chunk)
                    if print_progressbar:
                        done = int(bar_width * progress / file_size)
                        sys.stdout.write(
                            f"\r{disp_filename}\t[{'=' * done}{' ' * (bar_width-done)}]"
                        )
                        sys.stdout.flush()

            if print_progressbar:
                sys.stdout.write("\n")
                sys.stdout.flush()

        new_hash = get_file_hash("tmp/" + filename, hashlib.sha256())
        if (
            new_hash != "" and new_hash not in new_checksum.lower()
        ) and new_checksum != filename:
            logger.error(f"Checksum of {filename} does not match the one downloaded!")
            cleanup_files(filename)
            return False

        logger.info(f"Converting {filename} to .raw")

        # Convert to raw
        os.system(
            f"qemu-img convert {'-p' if print_progressbar else ''}\
                   -O raw tmp/{filename} tmp/{filename}.raw"
        )

        return True

    except requests.exceptions.HTTPError:
        logger.error(f"HTTP error while downloading {filename}")

    except requests.exceptions.ReadTimeout:
        logger.error(f"Download for {filename} timed out")

    except requests.exceptions.ConnectionError:
        logger.error(f"Unable to download {filename}")

    return False


def validate_raw_checksum(
    conn: openstack.connection.Connection, filename: str, version: any
) -> bool:
    """
    Compares the currently existing raw file's md5 checksum to the one on openstack.
    Returns true if they match

    Returns
    -------
    bool
        True if the MD5 hash of the local raw file is the same as the current one on openstack
        False otherwise
    """
    checksum_matches = False

    images = list(
        conn.image.images(
            name=version["image_name"],
            owner=conn.current_project_id,
            visibility=version["visibility"],
        )
    )

    for cur_img in images:
        if cur_img.checksum.lower() == get_file_hash(
            "tmp/" + filename + ".raw", hashlib.md5()
        ):
            checksum_matches = True

    if len(images) > 1:
        logger.warning(f"More than 1 image for {version['image_name']} already exists")

    return checksum_matches


def validate_checksum(
    version: any, filename: str, conn: openstack.connection.Connection, cloud: str
) -> str | None:
    """
    Validates checksums to check if an image requires updating

    Returns
    -------
    str
        A new checksum that's fetched online from version.checksum_url
        If checksum_url is not specified within version, the filename will be returned

    None
        If any critical error occures None is returned.
    """

    old_checksum = None
    if os.path.exists(
        f"checksums/{cloud}_{version['image_name'].replace(' ', '_')}_CHECKSUM"
    ):
        with open(
            f"checksums/{cloud}_{version['image_name'].replace(' ', '_')}_CHECKSUM",
            "r",
            encoding="utf-8",
        ) as f:
            old_checksum = f.read()
    elif version.get("checksum_url"):
        logger.info(
            f"Checksum file for {cloud}_{version['image_name']} does not exist. Creating one"
        )

    if not version.get("checksum_url"):
        logger.info(
            f"Checksum URL not specified for {version['image_name']}. Skipping checksum fetching"
        )
        return filename

    error = False
    new_checksum = old_checksum
    try:
        new_checksum = next(
            (
                line
                for line in requests.get(version["checksum_url"], timeout=5).text.split(
                    "\n"
                )
                if filename in line and not line.startswith("#")
            ),
            None,
        )

        if new_checksum is None:
            logger.error("New checksum could not be found in the list")
            new_checksum = old_checksum

    except requests.exceptions.HTTPError:
        error = True
        logger.error(
            f"HTTP error while downloading checksum for {version['image_name']}"
        )

    except requests.exceptions.ReadTimeout:
        error = True
        logger.error(f"Fetching the checksum for {version['image_name']} timed out")

    except requests.exceptions.ConnectionError:
        error = True
        logger.error(f"Unable to fetch new checksum for {version['image_name']}")

    if error:
        logger.error(f"An error occured validating {version['image_name']}")
        return None

    if new_checksum == old_checksum and validate_raw_checksum(conn, filename, version):

        current_images = list(
            conn.image.images(
                name=version["image_name"],
                owner=conn.current_project_id,
                visibility=version["visibility"],
                sort_key="created_at",
                sort_dir="desc",
            )
        )

        current_img_id = current_images[0].id if current_images else None

        if len(current_images) > 1:
            logger.warning(
                f"More than 1 image for {version['image_name']} already exists"
            )
        else:

            logger.info(f"{version['image_name']} already up to date")
            # Delete unused if they exist
            delete_unused_image(conn, version["image_name"], current_img_id)

        return None

    return new_checksum


def test_image_pinging(conn: openstack.connection.Connection, server_id: int) -> bool:
    """
    Test if the newly built image can be pinged by attaching it to a network
    Creates a router, interface and a floating ip which get cleaned up

    Returns
    -------
    bool
        True if the test server can be pinged properly
        False if an error is detected
    """

    disable_pinging = os.getenv("IMAGEBUILDER_DISABLE_PINGING")

    # the following test covers two cases
    # 1 - env var IMAGEBUILDER_DISABLE_PINGING is not defined (disable_pinging is None)
    # 2 - env var IMAGEBUILDER_DISABLE_PINGING is an empty string
    # in both cases, the result of the test is False
    if disable_pinging:
        logger.info("Skipping ping test...")
        return True

    public_id = conn.network.find_network("public", is_router_external=True).id

    try:
        floating_ip = conn.network.create_ip(floating_network_id=public_id)
    except openstack.exceptions.ConflictException as e:
        logger.error("Creation of floating ip failed")
        logger.error(str(e))
        return False

    port = next(conn.network.ports(device_id=server_id), None)

    if port is None:
        logger.error("No network port found")
        conn.network.delete_ip(floating_ip)
        return False

    floating_ip = conn.network.update_ip(floating_ip.id, port_id=port.id)

    logger.info(f"Testing pinging {floating_ip.floating_ip_address}")

    ping_result = os.system(
        f"timeout 360 bash -c \
        'while ! \
                ping -c 1 -W 1 {floating_ip.floating_ip_address} \
                > /dev/null 2>&1; \
                do sleep 1; \
            done'"
    )

    conn.network.delete_ip(floating_ip)

    if ping_result == 0:
        logger.info("Ping tests ok!")
        return True

    return False


def test_image(
        conn: openstack.connection.Connection,
        image: any,
        network: str,
        timeout: int,
        flavor: str) -> bool:
    """
    Test the newly created image by creating a test server and pinging it
    Returns
    -------
    bool
        True if the server can be created and pinged
        False if errors are detected
    """

    logger.info("Testing newly created image")

    secgroup = conn.network.find_security_group(
        "IMAGEBUILDER_PING_TEST", project_id=conn.current_project.id
    )

    if secgroup is None:
        secgroup = conn.network.create_security_group(name="IMAGEBUILDER_PING_TEST")
        logger.info(
            {
                "message": "Security group 'IMAGEBUILDER_PING_TEST' created",
                "security_group": secgroup,
            }
        )
        rule = conn.network.create_security_group_rule(
            security_group_id=secgroup.id,
            direction="ingress",
            ether_type="IPv4",
            protocol="icmp",
        )
        logger.info(
            {
                "message": "Created rule to allow ICMP ingress in security group",
                "rule": rule,
            }
        )
        rule = conn.network.create_security_group_rule(
            security_group_id=secgroup.id,
            direction="egress",
            ether_type="IPv4",
            protocol="icmp",
        )
        logger.info(
            {
                "message": "Created rule to allow ICMP egress in security group",
                "rule": rule,
            }
        )
    else:
        logger.info("Security group 'IMAGEBUILDER_PING_TEST' already exists.")

    try:
        test_server = conn.create_server(
            name=image.name + "_TESTSERVER",
            image=image,
            flavor=flavor,
            wait=True,
            auto_ip=False,
            network=network,
            timeout=timeout,
            security_groups=[secgroup.id],
        )
    except openstack.exceptions.SDKException as error:
        logger.error(
            {
                "message": f"Error creating server. {error}",
                "image_name": image.name,
                "network": network,
                "security_groups": secgroup.id,
                "server_name": image.name + "_TESTSERVER",
                "timeout": timeout,
            }
        )
        return False

    logger.info(f"Server with id ({test_server.id}) created")

    found_test_server = conn.get_server_by_id(test_server.id)

    if found_test_server["status"] == "ERROR":
        conn.delete_server(test_server.id)
        logger.error(
            {
                "message": f"Server {test_server.id} failed to start with image {image.id}!",
                "image_name": image.name
            }
        )
        return False

    # Test pinging
    ping_result = test_image_pinging(conn, test_server.id)

    conn.delete_server(test_server.id, wait=True)

    if not ping_result:
        logger.error(
            {
                "message": f"Server {test_server.id} failed to respond to ping!",
                "image_name": image.name
            }
        )
        return False

    logger.info("All tests ok!")

    return True


def cleanup_files(filename: str) -> None:
    """
    Delete image files that the program creates.
    """

    try:
        if os.path.exists(f"tmp/{filename}"):
            os.remove(f"tmp/{filename}")
    except OSError as error:
        logger.warning(f"Error removing file 'tmp/{filename}' from disk. {error}")

    try:
        if os.path.exists(f"tmp/{filename}.raw"):
            os.remove(f"tmp/{filename}.raw")
    except OSError as error:
        logger.warning(f"Error removing file 'tmp/{filename}' from disk. {error}")


def create_image(
    conn: openstack.connection.Connection,
    version: any,
    filename: str,
    new_checksum: str,
    network: str,
) -> any:
    """
    Creates an image

    Returns
    -------
    any:
        If everything is successful the newly created image is returned.
        If the image is already up to date or if there's an error None is returned
    """

    if "timeout" in version:
        timeout = version["timeout"]
    else:
        timeout = 600
    if "flavor" in version:
        flavor = version["flavor"]
    else:
        flavor = "standard.tiny"

    # Download image
    if not download_image(version["image_url"], filename, new_checksum):
        cleanup_files(filename)
        return None

    # Image downloaded, but is it the same as the current image on openstack?
    logger.info("Comparing downloaded image against openstack")

    if validate_raw_checksum(conn, filename, version):

        # Everything is ok so write the checksum pre-emptively
        if new_checksum != filename:

            img_name = version["image_name"].replace(" ", "_")
            with open(
                f"checksums/{conn.config.name}_{img_name}_CHECKSUM",
                "w",
                encoding="utf-8",
            ) as f:
                f.write(new_checksum)

        logger.info(
            f"Openstack already has the current version of {version['image_name']}"
        )

        current_img_id = next(
            conn.image.images(
                name=version["image_name"],
                owner=conn.current_project_id,
                visibility=version["visibility"],
            ),
            None,
        ).id

        delete_unused_image(conn, version["image_name"], current_img_id)
        return None

    logger.info("Uploading image")

    properties = version.get("properties", {})
    properties["description"] = "To find out which user to login with: ssh in as root."
    properties.setdefault("os_type", "linux")  # configures ephemeral drive formatting
    properties.setdefault(
        "hw_vif_multiqueue_enabled", True
    )  # set multiqueue to on by default

    new_image = conn.create_image(
        name=version["image_name"],
        filename="tmp/" + filename + ".raw",
        wait=True,
        visibility="private",
        allow_duplicates=True,
        disk_format="raw",
        container_format="bare",
        properties=properties,
    )

    if new_image is None:
        logger.error(
            {
                "message": "An error occured while creating the image",
                "image": version
            }
        )
        return None

    logger.info(f"Image with id ({new_image.id}) created")

    # Test image
    if not test_image(conn, new_image, network, timeout, flavor):
        return None

    return new_image


def delete_unused_image(
    conn: openstack.connection.Connection,
    name: str,
    skip: str = None,
    new_state: str = "community"
) -> dict:
    """
    Delete unused images but skip specified image if provided

    Returns
    -------
    dictionary with deleted, in error and in use images and their ids
    """

    result = {
        "deleted": {"count": 0, "ids": []},
        "in_use": {"count": 0, "ids": []},
        "errors": {"count": 0, "ids": []},
    }
    # Loop over existing images
    for img in conn.image.images(name=name, owner=conn.current_project_id):
        if img.id == skip:
            continue

        # Check if image is in use
        servers = list(conn.compute.servers(image=img.id, all_projects=True))
        volumes = list(conn.block_storage.volumes(image_id=img.id, all_projects=True))
        snapshots = list(
            conn.block_storage.snapshots(image_id=img.id, all_projects=True)
        )
        if len(servers) == 0 and len(volumes) == 0 and len(snapshots) == 0:
            logger.info(
                f"Image {img.id} not in use in a server, snapshot or volume, deleting..."
            )
            try:
                conn.delete_image(img.id)
                result["deleted"]["count"] += 1
                result["deleted"]["ids"].append(img.id)
            except openstack.exceptions.ConflictException as error:
                result["errors"]["count"] += 1
                result["errors"]["ids"].append(img.id)
                logger.error(
                    {
                        "message": "Error deleting image",
                        "image_id": img.id,
                        "error": error,
                    }
                )
        elif img.visibility != new_state:
            logger.info(
                f"Image {img.id} in use by a server, snapshot or volume"
                + f" setting it to 'hidden' and '{new_state}..."
            )
            conn.image.update_image(img.id, visibility=new_state)
            conn.image.update_image(img.id, os_hidden=True)
            result["in_use"]["count"] += 1
            result["in_use"]["ids"].append(img.id)
        else:
            logger.debug(
                f"Image {img.id} is in use by a server, snapshot or volume or it's a '{new_state}"
                + " image, not deleting it..."
            )
            result["in_use"]["count"] += 1
            result["in_use"]["ids"].append(img.id)

    return result


def process_current_images(
    input_data: dict,
    summary: dict,
    conn: openstack.connection.Connection
) -> None:
    """
    Main loop processing the current images, i.e., the images that we want to be public.
    """

    for version in input_data["current"]:
        version_name = version["image_name"]
        summary["current"][version_name] = {
            "state": None,
        }
        logger.debug(
            {
                "message": f"Checking current image '{version['image_name']}'...",
                "image": version,
            }
        )

        filename = version["image_url"].split("/")[-1]

        new_checksum = validate_checksum(version, filename, conn, cloud_name)

        if new_checksum is None:
            logger.debug(f"Image '{filename}' didn't change in remote source.")
            summary["current"][version_name]['state'] = "No remote change"
            continue

        logger.info(f"Downloading {version['image_name']}")

        new_image = create_image(conn, version, filename, new_checksum, network_name)

        if new_image is None:
            logger.debug("Image is already up to date or there was an error")
            summary["current"][version_name]['state'] = "Up-to-date or error creating image"
            continue

        # Everything is ok! Set visibility to what it should be
        logger.info(f"Setting image to {version['visibility']}")
        conn.image.update_image(new_image.id, visibility=version["visibility"])

        # Remove old ones
        result = delete_unused_image(conn, version_name, new_image.id)
        summary["current"][version_name]['state'] = "Ok"
        summary["current"][version_name]['old_images_delete_result'] = result
        summary["deleted_images"]["count"] += result["deleted"]["count"]
        summary["in_use_images"]["count"] += result["in_use"]["count"]
        summary["deleted_images"]["ids"] += result["deleted"]["ids"]
        summary["in_use_images"]["ids"] += result["in_use"]["ids"]

        if new_checksum != filename:
            with open(
                f"checksums/{cloud_name}_{version_name.replace(' ', '_')}_CHECKSUM",
                "w",
                encoding="utf-8",
            ) as f:
                f.write(new_checksum)

        logger.info(f"'{version_name}' has been successfully updated")


def process_deprecated_images(
    input_data: dict,
    summary: dict,
    conn: openstack.connection.Connection
) -> None:
    """
    Main loop processing the images to be deprecated, if present.
    """

    for version in input_data["deprecated"]:
        if "image_name" not in version:
            logger.warning(
                {
                    "message": "No image name found in image to delete",
                    "image": version
                }
            )
            continue
        version_name = version["image_name"]
        logger.debug(
            {
                "message": f"Checking deprecated image '{version['image_name']}'...",
                "image": version,
            }
        )
        filename = version["filename"]
        result = delete_unused_image(conn, version_name)
        summary["deprecated"][version_name] = result
        summary["deleted_images"]["count"] += result["deleted"]["count"]
        summary["in_use_images"]["count"] += result["in_use"]["count"]
        summary["deleted_images"]["ids"] += result["deleted"]["ids"]
        summary["in_use_images"]["ids"] += result["in_use"]["ids"]

        # It is deprecated so get rid of the files on disk
        try:
            if os.path.exists(
                f"checksums/{cloud_name}_{version_name.replace(' ', '_')}_CHECKSUM"
            ):
                logger.debug(
                    f"Image '{filename}' has been deprecated, deleting from local disk"
                )
                os.remove(
                    f"checksums/{cloud_name}_{version_name.replace(' ', '_')}_CHECKSUM"
                )
            else:
                logger.debug(f"Image '{filename}' not present in local disk")
        except OSError as error:
            logger.warning(
                f"Error removing image '{filename}' from local disk. {error}"
            )

        if version.get("filename"):
            cleanup_files(version["filename"])
    summary['duration'] = time.time() - summary['duration']
    summary['exit_code'] = logger.exit_code
    logger.info({"summary": summary})


def main() -> None:
    """
    Reads defined images from input.json, downloads new images and deprecates old ones
    """
    if cloud_name is None or network_name is None:

        missing = [
            x
            for x in ["IMAGEBUILDER_CLOUD", "IMAGEBUILDER_NETWORK"]
            if os.getenv(x) is None
        ]

        raise EnvironmentError(
            f"Environment variables are not set! ({', '.join(missing)})"
        )

    openstack.enable_logging(
        debug=(os.getenv("IMAGEBUILDER_OPENSTACK_DEBUG_LEVEL") is not None)
    )

    conn = openstack.connect(cloud=cloud_name)
    summary = {
        "duration": time.time(),
        "current": {},
        "deprecated": {},
        "deleted_images": {
            "count": 0,
            "ids": [],
        },
        "in_use_images": {
            "count": 0,
            "ids": [],
        },
        "errors": {
            "count": 0,
            "ids": [],
        },
    }

    # Load file from argv
    input_data = None
    input_file = os.getenv("IMAGEBUILDER_INPUT_FILE", "input.json")

    with open(input_file, "r", encoding="utf-8") as f:
        input_data = json.load(f)

    process_current_images(input_data, summary, conn)

    process_deprecated_images(input_data, summary, conn)


if __name__ == "__main__":
    main()
    sys.exit(logger.exit_code)