finally-a-monolithic-linux-hw-monitor/app/disks/drivers/smart.go

package drivers

import (
	"bytes"
	"encoding/json"
	"finally-a-monolithic-linux-hw-monitor/app/common"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
)

type SMARTDiskReport struct {
	Info  DiskInfo
	Usage DiskUsage
	SMART DiskSMART
}

// Internal JSON mapping matching smartctl 7.x+ output schemas
type smartctlJSON struct {
	Device struct {
		Protocol  string `json:"protocol"`
		ModelName string `json:"model_name"`
	} `json:"device"`
	Smartctl struct {
		ExitStatus int      `json:"exit_status"`
		Messages   []string `json:"messages"`
	} `json:"smartctl"`
	SmartStatus struct {
		Passed bool `json:"passed"`
	} `json:"smart_status"`

	// ATA / SATA Root Metrics
	PowerOnTime struct {
		Hours int `json:"hours"`
	} `json:"power_on_time"`
	PowerCycleCount int `json:"power_cycle_count"`
	Temperature     struct {
		Current int `json:"current"`
	} `json:"temperature"`
	ModelName    string `json:"model_name"`
	SerialNumber string `json:"serial_number"`
	Protocol     string `json:"protocol"`

	// NVMe-Specific Health Log Metrics
	NVMeHealthLog struct {
		PercentageUsed int `json:"percentage_used"`
		Temperature    int `json:"temperature"`
		PowerOnHours   int `json:"power_on_hours"`
		PowerCycles    int `json:"power_cycles"`
	} `json:"nvme_smart_health_information_log"`

	ATAMetrics struct {
		Table []struct {
			ID    int `json:"id"`
			Value int `json:"value"`
		} `json:"table"`
	} `json:"ata_smart_attributes"`
}

type SMARTDriver struct{}

func NewSMARTDriver() (*SMARTDriver, error) {
	r := SMARTDriver{}
	_, err := r.Read()
	return &r, err
}

// Helper function to convert raw byte counts to decimal Gigabytes (GB)
func formatBytesToGB(bytes uint64) string {
	gb := float64(bytes) / (1000 * 1000 * 1000)
	return fmt.Sprintf("%.2f GB", gb)
}

func (r *SMARTDriver) Render(term common.Terminal) error {
	d, err := r.Read()
	if err != nil {
		return err
	}

	disks := d.([]SMARTDiskReport)
	term.Subheader("Disks")

	for _, disk := range disks {
		term.Print(fmt.Sprintf("    Disk %s    ", disk.Info.DevicePath))
		term.Print(fmt.Sprintf("%-16s | %-16s", "Model Name", disk.Info.Model))
		term.Print(fmt.Sprintf("%-16s | %-16s", "Model Serial", disk.Info.Serial))
		term.Print(fmt.Sprintf("%-16s | %-16s", "Protocol", disk.Info.Protocol))
		var mp string
		if disk.Info.IsMounted {
			mp = disk.Info.MountPoint
		} else {
			mp = "Unmounted"
		}
		term.Print(fmt.Sprintf("%-16s | %-16s", "Mount Point", mp))

		// Formatted bytes in GB
		term.Print(fmt.Sprintf("%-16s | %-16s", "Total Capacity", formatBytesToGB(disk.Usage.TotalBytes)))
		term.Print(fmt.Sprintf("%-16s | %-16s", "Free Space", formatBytesToGB(disk.Usage.FreeBytes)))
		term.Print(fmt.Sprintf("%-16s | %-16s (%.2f%%)", "Used Space", formatBytesToGB(disk.Usage.UsedBytes), disk.Usage.UsagePercent))

		term.Print(fmt.Sprintf("%-16s | %t", "SMART Passed", disk.SMART.Passed))
		term.Print(fmt.Sprintf("%-16s | %d%%", "Wearout", disk.SMART.WearoutPercent))
		term.Print(fmt.Sprintf("%-16s | %d%%", "Health Remaining", disk.SMART.HealthRemaining))
		term.Print(fmt.Sprintf("%-16s | %d C", "Temperature", disk.SMART.TemperatureC))
		term.Print(fmt.Sprintf("%-16s | %d hrs", "Power-On Hours", disk.SMART.PowerOnHours))
	}

	return nil
}

func (r *SMARTDriver) Read() (any, error) {
	devicePaths, err := GetBlockDevices()
	if err != nil {
		return nil, err
	}

	disks := []SMARTDiskReport{}

	for _, devicePath := range devicePaths {
		disk, err := GetDiskReport(devicePath)
		if err != nil {
			continue
		}
		disks = append(disks, disk)
	}

	return disks, nil
}

func GetDiskReport(devicePath string) (SMARTDiskReport, error) {
	info, err := GetDiskInfo(devicePath)
	if err != nil {
		return SMARTDiskReport{}, fmt.Errorf("failed to get disk info for %s: %w", devicePath, err)
	}

	smart, err := GetDiskSMART(devicePath)
	if err != nil {
		return SMARTDiskReport{}, fmt.Errorf("failed to get SMART data for %s: %w", devicePath, err)
	}

	var usage DiskUsage
	if info.IsMounted {
		u, err := GetDiskUsage(info.MountPoint)
		if err == nil {
			usage = *u
		}
	}

	return SMARTDiskReport{
		Info:  *info,
		Usage: usage,
		SMART: *smart,
	}, nil
}

func extractProtocol(smartData *smartctlJSON) string {
	if smartData.Protocol != "" {
		return smartData.Protocol
	}
	return smartData.Device.Protocol
}

func extractModel(smartData *smartctlJSON) string {
	if smartData.ModelName != "" {
		return smartData.ModelName
	}
	return smartData.Device.ModelName
}

func GetDiskInfo(devicePath string) (*DiskInfo, error) {
	smartData, err := execSmartctl(devicePath)
	if err != nil {
		return nil, fmt.Errorf("failed to query device info: %w", err)
	}

	info := &DiskInfo{
		DevicePath: devicePath,
		Model:      extractModel(smartData),
		Serial:     smartData.SerialNumber,
		Protocol:   strings.ToUpper(extractProtocol(smartData)),
	}

	mountPoint, isMounted, err := findMountPoint(devicePath)
	if err != nil {
		return nil, fmt.Errorf("failed to check mount status: %w", err)
	}

	info.MountPoint = mountPoint
	info.IsMounted = isMounted

	return info, nil
}

func GetDiskUsage(path string) (*DiskUsage, error) {
	var stat syscall.Statfs_t
	if err := syscall.Statfs(path, &stat); err != nil {
		return nil, fmt.Errorf("statfs failed for path %s: %w", path, err)
	}

	total := stat.Blocks * uint64(stat.Bsize)
	free := stat.Bfree * uint64(stat.Bsize)
	used := total - free

	var percent float64
	if total > 0 {
		percent = (float64(used) / float64(total)) * 100.0
	}

	return &DiskUsage{
		TotalBytes:   total,
		FreeBytes:    free,
		UsedBytes:    used,
		UsagePercent: percent,
	}, nil
}

func GetDiskSMART(devicePath string) (*DiskSMART, error) {
	smartData, err := execSmartctl(devicePath)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch SMART data: %w", err)
	}

	smart := &DiskSMART{
		Passed: smartData.SmartStatus.Passed,
	}

	protocol := strings.ToUpper(extractProtocol(smartData))

	if protocol == "NVME" || strings.Contains(strings.ToLower(devicePath), "nvme") {
		log := smartData.NVMeHealthLog

		smart.TemperatureC = log.Temperature
		smart.PowerOnHours = log.PowerOnHours
		smart.PowerCycleCount = log.PowerCycles

		smart.WearoutPercent = log.PercentageUsed
		smart.HealthRemaining = 100 - log.PercentageUsed
		if smart.HealthRemaining < 0 {
			smart.HealthRemaining = 0
		}
		smart.HasWearData = true
	} else {
		smart.TemperatureC = smartData.Temperature.Current
		smart.PowerOnHours = smartData.PowerOnTime.Hours
		smart.PowerCycleCount = smartData.PowerCycleCount

		for _, attr := range smartData.ATAMetrics.Table {
			switch attr.ID {
			// Attributes that represent Health Remaining (100 -> 0)
			case 231, 177, 233, 169, 230, 173:
				smart.HealthRemaining = attr.Value
				smart.WearoutPercent = max(100-attr.Value, 0)
				smart.HasWearData = true

			// Attributes that represent Percent Used / Wearout (0 -> 100+)
			case 202:
				smart.WearoutPercent = attr.Value
				smart.HealthRemaining = max(100-attr.Value, 0)
				smart.HasWearData = true
			}

			if smart.HasWearData {
				break // Exits the for-loop once a matching attribute is found
			}
		}
	}

	return smart, nil
}

func GetBlockDevices() ([]string, error) {
	entries, err := os.ReadDir("/sys/block")
	if err != nil {
		return nil, err
	}

	var devices []string
	for _, entry := range entries {
		name := entry.Name()

		if strings.HasPrefix(name, "loop") || strings.HasPrefix(name, "ram") || strings.HasPrefix(name, "sr") {
			continue
		}

		devPath := filepath.Join("/dev", name)
		if _, err := os.Stat(devPath); err == nil {
			devices = append(devices, devPath)
		}
	}

	return devices, nil
}

func findMountPoint(devicePath string) (string, bool, error) {
	data, err := os.ReadFile("/proc/mounts")
	if err != nil {
		return "", false, err
	}

	lines := strings.Split(string(data), "\n")
	for _, line := range lines {
		fields := strings.Fields(line)
		if len(fields) >= 2 {
			if fields[0] == devicePath || strings.HasPrefix(fields[0], devicePath) {
				return fields[1], true, nil
			}
		}
	}

	return "", false, nil
}

func execSmartctl(devicePath string) (*smartctlJSON, error) {
	cmd := exec.Command("smartctl", "-a", "-j", devicePath)
	var out bytes.Buffer
	cmd.Stdout = &out

	_ = cmd.Run()

	var result smartctlJSON
	if err := json.Unmarshal(out.Bytes(), &result); err != nil {
		return nil, fmt.Errorf("failed to parse smartctl json output for %s: %w", devicePath, err)
	}

	return &result, nil
}