#!/usr/bin/python3
################################################################################
# sata:
# 	used to return information about the status of sata connections 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


# This contains the expected path(s), for 
# sata connectors on supported motherboards
g_motherboard_sata_addresses = {
	"X11SSH-CTF":[
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"X11SSM-F":[
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"H11SSL-i":[
		["pci-0000:43:00.2-ata-1","I-SATA0"],
		["pci-0000:43:00.2-ata-2","I-SATA1"],
		["pci-0000:43:00.2-ata-3","I-SATA2"],
		["pci-0000:43:00.2-ata-4","I-SATA3"],
		["pci-0000:43:00.2-ata-5","I-SATA4"],
		["pci-0000:43:00.2-ata-6","I-SATA5"],
		["pci-0000:43:00.2-ata-7","I-SATA6"],
		["pci-0000:43:00.2-ata-8","I-SATA7"],
		["pci-0000:44:00.0-ata-1","I-SATA0"],
		["pci-0000:44:00.0-ata-2","I-SATA1"]
	],
	"X11SPL-F":[
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"X11SPi-TF":[
		["pci-0000:00:11.5-ata-1","S-SATA0"],
		["pci-0000:00:11.5-ata-2","S-SATA1"],
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"EPC621D8A":[
		["pci-0000:00:11.5-ata-1","SSATA_0"],
		["pci-0000:00:11.5-ata-2","SSATA_1"],
		["pci-0000:00:11.5-ata-3","SSATA_2"],
		["pci-0000:00:11.5-ata-4","SSATA_3"],
		["pci-0000:00:11.5-ata-5","SSATA_4"],
		["pci-0000:00:11.5-ata-6","SSATA_5"],
		["pci-0000:00:17.0-ata-1","SATA_0"],
		["pci-0000:00:17.0-ata-2","SATA_1"],
		["pci-0000:00:17.0-ata-3","SATA_2"],
		["pci-0000:00:17.0-ata-4","SATA_3"],
		["pci-0000:00:17.0-ata-5","SATA_4"]
	],
	"X12SPL-F":[
		["pci-0000:00:11.5-ata-1","S-SATA0"],
		["pci-0000:00:11.5-ata-2","S-SATA1"],
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"H12SSL-i":[
		["pci-0000:4a:00.0-ata-1","I_SATA0"],
		["pci-0000:4a:00.0-ata-2","I_SATA1"],
		["pci-0000:4a:00.0-ata-3","I-SATA2"],
		["pci-0000:4a:00.0-ata-4","I-SATA3"],
		["pci-0000:4a:00.0-ata-5","I-SATA4"],
		["pci-0000:4a:00.0-ata-6","I-SATA5"],
		["pci-0000:4a:00.0-ata-7","I-SATA6"],
		["pci-0000:4a:00.0-ata-8","I-SATA7"]
	],
	"X11SPH-nCTF":[
		["pci-0000:00:11.5-ata-1","S-SATA0"],
		["pci-0000:00:11.5-ata-2","S-SATA1"],
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"X11SPH-nCTPF":[
		["pci-0000:00:11.5-ata-1","S-SATA0"],
		["pci-0000:00:11.5-ata-2","S-SATA1"],
		["pci-0000:00:17.0-ata-1","I-SATA0"],
		["pci-0000:00:17.0-ata-2","I-SATA1"],
		["pci-0000:00:17.0-ata-3","I-SATA2"],
		["pci-0000:00:17.0-ata-4","I-SATA3"],
		["pci-0000:00:17.0-ata-5","I-SATA4"],
		["pci-0000:00:17.0-ata-6","I-SATA5"],
		["pci-0000:00:17.0-ata-7","I-SATA6"],
		["pci-0000:00:17.0-ata-8","I-SATA7"]
	],
	"ME03-CE0-000":[
		["pci-0000:04:00.0-ata-1", "U2_P0_G3A0"],
		["pci-0000:04:00.0-ata-2", "U2_P0_G3A0"],
		["pci-0000:04:00.0-ata-3", "U2_P0_G3A0"],
		["pci-0000:04:00.0-ata-4", "U2_P0_G3A0"],
		# ["pci-0000:04:00.0-ata-5", "U2_P0_G3A0"],
		# ["pci-0000:04:00.0-ata-6", "U2_P0_G3A0"],
		# ["pci-0000:04:00.0-ata-7", "U2_P0_G3A0"],
		# ["pci-0000:04:00.0-ata-8", "U2_P0_G3A0"],
  
		# ["pci-0000:04:00.0-ata-1", "U2_P0_G3A1"],
		# ["pci-0000:04:00.0-ata-2", "U2_P0_G3A1"],
		# ["pci-0000:04:00.0-ata-3", "U2_P0_G3A1"],
		# ["pci-0000:04:00.0-ata-4", "U2_P0_G3A1"],
		["pci-0000:04:00.0-ata-5", "U2_P0_G3A1"],
		["pci-0000:04:00.0-ata-6", "U2_P0_G3A1"],
		["pci-0000:04:00.0-ata-7", "U2_P0_G3A1"],
		["pci-0000:04:00.0-ata-8", "U2_P0_G3A1"],
  
		["pci-0000:04:00.1-ata-1", "U2_P0_G3B0"],
		["pci-0000:04:00.1-ata-2", "U2_P0_G3B0"],
		["pci-0000:04:00.1-ata-3", "U2_P0_G3B0"],
		["pci-0000:04:00.1-ata-4", "U2_P0_G3B0"],
	],
	"MS03-6L0-000":[
  		["pci-0000:00:17.0-ata-1","SL_SATA1"],
		["pci-0000:00:17.0-ata-2","SL_SATA1"],
		["pci-0000:00:17.0-ata-3","SL_SATA1"],
		["pci-0000:00:17.0-ata-4","SL_SATA1"],
		["pci-0000:00:17.0-ata-5","SATA_4_5"],
		["pci-0000:00:17.0-ata-6","SATA_4_5"],

	],
	"MS73-HB0-000":[
  		["pci-0000:00:18.0-ata-1","SSATA2"],
		["pci-0000:00:18.0-ata-3","SSATA0"],
  
  		["pci-0000:00:17.0-ata-1","SL_SATA1"],
		["pci-0000:00:17.0-ata-2","SL_SATA1"],
		["pci-0000:00:17.0-ata-3","SL_SATA1"],
		["pci-0000:00:17.0-ata-4","SL_SATA1"],
  
  		["pci-0000:00:17.0-ata-5","SL_SATA2"],
		["pci-0000:00:17.0-ata-6","SL_SATA2"],
		["pci-0000:00:17.0-ata-7","SL_SATA2"],
		["pci-0000:00:17.0-ata-8","SL_SATA2"],
  
  		["pci-0000:00:18.0-ata-5","SL_SATA3"],
		["pci-0000:00:18.0-ata-6","SL_SATA3"],
		["pci-0000:00:18.0-ata-7","SL_SATA3"],
		["pci-0000:00:18.0-ata-8","SL_SATA3"],
  
  		# ["pci-0000:00:50.0-nvme-1","M2M_0"],
		# ["pci-0000:00:51.0-nvme-1","M2M_1"],
	],
	"MZ73-LM0-000":[
  		["pci-0000:65:00.0-ata-1","SLSAS_0"],
		["pci-0000:65:00.0-ata-2","SLSAS_0"],
		["pci-0000:65:00.0-ata-3","SLSAS_0"],
		["pci-0000:65:00.0-ata-4","SLSAS_0"],
  
	  	# ["pci-0000:61:00.0-nvme-1","M2_0"],
	],
 	"MC13-LE1-000":[
  		["pci-0000:0c:00.0-ata-1","SATA_0_1"],
		["pci-0000:0c:00.0-ata-2","SATA_0_1"],
		["pci-0000:0c:00.0-ata-3","SATA_2_3"],
		["pci-0000:0c:00.0-ata-4","SATA_2_3"],
	],
	"MH53-G40-000":[
		["pci-0000:64:00.0-ata-1","SATA_0_1"],
		["pci-0000:64:00.0-ata-2","SATA_0_1"],
		["pci-0000:64:00.0-ata-3","SATA_2_3"],
		["pci-0000:64:00.0-ata-4","SATA_2_3"],
	],
	"B550I AORUS PRO AX":[
		# ["pci-0000:01:00.1-ata-1", "S-SATA1"],
		# ["pci-0000:01:00.1-ata-2", "S-SATA2"],
		# ["pci-0000:01:00.1-ata-3", "S-SATA3"],
		# ["pci-0000:01:00.1-ata-4", "S-SATA4"],
  		# ["pci-0000:06:00.0-ata-1", "I-SATA1"],
		# ["pci-0000:06:00.0-ata-2", "I-SATA2"],
		# ["pci-0000:06:00.0-ata-3", "I-SATA3"],
		# ["pci-0000:06:00.0-ata-4", "I-SATA4"],
		# ["pci-0000:06:00.0-ata-5", "I-SATA5"],
	],
	"EC266D2I-2T/AQC":[
  		["pci-0000:00:17.0-ata-1", "OCU1"],
		["pci-0000:00:17.0-ata-2", "OCU1"],
		["pci-0000:00:17.0-ata-3", "OCU1"],
		["pci-0000:00:17.0-ata-4", "OCU1"],
  		["pci-0000:00:17.0-ata-5", "OCU2"],
		["pci-0000:00:17.0-ata-6", "OCU2"],
		["pci-0000:00:17.0-ata-7", "OCU2"],
		["pci-0000:00:17.0-ata-8", "OCU2"],
		# ["pci-0000:01:00.0-nvme-1", "M2_1"],
		# ["pci-0000:02:00.0-nvme-1", "M2_1"]
	],
"MW34-SP0-00":[
    ["pci-0000:00:17.0-ata-1", "SATA_0_1"],
    ["pci-0000:00:17.0-ata-2", "SATA_0_1"],
    ["pci-0000:00:17.0-ata-3", "SATA_2_3"],
    ["pci-0000:00:17.0-ata-4", "SATA_2_3"],
    ["pci-0000:00:17.0-ata-5", "SATA_4_5"],
    ["pci-0000:00:17.0-ata-6", "SATA_4_5"],
    ["pci-0000:00:17.0-ata-7", "SATA_6_7"],
    ["pci-0000:00:17.0-ata-8", "SATA_6_7"],
],
}

# Helper function to dynamically fetch SATA PCI addresses
def get_sata_pci_addresses():
    try:
        # Run the lshw command to get storage information
        lshw_result = subprocess.Popen(
            ["lshw", "-class", "storage"], stdout=subprocess.PIPE, universal_newlines=True
        ).stdout
    except Exception as e:
        print(f"Error running lshw: {e}")
        return []
    if lshw_result is None:
        return []

    pci_addresses = []
    sata_found = False  # Flag to track if we are in a SATA section

    bus_info_pattern = re.compile(r"bus info:\s+pci@(\S+)")  # Regex to capture PCI bus info
    description_pattern = re.compile(r"description:\s+SATA controller")  # To identify SATA sections

    for line in lshw_result:
        if description_pattern.search(line):
            sata_found = True
        elif line.strip() == "":  # Empty line indicates end of current section
            sata_found = False

        if sata_found:
            match = bus_info_pattern.search(line)
            if match:
                pci_addresses.append(match.group(1))  # Capture PCI address

    return pci_addresses

def ls_by_path(motherboard):
	try:
		ls_result = subprocess.Popen(
			["ls", "-l", "/dev/disk/by-path"], stdout=subprocess.PIPE, universal_newlines=True).stdout
	except:
		return False
	if ls_result is None:
		return False

	sata_connections = {}
	if motherboard == "B550I AORUS PRO AX":
        # Dynamically fetch SATA controller addresses
		sata_addresses = get_sata_pci_addresses()
		if not sata_addresses:
			print("No SATA controllers found dynamically for B550I AORUS PRO AX.")
			return sata_connections

		# Map SATA controllers to port numbers dynamically
		for line in ls_result:
			line = line.strip()
			for addr in sata_addresses:
				for port in range(1, 9):  # Assuming up to 8 ports
					path = f"pci-0000:{addr}-ata-{port}"
					if path in line:
						device = line.split(" -> ")[-1].strip().replace("../../", "")
						if not re.search(r'\d+$', device):  # Skip partition names
							if device not in sata_connections:
								sata_connections[device] = {
									"Path": path,
									"Connector": f"SATA-{port}",
									"Device": device,
									"Partitions": []
								}
						break
  
		return sata_connections

	else:
		if motherboard in g_motherboard_sata_addresses.keys():
			paths = g_motherboard_sata_addresses[motherboard]
			for line in ls_result:
				line = line.strip()
				for path in paths:
					# Match the full generated path directly with the line
					if path[0] in line:
						device = line.split(" -> ")[-1].strip().replace("../../", "")
						# Check if the device ends with a digit(s) using regex
						if not re.search(r'\d+$', device):
							# print(f"Matched path: {path[0]} -> {device}")
							if device not in sata_connections:
								sata_connections[device] = {
									"Path": path[0],
									"Connector": path[1],
									"Device": device,
									"Partitions": []
								}
						break

			return sata_connections
		# else:
		# 	for line in ls_result:
		# 		line = line.strip()
		# 		regex = re.search(r"(\S+-ata-\d+)\s->\s\W+(.*)", line)
		# 		if regex is not None:
		# 			device = regex.group(2).strip().replace("../../", "")
		# 			# print(f"Generic matched: {regex.group(1)} -> {device}")
		# 			if device not in sata_connections:
		# 				sata_connections[device] = {
		# 						"Path": regex.group(1),
		# 						"Connector": ("I-SATA" + regex.group(1).split('-ata-')[-1]),
		# 						"Device": device,
		# 						"Partitions": []
		# 					}
		else:
			# Fallback: Label SATA drives as SATA_0, SATA_1, SATA_2, ...
			port_counter = 0

			for line in ls_result:
				line = line.strip()
				match = re.search(r"(pci-0000:\S+-ata-(\d+))\s->\s\W+(.*)", line)
				if match:
					path = match.group(1)
					port = match.group(2)
					device = match.group(3).replace("../../", "")

					# Skip partitions
					if re.search(r'\d+$', device):
						continue

					if device not in sata_connections:
						sata_connections[device] = {
							"Path": path,
							"Connector": f"SATA_{port_counter}",
							"Device": device,
							"Partitions": []
						}
						port_counter += 1

			return sata_connections


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

	partitions = []
	for line in lsblk_result:
		regex = re.search(r"^(\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 get_motherboard_model():
	mobo_model = ""
	try:
		dmi_result = subprocess.Popen(["dmidecode","-t","2"],stdout=subprocess.PIPE,universal_newlines=True).stdout
	except:
		#print("ERROR: dmidecode is not installed")
		return False
	if dmi_result is None:
		return False
	for line in dmi_result:
		for field in g_motherboard_sata_addresses.keys():
			regex = re.search(r"^\sProduct Name:\s+({fld})".format(fld=field),line)
			if regex != None:
				mobo_model = field
				break
	return mobo_model


def main():
    motherboard = get_motherboard_model()
    sata = ls_by_path(motherboard)
    if sata:
        for device, connection in sata.items():
            partitions = lsblk(connection["Device"])
            if partitions:
                connection["Partitions"].extend(partitions)  # Append partitions

    if sata:
        output_str = json.dumps({"SATA Info": list(sata.values())}, indent=4)
    else:
        output_str = "{\"SATA Info\":[]}"

    # print output string to stdout
    print(output_str)


if __name__ == "__main__":
	main()