tyto/backend/internal/collectors/disk.go

package collectors

import (
	"bufio"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"tyto/internal/models"
)

type DiskCollector struct {
	procPath string
	mtabPath string
}

func NewDiskCollector(procPath, mtabPath string) *DiskCollector {
	return &DiskCollector{
		procPath: procPath,
		mtabPath: mtabPath,
	}
}

func (c *DiskCollector) Collect() (models.DiskStats, error) {
	stats := models.DiskStats{
		Mounts: []models.MountStats{},
		IO:     []models.DiskIO{},
	}

	// Get mount points
	mounts, err := c.getMounts()
	if err == nil && mounts != nil {
		stats.Mounts = mounts
	}

	// Get I/O stats
	io, err := c.getIOStats()
	if err == nil && io != nil {
		stats.IO = io
	}

	return stats, nil
}

func (c *DiskCollector) getMounts() ([]models.MountStats, error) {
	file, err := os.Open(c.mtabPath)
	if err != nil {
		// Fallback to /proc/mounts
		file, err = os.Open(filepath.Join(c.procPath, "mounts"))
		if err != nil {
			return nil, err
		}
	}
	defer file.Close()

	var mounts []models.MountStats
	scanner := bufio.NewScanner(file)
	seen := make(map[string]bool)

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if len(fields) < 3 {
			continue
		}

		device := fields[0]
		mountPoint := fields[1]
		fsType := fields[2]

		// Skip virtual filesystems and duplicates
		if !strings.HasPrefix(device, "/dev/") {
			continue
		}
		if seen[mountPoint] {
			continue
		}
		seen[mountPoint] = true

		// Skip certain filesystem types
		if fsType == "squashfs" || fsType == "overlay" {
			continue
		}

		var stat syscall.Statfs_t
		if err := syscall.Statfs(mountPoint, &stat); err != nil {
			continue
		}

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

		usedPercent := float64(0)
		if total > 0 {
			usedPercent = float64(used) / float64(total) * 100
		}

		mounts = append(mounts, models.MountStats{
			Device:      device,
			MountPoint:  mountPoint,
			Filesystem:  fsType,
			Total:       total,
			Used:        used,
			Available:   available,
			UsedPercent: usedPercent,
		})
	}

	return mounts, nil
}

func (c *DiskCollector) getIOStats() ([]models.DiskIO, error) {
	file, err := os.Open(filepath.Join(c.procPath, "diskstats"))
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var stats []models.DiskIO
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if len(fields) < 14 {
			continue
		}

		device := fields[2]

		// Only include real block devices (sd*, nvme*, vd*)
		if !strings.HasPrefix(device, "sd") &&
		   !strings.HasPrefix(device, "nvme") &&
		   !strings.HasPrefix(device, "vd") {
			continue
		}

		// Skip partitions (e.g., sda1, nvme0n1p1)
		if strings.ContainsAny(device[len(device)-1:], "0123456789") {
			lastChar := device[len(device)-1]
			if lastChar >= '0' && lastChar <= '9' {
				// Check if it's a partition
				if strings.Contains(device, "nvme") {
					if strings.Contains(device, "p") {
						continue
					}
				} else {
					continue
				}
			}
		}

		// Sectors read/written (field 5 and 9), sector = 512 bytes
		readSectors, _ := strconv.ParseUint(fields[5], 10, 64)
		writeSectors, _ := strconv.ParseUint(fields[9], 10, 64)

		stats = append(stats, models.DiskIO{
			Device:     device,
			ReadBytes:  readSectors * 512,
			WriteBytes: writeSectors * 512,
		})
	}

	return stats, nil
}