#!/usr/bin/python3
################################################################################
# network:
# 	used to return information about the network configuration in a .json
#   format. This is a helper script 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 re
import subprocess
import json
import sys
import os

g_dmi_fields = [
	"Bus Address",
	"Designation",
	"Type"
]

# ASRock Board 
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
}

# Gigabyte Board  
PCIE_SLOT_TYPE_GIGABYTE_ME03 = {
	"PCIE_1":"PCI Express 4 x16",
	"PCIE_2":"PCI Express 4 x8",
	"PCIE_3":"PCI Express 4 x16",
	"PCIE_4":"PCI Express 5 x16",
	"PCIE_5":"PCI Express 5 x16",
	"PCIE_6":"PCI Express 5 x16",
}

BA_LUT_MZ73 = {
	"PCIE_1": ["0000:a2:00.0"],
	"PCIE_2": ["0000:81:00.0", "0000:84:00.0"],  # Multiple potential bus addresses
	"PCIE_3": ["0000:01:00.0"],
	"PCIE_4": ["0000:41:00.0"],
}

PCIE_SLOT_LUT_B550I = {
	"J10":"PCIEX16",
	"J3700 M.2 Slot":"M2A_CPU",
	"J3708 PCIE x4 slot from Promontory":"M2B_SB",
}

BA_LUT_B550I = {
	# "0000:00:01.1":["0000:01:00.0", "0000:01:00.1"]
 	"PCIEX16":["0000:01:00.0", "0000:01:00.1"]
}

# Gigabyte MW34-SP0-00 — dmidecode bus addresses are wrong.
# Map slot names to root port addresses; endpoint bus numbers shift at each boot.
BA_LUT_MW34_NET = {
	"PCIE_1": ["0000:00:1d.0"],
	"PCIE_3": ["0000:00:01.0", "0000:00:01.1"],
	"PCIE_6": ["0000:00:01.0", "0000:00:01.1"],
}

_PCI_ADDR_RE_NET = re.compile(
	r'^[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9a-fA-F]$'
)

def _get_downstream_pci_addrs(root_port):
	"""Return PCI device addresses directly downstream of a root port via sysfs."""
	base = os.path.realpath(f"/sys/bus/pci/devices/{root_port}")
	if not os.path.isdir(base):
		return []
	return [e for e in os.listdir(base) if _PCI_ADDR_RE_NET.match(e)]


def get_motherboard_model_server_info():
	json_path = "/etc/45drives/server_info/server_info.json"
	if os.path.exists(json_path):
		f = open(json_path, "r")
		si = json.load(f)
		f.close()
		mobo_obj = {
			"Motherboard": [{
				"Manufacturer": si["Motherboard"]["Manufacturer"],
				"Product Name": si["Motherboard"]["Product Name"],
				"Serial Number": si["Motherboard"]["Serial Number"]
			}
			]
		}
		return json.dumps(mobo_obj)
	else:
		return "{\"Motherboard\":[{\"Manufacturer\": \"?\", \"Product Name\": \"?\", \"Serial Number\": \"?\"}]}"

def is_BCM57416_onboard_ethernet(device_name):
	return "BCM57416" in device_name

def get_ipv4(addr_info):
	if not addr_info:
		return "-"

	for addr in addr_info:
		if addr["family"] == "inet":
			return "{local}/{pre}".format(local=addr["local"],pre=str(addr["prefixlen"]))
	
	return "-"

def get_ipv6(addr_info):
	if not addr_info:
		return "-"

	for addr in addr_info:
		if addr["family"] == "inet6":
			return "{local}/{pre}".format(local=addr["local"],pre=str(addr["prefixlen"]))

	return "-"

def ip_addr():
	try:
		ipaddr_result = subprocess.Popen(
			["ip","-j","addr"],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		errorMessage = {"msg": "error running 'ip -j addr' command. "}
		print(json.dumps(errorMessage))
		exit(1)
	if ipaddr_result is None:
		exit(1)
	#ipaddr_json = json.loads(ipaddr_result.read())
	try:
		ipaddr_json = json.loads(ipaddr_result.read())
	except:
		errorMessage = {"msg": "error parsing json output of 'ip -j addr' command."}
		print(json.dumps(errorMessage))
		exit(1)
	
	j_ipaddr = []
	for connection in ipaddr_json:
		j_ipaddr.append({
			"Connection Name": connection["ifname"],
			"Connection State": connection["operstate"],
			"Type": connection["link_type"],
			"MAC": connection["address"] if "address" in connection.keys() else "-",
			"IPv4": get_ipv4(connection["addr_info"].copy()),
			"IPv6": get_ipv6(connection["addr_info"].copy())
		})
	return j_ipaddr


def get_network_info():
	#run ip -j addr
	j_connections = ip_addr()

	# LSHW COMMAND:
	# This runs the command: lshw -C network -businfo -quiet
	# this will allow us to determine which pci bus address is associated with a 
	# given connection. The results are stored as a list of tuples 
	# with the format of [(BUS ADDRESS 1, Connection Name 1),(BUS ADDRESS 2, Connection Name 2), ...]
	# the "network" option ensures that only pci devices associated with a given 
	# network connection are displayed. This is important in determining what pci network 
	# cards are installed in the system.
	try:
		lshw_result = subprocess.Popen(
			["lshw","-C","network","-businfo","-quiet"],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		return False
	if lshw_result is None:
		return False
	lshw = []
	for line in lshw_result:
		regex_lshw = re.search(r"^pci@(\S+)\s+(\S+)\s+(\S+)\s+(\S+)",line)
		if regex_lshw != None:
			lshw.append([regex_lshw.group(1),regex_lshw.group(2),regex_lshw.group(4)])

	# DMIDECODE COMMAND:
	# This runs the command: dmidecode -t 9
	# This lists all pci slot information. The results are stored in 
	# a list called "slot_entries". The g_dmi_fields list stores
	# the fields that we want to parse from the dmidecode output.
	# This list is then seperated into a list of dicts called "cards".
	# of these cards, we then search for cards with bus addresses
	# that match the bus addresses from the lshw command run prior. 
	# Because of the nature of the dmidecode results, we will only
	# see bus addresses in the form of 00:XX:00.0 the .0 results 
	# are the only ones that show up, the bus addresses ending in .1
	# do not. Therefore we have to add these to a list of dicts 
	# manually. This list is called "final_card_list". 
	# Lastly, we load in the netowrk_json_str and append the final_card_list
	# fields to the relevant connection, using the Connection Name 
	# as the method of determining the correct match. 
	try:
		dmi_result = subprocess.Popen(
			["dmidecode","-t","9"],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		return False
	if dmi_result is None:
		return False

	# get the relevant fields from the dmidecode output and store them in a list. 
	slot_entries = []
	for line in dmi_result:
		for field in g_dmi_fields:
			regex = re.search(r"^\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 = []
 
 	# Initialize a set to track unique designations
	used_designations = set()
 
	motherboard_info = json.loads(get_motherboard_model_server_info())
	motherboard_model = motherboard_info['Motherboard'][0].get('Product Name', '')
 
	for i in range(0,len(slot_entries),len(g_dmi_fields)):
		# cards.append(dict(slot_entries[i:i+len(g_dmi_fields)]))
		card = dict(slot_entries[i:i+len(g_dmi_fields)])

		if "Designation" in card:
			designation = card.get("Designation", "")
			if designation == 'N/A' or not designation:
				card = None
				continue
			if motherboard_model == 'B550I AORUS PRO AX':
				updated_designation = PCIE_SLOT_LUT_B550I[designation];
				actual_slot = updated_designation
			else:
				# Ensure that the designation suffix (_A, _B, etc.) is removed
				updated_designation = re.sub(r'_[A-Z]$', '', designation)  # This removes the trailing _A, _B, etc.
				actual_slot = updated_designation  # Use the updated designation for slot comparison

			if motherboard_model.startswith("ME03") and actual_slot in PCIE_SLOT_TYPE_GIGABYTE_ME03:
				card["Type"] = PCIE_SLOT_TYPE_GIGABYTE_ME03[actual_slot]
			else:
				card["Type"] = card.get("Type")  # Keep the original type if no match is found

			card['Designation'] = actual_slot  # Update the card with the cleaned-up slot designation
			# print(f"Updated card with cleaned designation and type: {card}")
			
			# Check if the designation has already been used
			if actual_slot not in used_designations:
				cards.append(card)  # Append only if the designation is unique
				used_designations.add(actual_slot)  # Mark the designation as used
			# else:
			# 	print(f"Skipping duplicate designation: {actual_slot}")

	for card in cards:
		if "Bus Address" in card.keys() and card["Bus Address"] in BA_EPC612D8A.keys():
			card["Bus Address"] = BA_EPC612D8A[card["Bus Address"]]
		
		if motherboard_model.startswith("MZ73") and "Designation" in card.keys() and card['Designation'] in BA_LUT_MZ73.keys():
			card["Bus Address"] = BA_LUT_MZ73[card["Designation"]]
   		
		if motherboard_model == 'B550I AORUS PRO AX' and "Designation" in card.keys() and card['Designation'] in BA_LUT_B550I.keys():
			card["Bus Address"] = BA_LUT_B550I[card["Designation"]]

		if motherboard_model.startswith("MW34") and "Designation" in card and card["Designation"] in BA_LUT_MW34_NET:
			root_ports = BA_LUT_MW34_NET[card["Designation"]]
			endpoints = []
			for rp in root_ports:
				endpoints.extend(_get_downstream_pci_addrs(rp))
			if endpoints:
				card["Bus Address"] = endpoints

	# for each card in the list of dictionaries, append 
	# dictionaries to a list that have the same Bus address
	# as those found in the list of connections found in 
	# the lshw command
	valid_cards = []
	for card in cards:
		matched_bus_address = None
		for hw_entry in lshw:
			device_name = hw_entry[2]
			
			# Exclude onboard ethernet from PCI matching
			if motherboard_model.startswith('MZ73') or motherboard_model.startswith('ROMED8'):
				if is_BCM57416_onboard_ethernet(device_name):
					continue

			# Check if Bus Address is a list or a single address
			# card_bus_addresses = card["Bus Address"] if isinstance(card["Bus Address"], list) else [card["Bus Address"]]
			if "Bus Address" not in card:
				continue  # Skip cards with no bus address

			card_bus_addresses = card["Bus Address"] if isinstance(card["Bus Address"], list) else [card["Bus Address"]]

			
			# Compare each address in card_bus_addresses to hw_entry[0]
			for card_bus_address in card_bus_addresses:
				
				# Normalize the formats for comparison (remove leading zeros, case insensitive)
				normalized_card_bus_address = card_bus_address.lower().lstrip('0')
				normalized_hw_entry_bus_address = hw_entry[0].lower().lstrip('0')

				if normalized_card_bus_address == normalized_hw_entry_bus_address:
					card["Connection Name"] = hw_entry[1]
					matched_bus_address = card_bus_address  # Store the matching bus address
					break  # Break inner loop if a match is found
			if matched_bus_address:
				card["Bus Address"] = matched_bus_address  # Set the card's bus address to the matching one
				valid_cards.append(card)
				break  # Break outer loop if a match is found

	# Copy the dictionaries for each card and modify the last digit of
	# both the Connection Name (ex enp24s0f0) and the bus address (ex. 0000:18:00.0)
	# to "enp24s0f1" and "0000:18:00.1" respectively. This is because we will
	# ALWAYS use pci cards with two ports, and never a single port according to Brett.  
	final_card_lst = []
	for card in valid_cards:
		card_duplicate = card.copy()
		
		# Check if the connection name ends with "0np0" and replace it with "1np1"
		if card_duplicate["Connection Name"].endswith("0np0"):
			card_duplicate["Connection Name"] = card_duplicate["Connection Name"][:-4] + "1np1"
		else:
			# Otherwise, just change the last digit from '0' to '1'
			card_duplicate["Connection Name"] = card_duplicate["Connection Name"][:-1] + "1"
		
		# Modify the Bus Address
		card_duplicate["Bus Address"] = card_duplicate["Bus Address"][:-1] + "1"
		
		# Append the original and duplicate card to the final list
		final_card_lst.append(card)
		final_card_lst.append(card_duplicate)

	# Add the relevant fields to the list of connections when we have a matching
	# Connection Name
	for obj in j_connections:
		matched = False  # Track if we matched the connection
		for card in final_card_lst:
			if obj["Connection Name"] == card["Connection Name"]:
				obj["Bus Address"] = card["Bus Address"]
				obj["Designation"] = card["Designation"]
				obj["PCI Type"] = card["Type"]

				# Modify the regex to capture the entire PCIE_X format
				# regex = re.search(r"(PCIE_\d+|(?:SLOT|PCIE)\d+)",card["Designation"])
				regex = re.search(r"(PCIE(?:X\d+|_\d+|\d+)|SLOT\d+)",card["Designation"])
				if regex is not None:
					obj["PCI Slot"] = regex.group(1)

	return j_connections

def main():
	network = get_network_info()
	result = []
	if network:
		for connection in network:
			result.append(
				{ 
					"connectionName": connection["Connection Name"], 
					"connectionState": connection["Connection State"], 
					"connectionType": connection["Type"], 
					"mac": connection["MAC"], 
					"ipv4": connection["IPv4"] if "IPv4" in connection.keys() else "-", 
					"ipv6": connection["IPv6"] if "IPv6" in connection.keys() else "-", 
					"pciSlot": connection["PCI Slot"] if "PCI Slot" in connection.keys() else "-", 
					"busAddress": connection["Bus Address"] if "Bus Address" in connection.keys() else "-"
				}
			)
		print(json.dumps(result,indent=4))
	else:
		print("Unable to provide network information.")
		print("Please ensure that the following programs are installed:")
		print("\tdmidecode, lshw")
		sys.exit(1)

if __name__ == "__main__":
	main()