#!/usr/bin/python

# This is meant to upgrade users to Hardy's developing kernel (2.6.24)
# Written by walkerk @ubuntuforums.org
# Do with it as you please...

import os, sys, socket, time

packages =  "linux linux-generic linux-headers-generic linux-image-generic linux-restricted-modules-generic" 
repository = "deb http://archive.ubuntu.com/ubuntu/ hardy main restricted"

# Prompt the user and return True/False
def query(prompt, retries=4, complaint="Yes or no, please!"):
    while True:
        ok = raw_input(prompt + " ")
        if ok in ("y", "ye", "yes"): return True
        if ok in ("n", "no", "nop", "nope"): return False

        retries = retries - 1
        if retries < 0: raise IOError, "refusenik user"
        print complaint


# Update /etc/apt/sources.list
def apt_sources(choice):
	if choice == 1: # Add the Gutsy repository
		fileObj = open("/etc/apt/sources.list.d/kernel.list", "w")
		fileObj.write(repository)
		fileObj.close()

		time.sleep(2)
		os.system("sudo apt-get update")

	elif choice == 2: # Remove the Hardy repository
		os.remove("/etc/apt/sources.list.d/kernel.list")		

		time.sleep(2)
		os.system("sudo apt-get update")

		clean_up = query("\n\033[1mWould you like to run autoclean to\nremove unnecessary files left behind?\033[0m")
		if clean_up == True: os.system("sudo apt-get autoremove")

	else:
		print "\033[1mAn error occured while trying to update sources.list\033[0m"
		sys.exit()

# Test network connectivity
def net_test():
	print "\nTesting your network connection..."	

	try:
		s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		s.connect(("www.google.com", 80))
		return True

	except socket.error:
		return False

# Determine if the user is using an nVidia graphics card
def nvidia():
	temp = os.popen("lspci")
	nvidia = temp.read()

	if nvidia.find("nVidia") != -1: return True
	else: return False

##########################
### Begin installation ###
##########################

if __name__ == "__main__":
	print "\n\033[1mUpgrade to Hardy's developing kernel (2.6.24)\033[0m"
	print "This has been tested on Gutsy but should work for others."
	print "Written by walkerk @ubuntuforums.org\nDo with it as you please...\n\n"

	install = query("Do you wish to continue the upgrade?")
	if install == False: sys.exit()

	# If an nVidia graphics card is present, prompt the user (Still working on this)
	if nvidia() == True:
		print "\n\033[1mAn nVidia graphics card has been detected...\033[0m"
		print "You will need to re-install the drivers (2 options)..."
		print "\n\033[1m[1]\033[0m Add the nvidia-glx-new package to the install (automated)"
		print "\033[1m[2]\033[0m Re-run the Envy script after installation (manually)\n"
		choice = raw_input("Please make your choice: ")

		if choice == "1":
			packages =  "linux linux-generic linux-headers-generic linux-image-generic linux-restricted-modules-generic nvidia-glx-new"
		elif choice == "2":
			packages =  "linux linux-generic linux-headers-generic linux-image-generic linux-restricted-modules-generic"
		else:
			print "\nInvalid choice. Aborting\n"
			sys.exit()

	# Testing network connection
	network = net_test()
	if network == False: 
		print "\nYou are not connected to the internet."
		print "You will not be able to download the packages. \033[1mAborting.\033[0m\n"
		sys.exit()

	# Download and install the packages
	apt_sources(1)
	os.system("sudo apt-get install " + packages)
	apt_sources(2)
	
	# Complete.. Reboot?
	print "\n\033[1mInstallation is complete. A reboot is necessary...\033[0m"
	reboot = query("Would you like to reboot now?")
	if reboot == True: os.system("sudo reboot")
	else: print "\n\033[1mYou will need to reboot before the changes take place. Done.\033[0m\n"