"""
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 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 -d\t\tForce a downgrade image; implicit when a more recent build exists in boot\n'
	print '\t -f\t\tSkip firmware update\n'
	print '\t -l\t\tPrint the associated change log\n'
	print '\t -L\t\tPrint the change log for the specified build without performing a download\n'
	print '\t -o\t\tOverwrite the existing boot directory with a clean auto-installation image\n'
	print '\t -s\t\tDownload the latest stable build; overrides the -b option\n'
	print '\t -v\t\tVerbose\n'

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

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

build = ""
build_path = LATEST
firmware = ""
base_path = os.getcwd() + '/'
verbose = False;
log = False
log_only = False
update_firmware = True
downgrade = False
overwrite = 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 == '-s':
		build = ""
		build_path = STABLE
	elif o == '-l':
		log = True
	elif o == '-L':
		log = True
		log_only = True
	elif o == '-v':
		verbose = True
	elif o == '-f':
		update_firmware = False
	elif o == '-d':
		downgrade = True
	elif o == '-o':
		overwrite = True
	elif o == '-h':
		usage()
		exit(0)
		
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"

#determine latest stable or development build number ------------------------
if build == "":
	
	try:
		url_file = urllib2.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:
			sys.exit('Error: Could not locate latest development build')
		else:
			sys.exit('Error: Could not locate latest stable build')

crc_filename = "os" + build + ".crc"
img_filename = "os" + build + ".img"

#print the change log -------------------------------------------------------
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--------------------------------------------------------------------------------'
		print "No change log available for build " + build
		print '--------------------------------------------------------------------------------\n'

	if log_only:
		sys.exit(0)

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

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

	os.rmdir(boot_path)

#create boot directory if needed --------------------------------------------
if not os.access(boot_path, 0777):

	print 'Creating boot directory at ' + base_path
	zip_filename = "olpc-auto.zip"	
	
	try:
		os.mkdir(boot_path, 0777);
		
		if verbose:
			print "\tdownloading " + zip_filename + "..."
		try:
			url_file = urllib2.urlopen(ai_path)
			ai_file = open(boot_path + zip_filename, 'wb')

			for line in url_file:
				ai_file.write(line)

			url_file.close()
			ai_file.close()
		except urllib2.HTTPError, e:
			sys.exit('Error: Could not download autoinstallation file (' + str(e.code) + ')')
		except:
			sys.exit('Error: Could not write ' + boot_path + zip_filename)
	
	except:
		sys.exit('Error: could not locate ' + base_path)
		
	try:
		zip_obj = zipfile.ZipFile(boot_path + zip_filename)

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

	if verbose:
		print "\tremoving " + zip_filename + "..."
	
		os.remove(boot_path + zip_filename)

#remove files from boot directory -------------------------------------------
if verbose:
	print "Removing files from " + boot_path

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)
		if verbose:
			print "\tremoving " + f + "..."
		os.remove(boot_path + f)

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

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

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

#determine the latet firmware version ---------------------------------------
if update_firmware:
	
	try:
		url_file = urllib2.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:
			sys.exit('Error: Could not locate latest firmware version')

rom_filename = firmware + ".rom"
md5_filename = firmware + ".rom.md5"

crc_filename = "os" + build + ".crc"
img_filename = "os" + build + ".img"

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

if verbose and update_firmware:
	print "Downloading latest firmware (" + firmware + ")"

if update_firmware:
	if verbose:
		print "\tdownloading " + rom_filename + "..."
	try:
		url_file = urllib2.urlopen(fw_path + rom_filename)
		rom_file = open(boot_path + rom_filename, 'wb')

		for line in url_file:
			rom_file.write(line)

		url_file.close()
		rom_file.close()
	except urllib2.HTTPError, e:
		sys.exit('Error: Could not download rom file (' + str(e.code) + ')')
	except:
		sys.exit('Error: Could not write ' + boot_path + rom_filename)


if verbose:
	if build_path == LATEST:
		print "Downloading latest development build (" + build + ")"
	elif build_path == STABLE:
		print "Downloading latest stable build (" + build + ")"
	else:
		print "Downloading build " + build	

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

if verbose:
	print "\tdownloading " + crc_filename + "..."
try:
	url_file = urllib2.urlopen(os_path + crc_filename)
	crc_file = open(boot_path + crc_filename, 'wb')

	for line in url_file:
		crc_file.write(line)

	url_file.close()
	crc_file.close()
except urllib2.HTTPError, e:
	sys.exit('Error: Could not download crc file (' + str(e.code) + ')')
except:
	sys.exit('Error: Could not write ' + boot_path + crc_filename)

#download the IMG file ------------------------------------------------------
if verbose:
	print "\tdownloading " + img_filename + "..."
try:
	url_file = urllib2.urlopen(os_path + img_filename)
	img_file = open(boot_path + img_filename, 'wb')

	for line in url_file:
		img_file.write(line)

	url_file.close()
	img_file.close()
except urllib2.HTTPError, e:
	sys.exit('Error: Could not download img file (' + str(e.code) + ')')
except:
	sys.exit('Error: Could not write ' + boot_path + img_filename)


if verbose:
	print 'Autoinstallation build ' + build + ' completed.\n'