#!/usr/bin/env python

# Copyright (c) 2012 OpenStack Foundation
# Copyright (c) 2010 Citrix Systems, Inc.
# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

# NOTE: XenServer still only supports Python 2.4 in it's dom0 userspace
# which means the Nova xenapi plugins must use only Python 2.4 features

"""Handle the uploading and downloading of images via Glance."""

import httplib
import md5
import socket
import urllib2

import pluginlib_nova
import utils


pluginlib_nova.configure_logging('glance')
logging = pluginlib_nova.logging
PluginError = pluginlib_nova.PluginError

SOCKET_TIMEOUT_SECONDS = 90


class RetryableError(Exception):
    pass


def _download_tarball_and_verify(request, staging_path):
    # NOTE(johngarbutt) By default, there is no timeout.
    # To ensure the script does not hang if we lose connection
    # to glance, we add this socket timeout.
    # This is here so there is no chance the timeout out has
    # been adjusted by other library calls.
    socket.setdefaulttimeout(SOCKET_TIMEOUT_SECONDS)

    try:
        response = urllib2.urlopen(request)
    except urllib2.HTTPError, error:
        raise RetryableError(error)
    except urllib2.URLError, error:
        raise RetryableError(error)
    except httplib.HTTPException, error:
        # httplib.HTTPException and derivatives (BadStatusLine in particular)
        # don't have a useful __repr__ or __str__
        raise RetryableError('%s: %s' % (error.__class__.__name__, error))

    url = request.get_full_url()
    logging.info("Reading image data from %s" % url)

    callback_data = {'bytes_read': 0}
    checksum = md5.new()

    def update_md5(chunk):
        callback_data['bytes_read'] += len(chunk)
        checksum.update(chunk)

    try:
        try:
            utils.extract_tarball(response, staging_path, callback=update_md5)
        except Exception, error:
            raise RetryableError(error)
    finally:
        bytes_read = callback_data['bytes_read']
        logging.info("Read %d bytes from %s", bytes_read, url)

    # Use ETag if available, otherwise X-Image-Meta-Checksum
    etag = response.info().getheader('etag', None)
    if etag is None:
        etag = response.info().getheader('x-image-meta-checksum', None)

    # Verify checksum using ETag
    checksum = checksum.hexdigest()

    if etag is None:
        msg = "No ETag found for comparison to checksum %(checksum)s"
        logging.info(msg % {'checksum': checksum})
    elif checksum != etag:
        msg = 'ETag %(etag)s does not match computed md5sum %(checksum)s'
        raise RetryableError(msg % {'checksum': checksum, 'etag': etag})
    else:
        msg = "Verified image checksum %(checksum)s"
        logging.info(msg % {'checksum': checksum})


def _download_tarball(sr_path, staging_path, image_id, glance_host,
                      glance_port, glance_use_ssl, extra_headers):
    """Download the tarball image from Glance and extract it into the staging
    area. Retry if there is any failure.
    """
    if glance_use_ssl:
        scheme = 'https'
    else:
        scheme = 'http'

    url = ("%(scheme)s://%(glance_host)s:%(glance_port)d/v1/images/"
           "%(image_id)s" % {'scheme': scheme, 'glance_host': glance_host,
                             'glance_port': glance_port,
                             'image_id': image_id})
    logging.info("Downloading %s" % url)

    request = urllib2.Request(url, headers=extra_headers)
    try:
        _download_tarball_and_verify(request, staging_path)
    except Exception:
        logging.exception('Failed to retrieve %(url)s' % {'url': url})
        raise


def _upload_tarball(staging_path, image_id, glance_host, glance_port,
                    glance_use_ssl, extra_headers, properties):
    """
    Create a tarball of the image and then stream that into Glance
    using chunked-transfer-encoded HTTP.
    """
    # NOTE(johngarbutt) By default, there is no timeout.
    # To ensure the script does not hang if we lose connection
    # to glance, we add this socket timeout.
    # This is here so there is no chance the timeout out has
    # been adjusted by other library calls.
    socket.setdefaulttimeout(SOCKET_TIMEOUT_SECONDS)

    if glance_use_ssl:
        scheme = 'https'
    else:
        scheme = 'http'

    url = '%s://%s:%s/v1/images/%s' % (scheme, glance_host, glance_port,
                                       image_id)
    logging.info("Writing image data to %s" % url)

    try:
        if glance_use_ssl:
            conn = httplib.HTTPSConnection(glance_host, glance_port)
        else:
            conn = httplib.HTTPConnection(glance_host, glance_port)
        conn.connect()
    except Exception, error:
        logging.exception('Failed to connect %(url)s' % {'url': url})
        raise RetryableError(error)

    try:
        validate_image_status_before_upload(conn, glance_host, glance_port,
                                            image_id, url, extra_headers)

        try:
            # NOTE(sirp): httplib under python2.4 won't accept
            # a file-like object to request
            conn.putrequest('PUT', '/v1/images/%s' % image_id)

            # NOTE(sirp): There is some confusion around OVF. Here's a summary
            # of where we currently stand:
            #   1. OVF as a container format is misnamed. We really should be
            #      using OVA since that is the name for the container format;
            #      OVF is the standard applied to the manifest file contained
            #      within.
            #   2. We're currently uploading a vanilla tarball. In order to be
            #      OVF/OVA compliant, we'll need to embed a minimal OVF
            #      manifest as the first file.

            # NOTE(dprince): In order to preserve existing Glance properties
            # we set X-Glance-Registry-Purge-Props on this request.
            headers = {
                'content-type': 'application/octet-stream',
                'transfer-encoding': 'chunked',
                'x-image-meta-is-public': 'False',
                'x-image-meta-status': 'queued',
                'x-image-meta-disk-format': 'vhd',
                'x-image-meta-container-format': 'ovf',
                'x-glance-registry-purge-props': 'False'}

            headers.update(**extra_headers)

            for key, value in properties.iteritems():
                header_key = "x-image-meta-property-%s" % key.replace('_', '-')
                headers[header_key] = str(value)

            for header, value in headers.iteritems():
                conn.putheader(header, value)
            conn.endheaders()
        except Exception, error:
            logging.exception('Failed to upload %(url)s' % {'url': url})
            raise RetryableError(error)

        callback_data = {'bytes_written': 0}

        def send_chunked_transfer_encoded(chunk):
            chunk_len = len(chunk)
            callback_data['bytes_written'] += chunk_len
            try:
                conn.send("%x\r\n%s\r\n" % (chunk_len, chunk))
            except Exception, error:
                logging.exception('Failed to upload when sending chunks')
                raise RetryableError(error)

        compression_level = properties.get('xenapi_image_compression_level')

        utils.create_tarball(
                None, staging_path, callback=send_chunked_transfer_encoded,
                compression_level=compression_level)

        send_chunked_transfer_encoded('')  # Chunked-Transfer terminator

        bytes_written = callback_data['bytes_written']
        logging.info("Wrote %d bytes to %s" % (bytes_written, url))

        resp = conn.getresponse()
        if resp.status == httplib.OK:
            return

        logging.error("Unexpected response while writing image data to %s: "
                      "Response Status: %i, Response body: %s"
                      % (url, resp.status, resp.read()))

        check_resp_status_and_retry(resp, image_id, glance_host, glance_port)

    finally:
        conn.close()


def check_resp_status_and_retry(resp, image_id, glance_host, glance_port):
    # Note(Jesse): This branch sorts errors into those that are permanent,
    # those that are ephemeral, and those that are unexpected.
    if resp.status in (httplib.BAD_REQUEST,                      # 400
                       httplib.UNAUTHORIZED,                     # 401
                       httplib.PAYMENT_REQUIRED,                 # 402
                       httplib.FORBIDDEN,                        # 403
                       httplib.NOT_FOUND,                        # 404
                       httplib.METHOD_NOT_ALLOWED,               # 405
                       httplib.NOT_ACCEPTABLE,                   # 406
                       httplib.PROXY_AUTHENTICATION_REQUIRED,    # 407
                       httplib.CONFLICT,                         # 409
                       httplib.GONE,                             # 410
                       httplib.LENGTH_REQUIRED,                  # 411
                       httplib.PRECONDITION_FAILED,              # 412
                       httplib.REQUEST_ENTITY_TOO_LARGE,         # 413
                       httplib.REQUEST_URI_TOO_LONG,             # 414
                       httplib.UNSUPPORTED_MEDIA_TYPE,           # 415
                       httplib.REQUESTED_RANGE_NOT_SATISFIABLE,  # 416
                       httplib.EXPECTATION_FAILED,               # 417
                       httplib.UNPROCESSABLE_ENTITY,             # 422
                       httplib.LOCKED,                           # 423
                       httplib.FAILED_DEPENDENCY,                # 424
                       httplib.UPGRADE_REQUIRED,                 # 426
                       httplib.NOT_IMPLEMENTED,                  # 501
                       httplib.HTTP_VERSION_NOT_SUPPORTED,       # 505
                       httplib.NOT_EXTENDED,                     # 510
                       ):
        raise PluginError("Got Permanent Error response [%i] while "
                          "uploading image [%s] to glance host [%s:%s]"
                          % (resp.status, image_id,
                             glance_host, glance_port))
    # NOTE(nikhil): Only a sub-set of the 500 errors are retryable. We
    # optimistically retry on 500 errors below.
    elif resp.status in (httplib.REQUEST_TIMEOUT,                # 408
                         httplib.INTERNAL_SERVER_ERROR,          # 500
                         httplib.BAD_GATEWAY,                    # 502
                         httplib.SERVICE_UNAVAILABLE,            # 503
                         httplib.GATEWAY_TIMEOUT,                # 504
                         httplib.INSUFFICIENT_STORAGE,           # 507
                         ):
        raise RetryableError("Got Ephemeral Error response [%i] while "
                             "uploading image [%s] to glance host [%s:%s]"
                             % (resp.status, image_id,
                                glance_host, glance_port))
    else:
        # Note(Jesse): Assume unexpected errors are retryable. If you are
        # seeing this error message, the error should probably be added
        # to either the ephemeral or permanent error list.
        raise RetryableError("Got Unexpected Error response [%i] while "
                             "uploading image [%s] to glance host [%s:%s]"
                             % (resp.status, image_id,
                                glance_host, glance_port))


def validate_image_status_before_upload(conn, glance_host, glance_port,
                                        image_id, url, extra_headers):
    try:
        # NOTE(nikhil): Attempt to determine if the Image has a status
        # of 'queued'. Because data will continued to be sent to Glance
        # until it has a chance to check the Image state, discover that
        # it is not 'active' and send back a 409. Hence, the data will be
        # unnecessarily buffered by Glance. This wastes time and bandwidth.
        # LP bug #1202785
        conn.request('HEAD', '/v1/images/%s' % image_id,
                     headers=extra_headers)
        head_resp = conn.getresponse()
        # NOTE(nikhil): read the response to re-use the conn object.
        body_data = head_resp.read(8192)
        if len(body_data) > 8:
            err_msg = ('Cannot upload data for image %(image_id)s as the '
                       'HEAD call had more than 8192 bytes of data in '
                       'the response body.' % {'image_id': image_id})
            raise PluginError("Got Permanent Error while uploading image "
                              "[%s] to glance host [%s:%s]. "
                              "Message: %s" % (image_id, glance_host,
                                               glance_port, err_msg))
        else:
            head_resp.read()

    except Exception, error:
        logging.exception('Failed to HEAD the image %(image_id)s while '
                          'checking image status before attempting to '
                          'upload %(url)s' % {'image_id': image_id,
                                              'url': url})
        raise RetryableError(error)

    if head_resp.status != httplib.OK:
        logging.error("Unexpected response while doing a HEAD call "
                      "to image %s , url = %s , Response Status: "
                      "%i" % (image_id, url, head_resp.status))

        check_resp_status_and_retry(head_resp, image_id, glance_host,
                                    glance_port)

    else:
        image_status = head_resp.getheader('x-image-meta-status')
        if image_status not in ('queued', ):
            err_msg = ('Cannot upload data for image %(image_id)s as the '
                       'image status is %(image_status)s' %
                       {'image_id': image_id, 'image_status': image_status})
            logging.exception(err_msg)
            raise PluginError("Got Permanent Error while uploading image "
                              "[%s] to glance host [%s:%s]. "
                              "Message: %s" % (image_id, glance_host,
                                               glance_port, err_msg))
        else:
            logging.info('Found image %(image_id)s in status '
                         '%(image_status)s. Attempting to '
                         'upload.' % {'image_id': image_id,
                                      'image_status': image_status})


def download_vhd(session, image_id, glance_host, glance_port, glance_use_ssl,
                 uuid_stack, sr_path, extra_headers):
    """Download an image from Glance, unbundle it, and then deposit the VHDs
    into the storage repository
    """
    staging_path = utils.make_staging_area(sr_path)
    try:
        # Download tarball into staging area and extract it
        _download_tarball(
            sr_path, staging_path, image_id, glance_host, glance_port,
            glance_use_ssl, extra_headers)

        # Move the VHDs from the staging area into the storage repository
        return utils.import_vhds(sr_path, staging_path, uuid_stack)
    finally:
        utils.cleanup_staging_area(staging_path)


def upload_vhd(session, vdi_uuids, image_id, glance_host, glance_port,
               glance_use_ssl, sr_path, extra_headers, properties):
    """Bundle the VHDs comprising an image and then stream them into Glance.
    """
    staging_path = utils.make_staging_area(sr_path)
    try:
        utils.prepare_staging_area(sr_path, staging_path, vdi_uuids)
        _upload_tarball(staging_path, image_id, glance_host, glance_port,
                        glance_use_ssl, extra_headers, properties)
    finally:
        utils.cleanup_staging_area(staging_path)


if __name__ == '__main__':
    utils.register_plugin_calls(download_vhd, upload_vhd)
