#!/usr/libexec/platform-python
import re
import subprocess
import os
import sys
import json

# supported CPU types
g_dmi_cpu_fields = {
    "Socket Designation":  ["CPU1", "CPU2"],
    "Version":  [
        "Intel(R) Xeon(R) Silver 4110 CPU @ 2.10GHz",
        "AMD EPYC 7281 16-Core Processor",
        "Intel(R) Xeon(R) CPU E3-1220 v6 @ 3.00GHz",
        "Intel(R) Xeon(R) Silver 4210 CPU @ 2.10GHz"
    ],
    "Current Speed": None,
    "Max Speed": None,
}

g_ipmitool_sensor_fields = {
    "CPU1 Temp": "(C)",
    "CPU2 Temp": "(C)",
    "PW Consumption": "(W)",
    "FAN1": "(RPM)",
    "FAN2": "(RPM)",
    "FAN3": "(RPM)",
    "FAN4": "(RPM)",
    "FAN5": "(RPM)",
    "FAN6": "(RPM)",
    "FANA": "(RPM)",
    "FANB": "(RPM)",
    "P1-DIMMA1 Temp": " (C)",
    "P1-DIMMB1 Temp": " (C)",
    "P1-DIMMD1 Temp": " (C)",
    "P1-DIMME1 Temp": " (C)",
    "P2-DIMMA1 Temp": " (C)",
    "P2-DIMMB1 Temp": " (C)",
    "P2-DIMMD1 Temp": " (C)",
    "P2-DIMME1 Temp": " (C)",
    "CPU 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)",
    "DIMMH1 Temp": " (C)"
}


def get_cpu_info():
	cpu = []
	try:
		dmi_result = subprocess.Popen(["dmidecode","-t","4"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		#print("ERROR: dmidecode is not installed")
		return False
	for line in dmi_result:
		for field in g_dmi_cpu_fields.keys():
			regex = re.search("^\s({fld}):\s+(.*)".format(fld=field),line)
			if regex != None:
				regex_group1_str = str(regex.group(1)).rstrip()
				regex_group2_str = str(regex.group(2)).rstrip()
				if g_dmi_cpu_fields[regex_group1_str] != None:
					if regex_group2_str in g_dmi_cpu_fields[regex_group1_str]:
						cpu.append((regex_group1_str,regex_group2_str))
					elif field == "Socket Designation" and regex_group2_str == "CPU":
						cpu.append((regex_group1_str,"CPU1"))
					else:
						cpu.append((regex_group1_str,regex_group2_str))
				else:
					cpu.append((regex_group1_str,regex_group2_str))
	cpu_dict = {
		"CPU":[]
	}

	cpu_count = len(cpu)//len(g_dmi_cpu_fields)
	for i in range(0,cpu_count):
		cpu_dict["CPU"].append(dict(cpu[i*len(g_dmi_cpu_fields):i*len(g_dmi_cpu_fields)+(len(g_dmi_cpu_fields))]))
	return cpu_dict["CPU"]


def get_sensor_readings():
    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:
                if regex.group(1) == "CPU Temp":
                    # ipmitool does not use the socket name
                    # we need to manually intervene and store it as
                    # CPU1 Temp
                    sensor_readings.append(
                        ("CPU1 Temp", regex.group(2) +
                         g_ipmitool_sensor_fields[regex.group(1)])
                    )
                else:
                    # All other readings are fine, use the field as captured.
                    sensor_readings.append((regex.group(1), regex.group(
                        2)+g_ipmitool_sensor_fields[regex.group(1)]))
    return dict(sensor_readings)


def main():
    cpus = get_cpu_info()
    sensor_readings = get_sensor_readings()

    if cpus:
        cpu_arr = []
        for cpu in cpus:
            cpu_arr.append(
                { 
                    'socket': cpu["Socket Designation"], 
                    'model': cpu["Version"],
                    'maxSpeed': cpu["Max Speed"],
                    'currentSpeed': cpu["Current Speed"],
                    'temperature': sensor_readings[cpu["Socket Designation"]+" Temp"] if (cpu["Socket Designation"]+" Temp") in sensor_readings.keys() else "-"
                },
            )
        result = {
            "cpus": cpu_arr
        }
    else:
        result = {
            "cpus": [
                { 
                    'socket': "UNKNOWN", 
                    'model': "UNKNOWN",
                    'maxSpeed': "-",
                    'currentSpeed': "-",
                    'temperature': "-"
                }
            ]
        }
    print(json.dumps(result, indent=4))

if __name__ == "__main__":
    main()
