Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 367f887

Browse files
arihant2matharihant2math
committed
more reg implementation
1 parent 3e291b3 commit 367f887

File tree

1 file changed

+286
-23
lines changed

1 file changed

+286
-23
lines changed

‎vm/src/stdlib/winreg.rs

Lines changed: 286 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,17 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
1212
mod winreg {
1313
use std::ffi::OsStr;
1414
use std::os::windows::ffi::OsStrExt;
15+
use std::ptr;
1516
use std::sync::Arc;
1617

17-
use crate::builtins::PyInt;
18+
use crate::builtins::{PyInt, PyTuple};
1819
use crate::common::lock::PyRwLock;
1920
use crate::function::FuncArgs;
2021
use crate::protocol::PyNumberMethods;
2122
use crate::types::AsNumber;
2223
use crate::{PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine};
2324

24-
use windows_sys::Win32::Foundation;
25+
use windows_sys::Win32::Foundation::{self, ERROR_MORE_DATA};
2526
use windows_sys::Win32::System::Registry;
2627

2728
use num_traits::ToPrimitive;
@@ -134,6 +135,7 @@ mod winreg {
134135
#[pymethod]
135136
fn Close(&self, vm: &VirtualMachine) -> PyResult<()> {
136137
let res = unsafe { Registry::RegCloseKey(*self.hkey.write()) };
138+
*self.hkey.write() = std::ptr::null_mut();
137139
if res == 0 {
138140
Ok(())
139141
} else {
@@ -167,6 +169,7 @@ mod winreg {
167169
#[pymethod(magic)]
168170
fn exit(zelf: PyRef<Self>, _args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
169171
let res = unsafe { Registry::RegCloseKey(*zelf.hkey.write()) };
172+
*zelf.hkey.write() = std::ptr::null_mut();
170173
if res == 0 {
171174
Ok(())
172175
} else {
@@ -252,22 +255,39 @@ mod winreg {
252255
// TODO: Computer name can be `None`
253256
#[pyfunction]
254257
fn ConnectRegistry(
255-
computer_name: String,
258+
computer_name: Option<String>,
256259
key: PyRef<PyHKEYObject>,
257260
vm: &VirtualMachine,
258-
) -> PyResult<()> {
259-
let wide_computer_name = to_utf16(computer_name);
260-
let res = unsafe {
261-
Registry::RegConnectRegistryW(
262-
wide_computer_name.as_ptr(),
263-
*key.hkey.read(),
264-
std::ptr::null_mut(),
265-
)
266-
};
267-
if res == 0 {
268-
Ok(())
261+
) -> PyResult<PyHKEYObject> {
262+
if let Some(computer_name) = computer_name {
263+
let mut ret_key = std::ptr::null_mut();
264+
let wide_computer_name = to_utf16(computer_name);
265+
let res = unsafe {
266+
Registry::RegConnectRegistryW(
267+
wide_computer_name.as_ptr(),
268+
*key.hkey.read(),
269+
&mut ret_key
270+
)
271+
};
272+
if res == 0 {
273+
Ok(PyHKEYObject::new(ret_key))
274+
} else {
275+
Err(vm.new_os_error(format!("error code: {}", res)))
276+
}
269277
} else {
270-
Err(vm.new_os_error(format!("error code: {}", res)))
278+
let mut ret_key = std::ptr::null_mut();
279+
let res = unsafe {
280+
Registry::RegConnectRegistryW(
281+
std::ptr::null_mut(),
282+
*key.hkey.read(),
283+
&mut ret_key
284+
)
285+
};
286+
if res == 0 {
287+
Ok(PyHKEYObject::new(ret_key))
288+
} else {
289+
Err(vm.new_os_error(format!("error code: {}", res)))
290+
}
271291
}
272292
}
273293

@@ -369,6 +389,144 @@ mod winreg {
369389
}
370390
}
371391

392+
// #[pyfunction]
393+
// fn EnumKey(key: PyRef<PyHKEYObject>, index: i32, vm: &VirtualMachine) -> PyResult<String> {
394+
// let mut tmpbuf = [0u16; 257];
395+
// let mut len = std::mem::sizeof(tmpbuf.len())/std::mem::sizeof(tmpbuf[0]);
396+
// let res = unsafe {
397+
// Registry::RegEnumKeyExW(
398+
// *key.hkey.read(),
399+
// index as u32,
400+
// tmpbuf.as_mut_ptr(),
401+
// &mut len,
402+
// std::ptr::null_mut(),
403+
// std::ptr::null_mut(),
404+
// std::ptr::null_mut(),
405+
// std::ptr::null_mut(),
406+
// )
407+
// };
408+
// if res != 0 {
409+
// return Err(vm.new_os_error(format!("error code: {}", res)));
410+
// }
411+
// let s = String::from_utf16(&tmpbuf[..len as usize])
412+
// .map_err(|e| vm.new_value_error(format!("UTF16 error: {}", e)))?;
413+
// Ok(s)
414+
// }
415+
416+
#[pyfunction]
417+
fn EnumValue(hkey: PyRef<PyHKEYObject>, index: u32, vm: &VirtualMachine) -> PyResult {
418+
// Query registry for the required buffer sizes.
419+
let mut ret_value_size: u32 = 0;
420+
let mut ret_data_size: u32 = 0;
421+
let hkey: *mut std::ffi::c_void = *hkey.hkey.read();
422+
let rc = unsafe {
423+
Registry::RegQueryInfoKeyW(
424+
hkey,
425+
ptr::null_mut(),
426+
ptr::null_mut(),
427+
ptr::null_mut(),
428+
ptr::null_mut(),
429+
ptr::null_mut(),
430+
ptr::null_mut(),
431+
ptr::null_mut(),
432+
&mut ret_value_size as *mut u32,
433+
&mut ret_data_size as *mut u32,
434+
ptr::null_mut(),
435+
ptr::null_mut(),
436+
)
437+
};
438+
if rc != 0 {
439+
return Err(vm.new_os_error(format!(
440+
"RegQueryInfoKeyW failed with error code {}",
441+
rc
442+
)));
443+
}
444+
445+
// Include room for null terminators.
446+
ret_value_size += 1;
447+
ret_data_size += 1;
448+
let mut buf_value_size = ret_value_size;
449+
let mut buf_data_size = ret_data_size;
450+
451+
// Allocate buffers.
452+
let mut ret_value_buf: Vec<u16> = vec![0; ret_value_size as usize];
453+
let mut ret_data_buf: Vec<u8> = vec![0; ret_data_size as usize];
454+
455+
// Loop to enumerate the registry value.
456+
loop {
457+
let mut current_value_size = ret_value_size;
458+
let mut current_data_size = ret_data_size;
459+
let rc = unsafe {
460+
Registry::RegEnumValueW(
461+
hkey,
462+
index,
463+
ret_value_buf.as_mut_ptr(),
464+
&mut current_value_size as *mut u32,
465+
ptr::null_mut(),
466+
{
467+
// typ will hold the registry data type.
468+
let mut t = 0u32;
469+
&mut t
470+
},
471+
ret_data_buf.as_mut_ptr(),
472+
&mut current_data_size as *mut u32,
473+
)
474+
};
475+
if rc == ERROR_MORE_DATA {
476+
// Double the buffer sizes.
477+
buf_data_size *= 2;
478+
buf_value_size *= 2;
479+
ret_data_buf.resize(buf_data_size as usize, 0);
480+
ret_value_buf.resize(buf_value_size as usize, 0);
481+
// Reset sizes for next iteration.
482+
ret_value_size = buf_value_size;
483+
ret_data_size = buf_data_size;
484+
continue;
485+
}
486+
if rc != 0 {
487+
return Err(vm.new_os_error(format!(
488+
"RegEnumValueW failed with error code {}",
489+
rc
490+
)));
491+
}
492+
493+
// At this point, current_value_size and current_data_size have been updated.
494+
// Retrieve the registry type.
495+
let mut reg_type: u32 = 0;
496+
unsafe {
497+
Registry::RegEnumValueW(
498+
hkey,
499+
index,
500+
ret_value_buf.as_mut_ptr(),
501+
&mut current_value_size as *mut u32,
502+
ptr::null_mut(),
503+
&mut reg_type as *mut u32,
504+
ret_data_buf.as_mut_ptr(),
505+
&mut current_data_size as *mut u32,
506+
)
507+
};
508+
509+
// Convert the registry value name from UTF‑16.
510+
let name_len = ret_value_buf
511+
.iter()
512+
.position(|&c| c == 0)
513+
.unwrap_or(ret_value_buf.len());
514+
let name = String::from_utf16(&ret_value_buf[..name_len])
515+
.map_err(|e| vm.new_value_error(format!("UTF16 conversion error: {}", e)))?;
516+
517+
// Slice the data buffer to the actual size returned.
518+
let data_slice = &ret_data_buf[..current_data_size as usize];
519+
let py_data = reg_to_py(vm, data_slice, reg_type)?;
520+
521+
// Return tuple (value_name, data, type)
522+
return Ok(vm.ctx.new_tuple(vec![
523+
vm.ctx.new_str(name).into(),
524+
py_data,
525+
vm.ctx.new_int(reg_type).into(),
526+
]).into());
527+
}
528+
}
529+
372530
#[pyfunction]
373531
fn FlushKey(key: PyRef<PyHKEYObject>, vm: &VirtualMachine) -> PyResult<()> {
374532
let res = unsafe { Registry::RegFlushKey(*key.hkey.read()) };
@@ -431,11 +589,11 @@ mod winreg {
431589
}
432590

433591
#[pyfunction]
434-
fn QueryInfoKey(key: PyRef<PyHKEYObject>, vm: &VirtualMachine) -> PyResult<()> {
592+
fn QueryInfoKey(key: PyRef<PyHKEYObject>, vm: &VirtualMachine) -> PyResult<PyRef<PyTuple>> {
435593
let key = *key.hkey.read();
436594
let mut lpcsubkeys: u32 = 0;
437595
let mut lpcvalues: u32 = 0;
438-
let lpftlastwritetime: *mut Foundation::FILETIME = std::ptr::null_mut();
596+
let mut lpftlastwritetime: Foundation::FILETIME = unsafe { std::mem::zeroed() };
439597
let err = unsafe {
440598
Registry::RegQueryInfoKeyW(
441599
key,
@@ -449,15 +607,16 @@ mod winreg {
449607
std::ptr::null_mut(),
450608
std::ptr::null_mut(),
451609
std::ptr::null_mut(),
452-
lpftlastwritetime,
610+
&mut lpftlastwritetime,
453611
)
454612
};
455613

456614
if err != 0 {
457-
Err(vm.new_os_error(format!("error code: {}", err)))
458-
} else {
459-
Ok(())
615+
return Err(vm.new_os_error(format!("error code: {}", err)));
460616
}
617+
let l: u64 = (lpftlastwritetime.dwHighDateTime as u64) << 32 | lpftlastwritetime.dwLowDateTime as u64;
618+
let tup: Vec<PyObjectRef> = vec![vm.ctx.new_int(lpcsubkeys).into(), vm.ctx.new_int(lpcvalues).into(), vm.ctx.new_int(l).into()];
619+
Ok(vm.ctx.new_tuple(tup))
461620
}
462621

463622
#[pyfunction]
@@ -482,7 +641,44 @@ mod winreg {
482641
Ok(())
483642
}
484643

485-
// TODO: QueryValueEx
644+
#[pyfunction]
645+
fn QueryValueEx(key: PyRef<PyHKEYObject>, name: String, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
646+
let wide_name = to_utf16(name);
647+
let mut buf_size = 0;
648+
let res = unsafe {
649+
Registry::RegQueryValueExW(
650+
*key.hkey.read(),
651+
wide_name.as_ptr(),
652+
std::ptr::null_mut(),
653+
std::ptr::null_mut(),
654+
std::ptr::null_mut(),
655+
&mut buf_size,
656+
)
657+
};
658+
// TODO: res == ERROR_MORE_DATA
659+
if res != 0 {
660+
return Err(vm.new_os_error(format!("error code: {}", res)));
661+
}
662+
let mut retBuf = Vec::with_capacity(buf_size as usize);
663+
let mut typ = 0;
664+
let res = unsafe {
665+
Registry::RegQueryValueExW(
666+
*key.hkey.read(),
667+
wide_name.as_ptr(),
668+
std::ptr::null_mut(),
669+
&mut typ,
670+
retBuf.as_mut_ptr(),
671+
&mut buf_size,
672+
)
673+
};
674+
// TODO: res == ERROR_MORE_DATA
675+
if res != 0 {
676+
return Err(vm.new_os_error(format!("error code: {}", res)));
677+
}
678+
let obj = reg_to_py(vm, retBuf.as_slice(), typ)?;
679+
Ok(obj)
680+
}
681+
486682
#[pyfunction]
487683
fn SaveKey(key: PyRef<PyHKEYObject>, file_name: String, vm: &VirtualMachine) -> PyResult<()> {
488684
let file_name = to_utf16(file_name);
@@ -531,6 +727,73 @@ mod winreg {
531727
}
532728
}
533729

730+
fn reg_to_py(vm: &VirtualMachine, ret_data: &[u8], typ: u32) -> PyResult {
731+
match typ {
732+
REG_DWORD => {
733+
// If there isn’t enough data, return 0.
734+
if ret_data.len() < std::mem::size_of::<u32>() {
735+
Ok(vm.ctx.new_int(0).into())
736+
} else {
737+
let val = u32::from_ne_bytes(ret_data[..4].try_into().unwrap());
738+
Ok(vm.ctx.new_int(val).into())
739+
}
740+
}
741+
REG_QWORD => {
742+
if ret_data.len() < std::mem::size_of::<u64>() {
743+
Ok(vm.ctx.new_int(0).into())
744+
} else {
745+
let val = u64::from_ne_bytes(ret_data[..8].try_into().unwrap());
746+
Ok(vm.ctx.new_int(val).into())
747+
}
748+
}
749+
REG_SZ | REG_EXPAND_SZ => {
750+
// Treat the data as a UTF-16 string.
751+
let u16_count = ret_data.len() / 2;
752+
let u16_slice = unsafe {
753+
std::slice::from_raw_parts(ret_data.as_ptr() as *const u16, u16_count)
754+
};
755+
// Only use characters up to the first NUL.
756+
let len = u16_slice.iter().position(|&c| c == 0).unwrap_or(u16_slice.len());
757+
let s = String::from_utf16(&u16_slice[..len])
758+
.map_err(|e| vm.new_value_error(format!("UTF16 error: {}", e)))?;
759+
Ok(vm.ctx.new_str(s).into())
760+
}
761+
REG_MULTI_SZ => {
762+
if ret_data.is_empty() {
763+
Ok(vm.ctx.new_list(vec![]).into())
764+
} else {
765+
let u16_count = ret_data.len() / 2;
766+
let u16_slice = unsafe {
767+
std::slice::from_raw_parts(ret_data.as_ptr() as *const u16, u16_count)
768+
};
769+
let mut strings: Vec<PyObjectRef> = Vec::new();
770+
let mut start = 0;
771+
for (i, &c) in u16_slice.iter().enumerate() {
772+
if c == 0 {
773+
// An empty string signals the end.
774+
if start == i {
775+
break;
776+
}
777+
let s = String::from_utf16(&u16_slice[start..i])
778+
.map_err(|e| vm.new_value_error(format!("UTF16 error: {}", e)))?;
779+
strings.push(vm.ctx.new_str(s).into());
780+
start = i + 1;
781+
}
782+
}
783+
Ok(vm.ctx.new_list(strings).into())
784+
}
785+
}
786+
// For REG_BINARY and any other unknown types, return a bytes object if data exists.
787+
_ => {
788+
if ret_data.is_empty() {
789+
Ok(vm.ctx.none())
790+
} else {
791+
Ok(vm.ctx.new_bytes(ret_data.to_vec()).into())
792+
}
793+
}
794+
}
795+
}
796+
534797
fn py2reg(value: PyObjectRef, typ: u32, vm: &VirtualMachine) -> PyResult<Option<Vec<u8>>> {
535798
match typ {
536799
REG_DWORD => {
@@ -570,7 +833,7 @@ mod winreg {
570833
fn SetValueEx(
571834
key: PyRef<PyHKEYObject>,
572835
value_name: String,
573-
reserved: u32,
836+
_reserved: u32,
574837
typ: u32,
575838
value: PyObjectRef,
576839
vm: &VirtualMachine,

0 commit comments

Comments
 (0)