#!/usr/bin/python
from __future__ import with_statement # This isn't required in Python 2.6, but must be
                                      # in the beginning and thus can't be conditional

import httplib
import urllib2
from time import strptime
from datetime import datetime, timedelta
from optparse import OptionParser
from os import environ
from BaseHTTPServer import BaseHTTPRequestHandler
import hashlib
import os
import re
import subprocess
import sys
import stat
import signal
import tempfile

current_release = 'lucid'
releasedict = {}
releasedict['hardy'] = []
releasedict['intrepid'] = []
releasedict['jaunty'] = []
releasedict['karmic'] = []
releasedict['lucid'] = []

#allreleases = ['hardy', 'intrepid']
allarchs = ['i386', 'amd64', 'lpia', 'armel']

default = {}
default['archs'] = ['i386', 'amd64']
default['variants'] = ['desktop', 'alternate', 'dvd']
default['build'] = 'current'
default['host'] = 'cdimage.ubuntu.com'
default['releases'] = ['hardy', 'lucid']
default['isoroot'] = environ['HOME'] + '/iso/'
default['do_release_check'] = True

# Hardy has no launchpadlib available
try:
    from launchpadlib.launchpad import Launchpad, STAGING_SERVICE_ROOT, EDGE_SERVICE_ROOT
    from launchpadlib.credentials import Credentials
except ImportError:
    print "Unable to import launchpadlib"
    default['do_release_check'] = False

zsync_binary = '/usr/bin/zsync'

class Flavor:
    def __init__(self, name, use_prefix=True, releases=default['releases'],
                 variants=default['variants'], archs=default['archs']):
        self.name = name
        self.prefix = None
        if use_prefix:
            self.prefix = name
        self.dir = dir
        self.releases = releases
        self.variants = variants
        self.archs = archs
        for r in releases:
            releasedict[r].append(self)
        flavors[name] = self

    def add_release(self, new):
        self.releases.append(new)

def my_log(x):
    print x

def get_launchpad_login():
    try:
        launchpad = Launchpad.login_anonymously("Ubuntu QA iso downloader", "edge")
    except AttributeError:
	    # fall back to more portable anonymous login
        launchpad = Launchpad.login("Ubuntu QA iso downloader", "", "", EDGE_SERVICE_ROOT)
    return launchpad

def get_devel_release():
    ''' Get the current development release from launchpad '''
    launchpad = get_launchpad_login()

    u = launchpad.projects['ubuntu']
    return u.current_series.name

def update_release():
    '''check launchpad to see if the current development release has
       changed and update configurations if that's the case'''
    global current_release

    lp_release = get_devel_release()
    if lp_release != current_release:
        default['releases'].append(lp_release)
        update_flavors(lp_release, current_release)

    current_release = lp_release

def update_flavors(new, current):
    '''Add new development release to each flavor that supports the
       current development release'''
    global releasedict

    releasedict[new] = []

    for flavor in releasedict[current]:
        releasedict[new].append(flavor)
        flavor.add_release(new)

## HTML Header, defines CSS and header image ##
def printHeader():
        print """
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.1//EN"
"http://www.w3.org/TR/xhtml11/DTD/xhtml11.dtd">
<html xmlns="http://www.w3.org/1999/xhtml" xml:lang="en">
<head>
  <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
  <title>Daily ISO Tracker</title>
  <style type="text/css">
    body { background: #FFFFFF; color: black; }
    a { text-decoration: none; }
    table.head { border-style: none none; }
    table.head td { text-align: center; padding-left: 10px; padding-right: 10px}
    table.body { background: #efe1c3; border-collapse: collapse; border-style: solid solid;
            border-width: 3px; margin-bottom: 3ex; empty-cells: show; }
    table.body th { text-align: left; border-style: none none dotted none;
               border-width: 1px; padding-right: 10px; }
    table.body td { text-align: left; border-style: none none dotted none;
               border-width: 1px; padding-right: 10px; }
    a { color: blue; }
    a.verified { color: green; font-weight: bold; }
    a.testing { color: blue; }
    pre { white-space: pre-wrap; }
  </style>
</head>
<body>
<table class="head">
        <tr><td width="30%"><a href="http://www.ubuntu.com/"><img src="http://www.ubuntu.com/themes/ubuntu07/images/ubuntulogo.png" border="0" hspace="0" vspace="0" alt="Ubuntu"></a></td>
        <td><h1>Daily ISO Tracker</h1></td></tr>
</table>
"""

## Print footer ##
def printFooter():
        time = datetime.datetime.utcnow()
        print "<p>Last Updated (UTC): %s by <a href=\"https://code.launchpad.net/~sbeattie/sru-tools/sru-buglist\">sru_buglist</a>" % time.ctime()
        print "written by <a href=\"mailto:sbeattie@ubuntu.com\">Steve Beattie</a>.</p>"
        print "</body></html>"

def dumpHeaders(headers):
	for header, value in headers:
		print header + ': ' + value

unit_names = {"year" : ("year", "years"),
	      "month" : ("month", "months"),
	      "week" : ("week", "weeks"),
	      "day" : ("day", "days"),
	      "hour" : ("hour", "hours"),
	      "minute" : ("minute", "minutes"),
	      "second" : ("second", "seconds")}

def seconds_in_units(seconds):
	"""
	Returns a tuple containing the most appropriate unit for the
	number of seconds supplied and the value in that units form.

		>>> seconds_in_units(7700)
		(2, 'hour')
	"""

	unit_limits = [("year", 365 * 24 * 3600),
		       ("month", 30 * 24 * 3600),
		       ("week", 7 * 24 * 3600),
		       ("day", 24 * 3600),
		       ("hour", 3600),
		       ("minute", 60)]

	for unit_name, limit in unit_limits:
		if seconds >= limit:
			amount = int(round(float(seconds) / limit))
			return amount, unit_name
	return seconds, "second"

def stringify_timedelta(td):
	"""
	Converts a timedelta into a nicely readable string.

	        >>> td = timedelta(days = 77, seconds = 5)
		>>> print readable_timedelta(td)
		two months
	"""
	seconds = td.days * 3600 * 24 + td.seconds
	amount, unit_name = seconds_in_units(seconds)

	str_unit = unit_names[unit_name][1]
	if amount == 1:
		str_unit = unit_names[unit_name][0]
	return "%d %s" % (amount, str_unit)

def parseDate(date):
	return datetime.strptime(date, "%a, %d %b %Y %H:%M:%S %Z")

def create_dir(path):
	# XXX: make this more defensive. Sigh.
	try:
		stat = os.stat(path)
	except OSError, e:
		os.makedirs(path)

def computepath(wanted, flavor, release, variant):
	path = '/'
	if flavor.prefix:
		path += flavor.prefix + '/'
	if release != current_release:
		path += release + '/'
	path += '%s/%s/' %(variantpaths[variant], wanted['build'])
	return path

def register_auth_handler(config):
    auth_handler = urllib2.HTTPBasicAuthHandler()
    auth_handler.add_password(realm = 'Daily ISO Mirror',
                              uri = 'http://%s' %(config.host),
                              user = config.username,
                              passwd = config.password)
    opener = urllib2.build_opener(auth_handler)
    urllib2.install_opener(opener)

class HeadRequest(urllib2.Request):
    def get_method(self):
        return "HEAD"

def get_status(uri, host):
    close = False
    result = False, None
    response = None

    try:
        response = urllib2.urlopen(HeadRequest("http://" + host + "/" + uri))
        if response.code == 200:
            info = response.info()
            modified = info.getheader('last-modified')
            age = datetime.utcnow() - parseDate(modified)
            result = True, stringify_timedelta(age)
        else:
            result = False, BaseHTTPRequestHandler.responses[response.code][0]
    except urllib2.HTTPError, e:
        result = False, BaseHTTPRequestHandler.responses[e.code][0]
    finally:
        if response:
            response.close()

    return result

def check_status(wanted, config):
	for count, image in enumerate(image_factory(wanted)):
		release, flavor, variant, arch = image
		isoname = release + '-' + variant + '-' + arch + '.iso'
		path = computepath(wanted, flavor, release, variant)
		uri = path + isoname
		#print 'Result for ' + uri + ': ' + str(response.status)
		rc, msg = get_status(uri + '.zsync', config.host)
		print 'Result for ' + isoname + '.zsync: ' + msg
		rc, msg = get_status(uri, config.host)
		print 'Result for ' + isoname + ': ' + msg

def _do_download(config, source, target):
	result = True
	conn = None
	out = None

	try:
	    conn = urllib2.urlopen('http://' + config.host + source)
	    out = open(target, 'w')
	    for data in conn:
	        out.write(data)
	except IOError:
	    result = False
	finally:
	    if conn:
	        conn.close()
	    if out:
	        out.close()
	return result

def _do_zsync(config, meta, src, target):
    command = [zsync_binary, '-k', meta, '-o', target,
               'http://%s/%s' %(config.host, src)]
    if config.username and config.password:
        command.insert(1, '-A')
        command.insert(2, '%s=%s:%s' %(config.host, config.username, config.password))
    if config.quiet:
        command.insert(1, '-q')
    if config.no_act or config.debug:
        if config.debug:
            print ' '.join(command)
        return True

    rc = 0
    interrupted = False
    try:
        rc = subprocess.call(command)
        # trap external SIGINT death
        if rc > 128 and rc % 128 == signal.SIGINT:
            interrupted = True
    except KeyboardInterrupt:
        interrupted = True
    finally:
        if os.path.exists(target + '.zs-old'):
            os.remove(target + '.zs-old')
        if (interrupted or rc != 0) and os.path.exists(target + '.part'):
            os.rename(target + '.part', target)
        if os.path.exists(target):
            os.chmod(target, stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH)
        if interrupted:
            raise KeyboardInterrupt
    return rc == 0

def do_zsync(config, destination, uri, isoname):
	zsync_uri = uri + '.zsync'
	zsync_iso = destination + '/' + isoname
	zsync_local = destination + '/'
	# hide zsync files when not mirroring tree structure
	if not config.mirror:
		zsync_local += '.'
	zsync_local += isoname + '.zsync'
	available, msg = get_status(zsync_uri, config.host)
	if not available:
		config._log("control file does not exist, skipping zsync: " + msg)
		return False

	# when config.build is set, we could be grabbing an iso older in
	# time than what we have now. zsync only updates it's local copy
	# of the meta data if the remote version is dated newer, so we
	# manually rsync the version we want.
	if config.build != default['build']:
		rc = _do_rsync(config, zsync_uri, zsync_local)

	return _do_zsync(config, zsync_local, zsync_uri, zsync_iso)
	

def _do_rsync(config, source, target):
	command = ['/usr/bin/rsync', '-zthP', 'rsync://%s/cdimage%s' %(config.host, source), target]
	if len(config.bwlimit) > 0:
		command.insert(2, '--bwlimit=' + config.bwlimit)
	if config.quiet:
		command.insert(2, '-q')
	if config.no_act or config.debug:
		if config.debug:
			print ' '.join(command)
		return True
	rc = subprocess.call(command)
	return rc == 0
	
def do_rsync(config, destination, uri, isoname):
	rsync_local = destination + '/' + isoname
	#rc = cmd(['/usr/bin/zsync', '-o', zsync_local, 'http://cdimages.ubuntu.com/' + zsync_uri])
	return _do_rsync(config, uri, rsync_local)

def lookup_hash(path, host):
	hash_table = {
		"SHA256" : hashlib.sha256,
		"SHA1"   : hashlib.sha1,
		"MD5"    : hashlib.md5,
	}
	for alg in hash_table.keys():
		uri = path + alg + 'SUMS'
		available, msg = get_status(uri, host)
		if available:
			return alg, hash_table[alg]
	return None, None

def do_verify(config, destination, hashsuffix, path, iso):
	alg, hashfunc  = lookup_hash(path, config.host)
	if not alg:
		return False
		
	uri = path + alg + "SUMS"
	if config.mirror:
		local = destination + '/' + alg + "SUMS"
	else:
		local = destination + '/' + alg + "SUMS" + hashsuffix
	rc = _do_download(config, uri, local)
	if not rc:
		return False
	
	digest = None
	with open(local) as f:
		for line in f:
			if iso in line:
				digest = re.split('\s+', line)[0]
				break
	
	if not digest:
		return False
	
	m = hashfunc()
	fd = os.open(destination + '/' + iso, os.O_RDONLY)
	data = os.read(fd, 16 * 256 * 256)
	while len(data):
		m.update(data)
		data = os.read(fd, 16 * 256 * 256)
	os.close(fd)
	result = m.hexdigest()
	#print digest + ' ' + iso
	#print result + ' ' + iso
	
	return digest == result

def iso_download(wanted, config):
	saved_cwd = os.getcwd()
	create_dir(wanted['isoroot'])
	os.chdir(wanted['isoroot'])

	for count, image in enumerate(image_factory(wanted)):
		release, flavor, variant, arch = image
		isoname = release + '-' + variant + '-' + arch + '.iso'
		hashsuffix = "." + release + '-' + variant
		path = computepath(wanted, flavor, release, variant)
		if config.mirror:
			destination = path[1:-1]
		else:
			destination = flavor.name
		create_dir(destination)
		uri = path + isoname
		#print 'Result for ' + uri + ': ' + str(response.status)

		config._log("Syncing %s %s %s %s" %(release, flavor.name, variant, arch))
		available, msg = get_status(uri, config.host)
		if not available:
			config._log("iso unavailable from server, skipping: " + msg)
			continue	

		if not (config.use_zsync and do_zsync(config, destination, uri, isoname)):
			do_rsync(config, destination, uri, isoname)
		if config.do_verify:
			config._log("Verifying " + destination + '/' + isoname)
			if not (config.debug or config.no_act):
				if do_verify(config, destination, hashsuffix, path, isoname):
					config._log("Verification succeeded")
				else:
					config._log("!!! Verification failed !!!")

	os.chdir(saved_cwd)

def dump_info(wanted):
	for count, image in enumerate(image_factory(wanted)):
		release, flavor, variant, arch = image
		print "%s %s %s %s %s" %(release, flavor.name, variant, arch, variantpaths[variant])

def image_factory(wanted):
	for release in releasedict.keys():
		if release not in wanted['releases']:
			continue
		for flavor in releasedict[release]:
			if flavor.name not in wanted['flavors']:
				continue
			for variant in flavor.variants:
				if variant not in wanted['variants']:
					continue
				for arch in flavor.archs:
					if arch not in wanted['archs']:
						continue
					yield (release, flavor, variant, arch)
					#print "%s %s %s %s %s" %(release, flavor.name, variant, arch, variantpaths[variant])

def read_config(pathname):
    '''Read the config file, using shell syntax for backwards
       compatibility.'''
    if not os.path.exists(pathname):
        return None

    shell_vars = ['ISOROOT', 'RELEASES', 'FLAVORS', 'VARIANTS', 'EXCLUDE',
                  'OPTS', 'QUIET', 'NO_ACT', 'VERIFY', 'BUILD', 'BASEURL',
                  'ARCHS', 'DLUSERNAME', 'DLPASSWORD', 'HOST']
    command = ". %s; " %(pathname)
    for var in shell_vars:
        command += 'if [ ! -z "$%s" ] ; then echo %s=$%s; fi ; ' %(var, var, var)
    p = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE)
    out, err = p.communicate()
    config = dict()
    for line in re.split('\n', out):
        l = re.split('=', line)
        if not len(l) == 2:
            continue
        if l[0] in ['RELEASES', 'FLAVORS', 'VARIANTS', 'EXCLUDE']:
            l[1] = l[1].lower()
            config[str.lower(l[0])] = re.split(' ', l[1])
        elif l[0] in ['BASEURL']:
            print "Warning! BASEURL is not a supported config option anymore, ignoring"
        else:
            config[str.lower(l[0])] = l[1]

    return config

def merge_config(cval, file, default, value):
    ret = None
    if len(cval) > 0:
        ret = cval
    elif file and value in file.keys() and len(file[value]) > 0:
        ret = file[value]
    elif value in default.keys():
        ret = default[value]
    return ret

def main():
    global releasedict, current_release, variantpaths, flavors

    config = {}
    wanted = {}
    flavors = {}
    variantpaths = {}
    variantpaths['desktop'] = 'daily-live'
    variantpaths['alternate'] = 'daily'
    variantpaths['jeos'] = 'daily'
    variantpaths['dvd'] = 'dvd'
    variantpaths['server'] = 'daily'
    variantpaths['addon'] = 'daily'
    variantpaths['usb'] = 'daily'
    variantpaths['netbook-remix'] = 'daily-live'
    variantpaths['netbook'] = 'daily-live'
    variantpaths['moblin-remix'] = 'daily-live'

    Flavor('ubuntu', use_prefix=False)
    Flavor('ubuntu-server', variants=['server'])
    Flavor('kubuntu')
    Flavor('kubuntu-kde4', releases=['hardy'], variants=['desktop', 'alternate'])
    Flavor('edubuntu', variants=['addon', 'dvd'])
    Flavor('xubuntu', variants=['desktop', 'alternate'])
    Flavor('ubuntustudio', variants=['alternate'])
    Flavor('jeos', releases=['hardy'], variants=['jeos'], archs=['i386'])
    Flavor('gobuntu', releases=['hardy'], variants=['alternate'])
    Flavor('mythbuntu', variants=['desktop'])
    Flavor('ubuntu-mid', variants=['usb'], archs=['lpia'])
    Flavor('ubuntu-netbook', variants=['netbook'], archs=['i386'])
    Flavor('kubuntu-netbook', variants=['netbook'], archs=['i386'])
    Flavor('ubuntu-moblin-remix', variants=['moblin-remix'], archs=['i386'])

    default['variants'] = variantpaths.keys()
    default['flavors'] = flavors.keys()

    parser = OptionParser()
    parser.add_option('-n', '--no-act', default=False, action='store_true', dest='no_act', help='compute everything but don\'t actually download')
    parser.add_option('-d', '--debug', default=False, action='store_true')
    parser.add_option('-q', '--quiet', default=False, action='store_true', help='suppress output except errors')
    parser.add_option('-P', '--progress', default=False, action='store_true', help='display a progress meter while downloading')
    parser.add_option('--only', '--flavor', action='append', default=[], dest='flavor', help='select specific flavor to download; multiple --flavor args can be passed')
    parser.add_option('--exclude', '--exclude-flavor', action='append', default=[], dest='exclude_flavor')
    parser.add_option('--release', '--only-release', action='append', default=[], dest='release')
    parser.add_option('--variant', '--only-variant', action='append', default=[], dest='variant')
    parser.add_option('--exclude-variant', action='append', default=[], dest='exclude_variant')
    parser.add_option('--arch', '--only-arch', action='append', default=[], dest='arch')
    parser.add_option('--bwlimit', default='', action="store", type="string", help='rate limit rsync download to BWLIMIT in KBytes/second')
    parser.add_option('--build', default=default['build'], action="store", type="string")
    parser.add_option('--no-verify', default=True, action="store_false", dest='do_verify', help='skip comparison against published checksums')
    parser.add_option('--no-zsync', default=True, action="store_false", dest='use_zsync', help="don't use zsync to download")
    #parser.add_option('--versions', default=False, action='store_true')
    parser.add_option('--isoroot', default='', action='store', type='string', help='local directory root to store isos; default is %s' %(default['isoroot']))
    parser.add_option('--config', default=environ['HOME'] + '/.dl-ubuntu-test-iso', action='store', type='string', help='choose alternate config file location, default is %default')
    parser.add_option('--host', default=default['host'], action='store', type='string', help='mirror to pull images from; default is %s' %(default['host']))
    parser.add_option('--mirror', default=False, action='store_true', dest='mirror', help='mirror tree structure from cdimages')
    parser.add_option('--no-release-check', default=default['do_release_check'], action='store_false', dest='do_release_check', help='skip checking launchpad for new release')
    parser.add_option('--user', '--username', default=None, action='store', dest='username', help='specify a user on the remote host')
    parser.add_option('--pass', '--password', default=None, action='store', dest='password', help='specify a password on the remote host')

    config, args = parser.parse_args()
    if len(args) != 0:
        parser.error("incorrect number of arguments")
    file_cfg =  read_config(config.config)

    wanted['flavors'] = merge_config(config.flavor, file_cfg, default, 'flavors')
    if len(config.exclude_flavor) > 0:
        wanted['flavors'] = filter(lambda x: x not in config.exclude_flavor, wanted['flavors'])

    wanted['variants'] = merge_config(config.variant, file_cfg, default, 'variants')
    if len(config.exclude_variant) > 0:
        wanted['variants'] = filter(lambda x: x not in config.exclude_variant, wanted['variants'])
    wanted['build'] = merge_config(config.build, file_cfg, default, 'build')
    wanted['isoroot'] = merge_config(config.isoroot, file_cfg, default, 'isoroot')

    wanted['archs'] = merge_config(config.arch, file_cfg, default, 'archs')

    if not config.no_act and file_cfg != None and 'no_act' in file_cfg.keys():
        config.no_act = ('true' == file_cfg['no_act'].lower())

    if config.do_verify and file_cfg != None and 'verify' in file_cfg.keys():
        config.no_act = ('true' == file_cfg['verify'].lower())

    if config.host == default['host'] and file_cfg != None and 'host' in file_cfg.keys():
        config.host = file_cfg['host']

    if not config.username and file_cfg != None and 'dlusername' in file_cfg.keys():
        config.username = file_cfg['dlusername']

    if not config.password and file_cfg != None and 'dlpassword' in file_cfg.keys():
        config.password = file_cfg['dlpassword']

    if (config.username or config.password):
        if (config.username and config.password):
            register_auth_handler(config)
        else:
            parser.error("Both a username and a password must be given,if one is given")

    if config.quiet:
        config._log = lambda x: x
    else:
        config._log = my_log

    if config.use_zsync and not os.path.exists(zsync_binary):
        config._log("Warning! zsync is not installed, falling back to rsync")
        config.use_zsync = False

    if config.use_zsync and len(config.bwlimit) > 0:
        config._log("Warning! zsync does not support bandwidth limiting")

    # need to do this as the last config option handler as it hits the network
    if config.do_release_check:
        update_release()
    wanted['releases'] = merge_config(config.release, file_cfg, default, 'releases')

    #print file_cfg
    #print config
    #print wanted
    #dump_info(wanted, config)
    #check_status(wanted, config)
    iso_download(wanted, config)

if __name__ == "__main__":
    main()
