arozos/mod/disk/diskfs/diskfs.go

package diskfs

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"regexp"
	"strings"

	"imuslab.com/arozos/mod/utils"
)

/*
	diskfs.go

	This module handle file system creation and formatting
*/

// Storage Device meta was generated by lsblk
// Partitions like sdX0
type PartitionMeta struct {
	Name       string `json:"name"`
	MajMin     string `json:"maj:min"`
	Rm         bool   `json:"rm"`
	Size       int64  `json:"size"`
	Ro         bool   `json:"ro"`
	Type       string `json:"type"`
	Mountpoint string `json:"mountpoint"`
}

// Block device, usually disk or rom, like sdX
type BlockDeviceMeta struct {
	Name       string          `json:"name"`
	MajMin     string          `json:"maj:min"`
	Rm         bool            `json:"rm"`
	Size       int64           `json:"size"`
	Ro         bool            `json:"ro"`
	Type       string          `json:"type"`
	Mountpoint string          `json:"mountpoint"`
	Children   []PartitionMeta `json:"children,omitempty"`
}

// A collection of information for lsblk output
type StorageDevicesMeta struct {
	Blockdevices []BlockDeviceMeta `json:"blockdevices"`
}

// Check if the file format driver is installed on this host
// if a format is supported, mkfs.(format) should be symlinked under /sbin
func FormatPackageInstalled(fsType string) bool {
	return utils.FileExists("/sbin/mkfs." + fsType)
}

// Create file system, support ntfs, ext4 and fat32 only
func FormatStorageDevice(fsType string, devicePath string) error {
	// Check if the filesystem type is supported
	switch fsType {
	case "ext4":
		// Format the device with the specified filesystem type
		cmd := exec.Command("sudo", "mkfs."+fsType, devicePath)
		output, err := cmd.CombinedOutput()
		if err != nil {
			return errors.New("unable to format device: " + string(output))
		}
		return nil

	case "vfat", "fat", "fat32":
		//Check if mkfs.fat exists
		if !FormatPackageInstalled("vfat") {
			return errors.New("unable to format device as fat (vfat). dosfstools not installed?")
		}

		// Format the device with the specified filesystem type
		cmd := exec.Command("sudo", "mkfs.vfat", devicePath)
		output, err := cmd.CombinedOutput()
		if err != nil {
			return errors.New("unable to format device: " + string(output))
		}
		return nil

	case "ntfs":
		//Check if ntfs-3g exists
		if !FormatPackageInstalled("ntfs") {
			return errors.New("unable to format device as ntfs: ntfs-3g not installed?")
		}

		//Format the drive
		cmd := exec.Command("sudo", "mkfs.ntfs", devicePath)
		output, err := cmd.CombinedOutput()
		if err != nil {
			return errors.New("unable to format device: " + string(output))
		}
		return nil

	default:
		return fmt.Errorf("unsupported filesystem type: %s", fsType)
	}
}

// List all the storage device in the system, set minSize to 0 for no filter
func ListAllStorageDevices() (*StorageDevicesMeta, error) {
	cmd := exec.Command("sudo", "lsblk", "-b", "--json")

	output, err := cmd.CombinedOutput()
	if err != nil {
		return nil, fmt.Errorf("lsblk error: %v", err)
	}

	var devices StorageDevicesMeta
	err = json.Unmarshal([]byte(output), &devices)
	return &devices, err
}

// Get block device (e.g. /dev/sdX) info
func GetBlockDeviceMeta(devicePath string) (*BlockDeviceMeta, error) {
	//Trim the /dev/ part of the device path
	deviceName := strings.TrimPrefix(devicePath, "/dev/")

	if len(deviceName) == 0 {
		return nil, errors.New("invalid device path given")
	}

	re := regexp.MustCompile(`\d+`)
	if re.MatchString(deviceName) {
		//This is a partition
		return nil, errors.New("given device path is a partition not a block device")
	}

	storageMeta, err := ListAllStorageDevices()
	if err != nil {
		return nil, err
	}

	for _, blockdevice := range storageMeta.Blockdevices {
		if blockdevice.Name == deviceName {
			return &blockdevice, nil
		}
	}

	return nil, errors.New("target block device not found")
}

// Get partition information (e.g. /dev/sdX1)
func GetPartitionMeta(devicePath string) (*PartitionMeta, error) {
	//Trim the /dev/ part of the device path
	deviceName := strings.TrimPrefix(devicePath, "/dev/")

	if len(deviceName) == 0 {
		return nil, errors.New("invalid device path given")
	}

	re := regexp.MustCompile(`\d+`)
	if !re.MatchString(deviceName) {
		//This is a partition
		return nil, errors.New("given device path is a block device not a partition")
	}

	storageMeta, err := ListAllStorageDevices()
	if err != nil {
		return nil, err
	}

	for _, blockdevice := range storageMeta.Blockdevices {
		if strings.Contains(deviceName, blockdevice.Name) {
			//Matching block device. Check for if there are a matching child
			for _, childPartition := range blockdevice.Children {
				if childPartition.Name == deviceName {
					return &childPartition, nil
				}
			}
		}
	}

	return nil, errors.New("target partition not found")
}

// Check if a device is mounted given the path name, like /dev/sdc
func DeviceIsMounted(devicePath string) (bool, error) {
	// Open the mountinfo file
	file, err := os.Open("/proc/mounts")
	if err != nil {
		return false, fmt.Errorf("error opening /proc/mounts: %v", err)
	}
	defer file.Close()

	// Scan the mountinfo file line by line
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Fields(line)
		if len(fields) >= 2 && fields[0] == devicePath {
			// Device is mounted
			return true, nil
		}
	}

	// Device is not mounted
	return false, nil
}

// UnmountDevice unmounts the specified device.
// Remember to use full path (e.g. /dev/md0) in the devicePath
func UnmountDevice(devicePath string) error {
	// Construct the bash command to unmount the device
	cmd := exec.Command("sudo", "bash", "-c", fmt.Sprintf("umount %s", devicePath))

	// Run the command
	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("error unmounting device: %v", err)
	}

	return nil
}

// Force Version of Unmount (dangerous)
// Remember to use full path (e.g. /dev/md0) in the devicePath
func ForceUnmountDevice(devicePath string) error {
	// Construct the bash command to unmount the device
	cmd := exec.Command("sudo", "bash", "-c", fmt.Sprintf("umount -l %s", devicePath))

	// Run the command
	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("error unmounting device: %v", err)
	}

	return nil
}

// DANGER: Wipe the whole disk given the disk path
func WipeDisk(diskPath string) error {
	// Unmount the disk
	isMounted, _ := DeviceIsMounted(diskPath)
	if isMounted {
		umountCmd := exec.Command("sudo", "umount", diskPath)
		if err := umountCmd.Run(); err != nil {
			return fmt.Errorf("error unmounting disk %s: %v", diskPath, err)
		}
	}

	// Wipe all filesystem signatures on the entire disk
	wipeCmd := exec.Command("sudo", "wipefs", "--all", "--force", diskPath)
	if err := wipeCmd.Run(); err != nil {
		return fmt.Errorf("error wiping filesystem signatures on %s: %v", diskPath, err)
	}

	return nil
}