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

Commitdb9774d

Browse files
committed
more reg implementation
1 parentd12db56 commitdb9774d

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
@@ -11,16 +11,17 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
1111
mod winreg{
1212
use std::ffi::OsStr;
1313
use std::os::windows::ffi::OsStrExt;
14+
use std::ptr;
1415
use std::sync::Arc;
1516

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

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

2627
use num_traits::ToPrimitive;
@@ -133,6 +134,7 @@ mod winreg {
133134
#[pymethod]
134135
fnClose(&self,vm:&VirtualMachine) ->PyResult<()>{
135136
let res =unsafe{Registry::RegCloseKey(*self.hkey.write())};
137+
*self.hkey.write() = std::ptr::null_mut();
136138
if res ==0{
137139
Ok(())
138140
}else{
@@ -166,6 +168,7 @@ mod winreg {
166168
#[pymethod(magic)]
167169
fnexit(zelf:PyRef<Self>,_args:FuncArgs,vm:&VirtualMachine) ->PyResult<()>{
168170
let res =unsafe{Registry::RegCloseKey(*zelf.hkey.write())};
171+
*zelf.hkey.write() = std::ptr::null_mut();
169172
if res ==0{
170173
Ok(())
171174
}else{
@@ -251,22 +254,39 @@ mod winreg {
251254
// TODO: Computer name can be `None`
252255
#[pyfunction]
253256
fnConnectRegistry(
254-
computer_name:String,
257+
computer_name:Option<String>,
255258
key:PyRef<PyHKEYObject>,
256259
vm:&VirtualMachine,
257-
) ->PyResult<()>{
258-
let wide_computer_name =to_utf16(computer_name);
259-
let res =unsafe{
260-
Registry::RegConnectRegistryW(
261-
wide_computer_name.as_ptr(),
262-
*key.hkey.read(),
263-
std::ptr::null_mut(),
264-
)
265-
};
266-
if res ==0{
267-
Ok(())
260+
) ->PyResult<PyHKEYObject>{
261+
ifletSome(computer_name) = computer_name{
262+
letmut ret_key = std::ptr::null_mut();
263+
let wide_computer_name =to_utf16(computer_name);
264+
let res =unsafe{
265+
Registry::RegConnectRegistryW(
266+
wide_computer_name.as_ptr(),
267+
*key.hkey.read(),
268+
&mut ret_key
269+
)
270+
};
271+
if res ==0{
272+
Ok(PyHKEYObject::new(ret_key))
273+
}else{
274+
Err(vm.new_os_error(format!("error code: {}", res)))
275+
}
268276
}else{
269-
Err(vm.new_os_error(format!("error code: {}", res)))
277+
letmut ret_key = std::ptr::null_mut();
278+
let res =unsafe{
279+
Registry::RegConnectRegistryW(
280+
std::ptr::null_mut(),
281+
*key.hkey.read(),
282+
&mut ret_key
283+
)
284+
};
285+
if res ==0{
286+
Ok(PyHKEYObject::new(ret_key))
287+
}else{
288+
Err(vm.new_os_error(format!("error code: {}", res)))
289+
}
270290
}
271291
}
272292

@@ -368,6 +388,144 @@ mod winreg {
368388
}
369389
}
370390

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

432590
#[pyfunction]
433-
fnQueryInfoKey(key:PyRef<PyHKEYObject>,vm:&VirtualMachine) ->PyResult<()>{
591+
fnQueryInfoKey(key:PyRef<PyHKEYObject>,vm:&VirtualMachine) ->PyResult<PyRef<PyTuple>>{
434592
let key =*key.hkey.read();
435593
letmut lpcsubkeys:u32 =0;
436594
letmut lpcvalues:u32 =0;
437-
let lpftlastwritetime:*mutFoundation::FILETIME = std::ptr::null_mut();
595+
letmutlpftlastwritetime:Foundation::FILETIME =unsafe{std::mem::zeroed()};
438596
let err =unsafe{
439597
Registry::RegQueryInfoKeyW(
440598
key,
@@ -448,15 +606,16 @@ mod winreg {
448606
std::ptr::null_mut(),
449607
std::ptr::null_mut(),
450608
std::ptr::null_mut(),
451-
lpftlastwritetime,
609+
&mutlpftlastwritetime,
452610
)
453611
};
454612

455613
if err !=0{
456-
Err(vm.new_os_error(format!("error code: {}", err)))
457-
}else{
458-
Ok(())
614+
returnErr(vm.new_os_error(format!("error code: {}", err)));
459615
}
616+
let l:u64 =(lpftlastwritetime.dwHighDateTimeasu64) <<32 | lpftlastwritetime.dwLowDateTimeasu64;
617+
let tup:Vec<PyObjectRef> =vec![vm.ctx.new_int(lpcsubkeys).into(), vm.ctx.new_int(lpcvalues).into(), vm.ctx.new_int(l).into()];
618+
Ok(vm.ctx.new_tuple(tup))
460619
}
461620

462621
#[pyfunction]
@@ -481,7 +640,44 @@ mod winreg {
481640
Ok(())
482641
}
483642

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

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

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp