Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit367f887

Browse files
committed
more reg implementation
1 parent3e291b3 commit367f887

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-
usecrate::builtins::PyInt;
18+
usecrate::builtins::{PyInt,PyTuple};
1819
usecrate::common::lock::PyRwLock;
1920
usecrate::function::FuncArgs;
2021
usecrate::protocol::PyNumberMethods;
2122
usecrate::types::AsNumber;
2223
usecrate::{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
fnClose(&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
fnexit(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
fnConnectRegistry(
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+
ifletSome(computer_name) = computer_name{
263+
letmut 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+
letmut 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+
fnEnumValue(hkey:PyRef<PyHKEYObject>,index:u32,vm:&VirtualMachine) ->PyResult{
418+
// Query registry for the required buffer sizes.
419+
letmut ret_value_size:u32 =0;
420+
letmut 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_sizeas*mutu32,
433+
&mut ret_data_sizeas*mutu32,
434+
ptr::null_mut(),
435+
ptr::null_mut(),
436+
)
437+
};
438+
if rc !=0{
439+
returnErr(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+
letmut buf_value_size = ret_value_size;
449+
letmut buf_data_size = ret_data_size;
450+
451+
// Allocate buffers.
452+
letmut ret_value_buf:Vec<u16> =vec![0; ret_value_sizeasusize];
453+
letmut ret_data_buf:Vec<u8> =vec![0; ret_data_sizeasusize];
454+
455+
// Loop to enumerate the registry value.
456+
loop{
457+
letmut current_value_size = ret_value_size;
458+
letmut 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_sizeas*mutu32,
465+
ptr::null_mut(),
466+
{
467+
// typ will hold the registry data type.
468+
letmut t =0u32;
469+
&mut t
470+
},
471+
ret_data_buf.as_mut_ptr(),
472+
&mut current_data_sizeas*mutu32,
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_sizeasusize,0);
480+
ret_value_buf.resize(buf_value_sizeasusize,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+
returnErr(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+
letmut reg_type:u32 =0;
496+
unsafe{
497+
Registry::RegEnumValueW(
498+
hkey,
499+
index,
500+
ret_value_buf.as_mut_ptr(),
501+
&mut current_value_sizeas*mutu32,
502+
ptr::null_mut(),
503+
&mut reg_typeas*mutu32,
504+
ret_data_buf.as_mut_ptr(),
505+
&mut current_data_sizeas*mutu32,
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_sizeasusize];
519+
let py_data =reg_to_py(vm, data_slice, reg_type)?;
520+
521+
// Return tuple (value_name, data, type)
522+
returnOk(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
fnFlushKey(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-
fnQueryInfoKey(key:PyRef<PyHKEYObject>,vm:&VirtualMachine) ->PyResult<()>{
592+
fnQueryInfoKey(key:PyRef<PyHKEYObject>,vm:&VirtualMachine) ->PyResult<PyRef<PyTuple>>{
435593
let key =*key.hkey.read();
436594
letmut lpcsubkeys:u32 =0;
437595
letmut lpcvalues:u32 =0;
438-
let lpftlastwritetime:*mutFoundation::FILETIME = std::ptr::null_mut();
596+
letmutlpftlastwritetime: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+
&mutlpftlastwritetime,
453611
)
454612
};
455613

456614
if err !=0{
457-
Err(vm.new_os_error(format!("error code: {}", err)))
458-
}else{
459-
Ok(())
615+
returnErr(vm.new_os_error(format!("error code: {}", err)));
460616
}
617+
let l:u64 =(lpftlastwritetime.dwHighDateTimeasu64) <<32 | lpftlastwritetime.dwLowDateTimeasu64;
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+
fnQueryValueEx(key:PyRef<PyHKEYObject>,name:String,vm:&VirtualMachine) ->PyResult<PyObjectRef>{
646+
let wide_name =to_utf16(name);
647+
letmut 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+
returnErr(vm.new_os_error(format!("error code: {}", res)));
661+
}
662+
letmut retBuf =Vec::with_capacity(buf_sizeasusize);
663+
letmut 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+
returnErr(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
fnSaveKey(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+
fnreg_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*constu16, 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*constu16, u16_count)
768+
};
769+
letmut strings:Vec<PyObjectRef> =Vec::new();
770+
letmut 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
fnpy2reg(value:PyObjectRef,typ:u32,vm:&VirtualMachine) ->PyResult<Option<Vec<u8>>>{
535798
match typ{
536799
REG_DWORD =>{
@@ -570,7 +833,7 @@ mod winreg {
570833
fnSetValueEx(
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)
[8]ページ先頭
©2009-2025 Movatter.jp