Got disk total size by lsblk

disk.go

package hardware

import (
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

/**
 * $ lsblk -l -b -p
 * NAME      MAJ:MIN RM         SIZE RO TYPE MOUNTPOINT
 * /dev/sda    8:0    0 500107862016  0 disk
 * /dev/sda1   8:1    0   4294967296  0 part /boot
 * /dev/sda2   8:2    0   4294967296  0 part [SWAP]
 * /dev/sda3   8:3    0         1024  0 part
 * /dev/sda5   8:5    0 107374182400  0 part /Data
 * /dev/sda6   8:6    0  53687091200  0 part /
 * /dev/sda7   8:7    0  32212254720  0 part
 * /dev/sda8   8:8    0  30719082496  0 part
 **/

const (
	diskTypeDisk = "disk"
	diskTypePart = "part"
)

type diskCore struct {
	Path string `json:"path"`
	Size int64  `json:"size"`
}

type diskLine struct {
	diskCore
	Type       string
	MountPoint string
}

type DiskPart struct {
	diskCore
	MountPoint string `json:"mountpoint"`
}
type DiskPartList []*DiskPart

type Disk struct {
	diskCore
	PartList DiskPartList `json:"part_list"`
}
type DiskList []*Disk

// RootTotal return the disk which has root total size(kB)
func (dlist DiskList) RootTotal() (int64, error) {
	d := dlist.GetByMountPoint("/")
	if d == nil {
		return 0, fmt.Errorf("not found disk by '/'")
	}
	return d.Size / 1024, nil
}

func (dlist DiskList) GetByMountPoint(mp string) *Disk {
	for _, d := range dlist {
		for _, p := range d.PartList {
			if p.MountPoint == mp {
				return d
			}
		}
	}
	return nil
}

func getDiskTotal() (int64, error) {
	data, err := getLsblkOutputs()
	if err != nil {
		return 0, err
	}
	dlist, err := decodeLsblkOutputs(data)
	if err != nil {
		return 0, err
	}
	return dlist.RootTotal()
}

func decodeLsblkOutputs(data []byte) (DiskList, error) {
	if len(data) == 0 {
		return nil, fmt.Errorf("invalid lsblk data")
	}
	var lines = strings.Split(string(data), "\n")
	var dlist DiskList
	var l = len(lines)
	var idx = 0
	for i := 1; i < l; i++ {
		if len(lines[i]) == 0 {
			continue
		}
		dl, err := decodeLsblkLine(lines[i])
		if err != nil {
			continue
		}
		if dl.Type == diskTypeDisk {
			dlist = append(dlist, &Disk{
				diskCore: dl.diskCore,
				PartList: DiskPartList{},
			})
			idx++
		} else {
			dlist[idx-1].PartList = append(dlist[idx-1].PartList, &DiskPart{
				diskCore:   dl.diskCore,
				MountPoint: dl.MountPoint,
			})
		}
	}

	if len(dlist) == 0 {
		return nil, fmt.Errorf("invalid lsblk data")
	}
	return dlist, nil
}

func decodeLsblkLine(line string) (*diskLine, error) {
	list := strings.Split(line, " ")
	list = filterEmptyItem(list)
	if len(list) < 6 {
		return nil, fmt.Errorf("invalid line: %s", line)
	}
	size, err := strconv.ParseInt(list[3], 10, 64)
	if err != nil {
		return nil, err
	}
	var mp string
	if len(list) == 7 {
		mp = list[6]
	}
	return &diskLine{
		diskCore: diskCore{
			Path: list[0],
			Size: size,
		},
		Type:       list[5],
		MountPoint: mp,
	}, nil
}

func getLsblkOutputs() ([]byte, error) {
	return exec.Command("/bin/sh", "-c", "lsblk -blp").CombinedOutput()
}

func filterEmptyItem(list []string) []string {
	var ret []string
	for _, v := range list {
		if len(v) == 0 {
			continue
		}
		ret = append(ret, v)
	}
	return ret
}