#!/usr/bin/env python

# 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.

#
# XenAPI plugin for managing glance images
#

import base64
import errno
import hmac
import httplib
import os
import os.path
import pickle
import sha
import subprocess
import time
import urlparse

import XenAPIPlugin

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

CHUNK_SIZE = 8192
KERNEL_DIR = '/boot/guest'
FILE_SR_PATH = '/var/run/sr-mount'


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


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


def _copy_kernel_vdi(dest, copy_args):
    vdi_uuid = copy_args['vdi_uuid']
    vdi_size = copy_args['vdi_size']
    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)
    f.close()
    of.close()
    logging.debug("Done. Filename: %s", filename)
    return filename


def put_vdis(session, args):
    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 = get_sr_path(session)
    #FIXME(sirp): writing to a temp file until Glance supports chunked-PUTs
    tmp_file = "%s.tar.gz" % os.path.join('/tmp', str(image_id))
    tar_cmd = ['tar', '-zcf', tmp_file, '--directory=%s' % sr_path]
    paths = ["%s.vhd" % vdi_uuid for vdi_uuid in vdi_uuids]
    tar_cmd.extend(paths)
    logging.debug("Bundling image with cmd: %s", tar_cmd)
    subprocess.call(tar_cmd)
    logging.debug("Writing to test file %s", tmp_file)
    put_bundle_in_glance(tmp_file, image_id, glance_host, glance_port)
    # FIXME(sirp): return anything useful here?
    return ""


def put_bundle_in_glance(tmp_file, image_id, glance_host, glance_port):
    size = os.path.getsize(tmp_file)
    basename = os.path.basename(tmp_file)

    bundle = open(tmp_file, 'r')
    try:
        headers = {
            'x-image-meta-store': 'file',
            'x-image-meta-is_public': 'True',
            'x-image-meta-type': 'raw',
            'x-image-meta-size': size,
            'content-length': size,
            'content-type': 'application/octet-stream',
         }
        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', '/images/%s' % image_id)

        for header, value in headers.iteritems():
            conn.putheader(header, value)
        conn.endheaders()

        chunk = bundle.read(CHUNK_SIZE)
        while chunk:
            conn.send(chunk)
            chunk = bundle.read(CHUNK_SIZE)

        res = conn.getresponse()
        #FIXME(sirp): should this be 201 Created?
        if res.status != httplib.OK:
            raise Exception("Unexpected response from Glance %i" % res.status)
    finally:
        bundle.close()


def get_sr_path(session):
    sr_ref = find_sr(session)

    if sr_ref is None:
        raise Exception('Cannot find SR to read VDI from')

    sr_rec = session.xenapi.SR.get_record(sr_ref)
    sr_uuid = sr_rec["uuid"]
    sr_path = os.path.join(FILE_SR_PATH, sr_uuid)
    return sr_path


#TODO(sirp): both objectstore and glance need this, should this be refactored
#into common lib
def find_sr(session):
    host = get_this_host(session)
    srs = session.xenapi.SR.get_all()
    for sr in srs:
        sr_rec = session.xenapi.SR.get_record(sr)
        if not ('i18n-key' in sr_rec['other_config'] and
                sr_rec['other_config']['i18n-key'] == 'local-storage'):
            continue
        for pbd in sr_rec['PBDs']:
            pbd_rec = session.xenapi.PBD.get_record(pbd)
            if pbd_rec['host'] == host:
                return sr
    return None


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