#!/usr/bin/env python3
###############################################################################
# dmap - used to create /etc/vdev_id.conf for 45Drives storage servers
#
# Copyright (C) 2020, Josh Boudreau <jboudreau@45drives.com>
#                     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 sys
import json
from optparse import OptionParser


g_quiet = False


###############################################################################
# Name: get_path_variables
# Args: None
# Desc: obtains and returns the path variables required by dmap.
#		conf_path: coresponds to where the vdev_id.conf file should be stored
#		dev_path: corresponds to the device path required when generating
#				  the ailiases found in vdev_id.conf.
###############################################################################
def get_path_variables():
	# get the alias config path, if it fails assume /etc
	# get the device path, if it fails assume /dev

	conf_path = os.getenv('ALIAS_CONFIG_PATH')
	if conf_path == None:
		log("No alias config path set in profile.d ... Defaulting to /etc")
		conf_path = "/etc"	
		
	dev_path = "/dev"

	return conf_path, dev_path


###############################################################################
# Name: trigger_udev
# Args: None
# Desc: runs two commands (udevadm trigger, udevadm settle), will only attempt
#       the second command it the first command returned a successful result.
###############################################################################
def trigger_udev():
	trigger_successful = False
	try:
		trigger_call = subprocess.run(["udevadm","trigger"])
		if trigger_call.returncode == 0:
			trigger_successful = True
	except OSError:
		log("Error triggering udevadm (udevadm trigger)")
		sys.exit(1)

	if trigger_successful:
		try:
			settle_call = subprocess.run(["udevadm","settle"])
			if settle_call.returncode != 0:
				raise OSError()
		except OSError:
			log("Error settling udevadm (udevadm settle)")
			sys.exit(1)			

###############################################################################
# Name: reset_map
# Args: config_path
# Desc: removes the existing vdev_id.conf file located in the directory
#       provided by config_path. Then udev is triggered, which removes
#		the drive ailiasing that existed on the machine.
###############################################################################
def reset_map(config_path):
	try:
		os.remove(config_path + "vdev_id.conf")
	except OSError:
		pass
	trigger_udev()
	log("Drive Aliasing reset")
	sys.exit(0)

###############################################################################
# Name: log
# Args: message
# Desc: outputs the message to stdout contingent on the global quiet flag.
###############################################################################
def log(message):
	if not g_quiet:
		print(message)

###############################################################################
# Name: check_root
# Args: None
# Desc: Ensures that dmap is running as root
###############################################################################
def check_root():
	root_test =	subprocess.run(["ls","/root"],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode
	if root_test:
		log("dmap must be run with root privileges")
		sys.exit(root_test)

###############################################################################
# Name: verify_vdev
# Args: None
# Desc: Checks to see if the proper vdev rules and scripts are present
#		and whether the script is executable. It will create them 
#		and ensure that they are setup to ensure that udevadm works as 
#		intended.
#		(/usr/lib/udev/rules.d/69-vdev.rules)
#		(/usr/lib/udev/vdev_id)
# 		returns true if files are present, returns false otherwise.
#
###############################################################################
def verify_vdev():
	VDEV_ID="https://raw.githubusercontent.com/45Drives/udev/master/vdev_id"
	rules_path="/usr/lib/udev/rules.d/69-vdev.rules"
	script_path="/usr/lib/udev/vdev_id"

	rules_test = os.path.exists(rules_path)
	script_test = os.path.exists(script_path)
	script_x_test = os.access(script_path,os.X_OK) if script_test else False

	if rules_test and script_test and script_x_test:
		return True
	else:
		log("valid vdev configuration not found. Resolving...")
	
	if not rules_test:
		log(rules_path +" not found: downloading...")
		rules_repo="https://raw.githubusercontent.com/45Drives/udev/master/69-vdev.rules"
		rv=subprocess.run(["curl","-o",rules_path,rules_repo],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode
		if rv:
			log("error downloading 69-vdev.rules from " + rules_repo)
		else:
			rules_test = os.path.exists(rules_path)
	
	if not script_test:
		log(script_path +" not found: downloading...")
		script_repo="https://raw.githubusercontent.com/45Drives/udev/master/vdev_id"
		rv=subprocess.run(["curl","-o",script_path,script_repo],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode
		if rv:
			log("error downloading vdev_id from " + script_repo)
		else:
			script_test = os.path.exists(script_path)

	if not script_x_test:
		if script_test:
			rv=subprocess.run(["chmod","+x",script_path],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode
			if rv:
				log("error making " + script_path + " executable")
			else:
				script_x_test = os.access(script_path,os.X_OK)

		else:
			log("cannot locate " + script_path)

	return rules_test and script_test and script_x_test


###############################################################################
# Name: get_server_info
# Args: none
# Desc: Returns a dict read from /etc/server_info/server_info.json. This file
#       is created by another python script /opt/tools/server_identifier. 
#       which automatically determines which 45Drives server is present. 
#       and provides us with the necessary data structure tol automatically
#       alias any server we make (see create_vdev_id())
###############################################################################
def get_server_info():
	try:
		return_code = subprocess.run(["/opt/tools/server_identifier"]).returncode
	except:
		print("dmap failed to get server information")
		sys.exit(1)
	if return_code != 0:
		print("dmap failed to get server information")
		sys.exit(1)
	else:
		# we should expect that there is a file that we can parse in
		# /opt/tools/server_info/server_info.json
		server = json.load(open("/etc/server_info/server_info.json","r"))
		return server


###############################################################################
# Name: create_vdev_id
# Args: server (dict)
#    Example of a server dict:
#    {
#        "Motherboard": {
#            "Manufacturer": "Supermicro",
#            "Product Name": "X11SPL-F",
#            "Serial Number": "ZM18AS011320"
#        },
#        "HBA": [
#            {
#                "Model": "SAS9305-16i",
#                "Adapter": "SAS3224",
#                "Bus Address": "0000:17:00.0",
#                "Drive Connections": 16,
#                "Kernel Driver": "mpt3sas",
#                "PCI Slot": 5
#            },
#            {
#                "Model": "SAS9305-24i",
#                "Adapter": "SAS3224",
#                "Bus Address": "0000:18:00.0",
#                "Drive Connections": 24,
#                "Kernel Driver": "mpt3sas",
#                "PCI Slot": 4
#            },
#            {
#                "Model": "SAS9305-24i",
#                "Adapter": "SAS3224",
#                "Bus Address": "0000:65:00.0",
#                "Drive Connections": 24,
#                "Kernel Driver": "mpt3sas",
#                "PCI Slot": 6
#            }
#        ],
#        "Hybrid": true,
#        "Serial": "13371337-1-1",
#        "Model": "Storinator-H32-S45-Enhanced",
#        "Alias Style": "H32",
#        "Chassis Size": "S45"
#    }
#    
# Desc: Using the server dict, we can alias all drives in the appropriate 
#       order based on the system information. The server dict is read in 
#       from /etc/server_info/server_info.json. The phy_order dict is used 
#       to assign the appropriate physical address. The keys correspond to 
#       those provided by the "/opt/tools/storcli64 show" command. 
#       The "Alias Style" key us used to reference the alias_template dict,
#       and the chassis size key stored in the 
#       alias_template[server["Alias Style"]] dict, will give us an array of 
#       ints used as indexes into the phy_order dict's arrays. 
###############################################################################
def create_vdev_id(server):
	phy_order = {
		"HBA 9405W-16i" : [9,11,13,15,  8,10,12,14,  1,3,5,7,  0,2,4,6],
		"SAS9305-16i" : [2,3,1,0,  6,7,5,4,  18,19,17,16,  22,23,21,20],
		"SAS9305-24i": [2,3,1,0,  6,7,5,4,  18,19,17,16,  22,23,21,20,  10,11,9,8,  14,15,13,12]
	}

	alias_template = {
		"H16":{
			"Q30":[15,23],
			"S45":[15,15,23],
			"XL60":[15,15,15,23]
		},
		"H32":{
			"Q30":[23,24],
			"S45":[15,23,24],
			"XL60":[15,15,23,24]
		},
		"STORINATOR":{
			"AV15":[15],
			"Q30":[15,15],
			"S45":[15,15,15],
			"XL60":[15,15,15],
			"F32":[16,16]	
		},
		"STORNADO":{
			"AV15":[16,16],
			"F32":[16,16]
		},
		"AV15-BASE":{
			"AV15":[0]
		},
		"?":{
			"?":[0]
		}
	}

	# Make sure that the user didn't make any manual edits that can throw dmap.
	if server["Alias Style"] not in alias_template.keys():
			print("/opt/tools/dmap: !! ERROR !! Invalid Alias Style set in /etc/server_info/server_info.json")
			print("                 Valid Options are: ", alias_template.keys())
			print("                 You can manually edit this file by setting \"Edit Mode\": true")
			print("                 in /etc/server_info/server_info.json along with any other parameters")
			sys.exit(1)
	elif server["Chassis Size"] not in alias_template[server["Alias Style"]].keys():
			print("/opt/tools/dmap: !! ERROR !! Invalid Chassis Size and Alias Style combination set in /etc/server_info/server_info.json")
			print("                 Valid Chassis Size Options for " + server["Alias Style"] + " are: ", alias_template[server["Alias Style"]].keys())
			print("                 You can manually edit this file by setting \"Edit Mode\": true")
			print("                 in /etc/server_info/server_info.json along with any other parameters")
			sys.exit(1)

	# We have a valid combination of alias styles and chassis sizes.
	vdev_id_str = "# This file was generated using dmap v1.8 (/opt/tools/dmap).\n"
	if len(server["HBA"]) > 0:
		# we have hba cards to inspect/alias
		for i in range(0,len(server["HBA"])):
			count = alias_template[server["Alias Style"]][server["Chassis Size"]][i]
			for j in range(0,count):
				if(j>=15 and server["Alias Style"] == "H32" and i==len(alias_template[server["Alias Style"]][server["Chassis Size"]])-2):
					# We are ailiasing a hybrid 32, which staggers the labels across the 2nd last and last rows.
					# But, the physical addresses need to increment by 1, while keeping the labels the same.
					vdev_id_str += (
						"alias {i}-{j} /dev/disk/by-path/pci-{addr}-sas-phy{p}-lun-0\n".format(
							i=i+1,j=j+1,addr=server["HBA"][i]["Bus Address"],p=phy_order[server["HBA"][i]["Model"]][j+1]
							)
						)
				else:
					# The default case for ailiasing hba cards.
					vdev_id_str += (
						"alias {i}-{j} /dev/disk/by-path/pci-{addr}-sas-phy{p}-lun-0\n".format(
							i=i+1,j=j+1,addr=server["HBA"][i]["Bus Address"],p=phy_order[server["HBA"][i]["Model"]][j]
							)
						)
	elif server["Alias Style"] == "AV15-BASE":
		# we are using the AV15-BASE ailiasing scheme.
		vdev_id_str += alias_av15_base()
	elif server["Alias Style"] == "?":
		# we don't know what kind of server this is. Likely a vm
		vdev_id_str += "# unable to map drives for unknown hardware configuration"

	return vdev_id_str

###############################################################################
# Name: alias_av15_base
# Args: none
# Desc: returns a string to to place within vdev_id.conf that will
#       alias drives in an av15 base model. The av15 base model does not 
#       use any hba cards. It instead uses two on board SAS connectors (1-1 to 1-8)
#       and 7 sata connectors (1-9 to 1-15). Using the output of lspci, we 
#       obtain the required bus addresses and generate the appropriate 
#       lines within the vdev_id.conf string. 
###############################################################################
def alias_av15_base():
	vdev_id_str = ""
	bus_address_sata = None
	bus_address_sas = None
	try:
		lspci_result = subprocess.Popen(["lspci"], stdout=subprocess.PIPE,
			universal_newlines=True).stdout
	except OSError:
		log("Error executing lspci.")
		sys.exit(1)

	for line in lspci_result:
		regex_sata = re.search("(\w\w:\w\w.\w).*Intel.*SATA Controller",line)
		regex_sas = re.search("(\w\w:\w\w.\w).*SAS3008",line)
		if regex_sata != None:
			bus_address_sata = "0000:"+regex_sata.group(1)
		if regex_sas != None:
			bus_address_sas = "0000:"+regex_sas.group(1)
		if bus_address_sata != None and bus_address_sas != None:
			break

	if bus_address_sata != None and bus_address_sas != None:
		for i in range(0,8):
			#alias the drives connected to the on board sas connectors
			vdev_id_str += "alias 1-{drive} /dev/disk/by-path/pci-{addr}-sas-phy{i}-lun-0\n".format(drive=i+1,addr=bus_address_sas,i=i)
		for i in range(0,7):
			#alias the drives connected to the on board sata sata connectors
			vdev_id_str += "alias 1-{drive} /dev/disk/by-path/pci-{addr}-ata-{i}\n".format(drive=i+9,addr=bus_address_sata,i=i+1)
	else:
		log("Error aliasing AV15-BASE")
		sys.exit(1)
	return vdev_id_str

###############################################################################
# Name: main (dmap)
# Args: (see parser)
# Desc: generates vdev_id.conf
###############################################################################
def main():
	# ensure that script has been run with root privilages
	check_root()
	parser = OptionParser()
	parser.add_option("-m","--no-udev",action="store_false",dest="trigger_udev",default=True,help="Creates map but doesnt trigger udev rules")
	parser.add_option("-s","--size",action="store",dest="sz",default=None,help="Specify chassis size")
	parser.add_option("-q","--quiet",action="store_true",dest="quiet",default=False,help="Quiet Mode")
	parser.add_option("-r","--reset-map",action="store_true",dest="reset_map",default=False,help="Resets the drive map")
	(options, args) = parser.parse_args()

	# apply the quiet flag to the global variable
	global g_quiet
	g_quiet = options.quiet

	# verify that vdev is configured
	if not verify_vdev():
		log("Invalid vdev configuration: run alias_setup.sh")
		sys.exit(1)

	# get the path variables required to determine config path and device path
	conf_path, dev_path = get_path_variables()
	
	# erase existing config file and reset map if specified by user (-r)
	if options.reset_map:
		reset_map(conf_path)

	server = get_server_info()
	vdev_id_conf = create_vdev_id(server)

	# write file to disk
	f = open(conf_path + "/vdev_id.conf","w")
	f.write(vdev_id_conf)
	f.close()

	#trigger udev
	if options.trigger_udev:
		trigger_udev()

	#print config file to stdout
	log("--------------------------------------------------------------------------------")
	log(conf_path + "/vdev_id.conf:")
	log("--------------------------------------------------------------------------------")
	log(vdev_id_conf)

if __name__ == "__main__":
    main()
