#!/usr/libexec/platform-python
################################################################################
# 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"]
	]
}

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

	sata_connections = []
	if(motherboard in g_motherboard_sata_addresses.keys()):
		paths = g_motherboard_sata_addresses[motherboard]	
		for line in ls_result:
			for path in paths:
				regex = re.search("{pth}(\.\d)?\s->\s\W+(.*)".format(pth=path[0]),line)
				if regex != None:
					sata_connections.append(
						{
							"Path": path[0],
							"Connector": path[1],
							"Device":regex.group(2)
						}
					)
		return sata_connections
	else:
		for line in ls_result:
			regex = re.search("(\S+)-ata-(\d)\s->\s\W+(.*)",line)
			if regex != None:
				sata_connections.append(
						{
							"Path": regex.group(1),
							"Connector":("I-SATA"+regex.group(2)),
							"Device":regex.group(3)
						}
					)
		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

	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 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
	for line in dmi_result:
		for field in g_motherboard_sata_addresses.keys():
			regex = re.search("^\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 connection in sata:
			partitions = lsblk(connection["Device"])
			if partitions:
				connection["Partitions"] = []
				for entry in partitions:
					connection["Partitions"].append(entry)

	if sata:
		output_str = "{\"SATA Info\":["
		for connection in sata:
			output_str += json.dumps(connection) + ","
		output_str = output_str[:-1] + "]}"
	else:
		output_str = "{\"SATA Info\":[]}"

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

if __name__ == "__main__":
    main()