#md5sum="ae545a6eb7fb8372d5b9d9d53336e43e"
from cisco import cli
from cisco import transfer 
from time import gmtime, strftime
import signal
import os
import string
import commands

# Host name and user credentials
username = "username"
password = "password\r"
hostname = "192.168.1.1"
vrf	 =  os.environ['POAP_VRF']

# POAP can use 3 modes to obtain the config file.
# - 'poap_static' - file is statically specified
# - 'poap_location' - CDP neighbor of interface on which DHCPDISCOVER arrived
#                     is used to derive the config file
# - 'poap_serial_number' - switch serial number is used to derive the config file

poap_config_file_mode = "poap_serial_number"

# Required space to copy config kickstart and system image in KB
required_space = 150000

####### Config File Infos #######

# Source file name of Config file
config_file_src = "poap.cfg"

# Temperory Destination file of Config file
config_file_dst = "poap_replay.cfg"

# indicates whether config file is copied 
config_copied = 0

# Destination file name for those lines in config which starts with hardware profile portmode or hardware profile tcam 
config_file_dst_first = "poap_1.cfg"

# Desination file name for those lines in config which does not match above criterea.
config_file_dst_second = "poap_2.cfg"

# Source path of Config file
config_path = "/"

# indicates whether first config file is empty or not 
emptyFirstFile = 1

# Destination Path 
destination_path = "/bootflash/"

####### System and Kickstart image info #######

# Source path of both System and Kickstart images 
image_path = "/"

# Source file name of System Image 
system_image_src = "n5000-uk9.5.2.1.N1.1.bin"

# Destination file name of System Image
system_image_dst = "n5k.s"

# indicates if System Image is copied
system_image_copied = 0 

# Source file name of Kickstart Image
kickstart_image_src = "n5000-uk9-kickstart.5.2.1.N1.1.bin"

# Destination file name of Kickstart Image
kickstart_image_dst = "n5k.k"

# indicates if Kickstart Image is copied
kickstart_image_copied = 0 

# Timeout info
config_timeout = 120 
system_timeout = 2100 
kickstart_timeout = 900  

# Log File name
try:
	poap_script_log = "/bootflash/%s_poap_%s_script.log" % (strftime("%Y%m%d%H%M%S", gmtime()), os.environ['POAP_PID']) 
except Exception as inst:
	print inst

# Log file handler
poap_script_log_handler = open(poap_script_log, "w+")

def poap_write_to_file (file_descriptor, info):
    file_descriptor.write(info)
    file_descriptor.flush()

def remove_file (filename) :
    try:
        os.remove(filename)
    except os.error:
        pass

def cleanup_files () : 
    global config_file_dst, config_file_dst_first, config_file_dst_second, system_image_dst, kickstart_image_dst, poap_script_log_handler
    poap_write_to_file(poap_script_log_handler, "\nINFO: delete all files")
    if config_copied == 1:
			remove_file("/bootflash/%s" % config_file_dst)
    remove_file("/bootflash/%s" % config_file_dst_first)
    remove_file("/bootflash/%s" % config_file_dst_second)
    if system_image_copied == 1:
    	remove_file("/bootflash/%s" % system_image_dst)
    if kickstart_image_copied == 1:
    	remove_file("/bootflash/%s" % kickstart_image_dst)
    remove_file("/bootflash/%s.tmp" % config_file_dst)
    remove_file("/bootflash/%s.tmp" % system_image_dst)
    remove_file("/bootflash/%s.tmp" % kickstart_image_dst)


def sig_handler_no_exit (signum, frame) : 
    global poap_script_log_handler
    poap_write_to_file(poap_script_log_handler, "\nINFO: SIGTERM Handler while configuring boot variables")


def sigterm_handler (signum, frame): 
    global poap_script_log_handler
    poap_write_to_file(poap_script_log_handler, "\nINFO: SIGTERM Handler") 
    cleanup_files()
    poap_script_log_handler.close()
    exit(1)

def removeFile (filename) :
    try:
    	os.remove(filename)
    except:
	pass

signal.signal(signal.SIGTERM, sigterm_handler)

# Procedure to split config file using global information
def splitConfigFile (): 
        global config_file_dst, config_file_dst_first, config_file_dst_second, emptyFirstFile, poap_script_log_handler 
	configFile 	  = open("/bootflash/%s" % config_file_dst, "r")
	configFile_first  = open("/bootflash/%s" % config_file_dst_first, "w+")
	configFile_second = open("/bootflash/%s" % config_file_dst_second, "w+")
        line = configFile.readline()
	while line != "": 
	    if not string.find(line, "hardware profile portmode", 0, 25) or not string.find(line, "hardware profile tcam", 0, 21):
		configFile_first.write(line)
		if emptyFirstFile is 1:
		    emptyFirstFile = 0
	    else: 
		configFile_second.write(line)
	    line = configFile.readline()

	configFile.close()
        removeFile("/bootflash/%s" % config_file_dst)
	configFile_first.close()
	if emptyFirstFile is 1:
      	    removeFile("/bootflash/%s" % config_file_dst_first) 

	configFile_second.close()

def verifyMD5sumofFile (md5given, filename) :
	md5calculated = commands.getstatusoutput("md5sum %s" % filename)
	md5calculated = md5calculated[1]
	md5calculated = md5calculated.split(" ")
	md5calculated = md5calculated[0]
	md5calculated = md5calculated.strip()

	if md5given == md5calculated:
		return True
	return False

def getMD5SumGiven (keyword, filename) :
	file = open("/bootflash/%s" % filename, "r")
	line = file.readline()
	while line != "":
		if not string.find(line, keyword, 0, len(keyword)) :
			line = line.split("=")
			line = line[1]
			line = line.strip()
			file.close()
			return line
		line = file.readline()
	file.close()
	return ""

def doCopyWithoutExit (protocol = "", host = "", source = "", dest = "", vrf = "management", login_timeout=10, user = "", password = "", dest_tmp = ""):
	if os.path.exists("/bootflash/%s" % dest_tmp): 
    		os.remove("/bootflash/%s" % dest_tmp)

	try:
		transfer(protocol, host, source, dest_tmp, vrf, login_timeout, user, password)
	except Exception as inst:
		poap_write_to_file(poap_script_log_handler, "\n Copy Failed: %s" % inst)
		return False
	dest_tmp = "%s%s" % (destination_path, dest_tmp)
	dest	 = "%s%s" % (destination_path, dest)
	os.rename(dest_tmp, dest)
	return True

def doCopy (protocol = "", host = "", source = "", dest = "", vrf = "management", login_timeout=10, user = "", password = "", dest_tmp = ""):
	if os.path.exists("/bootflash/%s" % dest_tmp): 
    		os.remove("/bootflash/%s" % dest_tmp)

	try:
		transfer(protocol, host, source, dest_tmp, vrf, login_timeout, user, password)
	except Exception as inst:
		poap_write_to_file(poap_script_log_handler, "\n Copy Failed: %s" % inst)
		cleanup_files()
		poap_script_log_handler.close()
		exit(1)
	dest_tmp = "%s%s" % (destination_path, dest_tmp)
	dest	 = "%s%s" % (destination_path, dest)
	os.rename(dest_tmp, dest)

def copyMd5Info (file_path, file_name):
	global username, hostname, poap_script_log_handler, password
	md5_file_name = "%s.md5" % file_name
	if os.path.exists("/bootflash/%s" % md5_file_name): 
		removeFile("/bootflash/%s" % md5_file_name)

	poap_write_to_file(poap_script_log_handler, "\nINFO: Starting Copy of MD5 File") 
	tmp_file = "%s.tmp" % md5_file_name 
	time = config_timeout
	src = "%s%s" % (file_path, md5_file_name)
	return doCopyWithoutExit ("ftp", hostname, src, md5_file_name, vrf, time, username, password, tmp_file)  
	 
# Procedure to copy config file using global information
def copyConfig ():
	global username, hostname, config_path, config_file_src, config_file_dst, config_timeout, poap_script_log_handler, emptyFirstFile, password
	org_file = config_file_dst
	if os.path.exists("/bootflash/%s" % org_file): 
		poap_write_to_file(poap_script_log_handler, "\nINFO: File already exists")
		emptyFirstFile = 0
		return

	poap_write_to_file(poap_script_log_handler, "\nINFO: Starting Copy of Config File") 
	tmp_file = "%s.tmp" % org_file
	time = config_timeout
	src = "%s%s" % (config_path, config_file_src)
	doCopy ("ftp", hostname, src, org_file, vrf, time, username, password, tmp_file)  
	config_copied = 1
	if copyMd5Info(config_path, config_file_src):
 		md5sumGiven = getMD5SumGiven("md5sum", "%s.md5" % config_file_src)
        	if md5sumGiven:
			if not verifyMD5sumofFile(md5sumGiven, "%s%s" % (destination_path, org_file)):
         			poap_write_to_file(poap_script_log_handler, "\n#### config file MD5 verification failed #####\n")
        			poap_script_log_handler.close()
	 			exit(1)	
        splitConfigFile() 
	poap_write_to_file(poap_script_log_handler, "\nINFO: Completed Copy of Config File") 

# Procedure to copy system image using global information
def copySystem ():
	global username, hostname, image_path, system_image_src, system_image_dst, system_timeout, poap_script_log_handler, password 
	poap_write_to_file(poap_script_log_handler, "\nINFO: Starting Copy of System Image")
	org_file = system_image_dst
	if os.path.exists("/bootflash/%s" % org_file): 
		poap_write_to_file(poap_script_log_handler, "\nINFO: File already exists")
		return
	tmp_file = "%s.tmp" % org_file
	time = system_timeout
	src = "%s%s" % (image_path, system_image_src)
	doCopy ("ftp", hostname, src, org_file, vrf, time, username, password, tmp_file)  
	system_image_copied = 1
	if copyMd5Info(image_path, system_image_src):
 		md5sumGiven = getMD5SumGiven("md5sum", "%s.md5" % system_image_src)
        	if md5sumGiven:
			if not verifyMD5sumofFile(md5sumGiven, "%s%s" % (destination_path, org_file)):
	         		poap_write_to_file(poap_script_log_handler, "\n#### System file MD5 verification failed #####\n")
       		 		poap_script_log_handler.close()
	 			exit(1)	
	poap_write_to_file(poap_script_log_handler, "\nINFO: Completed Copy of System Image" ) 

# Procedure to copy kickstart image using global information
def copyKickstart ():
	global username, hostname, image_path, kickstart_image_src, kickstart_image_dst, kickstart_timeout, poap_script_log_handler, password 
	poap_write_to_file(poap_script_log_handler, "\nINFO: Starting Copy of Kickstart Image")
	org_file = kickstart_image_dst 
	if os.path.exists("/bootflash/%s" % org_file): 
		poap_write_to_file(poap_script_log_handler, "\nINFO: File already exists")
		return
	tmp_file = "%s.tmp" % org_file 
	time = kickstart_timeout
	src = "%s%s" % (image_path, kickstart_image_src)
	doCopy ("ftp", hostname, src, org_file, vrf, time, username, password, tmp_file)  
	kickstart_image_copied = 1
	if copyMd5Info(image_path, kickstart_image_src):
 		md5sumGiven = getMD5SumGiven("md5sum", "%s.md5" % kickstart_image_src)
        	if md5sumGiven:
			if not verifyMD5sumofFile(md5sumGiven, "%s%s" % (destination_path, org_file)):
	         		poap_write_to_file(poap_script_log_handler, "\n#### Kickstart file MD5 verification failed #####\n")
        			poap_script_log_handler.close()
	 			exit(1)	

	poap_write_to_file(poap_script_log_handler, "\nINFO: Completed Copy of Kickstart Image") 

# Procedure to install both kickstart and system images
def installImages (): 
	global kickstart_image_dst, system_image_dst, poap_script_log_handler
	timeout = -1
	poap_write_to_file(poap_script_log_handler, "\n######### Copying the boot variables ##########")
	cli ("config terminal ; boot kickstart %s" % kickstart_image_dst)
	cli ("config terminal ; boot system %s" % system_image_dst)
	cli ("copy running-config startup-config")
	poap_write_to_file(poap_script_log_handler, "\n INFO: Configuration successful")

		
# Verify if free space is available to download config, kickstart and system
# images
def verifyfreespace (): 
    	global poap_script_log_handler, required_space
	s = os.statvfs("/bootflash/")
	freespace = (s.f_bavail * s.f_frsize) / 1024
        poap_write_to_file(poap_script_log_handler, "\n####The free space is %s##"  % freespace )

    	if required_space > freespace:
         	poap_write_to_file(poap_script_log_handler, "\n#### No enough space to copy the config, kickstart image and system image#####\n")
        	poap_script_log_handler.close()
        	exit(1)

# Procedure to set config_file based on switch serial number
def setSrcCfgFileNameSerial (): 
        global config_file_src, poap_script_log_handler
	if os.environ.has_key('POAP_SERIAL'):
		poap_write_to_file(poap_script_log_handler, "\n serial number %s" % os.environ['POAP_SERIAL'])
                config_file_src = "conf_%s.cfg" % os.environ['POAP_SERIAL'] 
	poap_write_to_file(poap_script_log_handler, "\n Selected conf file name : %s" % config_file_src)


# Procedure to set config_file_src
def setSrcCfgFileNameLocation():
 	global config_file_src, poap_script_log_handler, env
	startAppend = 0
	timeout = -1
	poap_write_to_file(poap_script_log_handler, "\nshow cdp neighbors interface %s" % os.environ['POAP_INTF'])
	cdpOutput = cli ("show cdp neighbors interface %s" % os.environ['POAP_INTF'])
	cdpOutArray = cdpOutput[1].split("\n")
	cdpRaw = cdpOutArray[7].split()
	cdpRawIntf = cdpOutArray[len(cdpOutArray) - 2].split()
	cdplist = cdpRaw[0].split('(')
	switchName = cdplist[0] 
	intfName   = cdpRawIntf[len(cdpRawIntf) - 1]
	config_file_src = "conf_%s_%s.cfg" % (switchName, intfName)
	config_file_src = string.replace(config_file_src, "/", "_")
	poap_write_to_file(poap_script_log_handler, "\nSelected conf file name : %s" % config_file_src)

if poap_config_file_mode == "poap_location": 
	#set source config file based on location
        setSrcCfgFileNameLocation()

elif poap_config_file_mode == "poap_serial_number": 
	#set source config file based on switch's serial number
        setSrcCfgFileNameSerial()


verifyfreespace()

# copy config file and images
copyConfig()

# copy config file and images 
copySystem()
copyKickstart()

signal.signal(signal.SIGTERM, sig_handler_no_exit)

# install images
installImages()

if emptyFirstFile is 0:
	cli ('copy bootflash:%s scheduled-config' % config_file_dst_first)
	poap_write_to_file(poap_script_log_handler, "\n######### Copying the first scheduled cfg done ##########")
	removeFile("/bootflash/%s" % config_file_dst_first)

cli ('copy bootflash:%s scheduled-config' % config_file_dst_second)
poap_write_to_file(poap_script_log_handler, "\n######### Copying the second scheduled cfg done ##########")
removeFile("/bootflash/%s" % config_file_dst_second)

poap_script_log_handler.close()
exit(0)