@@ -611,6 +611,7 @@ mod winreg {
611611 #[ pyfunction]
612612 fn QueryInfoKey ( key : HKEYArg , vm : & VirtualMachine ) -> PyResult < PyRef < PyTuple > > {
613613 let key = key. 0 ;
614+ vm. audit ( "winreg.QueryInfoKey" , || ( key as isize , ) ) ?;
614615 let info = host_winreg:: query_info_key_full ( key)
615616 . map_err ( |err| vm. new_os_error ( format ! ( "error code: {err}" ) ) ) ?;
616617 let tup: Vec < PyObjectRef > = vec ! [
@@ -622,20 +623,36 @@ mod winreg {
622623 }
623624
624625 #[ pyfunction]
625- fn QueryValue ( key : HKEYArg , sub_key : Option < String > , vm : & VirtualMachine ) -> PyResult < String > {
626+ fn QueryValue (
627+ key : HKEYArg ,
628+ sub_key : Option < PyRef < PyStr > > ,
629+ vm : & VirtualMachine ,
630+ ) -> PyResult < String > {
626631 let hkey = key. 0 ;
627632
633+ let sub_key = sub_key
634+ . as_ref ( )
635+ . map ( |s| {
636+ s. as_wtf8 ( )
637+ . to_wide_cstring ( )
638+ . map_err ( |_| nul_char_error ( vm) )
639+ } )
640+ . transpose ( ) ?;
641+ vm. audit ( "winreg.QueryValue" , || {
642+ let sub_key: PyObjectRef = sub_key. as_ref ( ) . map_or_else (
643+ || vm. ctx . none ( ) ,
644+ |s| vm. ctx . new_str ( Wtf8Buf :: from_wide ( s. as_slice ( ) ) ) . into ( ) ,
645+ ) ;
646+ ( hkey as isize , sub_key, vm. ctx . none ( ) )
647+ } ) ?;
648+
628649 if hkey == host_winreg:: HKEY_PERFORMANCE_DATA {
629650 return Err ( os_error_from_windows_code (
630651 vm,
631652 host_winreg:: ERROR_INVALID_HANDLE as i32 ,
632653 ) ) ;
633654 }
634655
635- let sub_key = sub_key
636- . map ( WideCString :: from_str)
637- . transpose ( )
638- . map_err ( |e| e. to_pyexception ( vm) ) ?;
639656 host_winreg:: query_default_value ( hkey, sub_key. as_deref ( ) )
640657 . map_err ( |err| err. to_pyexception ( vm) )
641658 }
0 commit comments