use std::mem;
use std::os::raw::{c_ulong, c_ulonglong};
use std::ptr;
use sysinfo::{get_current_pid, Pid, PidExt};
use wchar::wchar_t;
use widestring::U16CString;
use windows::core::{HRESULT, PCSTR, PCWSTR};
use windows::Win32::Foundation::{CloseHandle, HANDLE};
use windows::Win32::Storage::FileSystem::GetDriveTypeA;
use windows::Win32::Storage::InstallableFileSystems::{
FilterConnectCommunicationPort, FilterSendMessage,
};
use crate::driver_comm::DriveType::{
DriveCDRom, DriveFixed, DriveNoRootDir, DriveRamDisk, DriveRemote, DriveRemovable, DriveUnknown,
};
use crate::driver_comm::IrpMajorOp::{IrpCreate, IrpNone, IrpRead, IrpSetInfo, IrpWrite};
use crate::shared_def;
use crate::shared_def::ReplyIrp;
type BufPath = [wchar_t; 520];
#[derive(Debug)]
#[repr(C)]
struct DriverComMessage {
r#type: c_ulong,
pid: c_ulong,
gid: c_ulonglong,
path: BufPath,
}
#[repr(C)]
#[allow(dead_code)]
enum DriverComMessageType {
AddScanDirectory,
RemScanDirectory,
GetOps,
SetPid,
KillGid,
}
#[derive(Debug)]
#[repr(C)]
pub struct Driver {
handle: HANDLE,
}
impl Driver {
#[inline]
pub fn close_kernel_communication(&self) -> bool {
unsafe { CloseHandle(self.handle).as_bool() }
}
#[inline]
pub fn driver_set_app_pid(&self) -> Result<(), windows::core::Error> {
let buf = Self::string_to_commessage_buffer(r"\Device\harddiskVolume");
let mut get_irp_msg: DriverComMessage = DriverComMessage {
r#type: DriverComMessageType::SetPid as c_ulong,
pid: get_current_pid().unwrap().as_u32() as c_ulong,
gid: 140_713_315_094_899,
path: buf, };
let mut tmp: u32 = 0;
unsafe {
FilterSendMessage(
self.handle,
ptr::addr_of_mut!(get_irp_msg).cast::<std::ffi::c_void>(),
mem::size_of::<DriverComMessage>() as c_ulong,
None,
0,
std::ptr::addr_of_mut!(tmp),
)
}
}
#[inline]
pub fn open_kernel_driver_com() -> Result<Self, windows::core::Error> {
let com_port_name = U16CString::from_str("\\snFilter").unwrap().into_raw();
let handle;
unsafe {
handle = FilterConnectCommunicationPort(PCWSTR(com_port_name), 0, None, 0, None)?;
}
let res = Self { handle };
Ok(res)
}
#[inline]
pub fn get_irp(&self, vecnew: &mut Vec<u8>) -> Option<ReplyIrp> {
let mut get_irp_msg = Self::build_irp_msg(
DriverComMessageType::GetOps,
get_current_pid().unwrap(),
0,
"",
);
let mut tmp: u32 = 0;
unsafe {
FilterSendMessage(
self.handle,
ptr::addr_of_mut!(get_irp_msg).cast::<std::ffi::c_void>(),
mem::size_of::<DriverComMessage>() as c_ulong,
Option::from(vecnew.as_mut_ptr().cast::<std::ffi::c_void>()),
65536_u32,
ptr::addr_of_mut!(tmp).cast::<u32>(),
)
.expect("Cannot get driver message from driver");
}
if tmp != 0 {
let reply_irp: ReplyIrp;
unsafe {
reply_irp =
std::ptr::read_unaligned(vecnew.as_ptr().cast::<shared_def::ReplyIrp>());
}
return Some(reply_irp);
}
None
}
#[inline]
pub fn try_kill(&self, gid: c_ulonglong) -> Result<HRESULT, windows::core::Error> {
let mut killmsg = DriverComMessage {
r#type: DriverComMessageType::KillGid as c_ulong,
pid: 0, gid,
path: [0; 520],
};
let mut res: i32 = 0;
let mut res_size: u32 = 0;
unsafe {
FilterSendMessage(
self.handle,
ptr::addr_of_mut!(killmsg).cast::<std::ffi::c_void>(),
mem::size_of::<DriverComMessage>() as c_ulong,
Option::from(ptr::addr_of_mut!(res).cast::<std::ffi::c_void>()),
4_u32,
ptr::addr_of_mut!(res_size).cast::<u32>(),
)?;
}
Ok(HRESULT(res))
}
#[inline]
fn string_to_commessage_buffer(bufstr: &str) -> BufPath {
let temp = U16CString::from_str(bufstr).unwrap();
let mut buf: BufPath = [0; 520];
for (i, c) in temp.as_slice_with_nul().iter().enumerate() {
buf[i] = *c as wchar_t;
}
buf
}
#[inline]
fn build_irp_msg(
commsgtype: DriverComMessageType,
pid: Pid,
gid: u64,
path: &str,
) -> DriverComMessage {
DriverComMessage {
r#type: commsgtype as c_ulong, pid: pid.as_u32() as c_ulong,
gid,
path: Self::string_to_commessage_buffer(path),
}
}
}
#[repr(C)]
pub enum IrpMajorOp {
IrpNone,
IrpRead,
IrpWrite,
IrpSetInfo,
IrpCreate,
IrpCleanUp,
}
impl IrpMajorOp {
#[inline]
#[must_use]
pub fn from_byte(b: u8) -> Self {
match b {
1 => IrpRead,
2 => IrpWrite,
3 => IrpSetInfo,
4 | 5 => IrpCreate,
_ => IrpNone,
}
}
}
#[repr(C)]
pub enum DriveType {
DriveUnknown,
DriveNoRootDir,
DriveRemovable,
DriveFixed,
DriveRemote,
DriveCDRom,
DriveRamDisk,
}
impl DriveType {
#[inline]
#[must_use]
pub fn from_filepath(filepath: &str) -> Self {
let mut drive_type = 1u32;
if !filepath.is_empty() {
let drive_path = &filepath[..=filepath.find('\\').unwrap()];
unsafe {
drive_type = GetDriveTypeA(PCSTR(String::from(drive_path).as_ptr()));
}
}
match drive_type {
0 => DriveUnknown,
2 => DriveRemovable,
3 => DriveFixed,
4 => DriveRemote,
5 => DriveCDRom,
6 => DriveRamDisk,
_ => DriveNoRootDir,
}
}
}