#!/usr/bin/python # # Copyright 2016 Red Hat | Ansible # GNU General Public License v3.0+ (see COPYING or www.gnu.org/licenses/gpl-3.0.txt)

# Imports # ============================================================================

from __future__ import absolute_import, division, print_function __metaclass__ = type

import os import re import json import logging

from ansible.module_utils.docker_common import (

HAS_DOCKER_PY_2,
AnsibleDockerClient,
DockerBaseClass

)

from ansible.module_utils._text import to_native

try:

if HAS_DOCKER_PY_2:
    from docker.auth import resolve_repository_name
else:
    from docker.auth.auth import resolve_repository_name
from docker.utils.utils import parse_repository_tag

except ImportError:

# missing docker-py handled in docker_common
pass

import qb.ipc.stdio import qb.ipc.stdio.logging

# Globals # ============================================================================

logger = qb.ipc.stdio.logging.getLogger('qb_docker_image')

# Casses # ============================================================================

class ImageManager(DockerBaseClass):

'''
Adaptation and extension of the `ImageManager` class from Ansible's
`docker_image` module for QB's `qb_docker_image` module.
'''

# Construction
# ========================================================================

def __init__(self, client, results):

    super(ImageManager, self).__init__()

    self.client = client
    self.results = results
    parameters = self.client.module.params
    self.check_mode = self.client.check_mode

    self.archive_path = parameters.get('archive_path')
    self.container_limits = parameters.get('container_limits')
    self.dockerfile = parameters.get('dockerfile')
    self.force = parameters.get('force')
    self.load_path = parameters.get('load_path')
    self.name = parameters.get('name')
    self.nocache = parameters.get('nocache')
    self.path = parameters.get('path')
    self.pull = parameters.get('pull')
    self.repository = parameters.get('repository')
    self.rm = parameters.get('rm')
    self.state = parameters.get('state')
    self.tag = parameters.get('tag')
    self.http_timeout = parameters.get('http_timeout')
    self.push = parameters.get('push')
    self.buildargs = parameters.get('buildargs')

    # QB additions
    self.try_to_pull = parameters.get('try_to_pull')

    self.logger = qb.ipc.stdio.logging.getLogger(
        'qb_docker_image:ImageManager',
    )

    # If name contains a tag, it takes precedence over tag parameter.
    repo, repo_tag = parse_repository_tag(self.name)
    if repo_tag:
        self.name = repo
        self.tag = repo_tag

    if self.state in ['present', 'build']:
        self.present()
    elif self.state == 'absent':
        self.absent()

# END __init__

# Instance Methods
# ========================================================================

# Helpers
# ------------------------------------------------------------------------

def out(self, msg):
    '''
    Just proxies to :attr:`client.out`, which writes to the master QB
    process' STDOUT (if available).

    :param msg: String or dict with output, but should usually be a
                dict with Docker API/daemon output in this case; log
                messages using the :attr:`logger`.

    :return: None
    '''
    self.client.out(msg)

def warn(self, warning, **values):
    '''
    Append a warning message to the `warnings` array of :attr:`results`
    that will be returned to Ansible and log it as a warning.
    '''
    warning = str(warning)
    self.results['warnings'].append(warning.format(**values))
    self.logger.warning(warning, payload=values)

def fail(self, msg, **values):
    '''
    Proxy to :attr:`client.fail`.
    '''
    self.client.fail(msg, **values)

def append_action(self, msg, **values):
    '''
    Add a message to the `actions` array in :attr:`results` and log it
    (as `info`).

    I'm not sure what 'actions' are for, but they were here when I adapted
    it... maybe something to do with "check mode"?
    '''

    formatted = msg.format(**values)
    self.logger.info(formatted, payload=values)
    self.results['actions'].append(formatted)

def image_summary(self, image):
    return dict(
        Id          = image['Id'],
        RepoTags    = image['RepoTags'],
        Created     = image['Created'],
        Metadata    = image['Metadata'],
    )

# States
# ------------------------------------------------------------------------

def present(self):
    '''
    Handles state = 'present', which includes building, loading or pulling
    an image, depending on user provided parameters.

    :return: None
    '''

    self.logger.debug(
        "Starting state `present`...",
        payload = dict(
            name    = self.name,
            tag     = self.tag,
        )
    )

    existing_image = self.client.find_image(name=self.name, tag=self.tag)

    if existing_image:
        self.logger.info(
            "Found existing image `{find_name}` in local daemon",
            payload = dict(
                find_name   = "{}:{}".format(self.name, self.tag),
                **self.image_summary(existing_image)
            )
        )

    # Keep track of what images we get from where
    pulled_image = None
    built_image = None
    loaded_image = None

    if not existing_image or self.force:
        # Try to pull if we're not forcing (which means we want to
        # re-build/load regardless) and `self.try_to_pull` is `True`
        if not self.force and self.try_to_pull:
            self.append_action(
                'Tried to pull image `{name}:{tag}`',
                name    = self.name,
                tag     = self.tag
            )

            self.results['changed'] = True

            if not self.check_mode:
                pulled_image = self.client.try_pull_image(
                    self.name,
                    tag=self.tag
                )

                if pulled_image:
                    self.append_action(
                        'Pulled image `{name}:{tag}`',
                        name    = self.name,
                        tag     = self.tag
                    )
                    self.results['image'] = pulled_image
        # END if not self.force and self.try_to_pull:

        if pulled_image is None:
            if self.path:
                # Build the image
                if not os.path.isdir(self.path):
                    self.fail(
                        "Requested build path `{path}` could not be " +
                        "found or you do not have access.",
                        path    = self.path,
                    )

                image_name = self.name
                if self.tag:
                    image_name = "%s:%s" % (self.name, self.tag)

                self.logger.info(
                    "Building image `{image_name}`",
                    payload = dict(
                        image_name  = image_name,
                    )
                )

                self.append_action(
                    "Built image `{image_name}` from `{path}`",
                    image_name  = image_name,
                    path        = self.path,
                )

                self.results['changed'] = True

                if not self.check_mode:
                    built_image = self.build_image()
                    self.results['image'] = built_image

            elif self.load_path:

                # Load the image from an archive
                if not os.path.isfile(self.load_path):
                    self.fail(
                        "Error loading image `{name}`. " +
                        "Specified load path `{load_path}` does not exist.",
                        name        = self.name,
                        load_path   = self.load_path,
                    )

                image_name = self.name

                if self.tag:
                    image_name = "%s:%s" % (self.name, self.tag)

                self.append_action(
                    "Loaded image `{image_name}` from `{load_path}`",
                    image_name  = image_name,
                    load_path   = self.load_path,
                )

                self.results['changed'] = True

                if not self.check_mode:
                    loaded_image = self.load_image()
                    self.results['image'] = loaded_image

            else:
                # pull the image
                self.append_action(
                    'Pulled image `{name}:{tag}`',
                    name    = self.name,
                    tag     = self.tag,
                )

                self.results['changed'] = True

                if not self.check_mode:
                    pulled_image = self.client.pull_image(
                        self.name,
                        tag = self.tag,
                    )

                    self.results['image'] = pulled_image

                    if (
                        existing_image and
                        existing_image == self.results['image']
                    ):
                        self.results['changed'] = False
        # END if pulled_image is None:
    # END if not image or self.force:

    # Archive the image if we have an archive path
    if self.archive_path:
        self.archive_image(self.name, self.tag)

    # Tag the image to a repository if we have one
    if self.repository:
        self.tag_image(
            self.name,
            self.tag,
            self.repository,
            force   = self.force,
            push    = self.push,
        )

    # This is weird to me logically, but I'm attempting to stick to the
    # Ansible `docker_image` module behavior...
    # 
    # We only push to the default repository (Docker Hub) if we didn't
    # receive a `self.respository`. I guess this is for when you're using
    # another repo than Docker Hub that you provide as `repository` and
    # then it assumes you would never want to push to Docker Hub *too*.
    # 
    # OK, that kinda makes sense to me...
    # 
    elif self.push:
        if pulled_image is not None:
            # Regadless of anything, we never want to push an image we
            # just pulled... makes no sense.
            self.logger.debug(
                "Image was pulled from repo, not pushing",
                payload=self.image_summary(pulled_image)
            )

        else:
            # Ok, now we can look at pushing...
            if self.force:
                # We're forcing, so force the push
                self.logger.info(
                    "FORCING push of image `{name}:{tag}`...",
                    payload = dict(
                        name    = self.name,
                        tag     = self.tag,
                    )
                )
                self.push_image(self.name, self.tag)

            else:
                # We only want to push if we built or loaded an image
                if built_image is not None:
                    self.logger.info(
                        "Pushing built image `{name}:{tag}`",
                        payload = dict(
                            name    = self.name,
                            tag     = self.tag,
                            **self.image_summary(built_image)
                        )
                    )
                    self.push_image(self.name, self.tag)

                elif loaded_image is not None:
                    self.logger.info(
                        "Pushing loaded image `{name}:{tag}`",
                        payload = dict(
                            name    = self.name,
                            tag     = self.tag,
                            **self.image_summary(loaded_image)
                        )
                    )
                    self.push_image(self.name, self.tag)

                else:
                    self.logger.info(
                        "No image built or loaded, not pushing"
                    )

            # END if self.force / else
        # END if pulled_image is not None / else
    # END if self.repository / elif self.push

    # QB addition - set the existing image as the result. Not sure why
    # the Ansible version doesn't do this..?
    if not self.results['image'] and existing_image is not None:
        self.results['image'] = existing_image

    self.logger.debug(
        "State `present` done",
        payload = dict(
            results = self.results,
        )
    )
# END present()

def absent(self):
    '''
    Handles state = 'absent', which removes an image.

    :return None
    '''
    image = self.client.find_image(self.name, self.tag)
    if image:
        name = self.name
        if self.tag:
            name = "%s:%s" % (self.name, self.tag)
        if not self.check_mode:
            try:
                self.client.remove_image(name, force=self.force)
            except Exception as exc:
                self.fail("Error removing image %s - %s" % (name, str(exc)))

        self.results['changed'] = True
        self.append_action(
            "Removed image `{name}`",
            name = name,
        )
        self.results['image']['state'] = 'Deleted'

# Actions
# ------------------------------------------------------------------------

def archive_image(self, name, tag):
    '''
    Archive an image to a .tar file. Called when archive_path is passed.

    :param name - name of the image. Type: str
    :return None
    '''

    if not tag:
        tag = "latest"

    image = self.client.find_image(name=name, tag=tag)
    if not image:
        self.logger.info(
            "archive image: image {name}:{tag} not found",
            payload = dict(
                name    = name,
                tag     = tag,
            )
        )
        return

    image_name = "%s:%s" % (name, tag)

    self.append_action(
        'Archived image `{image_name}` to `{archive_path}`',
        image_name      = image_name,
        archive_path    = self.archive_path
    )

    self.results['changed'] = True

    if not self.check_mode:

        self.logger.info(
            "Getting archive of image `{image_name}`",
            payload = dict(
                image_name  = image_name
            )
        )

        try:
            image = self.client.get_image(image_name)
        except Exception as exc:
            self.fail(
                "Error getting image `%s` - %s" % (image_name, str(exc))
            )

        try:
            with open(self.archive_path, 'w') as fd:
                for chunk in image.stream(2048, decode_content=False):
                    fd.write(chunk)
        except Exception as exc:
            self.fail(
                "Error writing image archive `%s` - %s" % (
                    self.archive_path,
                    str(exc)
                )
            )

    image = self.client.find_image(name=name, tag=tag)
    if image:
        self.results['image'] = image

def push_image(self, name, tag=None):
    '''
    If the name of the image contains a repository path, then push the image.

    :param name Name of the image to push.
    :param tag Use a specific tag.
    :return: None
    '''

    repository = name
    if not tag:
        repository, tag = parse_repository_tag(name)
    registry, repo_name = resolve_repository_name(repository)

    self.logger.info(
        "push `{name}` to `{registry}/{repo_name}:{tag}`",
        payload = dict(
            name        = self.name,
            registry    = registry,
            repo_name   = repo_name,
            tag         = tag
        )
    )

    if registry:

        self.append_action(
            "Pushed image `{name}` to `{registry}/{repo_name}:{tag}`",
            name        = self.name,
            registry    = registry,
            repo_name   = repo_name,
            tag         = tag
        )

        self.results['changed'] = True

        if not self.check_mode:
            status = None
            try:
                for line in self.client.push(
                    repository,
                    tag     = tag,
                    stream  = True,
                    decode  = True
                ):
                    self.out(line)

                    if line.get('errorDetail'):
                        raise Exception(line['errorDetail']['message'])

                    status = line.get('status')

            except Exception as exc:
                if re.search('unauthorized', str(exc)):
                    if re.search('authentication required', str(exc)):
                        self.fail(
                            "Error pushing image %s/%s:%s - %s. Try logging into %s first." % (
                                registry,
                                repo_name,
                                tag,
                                str(exc),
                                registry
                            )
                        )
                    else:
                        self.fail(
                            "Error pushing image %s/%s:%s - %s. Does the repository exist?" % (
                                registry,
                                repo_name,
                                tag,
                                str(exc)
                            )
                        )

                self.fail(
                    "Error pushing image %s: %s" % (repository, str(exc))
                )

            self.results['image'] = self.client.find_image(
                name    = repository,
                tag     = tag
            )

            if not self.results['image']:
                self.results['image'] = dict()

            self.results['image']['push_status'] = status

def tag_image(self, name, tag, repository, force=False, push=False):
    '''
    Tag an image into a repository.

    :param name: name of the image. required.
    :param tag: image tag.
    :param repository: path to the repository. required.
    :param force: bool. force tagging, even it image already exists with the repository path.
    :param push: bool. push the image once it's tagged.
    :return: None
    '''
    repo, repo_tag = parse_repository_tag(repository)

    if not repo_tag:
        repo_tag = "latest"
        if tag:
            repo_tag = tag

    image = self.client.find_image(name=repo, tag=repo_tag)

    found = 'found' if image else 'not found'

    self.logger.info(
        "image `{repo}` was `{found}`",
        payload = dict(
            repo    = repo,
            found   = found
        )
    )

    if not image or force:
        self.logger.info(
            "tagging {name}:{tag} to {repo}:{repo_tag}",
            payload=dict(name=name, tag=tag, repo=repo, repo_tag=repo_tag)
        )

        self.results['changed'] = True

        append_action(
            "Tagged image {name}:{tag} to {repo}:{repo_tag}",
            name=name, tag=tag, repo=repo, repo_tag=repo_tag
        )

        if not self.check_mode:
            try:
                # Finding the image does not always work, especially running a localhost registry. In those
                # cases, if we don't set force=True, it errors.
                image_name = name
                if tag and not re.search(tag, name):
                    image_name = "%s:%s" % (name, tag)
                tag_status = self.client.tag(image_name, repo, tag=repo_tag, force=True)
                if not tag_status:
                    raise Exception("Tag operation failed.")
            except Exception as exc:
                self.fail("Error: failed to tag image - %s" % str(exc))
            self.results['image'] = self.client.find_image(name=repo, tag=repo_tag)
            if push:
                self.push_image(repo, repo_tag)

def build_image(self):
    '''
    Build an image

    :return: image dict
    '''
    params = dict(
        path=self.path,
        tag=self.name,
        rm=self.rm,
        nocache=self.nocache,
        # Docker Pythong client v3 doesn't support
        # stream=True,
        timeout=self.http_timeout,
        pull=self.pull,
        forcerm=self.rm,
        dockerfile=self.dockerfile,
        decode=True
    )
    build_output = []
    if self.tag:
        params['tag'] = "%s:%s" % (self.name, self.tag)
    if self.container_limits:
        params['container_limits'] = self.container_limits
    if self.buildargs:
        for key, value in self.buildargs.items():
            self.buildargs[key] = to_native(value)
        params['buildargs'] = self.buildargs

    self.logger.info(
        "Building",
        payload=params,
    )

    logs = self.client.build(**params)

    # self.logger.info("build result", payload=dict(result=result))

    for log in logs:

        self.out(log)

        if "stream" in log:
            build_output.append(log["stream"])

        if log.get('error'):
            if log.get('errorDetail'):
                errorDetail = log.get('errorDetail')
                self.fail(
                    "Error building %s - code: %s, message: %s, logs: %s" % (
                        self.name,
                        errorDetail.get('code'),
                        errorDetail.get('message'),
                        build_output
                    )
                )
            else:
                self.fail(
                    "Error building %s - message: %s, logs: %s" % (
                        self.name,
                        log.get('error'),
                        build_output
                    )
                )
    return self.client.find_image(name=self.name, tag=self.tag)

def load_image(self):
    '''
    Load an image from a .tar archive

    :return: image dict
    '''
    try:
        self.logger.info(
            "Opening image `{load_path}`",
            payload = dict(load_path=self.load_path)
        )

        image_tar = open(self.load_path, 'r')

    except Exception as exc:
        self.fail(
            "Error opening image `{load_path}` - `{error}`",
            load_path   = self.load_path,
            error       = str(exc)
        )

    try:
        self.logger.info(
            "Loading image from `{load_path}`",
            payload = dict(load_path=self.load_path)
        )

        self.client.load_image(image_tar)

    except Exception as exc:
        self.fail("Error loading image %s - %s" % (self.name, str(exc)))

    try:
        image_tar.close()
    except Exception as exc:
        self.fail("Error closing image %s - %s" % (self.name, str(exc)))

    return self.client.find_image(self.name, self.tag)