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"
	"unsafe"
)

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"` // Reported in Kelvin by smartctl NVMe JSON
		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
}

const BytesPerGByte = 1000 * 1000 * 1000

func formatBytesToGB(bytes uint64) string {
	gb := float64(bytes) / BytesPerGByte
	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))

		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 {
		mountPoints := strings.Split(info.MountPoint, ", ")
		var total, free, used uint64

		for _, mp := range mountPoints {
			u, err := GetDiskUsage(mp)
			if err == nil {
				total += u.TotalBytes
				free += u.FreeBytes
				used += u.UsedBytes
			}
		}

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

		usage = DiskUsage{
			TotalBytes:   total,
			FreeBytes:    free,
			UsedBytes:    used,
			UsagePercent: percent,
		}
	} else {
		u, err := GetRawBlockDeviceUsage(devicePath)
		if err == nil {
			usage = *u
		}
	}

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

func GetRawBlockDeviceUsage(devicePath string) (*DiskUsage, error) {
	file, err := os.Open(devicePath)
	if err != nil {
		return nil, fmt.Errorf("failed to open device %s: %w", devicePath, err)
	}
	defer file.Close()

	var size uint64
	const BLKGETSIZE64 = 0x80081272
	_, _, sysErr := syscall.Syscall(
		syscall.SYS_IOCTL,
		file.Fd(),
		uintptr(BLKGETSIZE64),
		uintptr(unsafe.Pointer(&size)),
	)

	if sysErr != 0 {
		return nil, fmt.Errorf("ioctl BLKGETSIZE64 failed: %v", sysErr)
	}

	return &DiskUsage{
		TotalBytes:   size,
		FreeBytes:    0,
		UsedBytes:    0,
		UsagePercent: 0.0,
	}, 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)),
	}

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

	info.MountPoint = strings.Join(mountPoints, ", ")
	info.IsMounted = isMounted

	return info, nil
}

func findAllMountPoints(devicePath string) ([]string, bool, error) {
	diskDevNum, err := getDeviceNumber(devicePath)
	if err != nil {
		return nil, false, err
	}

	// Map of device numbers (e.g. "8:0", "8:3", "252:6") belonging to this drive or its partitions/sub-devices
	associatedDevNums := getAssociatedDeviceNumbers(diskDevNum)

	data, err := os.ReadFile("/proc/mounts")
	if err != nil {
		return nil, false, err
	}

	var mountPoints []string
	seenMounts := make(map[string]bool)

	lines := strings.Split(string(data), "\n")
	for _, line := range lines {
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}

		mountedDev := fields[0]
		mountPoint := fields[1]

		// Skip non-device pseudo filesystems (sysfs, proc, tmpfs, etc.)
		if !strings.HasPrefix(mountedDev, "/dev/") {
			continue
		}

		mountedDevNum, err := getDeviceNumber(mountedDev)
		if err != nil {
			continue
		}

		if associatedDevNums[mountedDevNum] {
			if !seenMounts[mountPoint] {
				seenMounts[mountPoint] = true
				mountPoints = append(mountPoints, mountPoint)
			}
		}
	}

	return mountPoints, len(mountPoints) > 0, nil
}

func getDeviceNumber(devPath string) (string, error) {
	resolvedPath, err := filepath.EvalSymlinks(devPath)
	if err != nil {
		resolvedPath = devPath
	}

	var stat syscall.Stat_t
	if err := syscall.Stat(resolvedPath, &stat); err != nil {
		return "", err
	}

	rdev := uint64(stat.Rdev)
	major := (rdev >> 8) & 0xfff
	minor := (rdev & 0xff) | ((rdev >> 12) & 0xfff00)

	return fmt.Sprintf("%d:%d", major, minor), nil
}

func getAssociatedDeviceNumbers(diskDevNum string) map[string]bool {
	devNums := map[string]bool{diskDevNum: true}

	// Read /sys/dev/block/<major:minor> to locate sysfs device node
	sysDevPath, err := filepath.EvalSymlinks(filepath.Join("/sys/dev/block", diskDevNum))
	if err != nil {
		return devNums
	}

	// Recursively inspect sysfs entries for partition subdirectories and DM/LVM holders
	var collectDevs func(dir string)
	collectDevs = func(dir string) {
		entries, err := os.ReadDir(dir)
		if err != nil {
			return
		}

		for _, entry := range entries {
			devFile := filepath.Join(dir, entry.Name(), "dev")
			if devData, err := os.ReadFile(devFile); err == nil {
				num := strings.TrimSpace(string(devData))
				if num != "" && !devNums[num] {
					devNums[num] = true
					// Recurse into sysfs holders subdirectory to follow LVM / Device Mapper relationships
					collectDevs(filepath.Join(dir, entry.Name(), "holders"))
				}
			}
			if entry.Name() == "holders" {
				collectDevs(filepath.Join(dir, entry.Name()))
			}
		}
	}

	collectDevs(sysDevPath)
	return devNums
}

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

		tempC := log.Temperature
		if smartData.Temperature.Current > 0 {
			tempC = smartData.Temperature.Current
		} else if tempC > 200 {
			tempC -= 273
		}

		smart.TemperatureC = tempC
		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 {
			// 245: Intel DC
			case 231, 177, 233, 169, 230, 173:
				smart.HealthRemaining = attr.Value
				smart.WearoutPercent = max(100-attr.Value, 0)
				smart.HasWearData = true
			case 202, 245:
				smart.WearoutPercent = attr.Value
				smart.HealthRemaining = max(100-attr.Value, 0)
				smart.HasWearData = true
			}
			if smart.HasWearData {
				break
			}
		}
	}

	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") || strings.HasPrefix(name, "dm-") || strings.HasPrefix(name, "nbd") {
			continue
		}

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

	return devices, nil
}

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

	_ = cmd.Run()

	if out.Len() == 0 {
		return nil, fmt.Errorf("smartctl returned empty output for %s", devicePath)
	}

	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
}