#!/usr/bin/python

from __future__ import print_function
from future import standard_library
standard_library.install_aliases()
from builtins import map
from builtins import zip
from builtins import range
from invirt import database
import os
import sys
import subprocess
import random
import string
import tempfile
import urllib.request, urllib.parse, urllib.error
import math
import optparse as op

class InvirtImageException(Exception):
    pass

# verbosity = 0 means no output from the actual commands
# verbosity = 1 means only errors from the actual commands
# verbosity = 2 means all output from the actual commands
verbosity = 0

def getOutput():
    global verbosity
    return {
        'stdout': subprocess.PIPE if verbosity < 2 else None,
        'stderr': subprocess.PIPE if verbosity < 1 else None
        }

def lvcreate(name, size):
    lvc = subprocess.Popen(['lvcreate', '-L', size, '-n', name, 'xenvg'],
                           stderr=subprocess.PIPE,
                           stdout=getOutput()['stdout'])
    if not lvc.wait():
        return 0
    stderr = lvc.stderr.read()
    if 'already exists in volume group' in stderr:
        return 5
    else:
        if verbosity > 0:
            print(stderr)
        return 6

def lvrename(dest, src):
    lvr = subprocess.Popen(['lvchange', '-an', "xenvg/%s" % (src,)],
                           stderr=subprocess.PIPE,
                           stdout=getOutput()['stdout'])
    ret = lvr.wait()

    if ret:
       if verbosity > 0:
           print(lvr.stderr.read())
       return ret

    lvr = subprocess.Popen(['lvrename', "xenvg/%s" % (src,), "xenvg/%s" % (dest,)],
                           stderr=subprocess.PIPE,
                           stdout=getOutput()['stdout'])
    ret = lvr.wait()

    if ret:
        stderr = lvr.stderr.read()
        if not ('not found in volume group' in stderr):
            if verbosity > 0:
                print(stderr)
            return ret

    subprocess.Popen(['lvchange', '-ay', "xenvg/%s" % (dest,)],
                     stderr=subprocess.PIPE,
                     stdout=getOutput()['stdout'])
    ret = lvr.wait()

    if ret:
        if verbosity > 0:
           print(lvr.stderr.read())

    return ret


def lv_random(func, pattern, *args):
    """
    Run some LVM-related command, optionally with a random string in
    the LV name.
    
    func takes an LV name plus whatever's in *args and returns the
    return code of some LVM command, such as lvcreate or lvrename
    
    pattern can contain at most one '%s' pattern, which will be
    replaced by a 6-character random string.
    
    If pattern contains a '%s', the script will attempt to re-run
    itself if the error code indicates that the destination already
    exists
    """
    # Keep trying until it works
    while True:
        rand_string = ''.join(random.choice(string.ascii_letters) \
                                  for i in range(6))
        if '%s' in pattern:
            name = pattern % rand_string
        else:
            name = pattern
        ret = func(name, *args)
        if ret == 0:
            return name
        # 5 is the return code if the destination already exists
        elif '%s' not in pattern or ret != 5:
            raise InvirtImageException('E: Error running %s with args %s' % (func.__name__, args))

def lvcreate_random(pattern, size):
    """
    Creates an LV, optionally with a random string in the name.
    
    Call with a string formatting pattern with a single '%s' to use as
    a pattern for the name of the new LV.
    """
    return lv_random(lvcreate, pattern, size)

def lvrename_random(src, pattern):
    """
    Rename an LV to a new name with a random string incorporated.
    
    Call with a string formatting pattern with a single '%s' to use as
    a pattern for the name of the new LV
    """
    return lv_random(lvrename, pattern, src)

def fetch_image(cdrom):
    """
    Download a cdrom from a URI, shelling out to rsync if appropriate
    and otherwise trying to use urllib
    """
    full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
    temp_file = tempfile.mkstemp()[1]
    if verbosity > 0:
        print("Fetching image %s from %s to %s" % (cdrom.cdrom_id, full_uri, temp_file), file=sys.stderr)
    try:
        if full_uri.startswith('rsync://'):
            if subprocess.call(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
                               **getOutput()):
                raise InvirtImageException("E: Unable to download '%s'" % full_uri)
        else:
            # I'm not going to look for errors here, because I bet it'll
            # throw its own exceptions
            urllib.request.urlretrieve(full_uri, temp_file)
        return temp_file
    except:
        os.unlink(temp_file)
        raise

def copy_file(src, dest):
    """
    Copy a file from one location to another using dd
    """
    if subprocess.call(['dd', 'if=%s' % src, 'of=%s' % dest, 'bs=1M'],
                       **getOutput()):
        raise InvirtImageException('E: Unable to transfer %s into %s' % (src, dest))

def load_image(cdrom):
    """
    Update a cdrom image by downloading the latest version,
    transferring it into an LV, moving the old LV out of the way and
    the new LV into place
    """
    if cdrom.mirror_id is None:
        return
    try:
        temp_file = fetch_image(cdrom)
    except InvirtImageException as e:
        print('ERROR: %s.  Skipping.' % e, file=sys.stderr)
        return

    try:
        st_size = os.stat(temp_file).st_size
        if not st_size:
            print("Failed to fetch %s" % cdrom.cdrom_id, file=sys.stderr)
            return
        cdrom_size = '%sM' % math.ceil((float(st_size) / (1024 * 1024)))
        new_lv = lvcreate_random('image-new_%s_%%s' % cdrom.cdrom_id, cdrom_size)
        copy_file(temp_file, '/dev/xenvg/%s' % new_lv)
        lvrename_random('image_%s' % cdrom.cdrom_id, 'image-old_%s_%%s' % cdrom.cdrom_id)
        lvrename_random(new_lv, 'image_%s' % cdrom.cdrom_id)
        reap_images()
    finally:
        os.unlink(temp_file)

def reap_images():
    """
    Remove stale cdrom images that are no longer in use
    
    load_image doesn't attempt to remove the old image because it
    might still be in use. reap_images attempts to delete any LVs
    starting with 'image-old_', but ignores errors, in case they're
    still being used.
    """
    lvm_list = subprocess.Popen(['lvs', '-o', 'lv_name', '--noheadings'],
                               stdout=subprocess.PIPE,
                               stdin=subprocess.PIPE)
    lvm_list.wait()
    
    for lv in map(str.strip, lvm_list.stdout.read().splitlines()):
        if lv.startswith('image-old_'):
            subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
                            **getOutput())
            subprocess.call(['lvchange', '-a', 'n', '/dev/xenvg/%s' % lv],
                            **getOutput())
            subprocess.call(['lvchange', '-a', 'ey', '/dev/xenvg/%s' % lv],
                            **getOutput())
            subprocess.call(['lvremove', '--force', '/dev/xenvg/%s' % lv],
                            **getOutput())

def main():
    global verbosity
    
    database.connect()
    database.session.begin()

    usage = """%prog [options] --add [--cdrom] cdrom_id description mirror_id uri_suffix
       %prog [options] --add --mirror mirror_id uri_prefix

       %prog [options] --update [short_name1 [short_name2 ...]]
       %prog [options] --reap"""
    
    parser = op.OptionParser(usage=usage)
    parser.set_defaults(verbosity=0,
                        item='cdrom')
    
    parser.add_option('-a', '--add', action='store_const',
                      dest='action', const='add',
                      help='Add a new item to the database')
    
    parser.add_option('-u', '--update', action='store_const',
                      dest='action', const='update',
                      help='Update all cdrom images in the database with the latest version')
    parser.add_option('-r', '--reap', action='store_const',
                      dest='action', const='reap',
                      help='Reap stale cdrom images that are no longer in use')
    
    a_group = op.OptionGroup(parser, 'Adding new items')
    a_group.add_option('-c', '--cdrom', action='store_const',
                       dest='item', const='cdrom',
                       help='Add a new cdrom to the database')
    a_group.add_option('-m', '--mirror', action='store_const',
                       dest='item', const='mirror',
                       help='Add a new mirror to the database')
    parser.add_option_group(a_group)
    
    v_group = op.OptionGroup(parser, "Verbosity levels")
    v_group.add_option("-q", "--quiet", action='store_const',
                       dest='verbosity', const=0,
                       help='Show no output from commands this script runs (default)')
    v_group.add_option("-v", "--verbose", action='store_const',
                       dest='verbosity', const=1,
                       help='Show only errors from commands this script runs')
    v_group.add_option("--noisy", action='store_const',
                       dest='verbosity', const=2,
                       help='Show all output from commands this script runs')
    parser.add_option_group(v_group)
    
    (options, args) = parser.parse_args()
    verbosity = options.verbosity
    if options.action is None:
        print(parser.format_help())
    elif options.action == 'add':
        if options.item == 'cdrom':
            attrs = dict(list(zip(('cdrom_id', 'description', 'mirror_id', 'uri_suffix'),
                             args)))
            cdrom = database.CDROM(**attrs)
            database.session.add(cdrom)
            database.session.commit()
            
            load_image(cdrom)
        
        elif options.item == 'mirror':
            attrs = dict(list(zip(('mirror_id', 'uri_prefix'),
                             args)))
            mirror = database.Mirror(**attrs)
            database.session.add(mirror)
            database.session.commit()
    elif options.action == 'update':
        if len(args) > 0:
            images = [database.CDROM.query().get(arg) for arg in args]
        else:
            images = database.CDROM.query().all()
        for cdrom in images:
            if cdrom is not None:
                load_image(cdrom)
    elif options.action == 'reap':
        reap_images()

if __name__ == '__main__':
    main()