win-drives 0.1.0

use ntapi::ntioapi::*;
use winapi::{shared::ntdef::HANDLE, um::{winioctl::{DISK_GEOMETRY, DRIVE_LAYOUT_INFORMATION_EX, IOCTL_DISK_GET_DRIVE_GEOMETRY, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, IOCTL_DISK_GET_PARTITION_INFO_EX, IOCTL_STORAGE_GET_DEVICE_NUMBER, PARTITION_INFORMATION_EX, STORAGE_DEVICE_NUMBER}, winnt::{FILE_READ_ATTRIBUTES, FILE_SHARE_READ, FILE_SHARE_WRITE, GENERIC_READ, LARGE_INTEGER, SECURITY_QUALITY_OF_SERVICE, SYNCHRONIZE}}};
use core::{mem, ptr::null_mut};

use crate::{DriverError, ObjectAttributes};


pub fn open_handle(obj_name: heapless::String<20>) -> Result<HANDLE, DriverError> {
    let mut handle: *mut winapi::ctypes::c_void = null_mut();
    let mut obj_attr = ObjectAttributes::new();
    obj_attr.with_obj_name(obj_name);

    let sec_qos = SECURITY_QUALITY_OF_SERVICE {
        Length: mem::size_of::<SECURITY_QUALITY_OF_SERVICE>() as u32,
        ImpersonationLevel: 2,
        ContextTrackingMode: 0x1,
        EffectiveOnly: 1,
    };

    obj_attr.with_sec_qos(sec_qos);
    let mut obj_attr_raw = obj_attr.to_raw();
    let mut io_block: IO_STATUS_BLOCK = IO_STATUS_BLOCK::default();

    let status = unsafe { NtCreateFile(
        &mut handle,
        FILE_READ_ATTRIBUTES | GENERIC_READ | SYNCHRONIZE,
        &mut obj_attr_raw,
        &mut io_block,
        null_mut(),
        0,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        FILE_OPEN,
        FILE_NO_INTERMEDIATE_BUFFERING | FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT,
        null_mut(),
        0)
    };

    if status < 0 {
        return Err(DriverError::from(status))
    }

    Ok(handle)

}

pub fn geometry(handle: HANDLE) -> Result<DISK_GEOMETRY, DriverError> {
    let mut geometry: DISK_GEOMETRY = DISK_GEOMETRY::default();
    let mut io_block: IO_STATUS_BLOCK = IO_STATUS_BLOCK::default();

    let status = unsafe {
        NtDeviceIoControlFile(
            handle,
            null_mut(),
            None,
            null_mut(),
            &mut io_block,
            IOCTL_DISK_GET_DRIVE_GEOMETRY,
            null_mut(),
            0,
            &mut geometry as *mut _ as *mut _,
            size_of::<DISK_GEOMETRY>() as u32,
        )
    };

    if status < 0 {
        return Err(DriverError::from(status))
    }

    Ok(geometry)
}

pub fn partition_info(handle: HANDLE) -> Result<PARTITION_INFORMATION_EX, DriverError> {
    let mut info = PARTITION_INFORMATION_EX::default();
    let mut io_block = IO_STATUS_BLOCK::default();

    let status = unsafe { 
        NtDeviceIoControlFile(
            handle,
            null_mut(),
            None,
            null_mut(),
            &mut io_block,
            IOCTL_DISK_GET_PARTITION_INFO_EX,
            null_mut(),
            0,
            &mut info as *mut _ as *mut _,
            size_of::<PARTITION_INFORMATION_EX>() as u32
        )
    };

    if status < 0 {
        return Err(DriverError::from(status))
    }

    Ok(info)
}

pub fn drive_layout(handle: HANDLE) -> Result<DRIVE_LAYOUT_INFORMATION_EX, DriverError> {
    let mut info = DRIVE_LAYOUT_INFORMATION_EX::default();
    let mut io_block = IO_STATUS_BLOCK::default();

    let status = unsafe { 
        NtDeviceIoControlFile(
            handle,
            null_mut(),
            None,
            null_mut(),
            &mut io_block,
            IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
            null_mut(),
            0,
            &mut info as *mut _ as *mut _,
            size_of::<DRIVE_LAYOUT_INFORMATION_EX>() as u32
        )
    };

    if status < 0 {
        return Err(DriverError::from(status))
    }

    Ok(info)
}

pub fn device_number(handle: HANDLE) -> Result<STORAGE_DEVICE_NUMBER, DriverError> {
    let mut info = STORAGE_DEVICE_NUMBER::default();
    let mut io_block = IO_STATUS_BLOCK::default();

    let status = unsafe { 
        NtDeviceIoControlFile(
            handle,
            null_mut(),
            None,
            null_mut(),
            &mut io_block,
            IOCTL_STORAGE_GET_DEVICE_NUMBER,
            null_mut(),
            0,
            &mut info as *mut _ as *mut _,
            size_of::<STORAGE_DEVICE_NUMBER>() as u32,
        )
    };

    if status < 0 {
        return Err(DriverError::from(status))
    }

    Ok(info)
}

pub fn read_file(handle: HANDLE, buf: &mut [u8], offset: u64) -> Result<usize, DriverError> {
    let mut iosb: IO_STATUS_BLOCK = IO_STATUS_BLOCK::default();
    let mut byte_offset: LARGE_INTEGER = LARGE_INTEGER::default();
    unsafe { *byte_offset.QuadPart_mut() = offset as i64 };

    let status = unsafe {
        NtReadFile(
            handle,
            null_mut(),
            None,
            null_mut(),
            &mut iosb,
            buf.as_mut_ptr() as *mut winapi::ctypes::c_void,
            buf.len() as u32,
            &mut byte_offset,
            null_mut(),
        )
    };
    
    if status < 0 {
        return Err(status.into());
    }

    Ok(iosb.Information)
}