From eee91d8498c3c7260714cfbbfe5b2bb1a0216c81 Mon Sep 17 00:00:00 2001 From: yyc12345 Date: Wed, 15 Oct 2025 13:15:29 +0800 Subject: [PATCH] write ext shit --- Cargo.lock | 23 +++ wfassoc/Cargo.toml | 3 +- wfassoc/src/lib.rs | 314 ++++++++++++++++++++++++++++----------- wfassoc_exec/src/main.rs | 4 +- 4 files changed, 256 insertions(+), 88 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4431f6a..43715af 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -168,6 +168,12 @@ dependencies = [ "litrs", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "errno" version = "0.3.14" @@ -178,12 +184,28 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "hashbrown" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" + [[package]] name = "heck" version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "indexmap" +version = "2.11.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" +dependencies = [ + "equivalent", + "hashbrown", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.1" @@ -495,6 +517,7 @@ dependencies = [ name = "wfassoc" version = "0.1.0" dependencies = [ + "indexmap", "regex", "thiserror", "uuid", diff --git a/wfassoc/Cargo.toml b/wfassoc/Cargo.toml index 20b926b..00a2eea 100644 --- a/wfassoc/Cargo.toml +++ b/wfassoc/Cargo.toml @@ -8,7 +8,8 @@ license = "SPDX:MIT" [dependencies] thiserror = { workspace = true } -windows-sys = { version = "0.60.2", features = ["Win32_Security", "Win32_System_SystemServices"] } +windows-sys = { version = "0.60.2", features = ["Win32_Security", "Win32_System_SystemServices", "Win32_UI_Shell"] } winreg = { version = "0.55.0", features = ["transactions"] } +indexmap = "2.11.4" regex = "1.11.3" uuid = "1.18.1" diff --git a/wfassoc/src/lib.rs b/wfassoc/src/lib.rs index da276a4..964b006 100644 --- a/wfassoc/src/lib.rs +++ b/wfassoc/src/lib.rs @@ -4,14 +4,18 @@ #[cfg(not(target_os = "windows"))] compile_error!("Crate wfassoc is only supported on Windows."); +use regex::Regex; use std::ffi::OsStr; +use std::fmt::Display; use std::path::{Path, PathBuf}; +use std::str::FromStr; +use std::sync::LazyLock; use thiserror::Error as TeError; +use indexmap::{IndexMap, IndexSet}; use winreg::RegKey; use winreg::enums::{ HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE, KEY_READ, KEY_WRITE, }; -use winreg::transaction::Transaction; // region: Error Types @@ -22,10 +26,14 @@ pub enum WfError { NoPrivilege, #[error("error occurs when manipulating with Registry: {0}")] BadRegOper(#[from] std::io::Error), + #[error("given full path to application is invalid")] BadFullAppPath, - #[error("failed when casting path or OS string into string")] + #[error("failed when casting OS string into string")] BadOsStrCast, + + #[error("file extension {0} is already registered")] + DupExt(String), } /// The result type used in this crate. @@ -33,74 +41,6 @@ pub type WfResult = Result; // endregion -// region: Scope and View - -/// The scope where wfassoc will register and unregister application. -#[derive(Debug, Copy, Clone)] -pub enum Scope { - /// Scope for current user. - User, - /// Scope for all users under this computer. - System, -} - -/// The error occurs when cast View into Scope. -#[derive(Debug, TeError)] -#[error("hybrid view can not be cast into any scope")] -pub struct TryFromViewError {} - -impl TryFromViewError { - fn new() -> Self { - Self {} - } -} - -impl TryFrom for Scope { - type Error = TryFromViewError; - - fn try_from(value: View) -> Result { - match value { - View::User => Ok(Self::User), - View::System => Ok(Self::System), - View::Hybrid => Err(TryFromViewError::new()), - } - } -} - -impl Scope { - /// Check whether we have enough privilege when operating in current scope. - /// If we have, return true, otherwise false. - fn has_privilege(&self) -> bool { - // If we operate on System, and we do not has privilege, - // we think we do not have privilege, otherwise, - // there is no privilege required. - !matches!(self, Self::System if !has_privilege()) - } -} - -/// The view when wfassoc querying file extension association. -#[derive(Debug, Copy, Clone)] -pub enum View { - /// The view of current user. - User, - /// The view of system. - System, - /// Hybrid view of User and System. - /// It can be seen as that we use System first and then use User to override any existing items. - Hybrid, -} - -impl From for View { - fn from(value: Scope) -> Self { - match value { - Scope::User => Self::User, - Scope::System => Self::System, - } - } -} - -// endregion - // region: Utilities /// The println macro only works on Debug mode @@ -173,6 +113,20 @@ fn has_privilege() -> bool { is_member != 0 } +/// Notify Windows that some file associations are changed, and should refresh them. +/// This function must be called once you change any file associations. +fn notify_assoc_changed() -> () { + use windows_sys::Win32::UI::Shell::{SHCNE_ASSOCCHANGED, SHCNF_IDLIST, SHChangeNotify}; + unsafe { + SHChangeNotify( + SHCNE_ASSOCCHANGED as i32, + SHCNF_IDLIST, + Default::default(), + Default::default(), + ) + } +} + /// Try casting given &Path into &str. fn path_to_str(path: &Path) -> WfResult<&str> { path.to_str().ok_or(WfError::BadOsStrCast) @@ -185,29 +139,198 @@ fn osstr_to_str(osstr: &OsStr) -> WfResult<&str> { // endregion -// region: Registrar +// region: Types -/// The core registrar for register and unregister application. -pub struct Registrar { - /// The fully qualified path to the application. - full_path: PathBuf, +/// The token for access registered items in Program. +/// This is usually returned when you registering them. +pub type Token = usize; + +/// The scope where wfassoc will register and unregister application. +#[derive(Debug, Copy, Clone)] +pub enum Scope { + /// Scope for current user. + User, + /// Scope for all users under this computer. + System, } -impl Registrar { - /// Create a new registrar for following operations. - pub fn new(full_path: &Path) -> Self { - Self { - full_path: full_path.to_path_buf(), +/// The error occurs when cast View into Scope. +#[derive(Debug, TeError)] +#[error("hybrid View can not be cast into Scope")] +pub struct TryFromViewError {} + +impl TryFromViewError { + fn new() -> Self { + Self {} + } +} + +impl TryFrom for Scope { + type Error = TryFromViewError; + + fn try_from(value: View) -> Result { + match value { + View::User => Ok(Self::User), + View::System => Ok(Self::System), + View::Hybrid => Err(TryFromViewError::new()), } } } -impl Registrar { +impl Scope { + /// Check whether we have enough privilege when operating in current scope. + /// If we have, return true, otherwise false. + fn has_privilege(&self) -> bool { + // If we operate on System, and we do not has privilege, + // we think we do not have privilege, otherwise, + // there is no privilege required. + !matches!(self, Self::System if !has_privilege()) + } +} + +/// The view when wfassoc querying file extension association. +#[derive(Debug, Copy, Clone)] +pub enum View { + /// The view of current user. + User, + /// The view of system. + System, + /// Hybrid view of User and System. + /// It can be seen as that we use System first and then use User to override any existing items. + Hybrid, +} + +impl From for View { + fn from(value: Scope) -> Self { + match value { + Scope::User => Self::User, + Scope::System => Self::System, + } + } +} + +// endregion + +// region: File Extension + +/// The struct representing an file extension which must start with dot (`.`) +/// and followed by at least one arbitrary characters. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Ext { + /// The body of file extension (excluding dot). + body: String, +} + +impl Ext { + /// Create an new file extension. + pub fn new(raw: &str) -> Result { + Self::from_str(raw) + } + + /// Get the body part of file extension (excluding dot) + pub fn inner(&self) -> &str { + &self.body + } +} + +/// The error occurs when try parsing string into FileExt. +#[derive(Debug, TeError)] +#[error("given file extension name \"{inner}\" is invalid")] +pub struct ParseExtError { + inner: String +} + +impl ParseExtError { + fn new(inner: &str) -> Self { + Self { inner: inner.to_string() } + } +} + +impl Display for Ext { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, ".{}", self.body) + } +} + +impl FromStr for Ext { + type Err = ParseExtError; + + fn from_str(s: &str) -> Result { + static RE: LazyLock = LazyLock::new(|| Regex::new(r"^\.([^\.]+)$").unwrap()); + match RE.captures(s) { + Some(v) => Ok(Self { + body: v[1].to_string(), + }), + None => Err(ParseExtError::new(s)), + } + } +} + +// endregion + +// region: Program + +/// The struct representing a complete program for registration and unregistration. +pub struct Program { + /// The fully qualified path to the application. + full_path: PathBuf, + /// Optional default icon resource for overriding. + /// + /// TODO: Use specialized IconRc struct instead. + default_icon: Option, + /// Optional friendly app name for overriding. + /// + /// TODO: Use specialized StringRc for overriding. + friendly_app_name: Option, + + /// The collection holding all file extensions supported by this program. + exts: IndexSet, +} + +impl Program { + /// Create a new registrar for following operations. + /// + /// `full_path` is the fully qualified path to the application. + /// + /// `default_icon` is an optional icon resource replacing the default one + /// fetched from the first icon resource of your executable application. + /// + /// `friendly_app_name` also is an optional string or string resource replacing + /// the info fetched from executable application's version information. + pub fn new( + full_path: &Path, + default_icon: Option<&str>, + friendly_app_name: Option<&str>, + ) -> Self { + Self { + full_path: full_path.to_path_buf(), + default_icon: default_icon.map(|s| s.to_string()), + friendly_app_name: friendly_app_name.map(|s| s.to_string()), + exts: IndexSet::new(), + } + } + + /// Add file extension supported by this program. + pub fn add_ext(&mut self, ext: &Ext) -> WfResult<()> { + if self.exts.insert(ext.clone()) { + Ok(()) + } else { + Err(WfError::DupExt(ext.to_string())) + } + } +} + +impl Program { const APP_PATHS: &str = "Software\\Microsoft\\Windows\\CurrentVersion\\App Paths"; const APPLICATIONS: &str = "Software\\Classes\\Applications"; /// Register this application. pub fn register(&self, scope: Scope) -> WfResult<()> { + // Check privilege + if !scope.has_privilege() { + return Err(WfError::NoPrivilege); + } + // Fetch root key. let hk = RegKey::predef(match scope { Scope::User => HKEY_CURRENT_USER, @@ -216,7 +339,7 @@ impl Registrar { // Fetch file name and start in path. let file_name = self.extract_file_name()?; let start_in = self.extract_start_in()?; - + // Create App Paths subkey debug_println!("Adding App Paths subkey..."); let subkey_parent = hk.open_subkey_with_flags(Self::APP_PATHS, KEY_READ)?; @@ -230,13 +353,31 @@ impl Registrar { let subkey_parent = hk.open_subkey_with_flags(Self::APPLICATIONS, KEY_READ)?; let (subkey, _) = subkey_parent.create_subkey_with_flags(file_name, KEY_WRITE)?; // Write Applications values - subkey.set_value("FriendlyAppName", &"WoW!")?; + if let Some(default_icon) = &self.default_icon { + subkey.set_value("DefaultIcon", default_icon)?; + } + if let Some(friendly_app_name) = &self.friendly_app_name { + subkey.set_value("FriendlyAppName", friendly_app_name)?; + } + if !self.exts.is_empty() { + let (supported_types, _) = subkey.create_subkey_with_flags("SupportedTypes", KEY_WRITE)?; + for ext in &self.exts { + supported_types.set_value(ext.to_string(), &"")?; + } + } + // Okey + notify_assoc_changed(); Ok(()) } /// Unregister this application. pub fn unregister(&self, scope: Scope) -> WfResult<()> { + // Check privilege + if !scope.has_privilege() { + return Err(WfError::NoPrivilege); + } + // Fetch root key and file name. let hk = RegKey::predef(match scope { Scope::User => HKEY_CURRENT_USER, @@ -255,6 +396,7 @@ impl Registrar { subkey_parent.delete_subkey_all(file_name)?; // Okey + notify_assoc_changed(); Ok(()) } @@ -288,11 +430,12 @@ impl Registrar { } } -impl Registrar { +impl Program { /// Extract the file name part from full path to application, /// which was used in Registry path component. fn extract_file_name(&self) -> WfResult<&OsStr> { - // Get the file name part and make sure it is not empty + // Get the file name part and make sure it is not empty. + // Empty checker is CRUCIAL! self.full_path .file_name() .and_then(|p| if p.is_empty() { None } else { Some(p) }) @@ -303,6 +446,7 @@ impl Registrar { /// which basically is the stem of full path. fn extract_start_in(&self) -> WfResult<&OsStr> { // Get parent part and make sure it is not empty + // Empty checker is CRUCIAL! self.full_path .parent() .map(|p| p.as_os_str()) diff --git a/wfassoc_exec/src/main.rs b/wfassoc_exec/src/main.rs index 3728e02..1f352c7 100644 --- a/wfassoc_exec/src/main.rs +++ b/wfassoc_exec/src/main.rs @@ -2,7 +2,7 @@ use clap::{Parser, Subcommand}; use comfy_table::Table; use std::process; use thiserror::Error as TeError; -use wfassoc::{Error as WfError, FileExt, Scope, View}; +use wfassoc::{Error as WfError, Ext, Scope, View}; // region: Basic Types @@ -95,7 +95,7 @@ fn run_query(cli: Cli) -> Result<()> { ".kra", ".xcf", ".avif", ".qoi", ".apng", ".exr", ]; - for ext in exts.iter().map(|e| FileExt::new(e).unwrap()) { + for ext in exts.iter().map(|e| Ext::new(e).unwrap()) { if let Some(ext_assoc) = ext.query(View::Hybrid) { println!("{:?}", ext_assoc) }