import json
import subprocess

def getDevices():
    result = subprocess.run(['smartctl', '--scan-open', '-j'], stdout=subprocess.PIPE)
    result = result.stdout.decode('utf-8')

    devices = json.loads(result)

    return devices

def getAttributes(device):
    attributes = {}
    attributes["data"] = {}

    result = subprocess.run(['smartctl', device["name"], '-a', '-j'], stdout=subprocess.PIPE)
    result = result.stdout.decode('utf-8')
    data = json.loads(result)

    if "NVMe" == device["protocol"]:
        attributes["sector_size"] = data["nvme_namespaces"][0]["formatted_lba_size"]
        attributes["serial_number"] = data["serial_number"]
        attributes["type"] = "NVME"
        for attribute in data["nvme_smart_health_information_log"]:
            attributes["data"][attribute] = data["nvme_smart_health_information_log"][attribute]
        attributes["bytes_written"] = attributes["data"]["data_units_written"] * attributes["sector_size"] * 1000
    elif "ATA" == device["protocol"]:
        attributes["sector_size"] = data["physical_block_size"]
        attributes["serial_number"] = data["serial_number"]
        attributes["type"] = "ATA"
        for attribute in data["ata_smart_attributes"]["table"]:
            attributes["data"][attribute["name"]] = {}
            try:
                attributes["data"][attribute["name"]]["raw"] = int(attribute["raw"]["string"].split(' ')[0])
            except:
                attributes["data"][attribute["name"]]["raw"] = -1
            attributes["data"][attribute["name"]]["id"] = int(attribute["id"])
            attributes["data"][attribute["name"]]["value"] = int(attribute["value"])
            attributes["data"][attribute["name"]]["worst"] = int(attribute["worst"])
            attributes["data"][attribute["name"]]["thr"] = int(attribute["thresh"])
        try:
            attributes["bytes_written"] = attributes["data"]["Total_LBAs_Written"]["raw"] * attributes["sector_size"]
        except:
            attributes["bytes_written"] = -1
    else:
        pass

    return attributes

def getAllDeviceAttributes():
    devices = getDevices()
    attributes = {}

    for device in devices["devices"]:
        attributes[device["name"]] = getAttributes(device)

    return attributes

if __name__ == "__main__":
    smart = getAllDeviceAttributes()

    print(smart)

    for device in smart:
        print(smart[device]["serial_number"])