#!/usr/bin/python

"""
sugar-update.py

Created by Eben Eliason on 2007-08-28.
Please feel free to modify, extend, and redistribute this script
"""

import sys
import os
import getopt
import re
import urllib
import urllib2
import zipfile

#define help output ---------------------------------------------------------
def usage():
	print '\nUsage: sugar-update.py [-flLsv] [-b build] [directory]\n'
	print 'Description:\n'
	print '\tThis script will create an autoinstallation image boot directory within the'
	print '\tspecified directory.  If the boot directory already exists, it will only download'
	print '\tthe necessary image files (.img, .crc, and .rom).  Any of these files already'
	print '\tpresent within this boot directory will be deleted before the download begins.'
	print '\tIf no directory is specified, the current working directory will be used.\n'
	print '\tThis script has also has automatic support for downgrades, and will add the'
	print '\tnecessary force.os file to the boot directory if the build image already present'
	print '\tthere is newer than the build specified.\n'
	print 'Options:\n'
	print '\t -b\tNUM\tSpecify a build number; defaults to latest development build\n'
	print '\t -c\t\tCreate a clean auto-installation image, replacing existing boot directory\n'
	print '\t -d\t\tForce a downgrade image; implicit when a more recent build exists in boot\n'
	print '\t -f\t\tForce firmware update\n'
	print '\t -j\t\tUse joyride branch\n'
	print '\t -l\t\tPrint the associated change log\n'
	print '\t -L\t\tPrint the change log for the specified build and quit\n'
	print '\t -q\t\tQuiet; do not print progress output\n'
	print '\t -s\t\tDownload the latest stable build; overrides the -b option\n'

#check for valid arguments
try:
	opts,arg = getopt.getopt(sys.argv[1:], "b:cdfjlLqs")
except:
	usage()
	sys.exit(2)
		
#global constants -----------------------------------------------------------

LATEST = "latest"
STABLE = "LATEST-STABLE-BUILD"

build = ""
build_path = LATEST
firmware = ""
base_path = os.getcwd() + '/'
quiet = False;
log = False
log_only = False
update_firmware = False
downgrade = False
clean = False
joyride = False

if len(arg) > 0:
	base_path = arg[0].rstrip('/') + '/'
	
	#move base up one directory if boot directory was passed
	if base_path.endswith('boot/'):
		base_path = base_path[0:-5]

boot_path = base_path + 'boot/';	

#interpret arguments --------------------------------------------------------
for o, a in opts:
	if o == '-b':
		if build_path != STABLE:
			build = a
			build_path = 'build' + build
	elif o == '-c':
		clean = True
	elif o == '-d':
		downgrade = True
	elif o == '-f':
		update_firmware = True
	elif o == '-h':
		usage()
		exit(0)
	elif o == '-j':
		joyride = True
	elif o == '-l':
		log = True
	elif o == '-L':
		log = True
		log_only = True
	elif o == '-q':
		quiet = True
	elif o == '-s':
		build = ""
		build_path = STABLE

#this exception should be removed if the hieracrchies are fixed on olpc.downloads.redhat.com
if(build.find(".") > 0):
	if(joyride):
		os_path = "http://xs-dev.laptop.org/~cscott/olpc/streams/joyride/" + build_path + "/devel_jffs2/"
	else:
		os_path = "http://olpc.download.redhat.com/olpc/streams/development/" + build_path + "/"
else:
	if(joyride):
		os_path = "http://xs-dev.laptop.org/~cscott/olpc/streams/joyride/" + build_path + "/devel_jffs2/"
	else:
		os_path = "http://olpc.download.redhat.com/olpc/streams/development/" + build_path + "/devel_jffs2/"

fw_path = "http://dev.laptop.org/pub/firmware/LATEST/"
ai_path = "http://dev.laptop.org/git?p=users/cscott/autore;a=blob_plain;f=olpc-auto.zip;hb=HEAD"

#print a [message] and exit with [code]
def exit(message, code):
	print message
	sys.exit(code)

#verbose print: print [string] if not in quiet mode
def vprint(string):
	if not quiet:
		sys.stdout.write(string)
		sys.stdout.flush()
		
#progress update hook
def _reporthook(numblocks, blocksize, filesize, url=None):
	
	try:
		percent = min((numblocks*blocksize*100)/filesize, 100)
	except:
		percent = 100
	
	if(filesize > 0):
		
		if numblocks != 0:
			sys.stdout.write("\b"*4)
		
		sys.stdout.write("%3d%%" % percent)
	else:
		if numblocks != 0:
			sys.stdout.write("\b")
		
		n = numblocks%3
		
		if n == 0:
			sys.stdout.write("\\")
		elif n == 1:
			sys.stdout.write("/")
		elif n == 2:
			sys.stdout.write("-")
	
	sys.stdout.flush()


#download contents of [url] to [dst]
def geturl(url, dst):
	if sys.stdout.isatty() and not quiet:
		urllib.urlretrieve(url, dst, lambda nb, bs, fs, url=url: _reporthook(nb,bs,fs,url))
	else:
		urllib.urlretrieve(url, dst)
		
#determine latest stable or development build number ------------------------
if build == "":
	
	try:
		url_file = urllib.urlopen(os_path)
		
		for line in url_file:
			build,n = re.subn(r'.*os([\d.]+)\.img.*\n', r'\1', line)
			if n > 0:
				break
				
		url_file.close()
		
	except:
		if build_path == LATEST:
			exit('Error: Could not locate latest development build', 1)
		else:
			exit('Error: Could not locate latest stable build', 1)

#print the change log (when not in joyride) ---------------------------------
if log:
	try:
		
		url_file = urllib2.urlopen(os_path + "ChangeLog")
		
		print '\n--------------------------------------------------------------------------------'
		for line in url_file:
				print line.rstrip('\n')
		print '--------------------------------------------------------------------------------\n'
		
		url_file.close()	
	except urllib2.HTTPError, e:
		print '\n--------------------------------------------------------------------------------'
		if(joyride):
			print "No change log available for joyride build " + build
		else:
			print "No change log available for build " + build
		print '--------------------------------------------------------------------------------\n'

	if log_only:
		sys.exit(0)

#create a clean boot directory if specified ---------------------------------

if clean:
	
	try:
		for f in os.listdir(boot_path):
			os.remove(boot_path + f)

		os.rmdir(boot_path)
	
	except:
		pass

#create boot directory if needed --------------------------------------------
if not os.path.exists(boot_path):

	vprint('Creating boot directory at ' + base_path + "\n")
	zip_filename = "olpc-auto.zip"	
	
	try:
		os.mkdir(boot_path, 0777);
	except Exception, e:
		exit('Error: could not locate ' + base_path, 1)

	vprint("\tdownloading " + zip_filename + "...")
	geturl(ai_path, boot_path + zip_filename)
	print ""
	
	try:
		zip_obj = zipfile.ZipFile(boot_path + zip_filename)

		for filename in zip_obj.namelist():
			if filename.startswith('boot/'):
				vprint("\textracting " + filename + "...\n")
		
				f = open(base_path + filename, 'wb')
				f.write(zip_obj.read(filename))
				f.close()
	except:
		exit("Error: Could not extract files from " + zip_filename, 1)

	vprint("\tremoving " + zip_filename + "...\n")
	
	os.remove(boot_path + zip_filename)

#remove files from boot directory -------------------------------------------
vprint("Removing files from " + boot_path + "\n")

old_build = ""

for f in os.listdir(boot_path):
	if re.match(r'^os\d+\.', f):
		old_build = re.sub(r'^os(\d+)\..*', r'\1', f)
		vprint("\tremoving " + f + "...\n")
		os.remove(boot_path + f)

	if old_build < build and re.match(r'force\.os', f):
		vprint("\tremoving " + f + "...\n")
		os.remove(boot_path + f)

	if update_firmware and re.match(r'.*\.rom', f):
		vprint("\tremoving " + f + "...\n")
		os.remove(boot_path + f)

#force a downgrade when necessary -------------------------------------------

if old_build > build or downgrade:
	print "\tforcing downgrade..."
	downgrade_file = open(boot_path + "force.os", 'w')
	downgrade_file.write("The presence of this file forces a downgrade")
	downgrade_file.close()

#determine the latet firmware version ---------------------------------------
if update_firmware:
	
	try:
		url_file = urllib.urlopen(fw_path)
		
		for line in url_file:
			firmware,n = re.subn(r'.*href\s*=\s*"([^.]+)\.rom.*\n', r'\1', line)
			if n > 0:
				break
			
		url_file.close()
	
	except:
			exit('Error: Could not locate latest firmware version', 1)


rom_filename = firmware + ".rom"
md5_filename = firmware + ".rom.md5"
crc_filename = "os" + build + ".crc"
img_filename = "os" + build + ".img"

#download the ROM file ------------------------------------------------------

if update_firmware:
	vprint("Downloading latest firmware (" + firmware + ")\n")

if update_firmware:
	vprint("\tdownloading " + rom_filename + "...")
	geturl(fw_path + rom_filename, boot_path + rom_filename)
	print ""

#download the OS files ------------------------------------------------------

if build_path == LATEST:
	vprint("Downloading latest development build (" + build + ")\n")
elif build_path == STABLE:
	vprint("Downloading latest stable build (" + build + ")\n")
else:
	vprint("Downloading build " + build + "\n")

#download the CRC file ------------------------------------------------------

vprint("\tdownloading " + crc_filename + "...")
geturl(os_path + crc_filename, boot_path + crc_filename)
print ""

#download the IMG file ------------------------------------------------------
vprint("\tdownloading " + img_filename + "...")
geturl(os_path + img_filename, boot_path + img_filename)
print ""

vprint('Autoinstallation build ' + build + ' completed.\n\n')