#!/usr/libexec/platform-python

################################################################################
# serial45d - used to store serial, model and chassis information on 
#          45Drives storage servers. 
#
# Copyright (C) 2021, Mark Hooper   <mhooper@45drives.com>
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#   
################################################################################
import re
import subprocess
import os
import datetime

g_config_file_content = {
	"Motherboard":[],
	"CPU":[],
	"Chassis Type":[],
	"Product Type":[],
	"Product Name":[]
}

g_update_config = False

g_current_dir = os.path.dirname(os.path.realpath(__file__))

def outputSerialVersion():
	print("\n")
	print("-----------------------------------------------------------------------------------")
	print("                                                                                   ")
	print("                                $$\           $$\ $$\   $$\ $$$$$$$\  $$$$$$$\     ")
	print("                                \__|          $$ |$$ |  $$ |$$  ____| $$  __$$\    ")
	print("   $$$$$$$\  $$$$$$\   $$$$$$\  $$\  $$$$$$\  $$ |$$ |  $$ |$$ |      $$ |  $$ |   ")
	print("  $$  _____|$$  __$$\ $$  __$$\ $$ | \____$$\ $$ |$$$$$$$$ |$$$$$$$\  $$ |  $$ |   ")
	print("  \$$$$$$\  $$$$$$$$ |$$ |  \__|$$ | $$$$$$$ |$$ |\_____$$ |\_____$$\ $$ |  $$ |   ")
	print("   \____$$\ $$   ____|$$ |      $$ |$$  __$$ |$$ |      $$ |$$\   $$ |$$ |  $$ |   ")
	print("  $$$$$$$  |\$$$$$$$\ $$ |      $$ |\$$$$$$$ |$$ |      $$ |\$$$$$$  |$$$$$$$  |   ")
	print("  \_______/  \_______|\__|      \__| \_______|\__|      \__| \______/ \_______/    ")
	ver_file_path = f"{g_current_dir}/configs/version"
	if(os.path.exists(ver_file_path) and os.path.isfile(ver_file_path)):
		ver_file = open(ver_file_path,"r")
		version = ver_file.readline()
		ver_file.close()
		
		print("\n                                                                  Version: {v}".format(v=version))
		print("-----------------------------------------------------------------------------------\n")
	else:
		print("\n-----------------------------------------------------------------------------------\n")
	
		

##################################################################################
# getMotherboard()
# performs the command: dmidecode -t 2
# searches the output for the motherboard model and serial number and 
# returns them as strings.
##################################################################################
def getMotherboard():
	global g_update_config
	global g_config_file_content
	board_model = "unknown"
	board_serial = "unknown"
	board_manufacturer = "unknown"
	try:
		dmi_result = subprocess.Popen([f"{g_current_dir}/dependencies/dmidecode/dmidecode","-t","2"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		print("Error Executing dmidecode")
		exit(1)
	for line in dmi_result:
			regex_board_model = re.search("^\sProduct Name:\s+(.*)",line)
			regex_board_serial = re.search("^\sSerial Number:\s+(.*)",line)
			regex_board_manufacturer = re.search("^\sManufacturer:\s+(.*)",line)
			if regex_board_model != None:
				if regex_board_model.group(1) not in g_config_file_content["Motherboard"]:
					# board was not found in the config file, add it and set the 
					# g_update_config flag. 
					g_config_file_content["Motherboard"].append(regex_board_model.group(1))
					g_update_config = True
				board_model = regex_board_model.group(1)
			if regex_board_serial != None:
				board_serial = regex_board_serial.group(1)
			if regex_board_manufacturer != None:
				board_manufacturer = regex_board_manufacturer.group(1)
			if board_model != "unknown" and board_serial != "unknown" and board_manufacturer != "unknown":
				# all information has been assigned, break from the loop
				break
	return board_model,board_serial,board_manufacturer

##################################################################################
# getCPU()
# performs the command: dmidecode -t 4
# searches the output for the cpu model, and maintains a count of cpu's found. 
# Will also ensure that the config file is updated if a cpu is detected that
# was not already present in the config file by setting the g_update_config flag.
##################################################################################
def getCPU():
	global g_update_config
	global g_config_file_content
	cpu_model = "unknown"
	cpu_count = 0
	try:
		dmi_result = subprocess.Popen([f"{g_current_dir}/dependencies/dmidecode/dmidecode","-t","4"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		print("Error Executing dmidecode")
		exit(1)
	for line in dmi_result:
		regex_cpu = re.search("^\sVersion:\s+(.*)$",line)
		if regex_cpu != None:
			regex_cpu_version = str(regex_cpu.group(1)).rstrip()
			if regex_cpu_version not in g_config_file_content["CPU"]:
				# board was not found in the config file, add it and set the 
				# g_update_config flag. 
				g_config_file_content["CPU"].append(regex_cpu_version)
				g_update_config = True
			cpu_model = regex_cpu_version
			cpu_count += 1
		if cpu_count == 2:
			# we have found all possible entries, stop searching
			break
	return cpu_model, cpu_count

##################################################################################
# getHBA()
# performs the command: ...storcli64 show all
# counts the number of 16i and 24i hba cards and returns these as integers
##################################################################################
def getHBA():

	hba_models = {
		"SAS9305-16i":0,
		"SAS9305-24i":0
	}

	try:
		storcli64_result = subprocess.Popen(
			[f"{g_current_dir}/dependencies/storcli64/storcli64","show","all"],stdout=subprocess.PIPE,universal_newlines=True)
	except:
		print("Error running storcli64")
		exit(1)
	card_count = 0
	for line in storcli64_result.stdout:
		for model in hba_models.keys():
			# Model AdapterType VendId DevId SubVendId SubDevId PCIAddress 	
			regex = re.search("({mdl}).*00:\w\w:\w\w:\w\w\s+$".format(mdl=model),line)
			if regex != None:
				hba_models[model] += 1
				
	return hba_models["SAS9305-16i"], hba_models["SAS9305-24i"]


##################################################################################
# loadSystemConfigs()
# This function will load supported configs from an included file
# in ./configs/configs.txt
# This way, if we discover new configurations, we can save them by updating
# the config files.
##################################################################################
def loadSystemConfigs():
	config_file = open(f"{g_current_dir}/configs/config.txt","r")
	for line in config_file:
		for key in g_config_file_content.keys():
			regex = re.search("({k}):(.*)$".format(k=key),line)
			if regex != None:
				g_config_file_content[regex.group(1)].append(regex.group(2))

##################################################################################
# performFRUWrite()
# This function adjusts the remaining fru fields 
# and performs the final write operation by using
# ./dependencies/ipmi_linux/IPMICFG-Linux.x86_64 
##################################################################################
def performFRUWrite(fru_fields):
	print("\n")
	print("+----------------------------------------------------------------+")
	print("| Step 7: Perform Final Write Operation                          |")
	print("+----------------------------------------------------------------+")

	# get the current date for the manufacturing date
	date = datetime.datetime.now()
	fru_fields["board_manufacturing_date"] = date.strftime("%Y%m%d%H%M")

	# adjust the chassis_part_number and version fields
	# to indicate that the system is a hybrid
	if "H16" in fru_fields["product_name"]:
		fru_fields["chassis_part_number"] = "H16-" + fru_fields["chassis_type"]
		fru_fields["version"] = "Hybrid16(" + fru_fields["version"] + ")"
	elif "H32" in fru_fields["product_name"]:
		fru_fields["chassis_part_number"] = "H32-" + fru_fields["chassis_type"]
		fru_fields["version"] = "Hybrid32(" + fru_fields["version"] + ")"
	else:
		fru_fields["chassis_part_number"] = fru_fields["chassis_type"]
		fru_fields["version"] = "(" + fru_fields["version"] + ")"
	
	# assign the remaining fields.
	fru_fields["chassis_serial"] = fru_fields["serial"]
	fru_fields["part_number"] = fru_fields["chassis_type"]
	fru_fields["chassis_type"] = "17h"

	# prepare arguments for function call
	fru_args = [
		f"{g_current_dir}/dependencies/ipmi_linux/IPMICFG-Linux.x86_64",
		"-fru",
		"dmi",
		fru_fields["manufacturer_name"],
		fru_fields["product_name"],
		fru_fields["part_number"],
		fru_fields["version"],
		fru_fields["serial"],
		fru_fields["asset_tag"],
		fru_fields["board_manufacturing_date"],
		fru_fields["board_manufacturer"],
		fru_fields["board_model"],
		fru_fields["board_model"],
		fru_fields["board_serial"],
		fru_fields["chassis_type"],
		fru_fields["chassis_part_number"],
		fru_fields["chassis_serial"]
	]

	# perform the function call
	try:
		ipmi_result = subprocess.Popen(
			fru_args,stdout=subprocess.PIPE,universal_newlines=True)
	except:
		print("!!!! Error performing write !!!!")
		for i in range(0,len(fru_args)):
			print(fru_args[i],end=" ")
		exit(1)

	# output the result of the write operation
	print("\n!!!!   Performing Write  !!!!\n")
	print("\n Output from ipmicfg: ")
	for line in ipmi_result.stdout:
		print("\t" + line,end="")

	try:
		ipmitool_result = subprocess.Popen(
			["ipmitool","fru"],stdout=subprocess.PIPE,universal_newlines=True)
	except:
		print("Error executing \"ipmitool fru\"")
		exit(1)	

	ipmitool_fru_output = []
	# output the result of the write operation
	print("\n")
	print("+----------------------------------------------------------------+")
	print("| Step 8: Verify Output                                          |")
	print("+----------------------------------------------------------------+")
	print("\nOutput from ipmitool fru: ")
	for line in ipmitool_result.stdout:
		print("\t" + line,end="")
		ipmitool_fru_output.append(line)
	print("Setup Complete. If the results shown above appear incorrect") 
	print("run this program again, or contact 45Drives Support.")
	input("Press Enter to end program.") 
	return ipmitool_fru_output

##################################################################################
# countDrives()
# This function determines the total number of drives installed in the system
# and returns the count of HSS and SSD drives.
##################################################################################
def countDrives():
	ssd_count = 0
	hdd_count = 0
	devices = []
	try:
		ls_result = subprocess.Popen(
			["ls","-w","1","/sys/block/"],stdout=subprocess.PIPE,universal_newlines=True)
	except:
		print("Error Counting Drives")
		exit(1)
	for line in ls_result.stdout:
		regex_drive_count = re.search("^(sd.*)$",line)
		if regex_drive_count != None:
			devices.append(regex_drive_count.group(1))
	
	for dev in devices:
		rotational_path = "/sys/block/" + dev + "/queue/rotational"
		if os.path.isfile(rotational_path):
			rotational = open(rotational_path, mode='r')
			if bool(int(rotational.read(1))):
				hdd_count += 1
			else:
				ssd_count += 1
	return hdd_count, ssd_count

##################################################################################
# updateConfigFile()
# updates the config file in an instance where new hardware was detected, 
# or manually added by user.
##################################################################################
def updateConfigFile():
	config_file = open(f"{g_current_dir}/configs/config.txt","w")
	for key in g_config_file_content.keys():
		for entry in g_config_file_content[key]:
			config_file.write(key + ":" + entry + "\n")
	config_file.close()

##################################################################################
# getChassisType()
# gets the chassis type based on user input. returns a string. 
##################################################################################
def getChassisType():
	done = False
	while not done:
		print("\n")
		print("+----------------------------------------------------------------+")
		print("| Step 3: Select a Chassis Type                                  |")
		print("+----------------------------------------------------------------+")
		option = 0
		for product in g_config_file_content["Chassis Type"]:
			option += 1
			print("\t",option,": " + product)
		selection = input("Select an option (1 - " + str(option) + "): ")
		if selection.isnumeric() and int(selection) <= option and int(selection) > 0:
			return g_config_file_content["Chassis Type"][int(selection)-1]
		else:
			print("Invalid selection")

##################################################################################
# getProductType()
# Gets the product type selection from the user. Returns a string.
##################################################################################
def getProductType():
	done = False
	while not done:
		print("+----------------------------------------------------------------+")
		print("| Step 1: Select Product Type                                    |")
		print("+----------------------------------------------------------------+")
		option = 0
		for product in g_config_file_content["Product Type"]:
			option += 1
			print("\t",option,": " + product)
		selection = input("Select an option (1 - " + str(option) + "): ")
		if selection.isnumeric() and int(selection) <= option and int(selection) > 0:
			return g_config_file_content["Product Type"][int(selection)-1]
		else:
			print("Invalid selection")

##################################################################################
# getSerial()
# gets the serial number from the user. Performs basic check to ensure that
# numbers and hyphens are used. Typical job number is of format XXXXXXXX-YY-ZZ
##################################################################################
def getSerial():
	done = False
	selection = ""
	while not done:
		print("\n")
		print("+----------------------------------------------------------------+")
		print("| Step 4: Enter Product Serial                                   |")
		print("+----------------------------------------------------------------+")
		selection = input("Enter Product Serial Number (aka Job Number ex:104237-11-01): ")
		selCopy = selection
		selCopy = selCopy.split("-")
		if len(selCopy) == 3:
			#serial is of proper length
			if selCopy[0].isnumeric() and selCopy[2].isnumeric():
				# the first and last tokens are of proper format
				done = True
		elif len(selCopy) == 2:
			#serial is of improper length
			if selCopy[0].isnumeric():
				print("Serial number entered ("+selection+") doesn't match a full serial number format (ex. 123456-10-01).")
			if input("Are you sure that you want to use an incomplete serial number for this system? It can make support's job harder if you do.. (y/n):") == "y":
				done = True
		else:
			print("Invalid format used, try again")
	return selection


##################################################################################
# getHybridType()
# Lets the user specify if the product is a hybrid or not. 
##################################################################################
def getHybridType():
	done = False
	while not done:
		print("\n")
		print("+----------------------------------------------------------------+")
		print("| Step 2: Hybrid Check                                           |")
		print("+----------------------------------------------------------------+")

		valid_hybrid_strs = ["H32","H16","H8"]
		hybrid_str = ""
		# prompt user to see if system is a hybrid system.
		hybrid = input("Is this a hybrid machine? (y/n): ")
		if hybrid == "y":
			# we have a hybrid system, append an H to string
			hybrid_str += "H"
			# get type of hybrid from user.
			hybridCount = input("which kind of hybrid? (8/16/32): ")
			hybrid_str += hybridCount
			if hybrid_str in valid_hybrid_strs:
				return hybrid_str
			else:
				print("Invalid hybrid type \"{h}\". Try Again.".format(h=hybrid_str))
		elif hybrid == "n":
			return ""
		else:
			print("Invalid Selection")


##################################################################################
# getUserInput()
# Function that calls all user input functions outside of autoDetect().
# updates auto_detect_fields and fru_fields (declared in main()).
##################################################################################
def getUserInput(auto_detect_fields,fru_fields):
	product_type = getProductType()
	hybrid_type = getHybridType()
	chassis_type = getChassisType()
	serial = getSerial()
	asset_tag = "SELFSERIALIZED"

	auto_detect_fields["product_type"] = product_type
	auto_detect_fields["chassis_type"] = chassis_type
	auto_detect_fields["hybrid_type"] = hybrid_type
	fru_fields["chassis_type"] = chassis_type
	fru_fields["serial"] = serial
	fru_fields["asset_tag"] = asset_tag
	fru_fields["product_name"] = product_type

def manualSelect(auto_detect_fields,fru_fields):
	mobo_to_version_lut = {
		"X11SSH-CTF":"Base",
		"X11SPL-F": "Enhanced",
		"X11DPL-i": "Turbo",
		"X11SSM-F": "Base",
		"H11SSL-i": "Enhanced-AMD",
		"X11DPH-T": "Turbo",
		"X11SPi-TF": "Enhanced"
	}

	cpu_substr_to_version_lut = {
		"Bronze":"Base-B",
		"Silver":"Enhanced-S",
		"Gold":"Turbo-G"
	}

	done = False
	while not done:
		print("\n")
		print("+----------------------------------------------------------------+")
		print("| Step 5: Determine Unique Server Type                           |")
		print("+----------------------------------------------------------------+")
		final_product_str = "unknown"

		# build the product string using the input provided by the user.
		unique_product_str = auto_detect_fields["product_type"] + "-"
		if auto_detect_fields["hybrid_type"] != "":
			unique_product_str += auto_detect_fields["hybrid_type"] + "-"
		unique_product_str += auto_detect_fields["chassis_type"] + "-"

		#see if we can determine base, enhanced, turbo or enhanced-AMD type
		if auto_detect_fields["board_model"] in mobo_to_version_lut.keys():
			# we have a known motherboard model
			if auto_detect_fields["board_model"] == "X11SPL-F" or auto_detect_fields["board_model"] == "X11SPi-TF":
				# we need to see if the cpu is a gold, silver or bronze model
				for substr in cpu_substr_to_version_lut.keys():
					if substr in auto_detect_fields["cpu_model"]:
						# gold silver or bronze was found, we need to set the base, enhanced, and turbo accordingly
						unique_product_str_tmp = unique_product_str + cpu_substr_to_version_lut[substr]
						break
			else:
				unique_product_str_tmp = unique_product_str + mobo_to_version_lut[auto_detect_fields["board_model"]]
			if unique_product_str_tmp in g_config_file_content["Product Name"]:
				final_product_str = unique_product_str_tmp
			else:
				print("Product Configuration of \"{ps}\" is INVALID. \nRun Program again or contact Mark Hooper if issue persists.".format(ps=unique_product_str_tmp))
				print("Valid Product Configurations: ")
				for valid_product in g_config_file_content["Product Name"]:
					print("\t{vp}".format(vp=valid_product))
				input("Press ENTER to End Program.")
				exit(0)
			if final_product_str != "unknown":
				# we got valid input from user and assigned a value to final product string.
				auto_detect_fields["unique_product_str"] = final_product_str
				confirm = input("\nAuto-Detect determined that the product is a \"{pr}\", is this correct? (y/n):".format(pr=auto_detect_fields["unique_product_str"]))
				if confirm == "y":
					# user confirmed final product string to be correct. We're all done.
					fru_fields["product_name"] = auto_detect_fields["unique_product_str"]
					fru_fields["version"] = auto_detect_fields["unique_product_str"].split("-")[-1]
					return
				else:
					final_product_str = "unknown"
		
		# Unable to automatically determine product version, get 
		# this information from user. 
		versions = ["Base","Base-B","Enhanced","Enhanced-S","Turbo","Turbo-G","Enhanced-AMD"]
		print("Select a system version:")
		option = 0
		for version in versions:
			option += 1
			print("\t" + str(option) + ": " + version)
		selection = input("Select an option (1 - " + str(option) + "): ")
		if int(selection) <= option and int(selection) > 0:
			unique_product_str += versions[int(selection)-1]
			if unique_product_str in g_config_file_content["Product Name"]:
				final_product_str = unique_product_str
			else:
				print("Product Configuration of \"{ps}\" is INVALID. \nRun Program again or contact 45Drives Support if issue persists.".format(ps=unique_product_str))
				print("Valid Product Configurations: ")
				for valid_product in g_config_file_content["Product Name"]:
					print("\t{vp}".format(vp=valid_product))
				input("Press ENTER to End Program.")
				exit(0)
		if final_product_str != "unknown":
			# we got valid input from user and assigned a value to final product string.
			auto_detect_fields["unique_product_str"] = final_product_str
			confirm = input("\nAuto-Detect determined that the product is a \"{pr}\", is this correct? (y/n):".format(pr=auto_detect_fields["unique_product_str"]))
			if confirm == "y":
				# user confirmed final product string to be correct. We're all done.
				fru_fields["product_name"] = auto_detect_fields["unique_product_str"]
				fru_fields["version"] = auto_detect_fields["unique_product_str"].split("-")[-1]
				return

		# We need to try to get user unput again. Something was entered incorrectly. 
		print("Invalid information entered, Trying Again.\n")

def storePrevIpmitoolOutput():
	print("\n")
	print("+----------------------------------------------------------------+")
	print("| Step 6: Store backup                                           |")
	print("+----------------------------------------------------------------+")
	perform_backup = input("Would you like to store a backup of the previous serial data? (y/n): ")
	if perform_backup == "y":
		print("\nStoring output of 'ipmitool fru' command in backup file.")
		try:
			ipmitool_result = subprocess.Popen(
				["ipmitool","fru"],stdout=subprocess.PIPE,universal_newlines=True)
		except:
			print("Error executing \"ipmitool fru\"")
			exit(1)	

		ipmitool_fru_output = []
		for line in ipmitool_result.stdout:
			ipmitool_fru_output.append(line)

		backup_dir = f"{g_current_dir}/backup"
		if not os.path.exists(backup_dir):
			print("making directory to store backup files {bd}".format(bd=backup_dir))
			os.makedirs(backup_dir)

		date = datetime.datetime.now()
		timestamp = date.strftime("%Y-%m-%d_%H-%M-%S")
		backup_file_name = "{bd}/ipmitool_fru_backup_{ts}".format(bd=backup_dir,ts=timestamp)

		b_file = open(backup_file_name,"w")
		for line in ipmitool_fru_output:
			b_file.write(line)
		b_file.write("\n")
		b_file.close()
		print("Output stored successfully ({bfn})".format(bfn=backup_file_name))
	else:
		print("skipping...")

def check_root():
	root_test = subprocess.run(["ls","/root"],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode
	if root_test:
		print("serial45d must be run with root privileges.")
		exit(root_test)

def main():
	check_root()
	# Fields used in the IPMICFG-Linux.x86_64 write operation
	fru_fields = {
		"manufacturer_name":"45Drives",
		"product_name":"unknown",
		"part_number":"unknown",
		"version":"unknown",
		"serial":"unknown",
		"asset_tag":"unknown",
		"board_manufacturing_date":"unknown",
		"board_manufacturer":"unknown",
		"board_model":"unknown",
		"board_serial":"unknown",
		"chassis_type":"unknown",
		"chassis_part_number":"unknown",
		"chassis_serial":"unknown"
	}

	# Fields used to attempt to automatically detect the product
	auto_detect_fields = {
		"product_type":"unknown",
		"board_model":"unknown",
		"cpu_model":"unknown",
		"cpu_count":0,
		"hba_16i_count":0,
		"hba_24i_count":0,
		"hdd_count":0,
		"ssd_count":0,
		"chassis_type":"unknown",
		"unique_product_str":"unknown",
		"hybrid_type": ""
	}
	outputSerialVersion()
	
	# load config file information and store it in g_config_file_content
	loadSystemConfigs()

	# get system type, chassis size, and serial number from user.
	getUserInput(auto_detect_fields,fru_fields)

	# determine motherboard information and assign fields accordingly
	fru_fields["board_model"], fru_fields["board_serial"] , fru_fields["board_manufacturer"]= getMotherboard()
	auto_detect_fields["board_model"] = fru_fields["board_model"]

	# gather information and perform autodetection
	auto_detect_fields["cpu_model"], auto_detect_fields["cpu_count"] = getCPU()
	auto_detect_fields["hba_16i_count"], auto_detect_fields["hba_24i_count"] = getHBA()
	auto_detect_fields["hdd_count"], auto_detect_fields["ssd_count"] = countDrives()

	manualSelect(auto_detect_fields,fru_fields)

	storePrevIpmitoolOutput()

	# write final information using ipmicfg
	fru_write_output = performFRUWrite(fru_fields)

	# update the config file if we have discovered new information previously unknown. 
	if g_update_config:
		updateConfigFile()
		print("Config File has been updated.")


if __name__ == "__main__":
	main()
