#!/bin/env python
# encoding: iso-8859-1
"""
 Kernel updater
 Copyright (c) Phillip Berndt, 2006

 Updates the kernel to the newest version available.

 For instructions how to use this tool, execute "update-kernel -h".

"""
import urllib, sys, os, getopt, re, tarfile, math, shutil, gzip

class kernelVersions:
	"""
		Helper class for getting the kernel versions
	"""
	def newestVersion(self):
		"""
			Get the newest version's release number
		"""
		print " Downloading finger banner",
		sys.stdout.flush()
		try:
			data = urllib.urlopen("http://www.kernel.org/kdist/finger_banner").read()
			data = data.split("\n")[0].split(" ")[-1].split(".")
			if len(data) < 4:
				data.append(0)
		except:
			print "\r                              \r",
			return False
		print "\r                              \r",
		retVal = [ int(val) for val in data ] 
		return retVal

	def getDownloadLink(self, linkName):
		"""
			Get the link to a file based upon the link's text
		"""
		try:
			self.kernelMainpage
		except:
			print " Downloading kernel.org page",
			sys.stdout.flush()
			self.kernelMainpage = urllib.urlopen("http://www.kernel.org").read()
			print "\r                                    \r",
		link = re.search('"([^"]+)">%s<' % linkName, self.kernelMainpage)
		if link == None:
			print "Link not found"
			print "Maybe kernel.org changed it's layout?"
			print "Report this as a bug"
			sys.exit(0)

		return "http://www.kernel.org%s" % link.group(1)

	def currentVersion(self):
		"""
			Get the version number of the active kernel
		"""
		data = os.popen("uname -a").read()
		data = data.split(" ")[2].split(".")
		if len(data) < 4:
			data.append(0)
		return [ int(val) for val in data ] 

	def sourceVersion(self):
		"""
			Get the version number from /usr/src/linux
		"""
		try:
			makeFile = open("/usr/src/linux/Makefile", "r")
			makeOpts = {}		
			for opt in [ line.split(" = ")[0:2] for line in makeFile if len(line.split(" = ")) == 2 ]:
				makeOpts[opt[0]] = opt[1]
			
			try:
				makeOpts["EXTRAVERSION"]
			except:
				makeOpts["EXTRAVERSION"] = " 0"
			retVal = [int(makeOpts["VERSION"]), int(makeOpts["PATCHLEVEL"]), int(makeOpts["SUBLEVEL"]), int(makeOpts["EXTRAVERSION"][1:])]
			return retVal
		except:
			return False

class userInteraction:
	"""
		Output a text with a colored star before it inside TTYs (Gentoo style)
		Plus, this class provides you with the ability to ask questions
	"""
	def color(self, text, color):
		"""
			Output a text with color code "color", if possible
		"""
		if(sys.stdout.isatty()):
			print " \x1b[%d;01m*\x1b[39;00m %s" % (color, text)
		else:
			print " * ", text

	def colorMaker(self, color):
		"""
			A generator for color functions
			color must be a color code (ie. 32 for green)
		"""
		def tmp(text):
			self.color(text, color)
		return tmp

	def __init__(self):
		"""
			Initialize color functions
		"""
		self.green = self.colorMaker(32)
		self.red = self.colorMaker(31)
		self.yellow = self.colorMaker(33)

	def ask(self, text):
		"""
			Confront the user with a conditional question
		"""
		while 1:
			sys.stdout.write("   %s [y|n]: " % text)
			answer = sys.stdin.readline()[:-1]
			if answer == "n": return False
			if answer == "y": return True


def download(downloadFile, fileName):
	"""
		Function to download a file and show a status message
	"""
	if not os.access(fileName, os.F_OK):
		def status(sofar, size, total):
			percent = 1.0 * sofar * size / total * 100
			sys.stdout.write(" %d percent completed\r" % percent)
		urllib.urlretrieve(downloadFile, fileName, status)				

def handleAdditionalPatches(apply):
	"""
		Apply or remove additional patches
	"""
	additionalPatchDir = "/usr/src/patches/"
	try:
		patches = os.listdir(additionalPatchDir)
	except OSError:
		# Directory not found
		return
	if len(patches) is 0:
		return
	
	# Update patches
	if apply:
		for patch in filter(lambda x: x[-3:] == ".sh", patches):
			interactionGetter.green("Updating patches using %s" % patch)
			if os.system("cd /usr/src/patches; ./%s" % patch) is not 0:
				interactionGetter.yellow("Script execution failed")
	
	# Apply patches
	patches = os.listdir(additionalPatchDir)
	for patch in patches:
		oldDir = os.getcwd()
		os.chdir("/usr/src/linux")

		if patch[-4:] == ".bz2":
			catCommand = "bzcat"
		elif patch[-3:] == ".gz":
			catCommand = "gzcat"
		elif patch[-6:] == ".patch":
			catCommand = "cat"
		else:
			continue

		if apply:
			interactionGetter.green("Applying additional patch %s" % patch)

			patchApplied = False
			for pLevel in range(2):
				command = "%s %s | patch -p%d --dry-run 2>&1" % (catCommand, additionalPatchDir + patch, pLevel)
				result = os.popen(command, "r").read()
				if result.find("Perhaps you used the wrong -p") > -1:
					continue
				if result.find("Reversed (or previously applied)") > -1:
					interactionGetter.yellow("Already applied %s" % patch)
					patchApplied = True
					break
				if os.system(command + " -s &>/dev/null") is not 0:
					interactionGetter.red("Applying additional patch %s failed" % patch)
					sys.exit(1)

				command = "%s %s | patch -p%d %s" % (catCommand, additionalPatchDir + patch, pLevel, opts["silent"])
				if os.system(command) is not 0:
					interactionGetter.red("Applying additional patch %s failed IN REAL ENVIRONMENT" % patch)
					sys.exit(1)
				patchApplied = True

			if not patchApplied:
				interactionGetter.red("Failed to apply patch %s" % patch)
				if not interactionGetter.ask("Continue kernel compilation?"):
					sys.exit(1)
		else:
			interactionGetter.green("Removing additional patch %s" % patch)

			patchReversed = False
			for pLevel in range(2):
				command = "%s %s | patch -p%d -R --dry-run 2>&1" % (catCommand, additionalPatchDir + patch, pLevel)
				result = os.popen(command, "r").read()
				if result.find("Perhaps you used the wrong -p") > -1:
					continue
				if result.find("Unreversed patch") > -1:
					interactionGetter.yellow("Already reversed or never applied %s" % patch)
					patchReversed = True
					break
				if os.system(command + " -s &>/dev/null") is not 0:
					interactionGetter.red("Removing additional patch %s failed" % patch)
					sys.exit(1)

				command = "%s %s | patch -p%d -R %s" % (catCommand, additionalPatchDir + patch, pLevel, opts["silent"])
				if os.system(command) is not 0:
					interactionGetter.red("Removing additional patch %s failed IN REAL ENVIRONMENT" % patch)
					sys.exit(1)
				patchReversed = True

			if not patchReversed:
				interactionGetter.red("Failed to remove patch %s" % patch)
				sys.exit(1)

		os.chdir(oldDir)

if __name__ == "__main__":
	# Check arguments first
	args = {
		"a": "Automatically apply update (You'll still have to answer configuration questions manually)",
		"c": "Check for updates and return 1 if the kernel is out of date",
		"f": "Force rebuilding of the kernel",
		"h": "Display a help text on requirements and paths of this tool",
		"m": "Use menuconfig instead of silentoldconfig",
		"x": "Use xconfig instead of silentoldconfig",
		"r": "Automatically reboot after updating",
		"s": "Automatically shutdown after upgrading",
		"b": "Batch mode: Assume NO for all questions",
		"p:": "Apply (1) or unapply (0) patches to the sourcecode",
		"v": "Be verbose"
	}
	try:
		opts = getopt.getopt(sys.argv[1:], "".join(args.keys()))
	except:
		print "Kernel updater"
		print "Usage:"
		print " update-kernel %s" % " ".join([ "[-%s]" % arg for arg in args.keys() ])
		print
		for info in args.items():
			if len(info[0]) == 1:
				print " -%s    %s" % info
			else:
				print " -%s   %s" % info
		print
		sys.exit(0)
	opts = dict(opts[0])
	opts["silent"] = ""
	if not "-v" in opts:
		opts["silent"] = "-s &>/dev/null"

	# View help?
	if "-h" in opts:
		print "Kernel updater"
		helpText = """
			The kernel updater tool was initially intended to run on my pc only,
			so I fixed the locations for paths, etc.:

			 - /usr/src/linux
			   Must contain the sourcecode of the linux kernel. If the directory
			   does not exist, it will be created by this tool.
			   It will NOT be symlinked.
			   
			 - /usr/src/source
			   Will contain the source files for kernelarchives (patches, full source)

			 - /boot [Mountpoint!]
			   The kernel will be copied into this directory, named "bzImage".
			   Any existing "bzImage" file will be renamed to "bzImage_backup".

			 - /usr/src/patches
			   Every patch inside of this directory will be applied to the kernel
			   source and removed before sublevel updates. Useful for things like
			   swsusp2.
			   Every *.sh file in this directory will be executed before "apply".
			   You may use those file to update patches to new revisions
			   Every *.patch.(bz2|gz)? file in this directory will be applied
			   automatically.

			When updating to a new subversion, the old kernel source will be deleted
			and a new kernel will be downloaded. This tool tries to load the old
			configuration from /proc/config or /proc/config.gz.

			New kernel releases often contain new features. This tool uses the
			"silentoldconfig" make target, so that you'll be asked wheather to
			use the new kernel features or not. If a new kernel get's installed and
			you don't have a configuration yet I suggest that you use the menuconfig
			option instead.			

			Run "update-kernel -help" for command line options.
		"""
		print re.sub("\t\t\t", " ", helpText, 9999)
		print
		print
		sys.exit(0)
		
	# Check for requirements
	if not "-f" in opts and not os.access("/usr/sbin/module-rebuild", os.F_OK):
		print "You have to install module-rebuild to use update-kernel"
		print "In gentoo, try emerge module-rebuild"
		print
		sys.exit(1)
	

	# Create needed objects
	interactionGetter = userInteraction()
	kernelVersionGetter = kernelVersions()

	# If no output is wanted, close stdout
	if "-c" in opts:
		sys.stdout = open("/dev/null", "w")
	
	print "Kernel updater"

	# Apply or unapply patches
	if "-p" in opts:
		if opts["-p"] == "1":
			handleAdditionalPatches(True)
		elif opts["-p"] == "0":
			handleAdditionalPatches(False)
		else:
			print "Sorry, I did not understand '%s'" % opts["-p"]
		sys.exit(0)
	
	# Check version
	newestVersion   = kernelVersionGetter.newestVersion()
	currentVersion  = kernelVersionGetter.currentVersion()
	sourceVersion   = kernelVersionGetter.sourceVersion()
	prevVersion     = sourceVersion

	if newestVersion == False:
		# Not online
		interactionGetter.red("Not connected to the internet. Assuming that the kernel is up to date")
		newestVersion = currentVersion

	if sourceVersion == False:
		# No kernel available
		interactionGetter.yellow("No kernel source installed. Assuming source version 0.0.0.0")
		sourceVersion = [0, 0, 0, 0]

	sNewestVersion  = ".".join([str(a) for a in newestVersion])
	sSourceVersion = ".".join([str(a) for a in sourceVersion])
	sCurrentVersion = ".".join([str(a) for a in currentVersion])
	
	updateNeeded   = False
	interactionGetter.green("Newest version:        %s" % sNewestVersion)
	if newestVersion == sourceVersion:
		fcolor = interactionGetter.green
	else:
		fcolor = interactionGetter.yellow
		updateNeeded = True
	fcolor("Sourcecode version:    %s" % sSourceVersion)

	if newestVersion == currentVersion:
		fcolor = interactionGetter.green
	else:
		fcolor = interactionGetter.yellow
		updateNeeded = True
	fcolor("Active kernel version: %s" % sCurrentVersion)

	# If beeing run from a batch file, return now
	if "-c" in opts:
		if updateNeeded == True:
			sys.exit(1)
		else:
			sys.exit(0)
	
	# If no update is needed, exit
	if updateNeeded == False and not "-f" in opts:
		interactionGetter.green("Kernel is up to date")
		sys.exit(0)

	# Ask if the changelog shall be viewed
	if not "-a" in opts and not "-b" in opts and interactionGetter.ask("View Changelog?"):
		changeLogFile      = urllib.urlopen(kernelVersionGetter.getDownloadLink("Changelog")).read()
		changeLogChanges   = re.findall('^    \[.+$', changeLogFile, re.MULTILINE)
		try:
			less               = os.popen("less -F", "w")
			print >> less, "   Changes:"
			print >> less, "\n".join(changeLogChanges)
			less.close()
		except IOError:
			# Ignore, as this will occour when a user quits "less"
			pass
		
	# Ask if the user wants to update
	if not "-a" in opts and ("-b" in opts or not interactionGetter.ask("Update the system?")):
		sys.exit(1)
	
	# Check the difference between the versions
	# First Version and Subversion (i.e. 2.6)
	# I will not try to patch between those versions
	# but just download a new kernel
	if sourceVersion[0:2] != newestVersion[0:2] or sourceVersion[2] + 1 < newestVersion[2]:
		interactionGetter.yellow("Current kernel version is too old. Downloading newest kernel source...")
		downloadFile = kernelVersionGetter.getDownloadLink("F")
		fileName = "/usr/src/source/%s" % os.path.basename(downloadFile)
		# Ask the user
		if not "-a" in opts and ("-b" in opts or not interactionGetter.ask("WARNING: I will remove the old kernel source now. Proceed?")):
			sys.exit(1)
		# Create/Purge linux and source directories
		try:    os.makedirs("/usr/src/source")
		except: pass
		shutil.rmtree("/usr/src/linux", True)
		try:    os.makedirs("/usr/src/linux/_")
		except: pass
		# Download linux archive
		download(downloadFile, fileName)
		# Extract it
		interactionGetter.green("Extracting archive to /usr/src/linux")
		try:
			tarArchive = tarfile.open(fileName, "r:bz2")
			ilen = len(tarArchive.getnames())
			pos = 0
			for tarFile in tarArchive:
				tarArchive.extract(tarFile, "/usr/src/linux/_/")
				pos = pos + 1
				if math.fmod(pos, 10) == 0:
					percent = pos * 100.0 / ilen
					sys.stdout.write(" %d percent completed\r" % percent)
		except:
			interactionGetter.red("Failed to extract archive.")
			interactionGetter.red("Maybe the file is corrupted. Delete %s and try again." % fileName)
			sys.exit(1)
		# Put old config file into that directory
		try:
			shutil.copyfile("/proc/config", "/usr/src/linux/.config")
		except:
			try:
				config = open("/usr/src/linux/.config", "w")
				config.write(gzip.open("/proc/config.gz").read())
				config.close()
			except:
				interactionGetter.red("Failed to read old configuration. You'll have to configure everything manually")
		# Move the kernel source to the correct directory
		os.system("mv /usr/src/linux/_/linux*/* /usr/src/linux")
		shutil.rmtree("/usr/src/linux/_")
		# Apply additional patches
		handleAdditionalPatches(True)
		sourceVersion = newestVersion
			
	# Now, however, if there are differences left, reverse current extraversion patch
	if sourceVersion != newestVersion and sourceVersion[3] != 0:
		interactionGetter.green("Downloading source's extraversion patch")
		# Get the current filename based upon the newest link
		downloadFile = kernelVersionGetter.getDownloadLink(re.sub("\.0$", "", sNewestVersion))
		downloadFile = downloadFile.replace(re.sub("\.0$", "", sNewestVersion), sSourceVersion)
		fileName = "/usr/src/source/%s" % os.path.basename(downloadFile)
		# Download the patch
		download(downloadFile, fileName)
		# Reverse it
		if os.system("cd /usr/src/linux; bzcat %s | patch -R -p1 -t --dry-run %s" % (fileName, opts["silent"])):
			interactionGetter.red("Failed to reverse patch!")
			interactionGetter.red("This might be a corrupted download. Delete %s and try again." % fileName)
			sys.exit(1)
		if os.system("cd /usr/src/linux; bzcat %s | patch -R -p1 -t  %s" % (fileName, opts["silent"])):
			interactionGetter.red("Failed to reverse patch IN ACTION!")
			interactionGetter.red("Please remove /usr/src/linux and try again. Report this as a bug!")
			sys.exit(1)
		sourceVersion[3] = 0
		sSourceVersion = ".".join([str(a) for a in sourceVersion])
		interactionGetter.green("Reversed to version %s" % sSourceVersion)

	# Update to the newest sublevel
	if sourceVersion[2] != newestVersion[2]:
		# Remove additional patches
		handleAdditionalPatches(False)
		interactionGetter.green("Downloading patch to current sublevel")		
		# Get the current filename based upon the newest link
		downloadFile = kernelVersionGetter.getDownloadLink(re.sub("\.0$", "", sNewestVersion))
		sNewestVersionSublevel = "%s.%s.%s" % (newestVersion[0], newestVersion[1], newestVersion[2])
		downloadFile = downloadFile.replace(re.sub("\.0$", "", sNewestVersion), sNewestVersionSublevel)
		fileName = "/usr/src/source/%s" % os.path.basename(downloadFile)
		# Download the patch
		download(downloadFile, fileName)
		# Apply the patch
		interactionGetter.green("Applying patch to sublevel %s" % sNewestVersionSublevel)
		if os.system("cd /usr/src/linux; bzcat %s | patch -p1 -t --dry-run %s" % (fileName, opts["silent"])):
			interactionGetter.red("Failed to apply patch!")
			interactionGetter.red("This might be a corrupted download. Delete %s and try again." % fileName)
			sys.exit(1)
		if os.system("cd /usr/src/linux; bzcat %s | patch -p1 -t %s" % (fileName, opts["silent"])):
			interactionGetter.red("Failed to apply patch IN ACTION!")
			interactionGetter.red("Please remove /usr/src/linux and try again. Report this as a bug!")
			sys.exit(1)
		sourceVersion[2] = newestVersion[2]
		# Apply additional patches
		handleAdditionalPatches(True)

	# Update to the newest extraversion
	if sourceVersion != newestVersion:
		interactionGetter.green("Downloading patch to current extraversion")
		# Get the current filename based upon the newest link
		downloadFile = kernelVersionGetter.getDownloadLink(sNewestVersion)
		fileName = "/usr/src/source/%s" % os.path.basename(downloadFile)
		# Download the patch
		download(downloadFile, fileName)
		# Apply the patch
		interactionGetter.green("Applying patch to extraversion %s" % sNewestVersion)
		if os.system("cd /usr/src/linux; bzcat %s | patch -p1 -t --dry-run %s" % (fileName, opts["silent"])):
			interactionGetter.red("Failed to apply patch!")
			interactionGetter.red("This might be a corrupted download. Delete %s and try again." % fileName)
			sys.exit(1)
		if os.system("cd /usr/src/linux; bzcat %s | patch -p1 -t %s" % (fileName, opts["silent"])):
			interactionGetter.red("Failed to apply patch IN ACTION!")
			interactionGetter.red("Please remove /usr/src/linux and try again. Report this as a bug!")
			sys.exit(1)

	# Now the sourcecode should be up to date
	# Update the system
	if opts["silent"] != "":
		opts["silent"] = " &>/dev/null"
	os.chdir("/usr/src/linux")
	interactionGetter.green("Updating configuration...")
	command = "make silentoldconfig"
	if "-m" in opts:
		command = "make menuconfig"
	if "-x" in opts:
		command = "make xconfig"
	if os.system(command):
		interactionGetter.red("Failed to update configuration")
		sys.exit(0)
	interactionGetter.green("Compiling new kernel")
	if os.system("make %s" % opts["silent"]):
		interactionGetter.red("Failed to compile kernel")
		sys.exit(0)
	interactionGetter.green("Installing modules")
	if os.system("make modules_install %s" % opts["silent"]):
		interactionGetter.red("Failed to compile kernel")
		sys.exit(1)
	interactionGetter.green("Merging new kernel into system")
	arch = os.popen("uname -m").read().strip()
	if os.system("( ((mount | grep '/boot') || mount /boot) && (([ -e /boot/bzImage ] && cp /boot/bzImage /boot/bzImage_backup) || [ ! -e /boot/bzImage ]) && cp arch/%s/boot/bzImage /boot && umount /boot) %s" % (arch, opts["silent"])):
		interactionGetter.red("Failed to merge kernel")
		sys.exit(1)
	interactionGetter.green("Rebuilding needed modules")
	if prevVersion != newestVersion:
		if os.system("module-rebuild rebuild %s" % opts["silent"]):
			interactionGetter.red("Failed to rebuild modules")
			sys.exit(1)
	else:
		interactionGetter.green(" No rebuild required")

	# Now ask the user to reboot
	if "-s" in opts:
		os.system("shutdown -h now 'Shutdown after kernel update'")
	elif "-r" in opts or (not "-r" in opts and interactionGetter.ask("Finished. Shall I reboot?")):
		os.system("init 6")
	
	print
	print

	# Finished