#!/usr/bin/python
#
# This script helps the end user configure their drives (removable or otherwise)
# to be used with the toolkit
#
# Author: Dan Bracey 5/30/08
# Revision History:
# 6/12/08 -- 1.000 -- debracey -- Initial Creation

import os
import sys

bindir = os.path.abspath(os.path.dirname(os.path.realpath(sys.argv[0])))
libdir = os.path.join(bindir, "..", "python_lib")

sys.path.append(libdir)

import Internet2Lib       # Library functions
import Internet2Consts    # For constants
import readline           # Improve the usability of raw_input. Let you use those fancy "arrow" keys
import shutil             # For copying the files around
import commands           # For checking of the given drive is a volume group
import re                 # Regular Expressions

MINIMUM_SIZE = 10          # Minimum drive size in Gigabytes

# Gets the list of drives and their associated mount points
# @return A list of tuples -- (dev, mntPoint, size, Internal/USB/Floppy (0/1/2), %used, boolean formatted) (or None if error)
# Ex. [(/dev/sda, /media/sda, 70GB, 0, 2), ...]
def getDriveList():

    driveList = []

    # Check /proc/partitions to get the list of raw partitions
    (status, output) = commands.getstatusoutput("cat /proc/partitions")
    if status == 0:
        lines = output.split("\n")
        for line in lines:
            columns = line.split(None)
            if len(columns) < 4:
                continue

            major = columns[0]
            minor = columns[1]
            blocks = columns[2]
            name = columns[3]

            # major 8 == sd[1-15] disk
            # major 3 == hd[ab]* disk
            # major 22 == hd[cd]* disk
            # major 33 == hd[ef]* disk
            # major 34 == hd[gh]* disk
            # major 56 == hd[ij]* disk
            # major 57 == hd[kl]* disk
            # major 65 == sd[16-31]* disk
            # major 66 == sd[32-47]* disk
            # major 67 == sd[48-63]* disk
            # major 68 == sd[64-79]* disk
            # major 69 == sd[80-95]* disk
            # major 70 == sd[96-111]* disk
            # major 71 == sd[112-127]* disk
            # major 128 == sd[128-143]* disk
            # major 129 == sd[144-159]* disk
            # major 130 == sd[160-175]* disk
            # major 131 == sd[176-191]* disk
            # major 132 == sd[192-207]* disk
            # major 133 == sd[208-223]* disk
            # major 134 == sd[224-239]* disk
            # major 135 == sd[240-255]* disk
            # major 104 == cciss (HP Raid Card)

            if major == "major":
                continue

            if major == "8" or major == "3" or major == "22" or major == "33" or major == "34" or major == "56" or major == "57" or (int(major) >= 65 and int(major) <= 71) or (int(major) >= 128 and int(major) <= 135) or major == "104":

		# convert from bytes to gigabytes
                size = float(blocks) / 1024 / 1024

                dev_str = "/dev/"+name
                driveList.append({ 'name':dev_str, 'size':size })

    lvm_devs = []

    # Check LVM to get the list LVM partitions
    (status, output) = commands.getstatusoutput("lvs --noheadings -o lv_name,vg_name,size,devices --units G")
    if status == 0:
        lines = output.split("\n");
        for line in lines:
            columns = line.split(None)
            if len(columns) < 4:
                continue

            lv = columns[0]
            vg = columns[1]
            size = columns[2]
            device = columns[3]

            dev_str = "/dev/mapper/"+vg+"-"+lv
            if lv != "No" and lv != "LV":
	        m = re.match("(.*)([0-9]+)\([0-9]+\)", device)
	        if m:
		    # This will add /dev/sda2
		    lvm_devs.append(m.group(1)+m.group(2))
		    # This will add /dev/sda
		    lvm_devs.append(m.group(1))

	    	m = re.match("([0-9.]+)G", size)
	        if m:
		    size = float(m.group(1))

	            driveList.append({ 'name':dev_str, 'size':size })
    #end lvs checking

    # use blkid to query for the cdrom
    cdrom_devs = []
    (status, output) = commands.getstatusoutput("blkid -t TYPE=iso9660 -o device")
    if status == 0:
        lines = output.split("\n");
        for line in lines:
            cdrom_devs.append(line.rsplit('\n'))

    ext3_devs = []
    (status, output) = commands.getstatusoutput("blkid -t TYPE=ext3 -o device")
    if status == 0:
        lines = output.split("\n");
        for line in lines:
            ext3_devs.append(line.rsplit('\n')[0])

    devs_types = {}
    (status, output) = commands.getstatusoutput("blkid -o full")
    if status == 0:
        lines = output.split("\n");
        for line in lines:
            m = re.match('^([^:]*):.*TYPE="([^"]*)"', line)
            if m:
	        devs_types[m.group(1)] = m.group(2);

    retList   = []
    
    for drive in driveList:
        # Exclude any CD or floppy drives
        if drive["name"] in cdrom_devs or "cd" in drive["name"] or "floppy" in drive["name"] or ".hal-mtab-lock" in drive["name"] or "fd" in drive["name"]:
            continue

        # Exclude devices being used for LVM, if users really want to use them, it's not up to us to blow them away.
        if drive["name"] in lvm_devs:
	    continue

#        is_ext3 = False
#        if drive["name"] in ext3_devs:
#	    is_ext3 = True

        type = "unknown"
        try:
            type = devs_types[drive["name"]]
        except KeyError:
            type = "unknown"

        # Append the drive and the device
        tuple = (drive["name"], drive["size"], type)
        retList.append(tuple)
    # End loop over drives

    return retList
# End getDriveList
    
# Checks to see if the given partition/drive is USB or not
# @param partition The partition to check
# @return 1 if USB, 0 otherwise
def isDriveUSB(partition):
    command = "dmesg | grep removable | cut -d' ' -f7 | sort -u"
    (status, output) = commands.getstatusoutput(command)
    if status != 0 or output == None: 
        return False

    drives = output.split("\n")

    for drive in drives:
        drive   = drive.strip() # Just to be safe
        
        command = "dmesg | grep " + drive + ":." + drive + " | cut -d' ' -f3- | sort -u"
        parts   = Internet2Lib.runCommand(command)

        if parts == None:
            return 0

        for part in parts:
            part = part.strip() # Just to be safe
            
            # If what they gave us matches, it's USB (removable really)
            if partition == part:
                return 1
    
    return 0
## End isDriveUSB
def getMountPoint(dev):
    # Check /proc/partitions to get the list of raw partitions
    (status, output) = commands.getstatusoutput("cat /proc/mounts")
    if status == 0:
        lines = output.split("\n")
        for line in lines:
            columns = line.split(None)
            if len(columns) < 4:
                continue

            curr_dev    = columns[0]
            mount_point = columns[1]

            if curr_dev == dev:
                return mount_point

    return None
# End getMountPoint    

def readStoreInfo():
    try:
        fileHandle = open(Internet2Consts.SAVEINFO)
        fileHandle, fileLines = Internet2Lib.readFile(fileHandle)
        fileHandle.close()
    except IOError:
        return None

    version = None
    path    = None
    device  = None

    # Loop over file checking for the values we need
    for line in fileLines:
        columns = line.split(None)
        try:
             value = columns[1]
        except:
             value = None

        if columns[0] == "version": 
            version = value
        elif columns[0] == "path": 
            path = value
        elif columns[0] == "device": 
            device = value
    # End loop over file

    return { 'version':version, 'path':path, 'device':device }
# End readStoreInfo

# Removes old configuration files
# The purpose of this is to make sure we dont end up with duplicate sets of the config files, that could be bad
def removeOldConfigs(mount_point=None):
    if not mount_point:
        store_info = readStoreInfo()
        if store_info and store_info["path"]:
            mount_point = store_info["path"]

    if not mount_point:
        mount_point = os.path.join("/", "mnt", "store")

    nptools_path = os.path.join(mount_point, "NPTools")
    new_nptools_path = nptools_path + ".back"

    if os.path.exists(new_nptools_path):
        shutil.rmtree(new_nptools_path)
    if os.path.exists(nptools_path):
        shutil.move(nptools_path, new_nptools_path)

    return
# End removeOldConfigs

### MAIN ###
print "Welcome to the Internet2 pS-Performance Toolkit drive configuration program."
print "This program will prompt you to setup your harddrives/removable media so they can be used with the toolkit."
#print Internet2Consts.WHITE + "Note: You need to format/partition your drives before running this script." + Internet2Consts.NORMAL
print Internet2Consts.WHITE + "Default values are in []" + Internet2Consts.NORMAL

# Must be root
if not Internet2Lib.isRoot():
    print Internet2Consts.YELLOW + "You must run this script as root. Please rerun this script with root permissions." + Internet2Consts.NORMAL
    sys.exit(1)

# This actually mounts the drives
drives = getDriveList()

print "" # Need extra space

if not drives:
    print Internet2Consts.YELLOW + "Warning: No drives available to hold configuration files." + Internet2Consts.NORMAL
    print Internet2Consts.YELLOW + "Check your internal hard drive, or insert a thumb drive."  + Internet2Consts.NORMAL
    sys.exit(1)

# Sort the drives by size
drives.sort(cmp=lambda x,y: cmp(y[1], x[1]))

done     = False
count    = 0

while not done:
    saveNow  = False

    print "Please choose a location for where you would like to save the toolkit's configuration:"
    print "The minimum suggested size for the disk is %d G. You may select a smaller size disk, but that is not suggested." % MINIMUM_SIZE

    for i, drive in enumerate(drives):
        print "\t%d) %-40s Capacity: % 8.2f G Type: %s" % (i+1, drive[0], drive[1], drive[2])
    # End for

    # Print the exit option
    print "\t0. exit\n"

    # Process the user's input
    ans = raw_input("Please select a drive: ").strip()
    try:
        choice = int(ans)
    except ValueError, e:
        choice = -1 # Not valid

    if choice == 0:
        done = True
        break

    if choice < 0 or choice > len(drives):
        print Internet2Consts.YELLOW + "Error: Invalid selection\n" + Internet2Consts.NORMAL
        continue

    # Decrement the choice to represent the correct drive in the list
    count  = count+1
    choice = choice-1

    if drives[choice][1] < MINIMUM_SIZE:
        print "Drive " + drives[choice][0] + " is below the minimum support drive size of "+str(MINIMUM_SIZE)+" Gigabytes."
	ans = raw_input("Are you sure you'd like to use this drive? [no]: ").strip().upper()
	if not ans:
	    ans = "N"
	
	if ans[0] != "Y":
	    continue

    if drives[choice][2] != "ext3":
        print "To use drive " + drives[choice][0] + ", we will need to overwrite any existing data"
        print "that is there. You can opt out of using this drive by answering"
        print "no to the following question, but you will need to select one"
        print "drive to use. If you have a drive formatted as ext3 already, you"
        print "can use that drive without need to overwrite the existing data"

    ans = raw_input("Would you like to erase the existing data on drive "+drives[choice][0]+"? This will format the drive, removing any files currently on it. [no]: ").strip().upper()
    if not ans:
        ans = "N"

    if ans[0] != "Y" and drives[choice][2] == False:
        continue

    if ans[0] == "Y":
        # Check to see if a knoppix.sh exists. If so, the directories got
        # bind-mounted, and we need to stop that. XXX, we should check that the
        # stuff in SAVEINFO matches this. Otherwise, we blow away there
        # existing installation. This may not be a big deal since they're doing
        # so anyway...
        store_info = readStoreInfo()
        if store_info:
            if store_info["device"] and store_info["device"] == drives[choice][0]:
                print Internet2Consts.YELLOW + "Warning: this drive contains pS-Performance Toolkit data.\n\n" + Internet2Consts.NORMAL
                print "To use this drive, you will need to overwrite the data that is there. This will require a reboot."

                ans = raw_input("Would you like to overwrite the existing pS-Performance Toolkit data? [yes]: ").strip().upper()
                if not ans:
                    ans = "Y"

                if ans[0] == "Y":
                    removeOldConfigs()
                    status = os.system("/sbin/reboot");
                    sys.exit(1)
                else:
                    continue

        status = os.system("/sbin/mkfs.ext3 -F "+drives[choice][0]);
        if status != 0:
            print Internet2Consts.YELLOW + "Error: Couldn't format drive "+drives[choice][0]+"\n" + Internet2Consts.NORMAL
            continue

    # We don't want two sets of config files - that could mess up the boot
    # sequence
    removeOldConfigs()

    mount_point = os.path.join("/", "mnt", "store.new")
    nptools_directory = os.path.join(mount_point, "NPTools")
    scratch_file = os.path.join(nptools_directory, "scratch")
    scratch_file_size = 1000000000

    status = os.system("/bin/mkdir -p " + mount_point);
    Internet2Lib.forceMount(drives[choice][0], "/mnt/store.new")

    if not os.path.exists("/mnt/store.new/NPTools"):
        os.mkdir("/mnt/store.new/NPTools");
    elif not os.path.isdir("/mnt/store.new/NPTools"):
        print Internet2Consts.YELLOW + "Error: NPTools exists on this drive and is not a directory\n" + Internet2Consts.NORMAL
    # Create an empty 'scratch' file
    os.system("/bin/dd if=/dev/zero of="+scratch_file+" bs=1 count=1 seek="+scratch_file+"&> /dev/null")

    # Write the info file
    try:
        fileHandle = open(Internet2Consts.SAVEINFO, "w")
        saveFile   = []

        saveFile.append("device    "  + drives[choice][0])
        saveFile.append("path      /mnt/store.new")

        fileHandle, res = Internet2Lib.writeLines(fileHandle, saveFile)
        if not res:
            raise IOError()

        # Close the file handle
        fileHandle.close()


        print "Saving pS-Performance Toolkit customizations/logs to " + drives[choice][0] + "... " ,
        sys.stdout.flush()

        if not Internet2Lib.saveConfig():
            print Internet2Consts.RED + "failed." + Internet2Consts.NORMAL
        else:
            print Internet2Consts.GREEN + "sucessful." + Internet2Consts.NORMAL

    except IOError:
        print Internet2Consts.YELLOW + "Error: Cannot write drive information to save info file at: " + Internet2Consts.SAVEINFO + " check file/folder permissions and start this script again." + Internet2Consts.NORMAL
        sys.exit(1)

    print "Internet2 pS-Performance Toolkit data/customizations will now be saved to " + drives[choice][0];
    print "To change this configuration, rerun this script and select another drive."

    print "The changes you have made require a reboot. Would you like to reboot? [y] ",
    sys.stdout.flush() # Force display 
        
    # Get the user's input
    try:
        choice = sys.stdin.readline().strip()
    except ValueError, e:
        choice = -1 # Not valid
    print " "

    if choice == "" or choice[:1] == "y" or choice[:1] == "Y":
        os.system("/sbin/reboot");
        sys.exit(1)

    print "You have chosen not to reboot. Any further changes you make may not survive a reboot. Press enter to continue."
    sys.stdin.readline()

    done = True
# Done in that loop

# Exit success if they configured at least one drive
if count > 0:
    sys.exit(0)
else:
    sys.exit(1)
### END MAIN ###
