#!/usr/libexec/platform-python
################################################################################
# pci:
# 	used to return information about pci cards in a .json
#   format. This is a helper sctipt for use with the
#   cockpit-hardware package (https://github.com/45Drives/cockpit-hardware)
#
# Copyright (C) 2020, 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 subprocess
import re
import json

g_dmi_fields = [
	"Designation",
	"Type",
	"Current Usage",
	"ID",
	"Bus Address"
]



g_storcli64_fields = [
	"SAS9305-16i",
	"SAS9305-24i"
]

g_network_card_models = [
	"X540-AT2", 
	"XL710", 
	"XXV710", 
	"82599ES",
	"BCM57412",
	"MT27800"
]

g_sata_controllers = [
	"ASM1062"
]


def dmidecode():
	try:
		dmi_result = subprocess.Popen(
			["dmidecode","-t","9"],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		return False

	slot_entries = []
	for line in dmi_result:
		for field in g_dmi_fields:
			regex = re.search("^\s+({fld}):\s+(.*)".format(fld=field),line)
			if regex != None:
				slot_entries.append((regex.group(1),regex.group(2)))

	# Take the subset of fields in the list for each card and 
	# append them as a dictionary into a new list. 
	cards = []
	for i in range(0,len(slot_entries),len(g_dmi_fields)):
		cards.append(dict(slot_entries[i:i+len(g_dmi_fields)]))

	BA_EPC612D8A = {
		"ff00:16:00.0":"0000:17:00.0", #PCIE1
		"ff00:16:02.0":"0000:1c:00.0", #PCIE2
		"ff00:64:00.0":"0000:65:00.0", #PCIE4
		"ff00:64:02.0":"0000:66:00.0", #PCIE3
		"ff00:b2:00.0":"0000:b3:00.0", #PCIE6
		"ff00:b2:02.0":"0000:b4:00.0" #PCIE5
	}
	
	for slot in cards:
		if "Bus Address" in slot.keys() and slot["Bus Address"] in BA_EPC612D8A.keys():
			slot["Bus Address"] = BA_EPC612D8A[slot["Bus Address"]]
		if "Designation" in slot.keys():
			regex = r"(?:SLOT|PCIE)(\d+)"
			match = re.search(regex, slot["Designation"])
			if match != None:
				slot["ID"] = match.group(1)

			

	return cards


def storcli():
	try:
		storcli64_result = subprocess.Popen(
			["/opt/45drives/tools/storcli64","show","all"],stdout=subprocess.PIPE,universal_newlines=True)
	except:
		return False
	hba_cards = []
	hba_dict = {}
	for line in storcli64_result.stdout:
		for field in g_storcli64_fields:
			# Model AdapterType VendId DevId SubVendId SubDevId PCIAddress 	
			regex = re.search("({fld}).*(00:\w\w:\w\w:\w\w)\s+$".format(fld=field),line)
			if regex != None:
				hba_dict["Model"] = regex.group(1)
				hba_dict["Bus Address"] = "00" + regex.group(2)
				hba_dict["Bus Address"] = hba_dict["Bus Address"][:-3] + ".0"
				hba_cards.append(hba_dict.copy())
	
	return hba_cards

def lspci_hba():
	try:
		lspci_hba_result = subprocess.Popen(
			["lspci", "-d", "1000:*","-vv", "-i", "/opt/45drives/tools/pci.ids"],stdout=subprocess.PIPE,stderr=subprocess.STDOUT,universal_newlines=True).stdout.read()
	except:
		return False
	hba_cards = []
	hba_dict = {}
	rx_pci=re.compile(r"^(\w\w:\w\w\.\w).*\n.*(?:(?:(?:^\t).*\n)+^.*)?(9600-16i|9600-24i|SAS9305-16i|SAS9305-24i|HBA 9405W-16i|9361-16i|HBA 9400-16i|9361-24i).*\n",re.MULTILINE)
	for match in rx_pci.finditer(lspci_hba_result):
		hba_dict["Model"] = match.group(2)
		hba_dict["Bus Address"] = "0000:" + match.group(1)
		hba_cards.append(hba_dict.copy())

	return hba_cards

def network():
	try:
		network_result = subprocess.Popen(
			["/usr/share/cockpit/45drives-motherboard/helper_scripts/network"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		return False
	
	network_output = json.loads(network_result.read())

	network_cards = []
	for connection in network_output["Network Info"]:
		if "PCI Slot" in connection.keys():
			network_cards.append(connection)

	return network_cards

def getNetworkCardModel(busAddress):
	trimmedBusAddress = busAddress[5:]
	try:
		lspci_result = subprocess.Popen(
			["lspci"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		return "unknown"

	for line in lspci_result:
		regex_addr = re.search("^{addr}\s(.*)$".format(addr=trimmedBusAddress),line)
		if regex_addr != None:
			for model in g_network_card_models:
				regex_model = re.search("{mdl}".format(mdl=model),regex_addr.group(1))
				if regex_model != None:
					return model

	return "unknown"

def sata():
	try:
		lspci_result = subprocess.Popen(
			["lspci"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		return False

	sata = []
	sata_dict = {}
	for line in lspci_result:
		for field in g_sata_controllers:
			regex = re.search("^(\S+).*({fld}).*$".format(fld=field),line)
			if regex != None:
				sata_dict["Card Type"] = "Serial ATA Controller"
				sata_dict["Card Model"] = regex.group(2)
				sata_dict["Bus Address"] = "0000:" + regex.group(1)
				sata.append(sata_dict.copy())	

	# search through the sata list and append any partition information found. 
	try:
		ls_result = subprocess.Popen(
			["ls","-l","/dev/disk/by-path"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		return False

	drives = []
	
	for card in sata:
		for line in ls_result:
			drive_dict = {}
			regex = re.search("pci-({ba})-ata-(\d)\s->\s\W+(.*)".format(ba=card["Bus Address"]),line)
			if regex != None:
				drive_dict["Device"] = regex.group(3)
				drive_dict["Path"] = "pci-" + regex.group(1) + "-ata-" + regex.group(2)
				drive_dict["Partitions"] = lsblk(drive_dict["Device"])
				drives.append(drive_dict.copy())
		card["Connections"] = drives.copy()
	return sata

def lsblk(device):
	try:
		lsblk_result = subprocess.Popen(
			["lsblk","-l","/dev/"+device],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		return False

	partitions = []
	for line in lsblk_result:
		regex = re.search("^(\S+)\s+\S+\s+\S+\s+(\S+)\s+\S+\s+(\S+)(.*)$",line)
		if regex != None and regex.group(1) != "NAME":
			partitions.append(
					{
						"Name":regex.group(1),
						"Size":regex.group(2),
						"Type":regex.group(3),
						"Mount Point":regex.group(4).lstrip()
					}
				)
	return partitions

def main():
	pci_slots = dmidecode()
	hba_cards = lspci_hba()
	network_cards = network()
	sata_cards = sata()

	# Look for hba cards and pci slots with matching bus addresses.
	# and add a "Card Type" field to the list of pci slot dictionaries
	if pci_slots and hba_cards:
		for hba in hba_cards:
			for slot in pci_slots:
				if(hba["Bus Address"] == slot["Bus Address"]):
					slot["Card Type"] = "HBA"
					slot["Card Model"] = hba["Model"]
	
	# for each pci slot with an ID that corresponds to 
	# the "PCI Slot" field from the 
	# /usr/share/cockpit/45drives-system/helper_scripts/network script's
	# .json formatted output. Append this network connection 
	# to a list of connections for that specific PCI slot
	if pci_slots and network_cards:
		for slot in pci_slots:
			for card in network_cards:
				if ("ID" in slot.keys() and card["PCI Slot"] == slot["ID"]):
					slot["Card Type"] = "Network Card"
					slot["Card Model"] = getNetworkCardModel(str(slot["Bus Address"]))
					if "Connections" not in slot.keys():
						slot["Connections"] = []
					slot["Connections"].append(card)

	# for each pci slot, see if there are any sata_cards with
	# matching bus addresses. if so, append the following fields
	# to indicate the card type and model
	if pci_slots and sata_cards:
		for slot in pci_slots:
			for card in sata_cards:
				if(card["Bus Address"] == slot["Bus Address"]):
					slot["Card Type"] = card["Card Type"]
					slot["Card Model"] = card["Card Model"]
					if "Connections" not in slot.keys():
						slot["Connections"] = []
					slot["Connections"] = card["Connections"]
					# TODO: Get the partition information for any connected drives.

	# Create .json formatted output string
	output_str = "{ \"PCI Info\":["
	
	for slot in pci_slots:
		output_str += json.dumps(slot) + ","
	
	if pci_slots:
		output_str = output_str[:-1] + "]}"
	else:
		output_str = "{ \"PCI Info\":[]}"

	# print output string to stdout
	print(json.dumps(json.loads(output_str),indent=4))


if __name__ == "__main__":
    main()