image.py

"""
This module is used to interact with images on the DCAR
website.
"""
import json
import sys
import os
import tarfile
import shutil
import requests
import docker

from typing import Tuple
from datetime import timezone
from datetime import timedelta
from subprocess import run, PIPE
from time import sleep
from typing import Dict

from dateutil.parser import parse
from requests.exceptions import ProxyError, ReadTimeout

from dcar.request import Session
from dcar.utils import validate_argument, populate_image_name

def website_find_all(session) -> dict:
    """
    This function uses the DCAR website and python beautiful soup to collect
    metadata.
    """
    validate_argument("session", session, Session)

    image_details_dict: Dict[str, Dict] = {}

    vendor_level_request = session.get("/data/")
    vendors = vendor_level_request.json()

    # iterate through all vendor pages starting on https://dcar.dso.mil/
    for vendor in vendors['folders']:
        try:
            product_level_request = session.get("/data/{}".format(vendor))
            products = product_level_request.json()
        except ProxyError:
            print("Proxy Error Occurred, unable to pull vendor data for {}".format(vendor), flush=True, file=sys.stderr)
            continue

        # iterate through all product pages for each vendor page
        for product in products['folders']:
            try:
                image_level_request = session.get("/data/{}/{}".format(vendor, product))
                images = image_level_request.json()
            except ReadTimeout:
                print("Time Out Occurred, unable to pull image data for {}".format(product), flush=True, file=sys.stderr)
                continue
            except ProxyError:
                print("Proxy Error Occurred, unable to pull image data for {}".format(product), flush=True, file=sys.stderr)
                continue

            # iterate through all image pages for each product page
            for image in images['folders']:
                try:
                    image_detail_request = session.get("/data/{}/{}/{}".format(vendor, product, image))
                    image_details = image_detail_request.json()
                except ReadTimeout:
                    print("Time Out Occurred, unable to pull image data for {}".format(image), flush=True, file=sys.stderr)
                    continue
                except ProxyError:
                    print("Proxy Error Occurred, unable to pull image data for {}".format(image), flush=True, file=sys.stderr)
                    continue

                print("- found " + vendor + "/" + product + "/" + image, flush=True)

                try:
                    for build in image_details['builds']:
                        image_full_name = "{}/{}/{}:{}".format(vendor, product, image, build['tag'])

                        # Found instances of empty builds in API
                        if build['name'] == '':
                            print("Found invalid build for {}".format(image_full_name), flush=True, file=sys.stderr)
                            continue

                        build_date_string = build['buildDateAndNumber'].split("_")[0]
                        build_date = parse(build_date_string).replace(tzinfo=timezone.utc)

                        # DCAR image pages have multiple jenkins runs for each tag. For example
                        # ubi:8.1 may have two jenkins runs listed, we want to use the latest.
                        # only keep the latest jenkins run for a image.
                        if image_full_name in image_details_dict.keys():
                            image_date = parse(image_details_dict[image_full_name]['lastUpdate']).replace(tzinfo=timezone.utc)

                            # Drop builds that are older and the current image tag is approved
                            if (image_date >= build_date
                               and image_details_dict[image_full_name]['approval'] == 'approved'):
                                continue

                            # Drop if the build are older and neither are approved
                            if (image_date >= build_date
                               and image_details_dict[image_full_name]['approval'] != 'approved'
                               and build['approvalStatus'] != 'approved'):
                                continue

                        image_details_dict[image_full_name] = {
                            "vendor" : vendor,
                            "product" : product,
                            "image" : image,
                            "tag" : build['tag'],
                            "lastUpdate": build_date_string,
                            "approval": build['approvalStatus'],
                            "imageTar" : build['tar'],
                            "reportsSignatureTar" : build['allFilesTar'],
                            "buildNumber" : build['buildNumber'],
                            "sha" : build['sha'],
                            "publicKey" : build['publicKey']
                        }
                except KeyError as e:
                    print("Error {} key not found, likely from malformed build".format(e), flush=True, file=sys.stderr)
                except Exception as unkown_exception:
                    print("Error unkown exception {}".format(unkown_exception), flush=True, file=sys.stderr)

    return image_details_dict


def compare_image(image_info, acceptable_image_delta_min: int = 60) -> Tuple[str, bool]:
    """
    This function compares DCAR image data to a container image on a registry.
    """

    validate_argument("acceptable_image_delta_min", acceptable_image_delta_min, int)

    image_name = image_info[0]
    image_details = image_info[1]
    registry_user = os.environ["REGISTRY_USER"]
    registry_password = os.environ["REGISTRY_PASSWORD"]
    registry_image_name = populate_image_name(template_string=os.environ["REGISTRY_IMAGE_TEMPLATE"],
                                                   vendor=image_details["vendor"],
                                                   product=image_details["product"],
                                                   image=image_details["image"],
                                                   tag=image_details["tag"])

    # We assume that skopeo is installed, and we can executing it
    # through a subprocess. We'll use it to "inspect" the container registry's image
    # and find the image metadata to compare to metadata scrapped from DCAR.
    command_string = "skopeo inspect --creds=\"{}\":\"{}\" docker://{}".format(registry_user,
                                                                               registry_password,
                                                                               registry_image_name)
    # The line commented out below works for python 3.5 and lower. Command two lines down works for python 3.6 and up
    # command_result = run(command_string, shell=True, check=False, capture_output=True)
    command_result = run(command_string, shell=True, check=False, stdout=PIPE, stderr=PIPE)
    if command_result.returncode == 0:
        # TODO switch to manifest digests when available. Currently the
        # only way to compare an image from DCAR and an internal registry is by
        # comparing when the DCAR jenkins build ran and the creation date of
        # the image.
        container_image_created = parse(json.loads(command_result.stdout)['Created']).replace(tzinfo=timezone.utc)
        dcar_update = parse(image_details["lastUpdate"]).replace(tzinfo=timezone.utc)

        image_delta_time = dcar_update - container_image_created

        if timedelta(minutes=acceptable_image_delta_min) > image_delta_time:
            return (image_name, True)
    return (image_name, False)


def process_image(image_data: tuple,
                 max_retries: int = 3) -> Tuple[str, bool]:
    mirror_result = mirror_image(image_data, max_retries)
    cleanup_image(image_data)
    return mirror_result


def cleanup_image(image_data: tuple):
    image_name = image_data[0]
    image_details = image_data[1]
    registry_image_name = populate_image_name(template_string=os.environ["REGISTRY_IMAGE_TEMPLATE"],
                                                vendor=image_details["vendor"],
                                                product=image_details["product"],
                                                image=image_details["image"],
                                                tag=image_details["tag"])
    # Try and remove the local image with Docker
    if "PUSH_W_DOCKER" in os.environ:
        command_string = "docker rmi {}".format(registry_image_name)
        command_result = run(command_string, check=False, shell=True)
    else:
        pass

    # Get file and directory paths to clean up
    path = os.getcwd()
    absolute_dir = "{}/temp/{}/{}".format(path, image_name.split(":")[0], image_details["tag"])
    absolute_file = "{}{}-{}-reports-signature.tar.gz".format(absolute_dir,
                                                              image_details["image"],
                                                              image_details["tag"])

    # Remove tar files
    try:
        os.remove(absolute_file)
        print(f"Removed tar file: {absolute_file}")
    except:
        print(f"Failed to remove tar file: {absolute_file}")

    # Remove temp dir
    try:
        shutil.rmtree(absolute_dir)
        print(f"Removed temp dir: {absolute_dir}")
    except:
        print(f"Failed to remove tmp directory: {absolute_dir}")


def mirror_image(image_data: tuple,
                 max_retries: int = 3) -> Tuple[str, bool]:
    """
    This function takes DCAR image data and will download push an image to a container registry.
    """
    # Create neccessary variables
    image_name = image_data[0]
    image_details = image_data[1]
    registry_image_name = populate_image_name(template_string=os.environ["REGISTRY_IMAGE_TEMPLATE"],
                                                vendor=image_details["vendor"],
                                                product=image_details["product"],
                                                image=image_details["image"],
                                                tag=image_details["tag"])

    # Create session wrapper object
    session = Session(os.environ["DCAR_USERNAME"],
                      os.environ["DCAR_PASSWORD"],
                      os.environ["DCAR_TOTP_SEED"])

    registry_user = os.environ["REGISTRY_USER"]
    registry_password = os.environ["REGISTRY_PASSWORD"]

    # See if pushing with docker is an option
    if "PUSH_W_DOCKER" in os.environ:
        push_w_docker = True
        print("Pushing with docker selected")
    else:
        push_w_docker = False

    validate_argument("session", session, Session)
    validate_argument("registry_image_name", registry_image_name, str)
    validate_argument("registry_user", registry_user, str)
    validate_argument("registry_password", registry_password, str)
    validate_argument("image_name", image_name, str)
    validate_argument("image_details", image_details, dict)
    validate_argument("push_w_docker", push_w_docker, bool)
    validate_argument("max_retries", max_retries, int)

    print("\n- processing " + image_name, flush=True)

    # Create the download directory if it doesn't exist
    path = os.getcwd()
    absolute_dir = "{}/temp/{}/{}".format(path, image_name.split(":")[0], image_details["tag"])
    absolute_file = "{}{}-{}-reports-signature.tar.gz".format(absolute_dir,
                                                              image_details["image"],
                                                              image_details["tag"])

    try:
        os.makedirs(os.path.dirname(absolute_dir), exist_ok=True)
    except OSError:
        print("Creation of the directory {} failed".format(absolute_dir), flush=True, file=sys.stderr)
        return (image_name, False)

    try:
        # Download the tar.gz file
        with session.get(dcar_url=image_details["reportsSignatureTar"]["file"]) as response:
            signed_url = response.json()
            with requests.get(signed_url["signedUrl"], stream=True, verify=False) as response:
                with open(absolute_file, 'wb') as tar_download:
                    shutil.copyfileobj(response.raw, tar_download)
        print("\tdowloaded report", flush=True)
    except ReadTimeout:
        print("Time Out Error Occurred", flush=True, file=sys.stderr)
        return (image_name, False)
    except ProxyError:
        print("Proxy Error Occurred", flush=True, file=sys.stderr)
        return (image_name, False)
    except FileNotFoundError:
        print("Error saving file {}".format(absolute_file), flush=True, file=sys.stderr)
        return (image_name, False)
    except Exception as unkown_exception:
        print("Error unkown exception {}".format(unkown_exception), flush=True, file=sys.stderr)
        return (image_name, False)

    # The `reportsSignature.tar.gz` should be a tar file. If its not a tar file
    # then something went wrong, and we didn't download the exported container image.
    if not tarfile.is_tarfile(absolute_file):
        print("ERROR: {} is not a tarfile".format(absolute_file), flush=True, file=sys.stderr)
        return (image_name, False)

    tar_file = tarfile.open(absolute_file)
    tar_file.extractall(path=absolute_dir)

    # Check if the image tar is there. If not download the tar file separately
    image_tar_dict = absolute_dir + "/" + os.path.commonprefix(tar_file.getnames())
    image_tar = image_tar_dict + "/" + image_details["image"] + "-" + image_details["tag"] + ".tar"

    if not os.path.exists(image_tar):
        image_tar = image_tar_dict + "/" + image_details["image"] + "-" + image_details["buildNumber"] + ".tar"

    if not os.path.exists(image_tar):
        try:
            # Download the tar.gz file
            print("\timage tar was not present", flush=True)
            with session.get(dcar_url=image_details["imageTar"]["file"], stream=True) as response:
                signed_url = response.json()
                with requests.get(signed_url["signedUrl"], stream=True, verify=False) as response:
                    with open(image_tar, 'wb') as tar_download:
                        shutil.copyfileobj(response.raw, tar_download)
            print("\tdowloaded tar separately", flush=True)
        except ReadTimeout:
            print("Time Out Error Occurred", flush=True, file=sys.stderr)
            return (image_name, False)
        except ProxyError:
            print("Proxy Error Occurred", flush=True, file=sys.stderr)
            return (image_name, False)
        except FileNotFoundError:
            print("Error saving file {}".format(image_tar), flush=True, file=sys.stderr)
            return (image_name, False)
        except Exception as unkown_exception:
            print("Error unkown exception {}".format(unkown_exception), flush=True, file=sys.stderr)
            return (image_name, False)

    if not os.path.exists(image_tar):
        print("ERROR: Unable to download image tar {}".format(image_tar), flush=True, file=sys.stderr)
        return (image_name, False)

    # Next we need to figure out what kind of image format we have ...
    container_type = None
    tar_file = tarfile.open(image_tar)

    try:
        tar_file.getmember("oci-layout")
        container_type = "oci-archive"
    except KeyError:
        pass

    try:
        tar_file.getmember("manifest.json")
        container_type = "docker-archive"
    except KeyError:
        pass

    if container_type is None:
        print("ERROR: Failed to determine image type", flush=True, file=sys.stderr)
        return (image_name, False)

    # If its an OCI archive, extract into a directory
    if container_type == "oci-archive":
        print("\timage tar is an oci-archive", flush=True)
        image_dir = absolute_dir + "/" + image_details["image"] + "-" + image_details["tag"]
        tar_file = tarfile.open(image_tar)
        tar_file.extractall(path=image_dir)

    # We're going to use skopeo to move the image. Lets start
    # building the process's shell command based on the archive type
    # and where it is going to be stored.
    if container_type == "oci-archive":
        archive_portion_command_string = "oci:{}".format(image_dir)
    elif container_type == "docker-archive":
        archive_portion_command_string = "docker-archive:{}".format(image_tar)

    # The prefered behavior is to use skopeo to push docker images to the container registry.
    # If the registry is not compatible with skopeo, use skopeo to load the image into internal
    # storage, and then use docker to push to the container registry.
    if push_w_docker:
        docker_portion_command_string = "docker-daemon:{}".format(registry_image_name)
    else:
        docker_portion_command_string = "docker://{}".format(registry_image_name)

    for attempt in range(max_retries):
        if container_type == "docker-archive":
            command_string = "skopeo copy --dest-creds=\"{}\":\"{}\" {} {}".format(registry_user, registry_password,
                                                                                   archive_portion_command_string,
                                                                                   docker_portion_command_string)
            command_result = run(command_string, check=False, shell=True)
            if command_result.returncode == 0:
                break
            elif attempt >= max_retries-1:
                return (image_name, False)

        # sleep 10 seconds between retries
        sleep(10)

    # The preferred behavior is to use skopeo to push docker images to the container registry.
    # If the registry is not compatible with skopeo, use skopeo to load the image into internal
    # storage, and then use docker to push to the container registry.
    if push_w_docker:
        for attempt in range(max_retries):
            command_string = "docker push {}".format(registry_image_name)
            command_result = run(command_string, check=False, shell=True)
            if command_result.returncode == 0:
                break
            elif attempt >= max_retries-1:
                return (image_name, False)
            # sleep 10 seconds between retries
            sleep(10)

    return (image_name, True)