diff --git a/pkg/os/disk/api.go b/pkg/os/disk/api.go index 7b2e107f..0ca765a7 100644 --- a/pkg/os/disk/api.go +++ b/pkg/os/disk/api.go @@ -20,8 +20,13 @@ var ( ) const ( - IOCTL_STORAGE_GET_DEVICE_NUMBER = 0x2D1080 - IOCTL_STORAGE_QUERY_PROPERTY = 0x002d1400 + IOCTL_STORAGE_GET_DEVICE_NUMBER = 0x2D1080 + IOCTL_STORAGE_QUERY_PROPERTY = 0x002d1400 + IOCTL_DISK_SET_DISK_ATTRIBUTES = 0x0007c0f4 + IOCTL_DISK_GET_DISK_ATTRIBUTES = 0x000700f0 + IOCTL_DISK_GET_PARTITION_INFO_EX = 0x00070048 + IOCTL_DISK_GET_DRIVE_LAYOUT_EX = 0x00070050 + PARTITION_ENTRY_UNUSED_GUID = "00000000-0000-0000-0000-000000000000" ) // API declares the interface exposed by the internal API @@ -141,6 +146,29 @@ func (DiskAPI) IsDiskInitialized(diskNumber uint32) (bool, error) { return false, nil } +func (DiskAPI) IsDiskInitializedEx(diskPath string) (bool, error) { + disk, err := getDiskHandleFromPath(diskPath) + if err != nil { + return false, err + } + + partitionInfo := StoragePartitionInfo{} + partitionInfoSize := uint32(unsafe.Sizeof(partitionInfo)) + var size uint32 + + err = syscall.DeviceIoControl(disk, IOCTL_DISK_GET_PARTITION_INFO_EX, nil, 0, (*byte)(unsafe.Pointer(&partitionInfo)), partitionInfoSize, &size, nil) + if err != nil { + return false, fmt.Errorf("IOCTL_DISK_GET_PARTITION_INFO_EX failed: %v", err) + } + + if partitionInfo.PartitionStyle != 2 { + // Raw partition. + // TODO: Incorrect value is returned by syscall + return true, nil + } + return false, nil +} + func (DiskAPI) InitializeDisk(diskNumber uint32) error { cmd := fmt.Sprintf("Initialize-Disk -Number %d -PartitionStyle GPT", diskNumber) out, err := utils.RunPowershellCmd(cmd) @@ -163,6 +191,39 @@ func (DiskAPI) BasicPartitionsExist(diskNumber uint32) (bool, error) { return false, nil } +func (DiskAPI) BasicPartitionsExistEx(diskPath string) (bool, error) { + disk, err := getDiskHandleFromPath(diskPath) + if err != nil { + return false, err + } + + driveInfo := StorageDriveLayoutInfo{} + // driveInfo.PartitionEntry = make([]StoragePartitionInfo, driveInfo.PartitionCount*10) + // TODO: Converting any size array to go types + driveInfoSize := uint32(unsafe.Sizeof(driveInfo)) + var size uint32 + + // Without correct buffer size, errors with "The data area passed to a system call is too small." + err = syscall.DeviceIoControl(disk, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, nil, 0, (*byte)(unsafe.Pointer(&driveInfo)), driveInfoSize, &size, nil) + if err != nil { + + return false, fmt.Errorf("IOCTL_DISK_GET_DRIVE_LAYOUT_EX failed: %v", err) + } + + fmt.Printf("DriveInfo:\n%v\n", driveInfo) + if driveInfo.PartitionStyle == 1 && driveInfo.PartitionCount > 0 { + // GPT partition + return true, nil + } else if driveInfo.PartitionStyle == 0 && driveInfo.PartitionCount > 0 { + // MBR partition + return true, nil + } else { + // RAW partition + // Incorrect value returned by syscall for raw partition + return false, nil + } +} + func (DiskAPI) CreateBasicPartition(diskNumber uint32) error { cmd := fmt.Sprintf("New-Partition -DiskNumber %d -UseMaximumSize", diskNumber) out, err := utils.RunPowershellCmd(cmd) @@ -367,3 +428,79 @@ func (imp DiskAPI) GetDiskState(diskNumber uint32) (bool, error) { return !isOffline, nil } + +func (imp DiskAPI) SetDiskStateEx(diskPath string, isOffline bool, isReadOnly bool) error { + disk, err := getDiskHandleFromPath(diskPath) + if err != nil { + return err + } + + defer syscall.Close(disk) + attributes := SetStorageDeviceAttributes{} + attributesSize := uint32(unsafe.Sizeof(attributes)) + attributes.Version = attributesSize + attributes.Persist = true + attributes.Reserved1 = 0 + attributes.Reserved2 = 0 + var size uint32 + + if isOffline { + attributes.Attributes += 1 + } + if isReadOnly { + attributes.Attributes += 2 + } + err = syscall.DeviceIoControl(disk, IOCTL_DISK_SET_DISK_ATTRIBUTES, (*byte)(unsafe.Pointer(&attributes)), attributesSize, nil, 0, &size, nil) + // TODO: Errors with Access denied with write access to disk handle + if err != nil { + return fmt.Errorf("IOCTL_DISK_SET_DISK_ATTRIBUTES failed: %v", err) + } + return nil +} + +func (imp DiskAPI) GetDiskStateEx(diskPath string) (bool, bool, error) { + isDiskOnline := false + isDiskReadOnly := true + + disk, err := getDiskHandleFromPath(diskPath) + if err != nil { + return isDiskOnline, isDiskReadOnly, err + } + + defer syscall.Close(disk) + attributes := GetStorageDeviceAtrributes{} + attributesSize := uint32(unsafe.Sizeof(attributes)) + attributes.Version = attributesSize + var size uint32 + + err = syscall.DeviceIoControl(disk, IOCTL_DISK_GET_DISK_ATTRIBUTES, nil, 0, (*byte)(unsafe.Pointer(&attributes)), attributesSize, &size, nil) + if err != nil { + return isDiskOnline, isDiskReadOnly, fmt.Errorf("IOCTL_DISK_GET_DISK_ATTRIBUTES failed: %v", err) + } + + diskAttributes := attributes.Attributes + + if diskAttributes == 3 { + isDiskOnline = false + isDiskReadOnly = true + } else if diskAttributes == 2 { + isDiskOnline = true + isDiskReadOnly = true + } else if diskAttributes == 1 { + isDiskOnline = false + isDiskReadOnly = false + } else { + isDiskOnline = true + isDiskReadOnly = false + } + return isDiskOnline, isDiskReadOnly, nil +} + +func getDiskHandleFromPath(diskPath string) (syscall.Handle, error) { + handle, err := syscall.Open(diskPath, syscall.O_RDONLY|syscall.O_RDWR|syscall.O_NONBLOCK|syscall.O_SYNC, 7) + + if err != nil { + return 0, err + } + return handle, nil +} diff --git a/pkg/os/disk/types.go b/pkg/os/disk/types.go index 27bf2be1..98a984e8 100644 --- a/pkg/os/disk/types.go +++ b/pkg/os/disk/types.go @@ -1,5 +1,7 @@ package disk +import "golang.org/x/sys/windows" + type StorageDeviceNumber struct { DeviceType DeviceType DeviceNumber uint32 @@ -113,3 +115,76 @@ type Disk struct { Path string `json:"Path"` SerialNumber string `json:"SerialNumber"` } + +type SetStorageDeviceAttributes struct { + Version uint32 + Persist bool + Reserved1 uint32 + Attributes uint64 + AttributesMask uint64 + Reserved2 uint64 +} + +type GetStorageDeviceAtrributes struct { + Version uint32 + Reserved1 uint32 + Attributes uint64 +} + +type StoragePartitionStyle uint32 + +const ( + PartitionStyleMbr StoragePartitionStyle = iota + PartitionStyleGpt + PartitionStyleRaw +) + +type PartitionInfoMbr struct { + PartitionType byte + BootIndicator bool + RecognizedPartition bool + HiddenSectors uint32 + PartitionId windows.GUID +} + +type PartitionInfoGpt struct { + PartitionType windows.GUID + PartitionId windows.GUID + Attributes uint64 + Name [36]uint16 +} + +type StoragePartitionInfo struct { + PartitionStyle StoragePartitionStyle + StartingOffset int64 + PartitionLength int64 + PartitionNumber uint32 + RewritePartition bool + IsServicePartition bool + DummmyUnionName struct { + PartitionInfoMbr + PartitionInfoGpt + } +} + +type DriveLayoutInfoMbr struct { + Signature uint32 + CheckSum uint32 +} + +type DriveLayoutInfoGpt struct { + DiskId windows.GUID + StartingUsableOffset int64 + UsableLength int64 + MaxPartitionCount uint32 +} + +type StorageDriveLayoutInfo struct { + PartitionStyle StoragePartitionStyle + PartitionCount uint32 + DummmyUnionName struct { + DriveLayoutInfoMbr + DriveLayoutInfoGpt + } + PartitionEntry []StoragePartitionInfo +}