#!/usr/libexec/platform-python
################################################################################
# ram:
# 	used to return information about the ram configuration 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
import sys

g_dmi_fields = [
	"Size",
	"Form Factor",
	"Locator",
	"Type",
	"Manufacturer",
	"Serial Number"
]

g_ipmitool_sensor_fields = {
	"P1-DIMMA1 Temp": " (C)",
	"P1-DIMMB1 Temp": " (C)",
	"P1-DIMMC1 Temp": " (C)",
	"P1-DIMMD1 Temp": " (C)",
	"P1-DIMME1 Temp": " (C)",
	"P1-DIMMF1 Temp": " (C)",
	"P1-DIMMG1 Temp": " (C)",
	"P1-DIMMH1 Temp": " (C)",
	"P2-DIMMA1 Temp": " (C)",
	"P2-DIMMB1 Temp": " (C)",
	"P2-DIMMC1 Temp": " (C)",
	"P2-DIMMD1 Temp": " (C)",
	"P2-DIMME1 Temp": " (C)",
	"P2-DIMMF1 Temp": " (C)",
	"P2-DIMMG1 Temp": " (C)",
	"P2-DIMMH1 Temp": " (C)",
	"DIMMA1 Temp": " (C)",
	"DIMMA2 Temp": " (C)",
	"DIMMB1 Temp": " (C)",
	"DIMMB2 Temp": " (C)",
	"DIMMC1 Temp": " (C)",
	"DIMMC2 Temp": " (C)",
	"DIMMD1 Temp": " (C)",
	"DIMMD2 Temp": " (C)",
	"DIMME1 Temp": " (C)",
	"DIMME2 Temp": " (C)",
	"DIMMF1 Temp": " (C)",
	"DIMMF2 Temp": " (C)",
	"DIMMG1 Temp": " (C)",
	"DIMMG2 Temp": " (C)",
	"DIMMH1 Temp": " (C)",
	"DIMMH2 Temp": " (C)",
	"DDR4_A Temp": " (C)", #ASRockRack 
	"DDR4_B Temp": " (C)", #ASRockRack
	"DDR4_C Temp": " (C)", #ASRockRack
	"DDR4_D Temp": " (C)", #ASRockRack
	"DDR4_E Temp": " (C)", #ASRockRack
	"DDR4_F Temp": " (C)",  #ASRockRack
	"DIMMABCD Temp": " (C)", #X12SPL-F
	"DIMMEFGH Temp": " (C)", #X12SPL-F
	"P1_DIMMA~D Temp": " (C)", #X12DPi-N6
	"P1_DIMME~H Temp": " (C)", #X12DPi-N6
	"P2_DIMMA~D Temp": " (C)", #X12DPi-N6
	"P2_DIMME~H Temp": " (C)", #X12DPi-N6
}

def ipmitool_sensor():
	try:
		ipmitool_sensor_result = subprocess.Popen(
			["ipmitool","sensor"],stdout=subprocess.PIPE,stderr=subprocess.PIPE,universal_newlines=True).stdout
	except:
		print("Error running 'ipmitool sensor'")
		sys.exit(1)

	sensor_readings = []
	for line in ipmitool_sensor_result:
		for field in g_ipmitool_sensor_fields.keys():
			regex = re.search("^({fld})\s+\|\s+(\S+).*".format(fld=field),line)
			if regex != None:
				sensor_readings.append(
				(regex.group(1)[:-5],regex.group(2)+g_ipmitool_sensor_fields[regex.group(1)])
				)
	
	return dict(sensor_readings)

def dmidecode():
	try:
		dmi_result = subprocess.Popen(
			["dmidecode","-t","17"],stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		print("Error running 'dmidecode -t 17'")
		sys.exit(1)

	slot_entries = []
	slot_template = {
		"Size": "No Module Installed",
		"Form Factor": "-",
		"Locator" : "-",
		"Type": "-",
		"Manufacturer": "-",
		"Serial Number" : "-",
		"Temperature" : "-"
	}

	current_slot = None

	for line in dmi_result:
		if line.startswith("Handle "):
			if current_slot != None:
				slot_entries.append(current_slot.copy())
			current_slot = slot_template.copy()
		elif current_slot != None:
			for field in g_dmi_fields:
				regex = re.search("^\s+({fld}):\s+(.*)".format(fld=field),line)
				if regex != None and current_slot != None:
					current_slot[regex.group(1)] = regex.group(2)
	
	#append the last entry
	if current_slot != None:
		slot_entries.append(current_slot.copy())

	# check slot_entries for translatable locator tags (ASRockRack Motherboards, X12 and H12 boards)
	locator_tags = {
		"CPU1_DIMM_A1":["DDR4_A"],
		"CPU1_DIMM_A2":["DDR4_A"],
		"CPU1_DIMM_B1":["DDR4_B"],
		"CPU1_DIMM_C1":["DDR4_C"],
		"CPU1_DIMM_D1":["DDR4_D"],
		"CPU1_DIMM_D2":["DDR4_D"],
		"CPU1_DIMM_E1":["DDR4_E"],
		"CPU1_DIMM_F1":["DDR4_F"],
		"DIMMA1":["DIMMABCD","P1_DIMMA~D"],
		"DIMMB1":["DIMMABCD","P1_DIMMA~D"],
		"DIMMC1":["DIMMABCD","P1_DIMMA~D"],
		"DIMMD1":["DIMMABCD","P1_DIMMA~D"],
		"DIMME1":["DIMMEFGH","P1_DIMME~H"],
		"DIMMF1":["DIMMEFGH","P1_DIMME~H"],
		"DIMMG1":["DIMMEFGH","P1_DIMME~H"],
		"DIMMH1":["DIMMEFGH","P1_DIMME~H"],
		"P1-DIMMA1": ["P1_DIMMA~D"],
		"P1-DIMMB1": ["P1_DIMMA~D"],
		"P1-DIMMB2": ["P1_DIMMA~D"],
		"P1-DIMMC1": ["P1_DIMMA~D"],
		"P1-DIMMD1": ["P1_DIMMA~D"],
		"P2-DIMMA1": ["P2_DIMMA~D"],
		"P2-DIMMB1": ["P2_DIMMA~D"],
		"P2-DIMMB2": ["P2_DIMMA~D"],
		"P2-DIMMC1": ["P2_DIMMA~D"],
		"P2-DIMMD1": ["P2_DIMMA~D"],
		"P1-DIMME1": ["P1_DIMME~H"],
		"P1-DIMMF1": ["P1_DIMME~H"],
		"P1-DIMMG1": ["P1_DIMME~H"],
		"P1-DIMMH1": ["P1_DIMME~H"],
		"P2-DIMME1": ["P1_DIMME~H"],
		"P2-DIMMF1": ["P1_DIMME~H"],
		"P2-DIMMG1": ["P1_DIMME~H"],
		"P2-DIMMH1": ["P1_DIMME~H"]
	}

	for slot in slot_entries:
		if "Locator" in slot.keys() and slot["Locator"] in locator_tags.keys():
			slot["Temperature Keys"] = locator_tags[slot["Locator"]]

	return slot_entries

def main():
	ram = dmidecode()
	temps = ipmitool_sensor()

	if ram and temps:
		for module in ram:
			for key in g_dmi_fields:
				if key not in module.keys():
					module[key] = "-"
			if module["Locator"] in temps.keys():
				module["Temperature"] = temps[module["Locator"]]
			elif "Temperature Keys" in module.keys():
				for lookup_key in module["Temperature Keys"]:
					if lookup_key in temps.keys():
						module["Temperature"] = temps[lookup_key]
			if module["Temperature"] == "na (C)":
				module["Temperature"] = "-"
			if module["Manufacturer"] == "NO DIMM":
				module["Manufacturer"] = "-"
			if module["Serial Number"] == "NO DIMM":
				module["Serial Number"] = "-"

	result = []

	for slot in ram:
		result.append({ 
			"locator": slot["Locator"], 
			"type": slot["Type"], 
			"size": slot["Size"], 
			"manufacturer": slot["Manufacturer"], 
			"serialNumber": slot["Serial Number"], 
			"temperature": slot["Temperature"] 
			}
		)
	
	print(json.dumps(result,indent=4))


if __name__ == "__main__":
    main()