win-drives 0.1.0

use core::fmt::{self, Write};
use core::{mem::zeroed};
use core::{mem, ptr::{self, null_mut}};

use winapi::shared::guiddef::GUID;
use winapi::{shared::{ntdef::{OBJ_CASE_INSENSITIVE, OBJECT_ATTRIBUTES, UNICODE_STRING}, ntstatus::*}, um::{winioctl::*, winnt::SECURITY_QUALITY_OF_SERVICE}};

#[derive(Copy, Clone, Debug)]
pub enum DriverError {
    Permission,
    PathNotFound,
    NotFound,
    InvalidParameter,
    InvalidHandle,
    BufferTooSmall { needed: u32, got: usize },
    Unaligned { offset: u64, sector: u32 },
    NtStatus(i32)
}

impl core::fmt::Display for DriverError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            DriverError::Permission         => write!(f, "access denied"),
            DriverError::PathNotFound       => write!(f, "path not found"),
            DriverError::NotFound           => write!(f, "not found"),
            DriverError::InvalidParameter   => write!(f, "invalid parameter"),
            DriverError::InvalidHandle      => write!(f, "invalid handle"),
            DriverError::BufferTooSmall { needed, got } =>
                write!(f, "buffer too small: needed at least {needed}, got {got}"),
            DriverError::Unaligned { offset, sector } =>
                write!(f, "offset 0x{offset:X} not aligned to {sector}-byte sector"),
            DriverError::NtStatus(s) =>
                write!(f, "NTSTATUS 0x{:08X}", *s as u32),
        }
    }
}

impl core::error::Error for DriverError {}

impl From<i32> for DriverError {
    fn from(status: i32) -> Self {
        match status {
            STATUS_ACCESS_DENIED => DriverError::Permission,
            STATUS_OBJECT_NAME_NOT_FOUND => DriverError::NotFound,
            STATUS_OBJECT_PATH_NOT_FOUND => DriverError::PathNotFound,
            STATUS_INVALID_PARAMETER => DriverError::InvalidParameter,
            STATUS_INVALID_HANDLE => DriverError::InvalidHandle,
            _ => DriverError::NtStatus(status),
        }
    }
}

#[repr(u32)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum PartitionStyle {
    Gpt = PARTITION_STYLE_GPT,
    Mbr = PARTITION_STYLE_MBR,
    Raw = PARTITION_STYLE_RAW,
}

impl From<u32> for PartitionStyle {
    fn from(value: u32) -> Self {
        match value {
            PARTITION_STYLE_GPT => PartitionStyle::Gpt,
            PARTITION_STYLE_MBR => PartitionStyle::Mbr,
            PARTITION_STYLE_RAW => PartitionStyle::Raw,
            _ => PartitionStyle::Raw,
        }
    }
}

#[derive(Copy, Clone)]
pub struct PartitionInfo(PARTITION_INFORMATION_EX);

#[cfg(all(feature = "alloc", not(feature = "no-std")))]
pub type GptName = alloc::string::String;

#[cfg(feature = "no-std")]
pub type GptName = heapless::String<36>;

impl PartitionInfo {
    
    #[cfg(any(feature = "alloc", feature = "no-std"))]
    pub fn gpt_name(&self) -> Option<GptName> {
        if self.style() != PartitionStyle::Gpt {
            return None;
        }
        let gpt = unsafe { &self.0.u.Gpt() };

        #[cfg(feature = "no-std")]
        {
            let mut out = heapless::String::new();
            for c in core::char::decode_utf16(
                gpt.Name.iter().copied().take_while(|&u| u != 0)
            ) {
                let _ = out.push(c.unwrap_or('\u{FFFD}'));
            }
            if out.is_empty() { None } else { Some(out) }
        }

        #[cfg(not(feature = "no-std"))]
        {
            let s: alloc::string::String = core::char::decode_utf16(
                gpt.Name.iter().copied().take_while(|&u| u != 0)
            )
            .map(|c| c.unwrap_or('\u{FFFD}'))
            .collect();
            if s.is_empty() { None } else { Some(s) }
        }
    }

    /// The partition style: MBR, GPT, or RAW.
    #[inline]
    pub fn style(&self) -> PartitionStyle {
        PartitionStyle::from(self.0.PartitionStyle)
    }

    /// 1-based partition number. 0 means "the whole disk" entry.
    #[inline]
    pub const fn number(&self) -> u32 {
        self.0.PartitionNumber
    }

    /// Byte offset from the start of the disk.
    #[inline]
    pub fn starting_offset(&self) -> u64 {
        unsafe { *self.0.StartingOffset.QuadPart() as u64 }
    }

    /// Length of the partition in bytes.
    #[inline]
    pub fn length(&self) -> u64 {
        unsafe { *self.0.PartitionLength.QuadPart() as u64 }
    }

    /// Whether the partition needs to be rewritten.
    #[inline]
    pub const fn rewrite(&self) -> bool {
        self.0.RewritePartition != 0
    }

    /// The offset of the last byte of the partition (exclusive end).
    #[inline]
    pub fn end_offset(&self) -> u64 {
        self.starting_offset() + self.length()
    }
}

impl fmt::Debug for PartitionInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let p = &self.0;
        let style = PartitionStyle::from(p.PartitionStyle);

        let mut dbg = f.debug_struct("PartitionInfo");
        dbg.field("partition_style", &style)
           .field("partition_number", &p.PartitionNumber)
           .field("starting_offset", &unsafe { *p.StartingOffset.QuadPart() })
           .field("partition_length", &unsafe { *p.PartitionLength.QuadPart() })
           .field("rewrite_partition", &(p.RewritePartition != 0));

        match style {
            PartitionStyle::Mbr => {
                let mbr = unsafe { &p.u.Mbr() };
                dbg.field("mbr.boot_indicator", &mbr.BootIndicator)
                   .field("mbr.recognized_partition", &mbr.RecognizedPartition)
                   .field("mbr.hidden_sectors", &mbr.HiddenSectors)
                   .field("mbr.partition_type", &format_args!("0x{:02X}", mbr.PartitionType));
            }
            PartitionStyle::Gpt => {
                let gpt = unsafe { &p.u.Gpt() };
                dbg.field("gpt.partition_type", &GuidDebug(gpt.PartitionType))
                    .field("gpt.partition_id",   &GuidDebug(gpt.PartitionId))
                   .field("gpt.attributes", &format_args!("0x{:016X}", gpt.Attributes))
                   .field("gpt.name", &GptNameDisplay(&gpt.Name));
            }
            PartitionStyle::Raw => {}
        }

        dbg.finish()
    }
}

struct GptNameDisplay<'a>(&'a [u16; 36]);

impl fmt::Debug for GptNameDisplay<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let end = self.0.iter().position(|&u| u == 0).unwrap_or(self.0.len());

        f.write_char('"')?;
        for c in core::char::decode_utf16(self.0[..end].iter().copied()) {
            match c {
                Ok(ch) => f.write_char(ch)?,
                Err(_) => f.write_char('\u{FFFD}')?, // U+FFFD replacement char
            }
        }
        f.write_char('"')
    }
}

/// Newtype over a raw [`GUID`] with a human-readable `Debug` impl and
/// recognition of well-known GPT partition type GUIDs.
#[derive(Copy, Clone)]
pub struct Guid(pub GUID);

impl PartialEq for Guid {
    fn eq(&self, other: &Self) -> bool {
        let a = &self.0;
        let b = &other.0;
        a.Data1 == b.Data1
            && a.Data2 == b.Data2
            && a.Data3 == b.Data3
            && a.Data4 == b.Data4
    }
}

impl Eq for Guid {}

impl Guid {
    /// Returns a human-readable name for well-known GPT partition type GUIDs.
    ///
    /// Covers the identifiers defined by the UEFI spec that Windows actually
    /// uses on disks: EFI System, Microsoft Reserved, Basic Data, Windows
    /// Recovery, and BIOS Boot.
    #[inline]
    pub fn well_known_name(&self) -> Option<&'static str> {
        let g = &self.0;
        match (g.Data1, g.Data2, g.Data3, &g.Data4) {
            (0xC12A7328, 0xF81F, 0x11D2, b) if *b == [0xBA, 0x4B, 0x00, 0xA0, 0xC9, 0x3E, 0xC9, 0x3B] =>
                Some("EFI System"),
            (0xE3C9E316, 0x0B5C, 0x4DB8, b) if *b == [0x81, 0x7D, 0xF9, 0x2D, 0xF0, 0x02, 0x15, 0xAE] =>
                Some("Microsoft Reserved"),
            (0xEBD0A0A2, 0xB9E5, 0x4433, b) if *b == [0x87, 0xC0, 0x68, 0xB6, 0xB7, 0x26, 0x99, 0xC7] =>
                Some("Basic Data"),
            (0xDE94BBA4, 0x06D1, 0x4D40, b) if *b == [0xA1, 0x6A, 0xBF, 0xD5, 0x01, 0x79, 0xD6, 0xAC] =>
                Some("Windows Recovery"),
            (0x21686148, 0x6449, 0x6E6F, b) if *b == [0x74, 0x4E, 0x65, 0x64, 0x45, 0x46, 0x49, 0x00] =>
                Some("BIOS Boot"),
            _ => None,
        }
    }

    /// Formats the GUID in the canonical `{XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX}` form.
    #[inline]
    fn write_canonical(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let g = &self.0;
        write!(
            f,
            "{{{:08X}-{:04X}-{:04X}-{:02X}{:02X}-{:02X}{:02X}{:02X}{:02X}{:02X}{:02X}}}",
            g.Data1, g.Data2, g.Data3,
            g.Data4[0], g.Data4[1],
            g.Data4[2], g.Data4[3],
            g.Data4[4], g.Data4[5],
            g.Data4[6], g.Data4[7],
        )
    }
}

impl core::fmt::Debug for Guid {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self.well_known_name() {
            Some(name) => write!(f, "{} (", name).and_then(|_| {
                self.write_canonical(f)?;
                write!(f, ")")
            }),
            None => self.write_canonical(f),
        }
    }
}

impl From<GUID> for Guid {
    #[inline]
    fn from(value: GUID) -> Self {
        Self(value)
    }
}

impl From<Guid> for GUID {
    #[inline]
    fn from(value: Guid) -> Self {
        value.0
    }
}

pub struct GuidDebug(pub GUID);

impl fmt::Debug for GuidDebug {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let g = &self.0;
        write!(
            f,
            "{{{:08X}-{:04X}-{:04X}-{:02X}{:02X}-{:02X}{:02X}{:02X}{:02X}{:02X}{:02X}}}",
            g.Data1, g.Data2, g.Data3,
            g.Data4[0], g.Data4[1],
            g.Data4[2], g.Data4[3],
            g.Data4[4], g.Data4[5],
            g.Data4[6], g.Data4[7],
        )
    }
}

#[derive(Copy, Clone)]
pub struct DiskGeometry(DISK_GEOMETRY);

impl core::fmt::Debug for DiskGeometry {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let g = &self.0;
        f.debug_struct("DiskGeometry")
            .field("cylinders", &self.cylinders())
            .field("tracks_per_cylinder", &self.tracks_per_cylinder())
            .field("sectors_per_track", &(g.SectorsPerTrack as u64))
            .field("bytes_per_sector", &(g.BytesPerSector as u64))
            .field("media_type", &g.MediaType)
            .field("total_sectors", &self.sectors())
            .field("total_bytes", &self.size())
            .finish()
    }
}

impl DiskGeometry {

    #[inline]
    pub fn size(&self) -> u64 {
        self.sectors() * self.0.BytesPerSector as u64
    }

    #[inline]
    pub fn cylinders(&self) -> u64 {
        unsafe { *self.0.Cylinders.QuadPart() as u64 } 
    }

    #[inline]
    pub const fn media_type(&self) -> u32 {
        self.0.MediaType
    }

    #[inline]
    pub const fn bytes_per_sector(&self) -> u64 {
        self.0.BytesPerSector as u64
    }

    #[inline]
    pub const fn tracks_per_cylinder(&self) -> u32 {
        self.0.TracksPerCylinder as u32
    }

    #[inline]
    pub const fn sectors_per_track(&self) -> u32 {
        self.0.SectorsPerTrack as u32
    }

    #[inline]
    pub fn sectors(&self) -> u64 {
        self.cylinders() * self.tracks_per_cylinder() as u64 * self.sectors_per_track() as u64
    }
}

impl From<DISK_GEOMETRY> for DiskGeometry {
    fn from(value: DISK_GEOMETRY) -> Self {
        DiskGeometry(value)
    }
}

impl From<PARTITION_INFORMATION_EX> for PartitionInfo {
    fn from(value: PARTITION_INFORMATION_EX) -> Self {
        Self(value)
    }
}

pub struct ObjectAttributes {
    obj_name: heapless::String<20>,
    obj_name_u16: heapless::Vec<u16, 40>,
    obj_name_uc: UNICODE_STRING,
    qos: Option<SECURITY_QUALITY_OF_SERVICE>
}

impl ObjectAttributes {
    pub const fn new() -> Self {
        Self {
            obj_name: heapless::String::new(),
            obj_name_u16: heapless::Vec::new(),
            obj_name_uc: unsafe { zeroed() },
            qos: None
        }
    }

    pub fn with_obj_name(&mut self, obj_name: heapless::String<20>) {
        self.obj_name = obj_name;
        self.obj_name_u16 = self.obj_name
            .encode_utf16()
            .collect::<heapless::Vec<_, 40>>();

        self.obj_name_uc = UNICODE_STRING {
            Length: self.obj_name_u16.len() as u16 * 2,
            MaximumLength: self.obj_name_u16.len() as u16 * 2 + 2,
            Buffer: self.obj_name_u16.as_mut_ptr(),
        };
    }

    pub fn with_sec_qos(&mut self, qos: SECURITY_QUALITY_OF_SERVICE) {
        self.qos = Some(qos);
    }

    pub fn to_raw(&mut self) -> OBJECT_ATTRIBUTES {

        let sec_qos = match self.qos.as_mut() {
            Some(sec_qos) => sec_qos as *mut _ as *mut _,
            None => null_mut(),
        };

        OBJECT_ATTRIBUTES {
            Length: mem::size_of::<OBJECT_ATTRIBUTES>() as u32,
            RootDirectory: null_mut(),
            ObjectName: &mut self.obj_name_uc,
            Attributes: OBJ_CASE_INSENSITIVE,
            SecurityDescriptor: ptr::null_mut(),
            SecurityQualityOfService: sec_qos,
        }
    }
    
}

pub struct DriveLayout(DRIVE_LAYOUT_INFORMATION_EX);

impl From<DRIVE_LAYOUT_INFORMATION_EX> for DriveLayout {
    fn from(value: DRIVE_LAYOUT_INFORMATION_EX) -> Self {
        Self(value)
    }
}

impl DriveLayout {
    pub const fn partition_count(&self) -> u32 {
        self.0.PartitionCount
    }

    #[inline]
    pub fn style(&self) -> PartitionStyle {
        PartitionStyle::from(self.0.PartitionStyle)
    }

    pub fn partitions(&self) -> &[PARTITION_INFORMATION_EX] {
        unsafe {
            core::slice::from_raw_parts(
                self.0.PartitionEntry.as_ptr(),
                self.0.PartitionCount as usize,
            )
        }
    }
}

/// The Windows `DEVICE_TYPE` values relevant to disk enumeration.
///
/// See the `FILE_DEVICE_*` constants defined in `winapi::um::winioctl`
/// and the Windows SDK header `ntddk.h` for the full list.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
pub enum DeviceType {
    /// `FILE_DEVICE_BEEP`
    Beep = 0x0000_0001,
    /// `FILE_DEVICE_CD_ROM`
    CdRom = 0x0000_0002,
    /// `FILE_DEVICE_CD_ROM_FILE_SYSTEM`
    CdRomFileSystem = 0x0000_0003,
    /// `FILE_DEVICE_CONTROLLER`
    Controller = 0x0000_0004,
    /// `FILE_DEVICE_DATALINK`
    DataLink = 0x0000_0005,
    /// `FILE_DEVICE_DFS`
    Dfs = 0x0000_0006,
    /// `FILE_DEVICE_DISK`
    Disk = 0x0000_0007,
    /// `FILE_DEVICE_DISK_FILE_SYSTEM`
    DiskFileSystem = 0x0000_0008,
    /// `FILE_DEVICE_FILE_SYSTEM`
    FileSystem = 0x0000_0009,
    /// `FILE_DEVICE_TAPE`
    Tape = 0x0000_001F,
    /// `FILE_DEVICE_TAPE_FILE_SYSTEM`
    TapeFileSystem = 0x0000_0020,
    /// `FILE_DEVICE_VIRTUAL_DISK`
    VirtualDisk = 0x0000_0024,
    /// Any device type not covered above.
    Other(u32),
}

impl From<u32> for DeviceType {
    fn from(value: u32) -> Self {
        match value {
            0x0000_0001 => Self::Beep,
            0x0000_0002 => Self::CdRom,
            0x0000_0003 => Self::CdRomFileSystem,
            0x0000_0004 => Self::Controller,
            0x0000_0005 => Self::DataLink,
            0x0000_0006 => Self::Dfs,
            0x0000_0007 => Self::Disk,
            0x0000_0008 => Self::DiskFileSystem,
            0x0000_0009 => Self::FileSystem,
            0x0000_001F => Self::Tape,
            0x0000_0020 => Self::TapeFileSystem,
            0x0000_0024 => Self::VirtualDisk,
            other => Self::Other(other),
        }
    }
}

impl From<DeviceType> for u32 {
    fn from(value: DeviceType) -> Self {
        match value {
            DeviceType::Beep => 0x0000_0001,
            DeviceType::CdRom => 0x0000_0002,
            DeviceType::CdRomFileSystem => 0x0000_0003,
            DeviceType::Controller => 0x0000_0004,
            DeviceType::DataLink => 0x0000_0005,
            DeviceType::Dfs => 0x0000_0006,
            DeviceType::Disk => 0x0000_0007,
            DeviceType::DiskFileSystem => 0x0000_0008,
            DeviceType::FileSystem => 0x0000_0009,
            DeviceType::Tape => 0x0000_001F,
            DeviceType::TapeFileSystem => 0x0000_0020,
            DeviceType::VirtualDisk => 0x0000_0024,
            DeviceType::Other(v) => v,
        }
    }
}

pub struct DeviceNumber(STORAGE_DEVICE_NUMBER);

impl core::fmt::Debug for DeviceNumber {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("DeviceNumber")
            .field("device_type", &self.device_type())
            .field("device_number", &self.device_number())
            .field("partition_number", &self.partition_number())
            .finish()
    }
}

impl From<STORAGE_DEVICE_NUMBER> for DeviceNumber {
    fn from(value: STORAGE_DEVICE_NUMBER) -> Self {
        Self(value)
    }
}

impl DeviceNumber {
    pub fn device_type(&self) -> DeviceType {
        self.0.DeviceType.into()
    }

    pub const fn device_number(&self) -> u32 {
        self.0.DeviceNumber
    }

    pub const fn partition_number(&self) -> u32 {
        self.0.PartitionNumber
    }
}