rustpython_vm/stdlib/
nt.rs

1use crate::{builtins::PyModule, PyRef, VirtualMachine};
2
3pub use module::raw_set_handle_inheritable;
4
5pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
6    let module = module::make_module(vm);
7    super::os::extend_module(vm, &module);
8    module
9}
10
11#[pymodule(name = "nt", with(super::os::_os))]
12pub(crate) mod module {
13    use crate::{
14        builtins::{PyDictRef, PyListRef, PyStrRef, PyTupleRef},
15        common::{crt_fd::Fd, os::last_os_error, suppress_iph},
16        convert::ToPyException,
17        function::{Either, OptionalArg},
18        ospath::OsPath,
19        stdlib::os::{errno_err, DirFd, FollowSymlinks, SupportFunc, TargetIsDirectory, _os},
20        PyResult, TryFromObject, VirtualMachine,
21    };
22    use libc::intptr_t;
23    use std::{
24        env, fs, io,
25        mem::MaybeUninit,
26        os::windows::ffi::{OsStrExt, OsStringExt},
27    };
28    use windows_sys::Win32::{
29        Foundation::{self, INVALID_HANDLE_VALUE},
30        Storage::FileSystem,
31        System::{Console, Threading},
32    };
33
34    #[pyattr]
35    use libc::{O_BINARY, O_TEMPORARY};
36
37    #[pyfunction]
38    pub(super) fn access(path: OsPath, mode: u8, vm: &VirtualMachine) -> PyResult<bool> {
39        let attr = unsafe { FileSystem::GetFileAttributesW(path.to_widecstring(vm)?.as_ptr()) };
40        Ok(attr != FileSystem::INVALID_FILE_ATTRIBUTES
41            && (mode & 2 == 0
42                || attr & FileSystem::FILE_ATTRIBUTE_READONLY == 0
43                || attr & FileSystem::FILE_ATTRIBUTE_DIRECTORY != 0))
44    }
45
46    #[derive(FromArgs)]
47    pub(super) struct SymlinkArgs {
48        src: OsPath,
49        dst: OsPath,
50        #[pyarg(flatten)]
51        target_is_directory: TargetIsDirectory,
52        #[pyarg(flatten)]
53        _dir_fd: DirFd<{ _os::SYMLINK_DIR_FD as usize }>,
54    }
55
56    #[pyfunction]
57    pub(super) fn symlink(args: SymlinkArgs, vm: &VirtualMachine) -> PyResult<()> {
58        use std::os::windows::fs as win_fs;
59        let dir = args.target_is_directory.target_is_directory
60            || args
61                .dst
62                .as_path()
63                .parent()
64                .and_then(|dst_parent| dst_parent.join(&args.src).symlink_metadata().ok())
65                .map_or(false, |meta| meta.is_dir());
66        let res = if dir {
67            win_fs::symlink_dir(args.src.path, args.dst.path)
68        } else {
69            win_fs::symlink_file(args.src.path, args.dst.path)
70        };
71        res.map_err(|err| err.to_pyexception(vm))
72    }
73
74    #[pyfunction]
75    fn set_inheritable(fd: i32, inheritable: bool, vm: &VirtualMachine) -> PyResult<()> {
76        let handle = Fd(fd).to_raw_handle().map_err(|e| e.to_pyexception(vm))?;
77        set_handle_inheritable(handle as _, inheritable, vm)
78    }
79
80    #[pyattr]
81    fn environ(vm: &VirtualMachine) -> PyDictRef {
82        let environ = vm.ctx.new_dict();
83
84        for (key, value) in env::vars() {
85            environ.set_item(&key, vm.new_pyobj(value), vm).unwrap();
86        }
87        environ
88    }
89
90    #[pyfunction]
91    fn chmod(
92        path: OsPath,
93        dir_fd: DirFd<0>,
94        mode: u32,
95        follow_symlinks: FollowSymlinks,
96        vm: &VirtualMachine,
97    ) -> PyResult<()> {
98        const S_IWRITE: u32 = 128;
99        let [] = dir_fd.0;
100        let metadata = if follow_symlinks.0 {
101            fs::metadata(&path)
102        } else {
103            fs::symlink_metadata(&path)
104        };
105        let meta = metadata.map_err(|err| err.to_pyexception(vm))?;
106        let mut permissions = meta.permissions();
107        permissions.set_readonly(mode & S_IWRITE == 0);
108        fs::set_permissions(&path, permissions).map_err(|err| err.to_pyexception(vm))
109    }
110
111    // cwait is available on MSVC only (according to CPython)
112    #[cfg(target_env = "msvc")]
113    extern "C" {
114        fn _cwait(termstat: *mut i32, procHandle: intptr_t, action: i32) -> intptr_t;
115    }
116
117    #[cfg(target_env = "msvc")]
118    #[pyfunction]
119    fn waitpid(pid: intptr_t, opt: i32, vm: &VirtualMachine) -> PyResult<(intptr_t, i32)> {
120        let mut status = 0;
121        let pid = unsafe { suppress_iph!(_cwait(&mut status, pid, opt)) };
122        if pid == -1 {
123            Err(errno_err(vm))
124        } else {
125            Ok((pid, status << 8))
126        }
127    }
128
129    #[cfg(target_env = "msvc")]
130    #[pyfunction]
131    fn wait(vm: &VirtualMachine) -> PyResult<(intptr_t, i32)> {
132        waitpid(-1, 0, vm)
133    }
134
135    #[pyfunction]
136    fn kill(pid: i32, sig: isize, vm: &VirtualMachine) -> PyResult<()> {
137        let sig = sig as u32;
138        let pid = pid as u32;
139
140        if sig == Console::CTRL_C_EVENT || sig == Console::CTRL_BREAK_EVENT {
141            let ret = unsafe { Console::GenerateConsoleCtrlEvent(sig, pid) };
142            let res = if ret == 0 { Err(errno_err(vm)) } else { Ok(()) };
143            return res;
144        }
145
146        let h = unsafe { Threading::OpenProcess(Threading::PROCESS_ALL_ACCESS, 0, pid) };
147        if h == 0 {
148            return Err(errno_err(vm));
149        }
150        let ret = unsafe { Threading::TerminateProcess(h, sig) };
151        let res = if ret == 0 { Err(errno_err(vm)) } else { Ok(()) };
152        unsafe { Foundation::CloseHandle(h) };
153        res
154    }
155
156    #[pyfunction]
157    fn get_terminal_size(
158        fd: OptionalArg<i32>,
159        vm: &VirtualMachine,
160    ) -> PyResult<_os::PyTerminalSize> {
161        let (columns, lines) = {
162            let stdhandle = match fd {
163                OptionalArg::Present(0) => Console::STD_INPUT_HANDLE,
164                OptionalArg::Present(1) | OptionalArg::Missing => Console::STD_OUTPUT_HANDLE,
165                OptionalArg::Present(2) => Console::STD_ERROR_HANDLE,
166                _ => return Err(vm.new_value_error("bad file descriptor".to_owned())),
167            };
168            let h = unsafe { Console::GetStdHandle(stdhandle) };
169            if h == 0 {
170                return Err(vm.new_os_error("handle cannot be retrieved".to_owned()));
171            }
172            if h == INVALID_HANDLE_VALUE {
173                return Err(errno_err(vm));
174            }
175            let mut csbi = MaybeUninit::uninit();
176            let ret = unsafe { Console::GetConsoleScreenBufferInfo(h, csbi.as_mut_ptr()) };
177            let csbi = unsafe { csbi.assume_init() };
178            if ret == 0 {
179                return Err(errno_err(vm));
180            }
181            let w = csbi.srWindow;
182            (
183                (w.Right - w.Left + 1) as usize,
184                (w.Bottom - w.Top + 1) as usize,
185            )
186        };
187        Ok(_os::PyTerminalSize { columns, lines })
188    }
189
190    #[cfg(target_env = "msvc")]
191    extern "C" {
192        fn _wexecv(cmdname: *const u16, argv: *const *const u16) -> intptr_t;
193    }
194
195    #[cfg(target_env = "msvc")]
196    #[pyfunction]
197    fn execv(
198        path: PyStrRef,
199        argv: Either<PyListRef, PyTupleRef>,
200        vm: &VirtualMachine,
201    ) -> PyResult<()> {
202        use std::iter::once;
203
204        let make_widestring =
205            |s: &str| widestring::WideCString::from_os_str(s).map_err(|err| err.to_pyexception(vm));
206
207        let path = make_widestring(path.as_str())?;
208
209        let argv = vm.extract_elements_with(argv.as_ref(), |obj| {
210            let arg = PyStrRef::try_from_object(vm, obj)?;
211            make_widestring(arg.as_str())
212        })?;
213
214        let first = argv
215            .first()
216            .ok_or_else(|| vm.new_value_error("execv() arg 2 must not be empty".to_owned()))?;
217
218        if first.is_empty() {
219            return Err(
220                vm.new_value_error("execv() arg 2 first element cannot be empty".to_owned())
221            );
222        }
223
224        let argv_execv: Vec<*const u16> = argv
225            .iter()
226            .map(|v| v.as_ptr())
227            .chain(once(std::ptr::null()))
228            .collect();
229
230        if (unsafe { suppress_iph!(_wexecv(path.as_ptr(), argv_execv.as_ptr())) } == -1) {
231            Err(errno_err(vm))
232        } else {
233            Ok(())
234        }
235    }
236
237    #[pyfunction]
238    fn _getfinalpathname(path: OsPath, vm: &VirtualMachine) -> PyResult {
239        let real = path
240            .as_ref()
241            .canonicalize()
242            .map_err(|e| e.to_pyexception(vm))?;
243        path.mode.process_path(real, vm)
244    }
245
246    #[pyfunction]
247    fn _getfullpathname(path: OsPath, vm: &VirtualMachine) -> PyResult {
248        let wpath = path.to_widecstring(vm)?;
249        let mut buffer = vec![0u16; Foundation::MAX_PATH as usize];
250        let ret = unsafe {
251            FileSystem::GetFullPathNameW(
252                wpath.as_ptr(),
253                buffer.len() as _,
254                buffer.as_mut_ptr(),
255                std::ptr::null_mut(),
256            )
257        };
258        if ret == 0 {
259            return Err(errno_err(vm));
260        }
261        if ret as usize > buffer.len() {
262            buffer.resize(ret as usize, 0);
263            let ret = unsafe {
264                FileSystem::GetFullPathNameW(
265                    wpath.as_ptr(),
266                    buffer.len() as _,
267                    buffer.as_mut_ptr(),
268                    std::ptr::null_mut(),
269                )
270            };
271            if ret == 0 {
272                return Err(errno_err(vm));
273            }
274        }
275        let buffer = widestring::WideCString::from_vec_truncate(buffer);
276        path.mode.process_path(buffer.to_os_string(), vm)
277    }
278
279    #[pyfunction]
280    fn _getvolumepathname(path: OsPath, vm: &VirtualMachine) -> PyResult {
281        let wide = path.to_widecstring(vm)?;
282        let buflen = std::cmp::max(wide.len(), Foundation::MAX_PATH as usize);
283        let mut buffer = vec![0u16; buflen];
284        let ret = unsafe {
285            FileSystem::GetVolumePathNameW(wide.as_ptr(), buffer.as_mut_ptr(), buflen as _)
286        };
287        if ret == 0 {
288            return Err(errno_err(vm));
289        }
290        let buffer = widestring::WideCString::from_vec_truncate(buffer);
291        path.mode.process_path(buffer.to_os_string(), vm)
292    }
293
294    #[pyfunction]
295    fn _path_splitroot(path: OsPath, vm: &VirtualMachine) -> PyResult<(String, String)> {
296        let orig: Vec<_> = path.path.encode_wide().collect();
297        if orig.is_empty() {
298            return Ok(("".to_owned(), "".to_owned()));
299        }
300        let backslashed: Vec<_> = orig
301            .iter()
302            .copied()
303            .map(|c| if c == b'/' as u16 { b'\\' as u16 } else { c })
304            .chain(std::iter::once(0)) // null-terminated
305            .collect();
306
307        fn from_utf16(wstr: &[u16], vm: &VirtualMachine) -> PyResult<String> {
308            String::from_utf16(wstr).map_err(|e| vm.new_unicode_decode_error(e.to_string()))
309        }
310
311        let wbuf = windows::core::PCWSTR::from_raw(backslashed.as_ptr());
312        let (root, path) = match unsafe { windows::Win32::UI::Shell::PathCchSkipRoot(wbuf) } {
313            Ok(end) => {
314                assert!(!end.is_null());
315                let len: usize = unsafe { end.as_ptr().offset_from(wbuf.as_ptr()) }
316                    .try_into()
317                    .expect("len must be non-negative");
318                assert!(
319                    len < backslashed.len(), // backslashed is null-terminated
320                    "path: {:?} {} < {}",
321                    std::path::PathBuf::from(std::ffi::OsString::from_wide(&backslashed)),
322                    len,
323                    backslashed.len()
324                );
325                (from_utf16(&orig[..len], vm)?, from_utf16(&orig[len..], vm)?)
326            }
327            Err(_) => ("".to_owned(), from_utf16(&orig, vm)?),
328        };
329        Ok((root, path))
330    }
331
332    #[pyfunction]
333    fn _getdiskusage(path: OsPath, vm: &VirtualMachine) -> PyResult<(u64, u64)> {
334        use FileSystem::GetDiskFreeSpaceExW;
335
336        let wpath = path.to_widecstring(vm)?;
337        let mut _free_to_me: u64 = 0;
338        let mut total: u64 = 0;
339        let mut free: u64 = 0;
340        let ret =
341            unsafe { GetDiskFreeSpaceExW(wpath.as_ptr(), &mut _free_to_me, &mut total, &mut free) };
342        if ret != 0 {
343            return Ok((total, free));
344        }
345        let err = io::Error::last_os_error();
346        if err.raw_os_error() == Some(Foundation::ERROR_DIRECTORY as i32) {
347            if let Some(parent) = path.as_ref().parent() {
348                let parent = widestring::WideCString::from_os_str(parent).unwrap();
349
350                let ret = unsafe {
351                    GetDiskFreeSpaceExW(parent.as_ptr(), &mut _free_to_me, &mut total, &mut free)
352                };
353
354                return if ret == 0 {
355                    Err(errno_err(vm))
356                } else {
357                    Ok((total, free))
358                };
359            }
360        }
361        Err(err.to_pyexception(vm))
362    }
363
364    #[pyfunction]
365    fn get_handle_inheritable(handle: intptr_t, vm: &VirtualMachine) -> PyResult<bool> {
366        let mut flags = 0;
367        if unsafe { Foundation::GetHandleInformation(handle as _, &mut flags) } == 0 {
368            Err(errno_err(vm))
369        } else {
370            Ok(flags & Foundation::HANDLE_FLAG_INHERIT != 0)
371        }
372    }
373
374    pub fn raw_set_handle_inheritable(handle: intptr_t, inheritable: bool) -> std::io::Result<()> {
375        let flags = if inheritable {
376            Foundation::HANDLE_FLAG_INHERIT
377        } else {
378            0
379        };
380        let res = unsafe {
381            Foundation::SetHandleInformation(handle as _, Foundation::HANDLE_FLAG_INHERIT, flags)
382        };
383        if res == 0 {
384            Err(last_os_error())
385        } else {
386            Ok(())
387        }
388    }
389
390    #[pyfunction]
391    fn listdrives(vm: &VirtualMachine) -> PyResult<PyListRef> {
392        use windows_sys::Win32::Foundation::ERROR_MORE_DATA;
393
394        let mut buffer = [0u16; 256];
395        let len =
396            unsafe { FileSystem::GetLogicalDriveStringsW(buffer.len() as _, buffer.as_mut_ptr()) };
397        if len == 0 {
398            return Err(errno_err(vm));
399        }
400        if len as usize >= buffer.len() {
401            return Err(std::io::Error::from_raw_os_error(ERROR_MORE_DATA as _).to_pyexception(vm));
402        }
403        let drives: Vec<_> = buffer[..(len - 1) as usize]
404            .split(|&c| c == 0)
405            .map(|drive| vm.new_pyobj(String::from_utf16_lossy(drive)))
406            .collect();
407        Ok(vm.ctx.new_list(drives))
408    }
409
410    #[pyfunction]
411    fn set_handle_inheritable(
412        handle: intptr_t,
413        inheritable: bool,
414        vm: &VirtualMachine,
415    ) -> PyResult<()> {
416        raw_set_handle_inheritable(handle, inheritable).map_err(|e| e.to_pyexception(vm))
417    }
418
419    #[pyfunction]
420    fn mkdir(
421        path: OsPath,
422        mode: OptionalArg<i32>,
423        dir_fd: DirFd<{ _os::MKDIR_DIR_FD as usize }>,
424        vm: &VirtualMachine,
425    ) -> PyResult<()> {
426        let mode = mode.unwrap_or(0o777);
427        let [] = dir_fd.0;
428        let _ = mode;
429        let wide = path.to_widecstring(vm)?;
430        let res = unsafe { FileSystem::CreateDirectoryW(wide.as_ptr(), std::ptr::null_mut()) };
431        if res == 0 {
432            return Err(errno_err(vm));
433        }
434        Ok(())
435    }
436
437    pub(crate) fn support_funcs() -> Vec<SupportFunc> {
438        Vec::new()
439    }
440}
rustpython_vm/stdlib/nt.rs

rustpython_vm/stdlib/
nt.rs
