#!/usr/bin/env python

# Copyright (c) 2012 Openstack, LLC
# 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.

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

import cPickle as pickle
import httplib
try:
    import json
except ImportError:
    import simplejson as json
import md5
import os
import os.path
import shutil

import urllib2
import XenAPIPlugin

import utils

#FIXME(sirp): should this use pluginlib from 5.6?
from pluginlib_nova import *
configure_logging('glance')

KERNEL_DIR = '/boot/guest'


class RetryableError(Exception):
    pass


def _copy_kernel_vdi(dest, copy_args):
    vdi_uuid = copy_args['vdi_uuid']
    vdi_size = copy_args['vdi_size']
    cached_image = copy_args['cached-image']
    logging.debug("copying kernel/ramdisk file from %s to /boot/guest/%s",
                  dest, vdi_uuid)
    filename = KERNEL_DIR + '/' + vdi_uuid
    #make sure KERNEL_DIR exists, otherwise create it
    if not os.path.isdir(KERNEL_DIR):
        logging.debug("Creating directory %s", KERNEL_DIR)
        os.makedirs(KERNEL_DIR)
    #read data from /dev/ and write into a file on /boot/guest
    of = open(filename, 'wb')
    f = open(dest, 'rb')
    #copy only vdi_size bytes
    data = f.read(vdi_size)
    of.write(data)
    if cached_image:
        #create a cache file. If caching is enabled, kernel images do not have
        #to be fetched from glance.
        cached_image = KERNEL_DIR + '/' + cached_image
        logging.debug("copying kernel/ramdisk file from %s to /boot/guest/%s",
                      dest, cached_image)
        cache_file = open(cached_image, 'wb')
        cache_file.write(data)
        cache_file.close()
        logging.debug("Done. Filename: %s", cached_image)

    f.close()
    of.close()
    logging.debug("Done. Filename: %s", filename)
    return filename


def _download_tarball_and_verify(request, staging_path):
    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 % locals())
    elif checksum != etag:
        msg = 'ETag %(etag)s does not match computed md5sum %(checksum)s'
        raise RetryableError(msg % locals())
    else:
        msg = "Verified image checksum %(checksum)s"
        logging.info(msg % locals())


def _download_tarball(sr_path, staging_path, image_id, glance_host,
                      glance_port, auth_token):
    """Download the tarball image from Glance and extract it into the staging
    area. Retry if there is any failure.
    """
    # Build request headers
    headers = {}
    if auth_token:
        headers['x-auth-token'] = auth_token

    url = ("http://%(glance_host)s:%(glance_port)d/v1/images/"
           "%(image_id)s" % locals())
    logging.info("Downloading %s" % url)

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


def _upload_tarball(staging_path, image_id, glance_host, glance_port,
                    auth_token, properties):
    """
    Create a tarball of the image and then stream that into Glance
    using chunked-transfer-encoded HTTP.
    """
    url = 'http://%s:%s/v1/images/%s' % (glance_host, glance_port, image_id)
    logging.info("Writing image data to %s" % url)
    conn = httplib.HTTPConnection(glance_host, glance_port)

    # 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'}

    # If we have an auth_token, set an x-auth-token header
    if auth_token:
        headers['x-auth-token'] = auth_token

    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()

    callback_data = {'bytes_written': 0}

    def send_chunked_transfer_encoded(chunk):
        chunk_len = len(chunk)
        callback_data['bytes_written'] += chunk_len
        conn.send("%x\r\n%s\r\n" % (chunk_len, chunk))

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

    conn.send("0\r\n\r\n")  # 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:
        logging.error("Unexpected response while writing image data to %s: "
                      "Response Status: %i, Response body: %s"
                      % (url, resp.status, resp.read()))
        raise Exception("Unexpected response [%i] while uploading image [%s] "
                        "to glance host [%s:%s]"
                        % (resp.status, image_id, glance_host, glance_port))
    conn.close()


def create_kernel_ramdisk(session, args):
    """Creates a copy of the kernel/ramdisk image if it is present in the
    cache. If the image is not present in the cache, it does nothing.
    """
    cached_image = exists(args, 'cached-image')
    image_uuid = exists(args, 'new-image-uuid')
    cached_image_filename = KERNEL_DIR + '/' + cached_image
    filename = KERNEL_DIR + '/' + image_uuid

    if os.path.isfile(cached_image_filename):
        shutil.copyfile(cached_image_filename, filename)
        logging.debug("Done. Filename: %s", filename)
    else:
        filename = ""
        logging.debug("Cached kernel/ramdisk image not found")
    return filename


def download_vhd(session, args):
    """Download an image from Glance, unbundle it, and then deposit the VHDs
    into the storage repository
    """
    params = pickle.loads(exists(args, 'params'))
    image_id = params["image_id"]
    glance_host = params["glance_host"]
    glance_port = params["glance_port"]
    uuid_stack = params["uuid_stack"]
    sr_path = params["sr_path"]
    auth_token = params["auth_token"]

    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,
            auth_token)

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

    # Right now, it's easier to return a single string via XenAPI,
    # so we'll json encode the list of VHDs.
    return json.dumps(vdi_list)


def upload_vhd(session, args):
    """Bundle the VHDs comprising an image and then stream them into Glance.
    """
    params = pickle.loads(exists(args, 'params'))
    vdi_uuids = params["vdi_uuids"]
    image_id = params["image_id"]
    glance_host = params["glance_host"]
    glance_port = params["glance_port"]
    sr_path = params["sr_path"]
    auth_token = params["auth_token"]
    properties = params["properties"]

    staging_path = utils.make_staging_area(sr_path)
    try:
        utils.prepare_staging_area_for_upload(sr_path, staging_path, vdi_uuids)
        _upload_tarball(staging_path, image_id, glance_host, glance_port,
                        auth_token, properties)
    finally:
        utils.cleanup_staging_area(staging_path)

    return ""  # Nothing useful to return on an upload


def copy_kernel_vdi(session, args):
    vdi = exists(args, 'vdi-ref')
    size = exists(args, 'image-size')
    cached_image = optional(args, 'cached-image')
    #Use the uuid as a filename
    vdi_uuid = session.xenapi.VDI.get_uuid(vdi)
    copy_args = {'vdi_uuid': vdi_uuid,
                 'vdi_size': int(size),
                 'cached-image': cached_image}
    filename = with_vdi_in_dom0(session, vdi, False,
                                lambda dev:
                               _copy_kernel_vdi('/dev/%s' % dev, copy_args))
    return filename


def remove_kernel_ramdisk(session, args):
    """Removes kernel and/or ramdisk from dom0's file system"""
    kernel_file = optional(args, 'kernel-file')
    ramdisk_file = optional(args, 'ramdisk-file')
    if kernel_file:
        os.remove(kernel_file)
    if ramdisk_file:
        os.remove(ramdisk_file)
    return "ok"


if __name__ == '__main__':
    XenAPIPlugin.dispatch({'upload_vhd': upload_vhd,
                           'download_vhd': download_vhd,
                           'copy_kernel_vdi': copy_kernel_vdi,
                           'create_kernel_ramdisk': create_kernel_ramdisk,
                           'remove_kernel_ramdisk': remove_kernel_ramdisk})
