use crate::windows_utils::string::EasyPCWSTR;
use arbitrary::Arbitrary;
use eyre::ensure;
use facet::Facet;
use std::fmt;
use std::str::FromStr;
use windows::Win32::System::WindowsProgramming::DRIVE_REMOTE;
#[derive(Clone, PartialEq, Eq, Debug, Facet, Arbitrary)]
#[facet(transparent)]
pub struct DriveLetterPattern(pub String);
impl Default for DriveLetterPattern {
fn default() -> Self {
DriveLetterPattern("*".to_string())
}
}
impl DriveLetterPattern {
pub fn into_drive_letters(&self) -> eyre::Result<Vec<char>> {
let input = self.as_ref().trim();
if input == "*" {
return get_available_drives();
}
let mut rtn = Vec::new();
for (i, char) in input.chars().enumerate() {
let skippable = char.is_whitespace() || char == ',' || char == ';';
if skippable {
continue;
}
ensure!(
char.is_ascii_alphabetic(),
"Invalid drive letter character at position {i}: '{char}'"
);
rtn.push(char.to_ascii_uppercase());
}
ensure!(!rtn.is_empty(), "No drive letters found in: '{}'", input);
Ok(rtn)
}
}
impl fmt::Display for DriveLetterPattern {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
impl FromStr for DriveLetterPattern {
type Err = eyre::Report;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let s = s.trim();
ensure!(!s.is_empty(), "empty drive letter pattern");
Ok(DriveLetterPattern(s.to_string()))
}
}
impl AsRef<str> for DriveLetterPattern {
fn as_ref(&self) -> &str {
&self.0
}
}
fn get_available_drives() -> eyre::Result<Vec<char>> {
use windows::Win32::Storage::FileSystem::GetDriveTypeW;
use windows::Win32::Storage::FileSystem::GetLogicalDrives;
let drives_bitmask = unsafe { GetLogicalDrives() };
let mut available_drives = Vec::new();
for i in 0..26 {
if (drives_bitmask & (1 << i)) != 0 {
let idx = u8::try_from(i).unwrap_or_default();
let drive_letter = (b'A' + idx) as char;
let drive_root = format!("{drive_letter}:\\");
let drive_root = drive_root.easy_pcwstr()?;
let drive_type = unsafe { GetDriveTypeW(drive_root.as_ref()) };
if should_enumerate_drive_type(drive_type) {
available_drives.push(drive_letter);
}
}
}
ensure!(!available_drives.is_empty(), "No drives found on system");
Ok(available_drives)
}
fn should_enumerate_drive_type(drive_type: u32) -> bool {
drive_type != DRIVE_REMOTE
}
#[cfg(test)]
mod tests {
use super::DRIVE_REMOTE;
use super::DriveLetterPattern;
use super::should_enumerate_drive_type;
#[test]
fn parses_explicit_drive_letters_with_separators() -> eyre::Result<()> {
let drive_letters = DriveLetterPattern("C,D;E F".to_string()).into_drive_letters()?;
assert_eq!(drive_letters, vec!['C', 'D', 'E', 'F']);
Ok(())
}
#[test]
fn rejects_non_alphabetic_drive_letters() {
let error = DriveLetterPattern("C1".to_string())
.into_drive_letters()
.unwrap_err();
assert!(
error
.to_string()
.contains("Invalid drive letter character at position 1: '1'")
);
}
#[test]
fn excludes_remote_drive_types_from_wildcard_enumeration() {
assert!(!should_enumerate_drive_type(DRIVE_REMOTE));
assert!(should_enumerate_drive_type(3));
}
}