golang实现windows获取加密盘符的总大小

golang实现windows获取加密盘符的总大小

package main

import (
	"fmt"
	"syscall"
	"unsafe"
)

type PartitionStyle uint32

const (
	IOCTL_DISK_GET_DRIVE_LAYOUT_EX         = 0x00070050
	FILE_DEVICE_MASS_STORAGE        uint32 = 0x0000002d
	IOCTL_STORAGE_BASE              uint32 = FILE_DEVICE_MASS_STORAGE
	FILE_ANY_ACCESS                 uint16 = 0
	FILE_SPECIAL_ACCESS             uint16 = FILE_ANY_ACCESS
	FILE_READ_ACCESS                uint16 = 0x0001
	FILE_WRITE_ACCESS               uint16 = 0x0002
	METHOD_BUFFERED                 uint8  = 0
	METHOD_IN_DIRECT                uint8  = 1
	METHOD_OUT_DIRECT               uint8  = 2
	METHOD_NEITHER                  uint8  = 3
	IOCTL_STORAGE_GET_DEVICE_NUMBER uint32 = (IOCTL_STORAGE_BASE << 16) | uint32(FILE_ANY_ACCESS<<14) | uint32(0x0420<<2) | uint32(METHOD_BUFFERED)

	PartitionStyleMbr PartitionStyle = 0
	PartitionStyleGpt PartitionStyle = 1
	PartitionStyleRaw PartitionStyle = 2
	FILE_DEVICE_DISK  uint32         = 0x7
)

type GUID struct {
	Data1 uint32
	Data2 uint16
	Data3 uint16
	Data4 [8]byte
}

type DRIVE_LAYOUT_INFORMATION_GPT struct {
	DiskId               GUID
	StartingUsableOffset uint64
	UsableLength         uint64
	MaxPartitionCount    uint32
}

type PARTITION_INFORMATION_MBR struct {
	PartitionType       byte
	BootIndicator       bool
	RecognizedPartition bool
	HiddenSectors       uint32
	PartitionId         GUID
}

type PARTITION_INFORMATION_GPT struct {
	PartitionType GUID
	PartitionId   GUID
	Attributes    uint64
	Name          [36]uint16
}

type PARTITION_INFORMATION_EX struct {
	PartitionStyle   PartitionStyle
	StartingOffset   int64
	PartitionLength  int64
	DeviceNumber     int32
	RewritePartition bool
	Rev01            bool
	Rev02            bool
	Rev03            bool
	PartitionInfo    [112]byte
}
type DRIVE_LAYOUT_INFORMATION_MBR struct {
	Signature uint32
	CheckSum  uint32
}

type DRIVE_LAYOUT_INFORMATION_EX_HEADER struct {
	PartitionStyle PartitionStyle
	PartitionCount uint32
}

func getDiskHandleByNum(num uint32) (syscall.Handle, error) {
	diskName := fmt.Sprintf(`\\.\PhysicalDrive%d`, num)
	disk, _ := syscall.UTF16PtrFromString(diskName)
	handle, err := syscall.CreateFile(
		disk,
		syscall.GENERIC_READ,
		syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE,
		nil,
		syscall.OPEN_EXISTING,
		0,
		0,
	)
	return handle, err
}
func GetSizeOf_DRIVE_LAYOUT_INFORMATION() int {
	a := unsafe.Sizeof(DRIVE_LAYOUT_INFORMATION_GPT{})
	b := unsafe.Sizeof(DRIVE_LAYOUT_INFORMATION_MBR{})
	if a > b {
		return int(a)
	} else {
		return int(b)
	}
}
func getAllPartitionInfo(diskHandle syscall.Handle) ([]byte, error) {
	var bytesReturned uint32
	buffer := make([]byte, 4096)
	err := syscall.DeviceIoControl(diskHandle, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, nil, 0, &buffer[0], uint32(len(buffer)), &bytesReturned, nil)
	if err != nil {
		return nil, err
	}
	return buffer, nil
}

type STORAGE_DEVICE_NUMBER struct {
	DeviceType      uint32
	DeviceNumber    uint32
	PartitionNumber uint32
}

// formatFileSize 将文件大小转换为易读的格式
func FormatFileSize(size int64) string {
	// 定义文件大小单位
	units := []string{"B", "KB", "MB", "GB", "TB", "PB"}

	// 处理文件大小为0的情况
	if size == 0 {
		return "0 B"
	}

	// 计算文件大小所在单位的索引
	unitIndex := 0
	for size >= 1024 && unitIndex < len(units)-1 {
		size /= 1024
		unitIndex++
	}

	// 格式化文件大小
	return fmt.Sprintf("%d%s", size, units[unitIndex])
}
func GetDriveBasicInfo(drive string) (STORAGE_DEVICE_NUMBER, error) {
	var disk_num STORAGE_DEVICE_NUMBER
	var err error
	filepath, _ := syscall.UTF16PtrFromString(`\\.\` + drive + ":")
	handle, err := syscall.CreateFile(
		filepath,
		syscall.GENERIC_READ,
		syscall.FILE_SHARE_READ|syscall.FILE_SHARE_WRITE,
		nil,
		syscall.OPEN_EXISTING,
		0,
		0)
	if ^uintptr(0) == uintptr(handle) {
		fmt.Printf("CreateFile() failed, errmsg = %s\n", err.Error())
		return disk_num, nil
	}
	var size uint32 = uint32(unsafe.Sizeof(disk_num))
	var ret_size uint32 = 0
	var outbuf *byte = (*byte)(unsafe.Pointer(&disk_num))
	syscall.DeviceIoControl(
		handle,
		IOCTL_STORAGE_GET_DEVICE_NUMBER,
		nil, 0,
		outbuf, size,
		&ret_size, nil)
	syscall.CloseHandle(handle)
	return disk_num, nil
}
func GetDriveTotal(drive string) int64 {
	dinfo, err := GetDriveBasicInfo(drive)
	if err != nil {
		fmt.Println("dinfo", dinfo, err)
		return 0
	}
	DeviceNumber := dinfo.DeviceNumber
	disk, err := getDiskHandleByNum(DeviceNumber)
	if err != nil {
		if err == syscall.ERROR_FILE_NOT_FOUND {
			// 物理磁盘号不存在，结束枚举
			fmt.Println("err", err)
			return 0
		}
	}
	defer syscall.CloseHandle(disk)

	data, err := getAllPartitionInfo(disk)
	if err != nil {
		fmt.Errorf("Failed to get partition info: %v\n", err)
		return 0
	}
	header := (*DRIVE_LAYOUT_INFORMATION_EX_HEADER)(unsafe.Pointer(&data[0]))
	next := data[int(unsafe.Sizeof(*header)):]
	entryOffset := GetSizeOf_DRIVE_LAYOUT_INFORMATION()
	entryData := next[entryOffset:]
	entrySize := unsafe.Sizeof(PARTITION_INFORMATION_EX{})
	for i := 0; i < int(header.PartitionCount); i++ {
		if len(entryData) < int(entrySize) {
			break
		}
		partitionEntry := (*PARTITION_INFORMATION_EX)(unsafe.Pointer(&entryData[0]))
		entryData = entryData[entrySize:]
		if partitionEntry.DeviceNumber == int32(dinfo.PartitionNumber) {
			return partitionEntry.PartitionLength
		}
	}
	return 0
}

func main() {
	total := GetDriveTotal("C")
	fmt.Println("total", total, FormatFileSize(total))

}