From 65ed30b2ab9d837b80e76a6bd81c75f8a5548205 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Tue, 25 Feb 2020 03:51:52 +0000 Subject: [PATCH 1/4] Add winreg module --- Cargo.lock | 10 +++ vm/Cargo.toml | 3 + vm/src/stdlib/mod.rs | 3 + vm/src/stdlib/winreg.rs | 170 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 186 insertions(+) create mode 100644 vm/src/stdlib/winreg.rs diff --git a/Cargo.lock b/Cargo.lock index 50b3be9a1..2f1ebfa1f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1557,6 +1557,7 @@ dependencies = [ "volatile", "wasm-bindgen", "winapi", + "winreg", ] [[package]] @@ -2429,3 +2430,12 @@ name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "winreg" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2986deb581c4fe11b621998a5e53361efe6b48a151178d0cd9eeffa4dc6acc9" +dependencies = [ + "winapi", +] diff --git a/vm/Cargo.toml b/vm/Cargo.toml index f0e9a89e5..cafa8d357 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -94,6 +94,9 @@ rustyline = "6.0" [target.'cfg(not(any(target_arch = "wasm32", target_os = "redox")))'.dependencies] dns-lookup = "1.0" +[target.'cfg(windows)'.dependencies] +winreg = "0.6" + [target."cfg(windows)".dependencies.winapi] version = "0.3" features = ["winsock2", "handleapi", "ws2def", "std", "winbase"] diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs index ea305c19b..6c7cc8ec2 100644 --- a/vm/src/stdlib/mod.rs +++ b/vm/src/stdlib/mod.rs @@ -52,6 +52,8 @@ pub mod signal; mod subprocess; #[cfg(windows)] mod winapi; +#[cfg(windows)] +mod winreg; #[cfg(not(target_arch = "wasm32"))] mod zlib; @@ -135,6 +137,7 @@ pub fn get_module_inits() -> HashMap { #[cfg(windows)] { modules.insert("_winapi".to_owned(), Box::new(winapi::make_module)); + modules.insert("winreg".to_owned(), Box::new(winreg::make_module)); } modules diff --git a/vm/src/stdlib/winreg.rs b/vm/src/stdlib/winreg.rs new file mode 100644 index 000000000..849016828 --- /dev/null +++ b/vm/src/stdlib/winreg.rs @@ -0,0 +1,170 @@ +#![allow(non_snake_case)] + +use std::cell::{Ref, RefCell}; + +use super::os; +use crate::function::OptionalArg; +use crate::obj::objstr::PyStringRef; +use crate::obj::objtype::PyClassRef; +use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject}; +use crate::VirtualMachine; + +use winreg::RegKey; + +#[pyclass] +#[derive(Debug)] +struct PyHKEY { + key: RefCell, +} +type PyHKEYRef = PyRef; + +impl PyValue for PyHKEY { + fn class(vm: &VirtualMachine) -> PyClassRef { + vm.class("winreg", "HKEYType") + } +} + +#[pyimpl] +impl PyHKEY { + fn new(key: RegKey) -> Self { + Self { + key: RefCell::new(key), + } + } + + fn key(&self) -> Ref { + self.key.borrow() + } + + #[pymethod] + fn Close(&self) { + let null_key = RegKey::predef(0 as winreg::HKEY); + let key = self.key.replace(null_key); + drop(key); + } + #[pymethod] + fn Detach(&self) -> usize { + let null_key = RegKey::predef(0 as winreg::HKEY); + let key = self.key.replace(null_key); + let handle = key.raw_handle(); + std::mem::forget(key); + handle as usize + } + + #[pymethod(magic)] + fn bool(&self) -> bool { + !self.key().raw_handle().is_null() + } + #[pymethod(magic)] + fn enter(zelf: PyRef) -> PyRef { + zelf + } + #[pymethod(magic)] + fn exit(&self, _cls: PyObjectRef, _exc: PyObjectRef, _tb: PyObjectRef) { + self.Close(); + } +} + +enum Hkey { + PyHKEY(PyHKEYRef), + Constant(winreg::HKEY), +} +impl TryFromObject for Hkey { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + obj.downcast() + .map(Self::PyHKEY) + .or_else(|o| usize::try_from_object(vm, o).map(|i| Self::Constant(i as winreg::HKEY))) + } +} +impl Hkey { + fn with_key(&self, f: impl FnOnce(&RegKey) -> R) -> R { + match self { + Self::PyHKEY(py) => f(&py.key()), + Self::Constant(hkey) => { + let k = RegKey::predef(*hkey); + let res = f(&k); + std::mem::forget(k); + res + } + } + } +} + +fn winreg_OpenKey( + key: Hkey, + subkey: Option, + reserved: OptionalArg, + access: OptionalArg, + vm: &VirtualMachine, +) -> PyResult { + let reserved = reserved.unwrap_or(0); + let access = access.unwrap_or(winreg::enums::KEY_READ); + if reserved != 0 { + // RegKey::open_subkey* doesn't have a reserved param, so this'll do + return Err(vm.new_value_error("reserved param must be 0".to_owned())); + } + + let subkey = subkey.as_ref().map_or("", |s| s.as_str()); + let key = key + .with_key(|k| k.open_subkey_with_flags(subkey, access)) + .map_err(|e| os::convert_io_error(vm, e))?; + + Ok(PyHKEY::new(key)) +} + +fn winreg_QueryValue( + key: Hkey, + subkey: Option, + vm: &VirtualMachine, +) -> PyResult { + let subkey = subkey.as_ref().map_or("", |s| s.as_str()); + key.with_key(|k| k.get_value(subkey)) + .map_err(|e| os::convert_io_error(vm, e)) +} + +fn winreg_QueryValueEx( + key: Hkey, + subkey: Option, + vm: &VirtualMachine, +) -> PyResult<(usize, Vec)> { + let subkey = subkey.as_ref().map_or("", |s| s.as_str()); + key.with_key(|k| k.get_raw_value(subkey)) + .map(|regval| (regval.vtype as usize, regval.bytes)) + .map_err(|e| os::convert_io_error(vm, e)) +} + +fn winreg_CloseKey(key: PyHKEYRef) { + key.Close(); +} + +pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { + let ctx = &vm.ctx; + let hkey_type = PyHKEY::make_class(ctx); + let module = py_module!(vm, "winreg", { + "HKEYType" => hkey_type, + "OpenKey" => ctx.new_function(winreg_OpenKey), + "QueryValue" => ctx.new_function(winreg_QueryValue), + "QueryValueEx" => ctx.new_function(winreg_QueryValueEx), + "CloseKey" => ctx.new_function(winreg_CloseKey), + }); + + macro_rules! add_hkey_constants { + ($($name:ident),*) => { + extend_module!(vm, module, { + $((stringify!($name)) => ctx.new_int(winreg::enums::$name as usize)),* + }) + }; + } + + add_hkey_constants!( + HKEY_CLASSES_ROOT, + HKEY_CURRENT_USER, + HKEY_LOCAL_MACHINE, + HKEY_USERS, + HKEY_PERFORMANCE_DATA, + HKEY_CURRENT_CONFIG, + HKEY_DYN_DATA + ); + + module +} From a79dfcd0274298d80f7512bfd713c18b386b43ce Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Tue, 25 Feb 2020 11:39:51 -0600 Subject: [PATCH 2/4] Enable importing winreg in importlib._boostrap_external --- Lib/importlib/_bootstrap_external.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Lib/importlib/_bootstrap_external.py b/Lib/importlib/_bootstrap_external.py index 18eb43f2a..d60cb7691 100644 --- a/Lib/importlib/_bootstrap_external.py +++ b/Lib/importlib/_bootstrap_external.py @@ -1605,8 +1605,7 @@ def _setup(_bootstrap_module): setattr(self_module, '_weakref', weakref_module) # Directly load the winreg module (needed during bootstrap). - # XXX RustPython TODO: winreg module - if builtin_os == 'nt' and False: + if builtin_os == 'nt': winreg_module = _bootstrap._builtin_from_name('winreg') setattr(self_module, '_winreg', winreg_module) From 2f617feb107005e6733be7b7cb0f02a0b4986788 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Fri, 28 Feb 2020 11:45:08 -0600 Subject: [PATCH 3/4] Convert registry values to Python values --- Cargo.lock | 4 +- vm/Cargo.toml | 2 +- vm/src/stdlib/winreg.rs | 115 ++++++++++++++++++++++++++++++++++++++-- 3 files changed, 113 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2f1ebfa1f..939e39baa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2433,9 +2433,9 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" [[package]] name = "winreg" -version = "0.6.2" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2986deb581c4fe11b621998a5e53361efe6b48a151178d0cd9eeffa4dc6acc9" +checksum = "0120db82e8a1e0b9fb3345a539c478767c0048d842860994d96113d5b667bd69" dependencies = [ "winapi", ] diff --git a/vm/Cargo.toml b/vm/Cargo.toml index cafa8d357..be6a88622 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -95,7 +95,7 @@ rustyline = "6.0" dns-lookup = "1.0" [target.'cfg(windows)'.dependencies] -winreg = "0.6" +winreg = "0.7" [target."cfg(windows)".dependencies.winapi] version = "0.3" diff --git a/vm/src/stdlib/winreg.rs b/vm/src/stdlib/winreg.rs index 849016828..6259aa83a 100644 --- a/vm/src/stdlib/winreg.rs +++ b/vm/src/stdlib/winreg.rs @@ -1,6 +1,8 @@ #![allow(non_snake_case)] use std::cell::{Ref, RefCell}; +use std::convert::TryInto; +use std::io; use super::os; use crate::function::OptionalArg; @@ -9,7 +11,8 @@ use crate::obj::objtype::PyClassRef; use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject}; use crate::VirtualMachine; -use winreg::RegKey; +use winapi::shared::winerror; +use winreg::{enums::RegType, RegKey, RegValue}; #[pyclass] #[derive(Debug)] @@ -126,15 +129,115 @@ fn winreg_QueryValueEx( key: Hkey, subkey: Option, vm: &VirtualMachine, -) -> PyResult<(usize, Vec)> { +) -> PyResult<(PyObjectRef, usize)> { let subkey = subkey.as_ref().map_or("", |s| s.as_str()); key.with_key(|k| k.get_raw_value(subkey)) - .map(|regval| (regval.vtype as usize, regval.bytes)) + .map_err(|e| os::convert_io_error(vm, e)) + .and_then(|regval| { + let ty = regval.vtype.clone() as usize; + Ok((reg_to_py(regval, vm)?, ty)) + }) +} + +fn winreg_EnumKey(key: Hkey, index: u32, vm: &VirtualMachine) -> PyResult { + key.with_key(|k| k.enum_keys().nth(index as usize)) + .unwrap_or_else(|| { + Err(io::Error::from_raw_os_error( + winerror::ERROR_NO_MORE_ITEMS as i32, + )) + }) .map_err(|e| os::convert_io_error(vm, e)) } -fn winreg_CloseKey(key: PyHKEYRef) { - key.Close(); +fn winreg_EnumValue( + key: Hkey, + index: u32, + vm: &VirtualMachine, +) -> PyResult<(String, PyObjectRef, usize)> { + key.with_key(|k| k.enum_values().nth(index as usize)) + .unwrap_or_else(|| { + Err(io::Error::from_raw_os_error( + winerror::ERROR_NO_MORE_ITEMS as i32, + )) + }) + .map_err(|e| os::convert_io_error(vm, e)) + .and_then(|(name, value)| { + let ty = value.vtype.clone() as usize; + Ok((name, reg_to_py(value, vm)?, ty)) + }) +} + +fn winreg_CloseKey(key: Hkey) { + match key { + Hkey::PyHKEY(py) => py.Close(), + Hkey::Constant(hkey) => drop(RegKey::predef(hkey)), + } +} + +fn reg_to_py(value: RegValue, vm: &VirtualMachine) -> PyResult { + macro_rules! bytes_to_int { + ($int:ident, $f:ident, $name:ident) => {{ + let i = if value.bytes.is_empty() { + Ok(0 as $int) + } else { + (&*value.bytes).try_into().map($int::$f).map_err(|_| { + vm.new_value_error(format!("{} value is wrong length", stringify!(name))) + }) + }; + i.map(|i| vm.new_int(i)) + }}; + }; + let bytes_to_wide = |b: &[u8]| -> Option<&[u16]> { + if b.len() % 2 == 0 { + Some(unsafe { std::slice::from_raw_parts(b.as_ptr().cast(), b.len() / 2) }) + } else { + None + } + }; + match value.vtype { + RegType::REG_DWORD => bytes_to_int!(u32, from_ne_bytes, REG_DWORD), + RegType::REG_DWORD_BIG_ENDIAN => bytes_to_int!(u32, from_be_bytes, REG_DWORD_BIG_ENDIAN), + RegType::REG_QWORD => bytes_to_int!(u64, from_ne_bytes, REG_DWORD), + // RegType::REG_QWORD_BIG_ENDIAN => bytes_to_int!(u64, from_be_bytes, REG_DWORD_BIG_ENDIAN), + RegType::REG_SZ | RegType::REG_EXPAND_SZ => { + let wide_slice = bytes_to_wide(&value.bytes).ok_or_else(|| { + vm.new_value_error("REG_SZ string doesn't have an even byte length".to_owned()) + })?; + let nul_pos = wide_slice + .iter() + .position(|w| *w == 0) + .unwrap_or_else(|| wide_slice.len()); + let s = String::from_utf16_lossy(&wide_slice[..nul_pos]); + Ok(vm.new_str(s)) + } + RegType::REG_MULTI_SZ => { + if value.bytes.is_empty() { + return Ok(vm.ctx.new_list(vec![])); + } + let wide_slice = bytes_to_wide(&value.bytes).ok_or_else(|| { + vm.new_value_error( + "REG_MULTI_SZ string doesn't have an even byte length".to_owned(), + ) + })?; + let wide_slice = if let Some((0, rest)) = wide_slice.split_last() { + rest + } else { + wide_slice + }; + let strings = wide_slice + .split(|c| *c == 0) + .map(|s| vm.new_str(String::from_utf16_lossy(s))) + .collect(); + Ok(vm.ctx.new_list(strings)) + } + RegType::REG_BINARY | _ => { + if value.bytes.is_empty() { + Ok(vm.get_none()) + } else { + Ok(vm.ctx.new_bytes(value.bytes)) + } + } + } } pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { @@ -145,6 +248,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "OpenKey" => ctx.new_function(winreg_OpenKey), "QueryValue" => ctx.new_function(winreg_QueryValue), "QueryValueEx" => ctx.new_function(winreg_QueryValueEx), + "EnumKey" => ctx.new_function(winreg_EnumKey), + "EnumValue" => ctx.new_function(winreg_EnumValue), "CloseKey" => ctx.new_function(winreg_CloseKey), }); From eec333399fa295e2b6d73d2eefbed5a5ed578a30 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Tue, 3 Mar 2020 11:32:00 -0600 Subject: [PATCH 4/4] Add more winreg constants --- vm/src/stdlib/winreg.rs | 45 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 41 insertions(+), 4 deletions(-) diff --git a/vm/src/stdlib/winreg.rs b/vm/src/stdlib/winreg.rs index 6259aa83a..e08eabc8c 100644 --- a/vm/src/stdlib/winreg.rs +++ b/vm/src/stdlib/winreg.rs @@ -253,22 +253,59 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "CloseKey" => ctx.new_function(winreg_CloseKey), }); - macro_rules! add_hkey_constants { - ($($name:ident),*) => { + macro_rules! add_constants { + (hkey, $($name:ident),*$(,)?) => { extend_module!(vm, module, { $((stringify!($name)) => ctx.new_int(winreg::enums::$name as usize)),* }) }; + (winnt, $($name:ident),*$(,)?) => { + extend_module!(vm, module, { + $((stringify!($name)) => ctx.new_int(winapi::um::winnt::$name)),* + }) + }; } - add_hkey_constants!( + add_constants!( + hkey, HKEY_CLASSES_ROOT, HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE, HKEY_USERS, HKEY_PERFORMANCE_DATA, HKEY_CURRENT_CONFIG, - HKEY_DYN_DATA + HKEY_DYN_DATA, + ); + add_constants!( + winnt, + // access rights + KEY_ALL_ACCESS, + KEY_WRITE, + KEY_READ, + KEY_EXECUTE, + KEY_QUERY_VALUE, + KEY_SET_VALUE, + KEY_CREATE_SUB_KEY, + KEY_ENUMERATE_SUB_KEYS, + KEY_NOTIFY, + KEY_CREATE_LINK, + KEY_WOW64_64KEY, + KEY_WOW64_32KEY, + // value types + REG_BINARY, + REG_DWORD, + REG_DWORD_LITTLE_ENDIAN, + REG_DWORD_BIG_ENDIAN, + REG_EXPAND_SZ, + REG_LINK, + REG_MULTI_SZ, + REG_NONE, + REG_QWORD, + REG_QWORD_LITTLE_ENDIAN, + REG_RESOURCE_LIST, + REG_FULL_RESOURCE_DESCRIPTOR, + REG_RESOURCE_REQUIREMENTS_LIST, + REG_SZ, ); module