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

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

package mainimport ("fmt""syscall""unsafe"
)type PartitionStyle uint32const (IOCTL_DISK_GET_DRIVE_LAYOUT_EX         = 0x00070050FILE_DEVICE_MASS_STORAGE        uint32 = 0x0000002dIOCTL_STORAGE_BASE              uint32 = FILE_DEVICE_MASS_STORAGEFILE_ANY_ACCESS                 uint16 = 0FILE_SPECIAL_ACCESS             uint16 = FILE_ANY_ACCESSFILE_READ_ACCESS                uint16 = 0x0001FILE_WRITE_ACCESS               uint16 = 0x0002METHOD_BUFFERED                 uint8  = 0METHOD_IN_DIRECT                uint8  = 1METHOD_OUT_DIRECT               uint8  = 2METHOD_NEITHER                  uint8  = 3IOCTL_STORAGE_GET_DEVICE_NUMBER uint32 = (IOCTL_STORAGE_BASE << 16) | uint32(FILE_ANY_ACCESS<<14) | uint32(0x0420<<2) | uint32(METHOD_BUFFERED)PartitionStyleMbr PartitionStyle = 0PartitionStyleGpt PartitionStyle = 1PartitionStyleRaw PartitionStyle = 2FILE_DEVICE_DISK  uint32         = 0x7
)type GUID struct {Data1 uint32Data2 uint16Data3 uint16Data4 [8]byte
}type DRIVE_LAYOUT_INFORMATION_GPT struct {DiskId               GUIDStartingUsableOffset uint64UsableLength         uint64MaxPartitionCount    uint32
}type PARTITION_INFORMATION_MBR struct {PartitionType       byteBootIndicator       boolRecognizedPartition boolHiddenSectors       uint32PartitionId         GUID
}type PARTITION_INFORMATION_GPT struct {PartitionType GUIDPartitionId   GUIDAttributes    uint64Name          [36]uint16
}type PARTITION_INFORMATION_EX struct {PartitionStyle   PartitionStyleStartingOffset   int64PartitionLength  int64DeviceNumber     int32RewritePartition boolRev01            boolRev02            boolRev03            boolPartitionInfo    [112]byte
}
type DRIVE_LAYOUT_INFORMATION_MBR struct {Signature uint32CheckSum  uint32
}type DRIVE_LAYOUT_INFORMATION_EX_HEADER struct {PartitionStyle PartitionStylePartitionCount 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 uint32buffer := 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      uint32DeviceNumber    uint32PartitionNumber uint32
}// formatFileSize 将文件大小转换为易读的格式
func FormatFileSize(size int64) string {// 定义文件大小单位units := []string{"B", "KB", "MB", "GB", "TB", "PB"}// 处理文件大小为0的情况if size == 0 {return "0 B"}// 计算文件大小所在单位的索引unitIndex := 0for size >= 1024 && unitIndex < len(units)-1 {size /= 1024unitIndex++}// 格式化文件大小return fmt.Sprintf("%d%s", size, units[unitIndex])
}
func GetDriveBasicInfo(drive string) (STORAGE_DEVICE_NUMBER, error) {var disk_num STORAGE_DEVICE_NUMBERvar err errorfilepath, _ := 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 = 0var 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.DeviceNumberdisk, 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))}