1"""Windows-specific implementation of process utilities.
2
3This file is only meant to be imported by process.py, not by end-users.
4"""
5
6#-----------------------------------------------------------------------------
7#  Copyright (C) 2010-2011  The IPython Development Team
8#
9#  Distributed under the terms of the BSD License.  The full license is in
10#  the file COPYING, distributed as part of this software.
11#-----------------------------------------------------------------------------
12
13#-----------------------------------------------------------------------------
14# Imports
15#-----------------------------------------------------------------------------
16from __future__ import print_function
17
18# stdlib
19import os
20import sys
21import ctypes
22
23from ctypes import c_int, POINTER
24from ctypes.wintypes import LPCWSTR, HLOCAL
25from subprocess import STDOUT
26
27# our own imports
28from ._process_common import read_no_interrupt, process_handler, arg_split as py_arg_split
29from . import py3compat
30from .encoding import DEFAULT_ENCODING
31
32#-----------------------------------------------------------------------------
33# Function definitions
34#-----------------------------------------------------------------------------
35
36class AvoidUNCPath(object):
37    """A context manager to protect command execution from UNC paths.
38
39    In the Win32 API, commands can't be invoked with the cwd being a UNC path.
40    This context manager temporarily changes directory to the 'C:' drive on
41    entering, and restores the original working directory on exit.
42
43    The context manager returns the starting working directory *if* it made a
44    change and None otherwise, so that users can apply the necessary adjustment
45    to their system calls in the event of a change.
46
47    Examples
48    --------
49    ::
50        cmd = 'dir'
51        with AvoidUNCPath() as path:
52            if path is not None:
53                cmd = '"pushd %s &&"%s' % (path, cmd)
54            os.system(cmd)
55    """
56    def __enter__(self):
57        self.path = py3compat.getcwd()
58        self.is_unc_path = self.path.startswith(r"\\")
59        if self.is_unc_path:
60            # change to c drive (as cmd.exe cannot handle UNC addresses)
61            os.chdir("C:")
62            return self.path
63        else:
64            # We return None to signal that there was no change in the working
65            # directory
66            return None
67
68    def __exit__(self, exc_type, exc_value, traceback):
69        if self.is_unc_path:
70            os.chdir(self.path)
71
72
73def _find_cmd(cmd):
74    """Find the full path to a .bat or .exe using the win32api module."""
75    try:
76        from win32api import SearchPath
77    except ImportError:
78        raise ImportError('you need to have pywin32 installed for this to work')
79    else:
80        PATH = os.environ['PATH']
81        extensions = ['.exe', '.com', '.bat', '.py']
82        path = None
83        for ext in extensions:
84            try:
85                path = SearchPath(PATH, cmd, ext)[0]
86            except:
87                pass
88        if path is None:
89            raise OSError("command %r not found" % cmd)
90        else:
91            return path
92
93
94def _system_body(p):
95    """Callback for _system."""
96    enc = DEFAULT_ENCODING
97    for line in read_no_interrupt(p.stdout).splitlines():
98        line = line.decode(enc, 'replace')
99        print(line, file=sys.stdout)
100    for line in read_no_interrupt(p.stderr).splitlines():
101        line = line.decode(enc, 'replace')
102        print(line, file=sys.stderr)
103
104    # Wait to finish for returncode
105    return p.wait()
106
107
108def system(cmd):
109    """Win32 version of os.system() that works with network shares.
110
111    Note that this implementation returns None, as meant for use in IPython.
112
113    Parameters
114    ----------
115    cmd : str or list
116      A command to be executed in the system shell.
117
118    Returns
119    -------
120    None : we explicitly do NOT return the subprocess status code, as this
121    utility is meant to be used extensively in IPython, where any return value
122    would trigger :func:`sys.displayhook` calls.
123    """
124    # The controller provides interactivity with both
125    # stdin and stdout
126    #import _process_win32_controller
127    #_process_win32_controller.system(cmd)
128
129    with AvoidUNCPath() as path:
130        if path is not None:
131            cmd = '"pushd %s &&"%s' % (path, cmd)
132        return process_handler(cmd, _system_body)
133
134def getoutput(cmd):
135    """Return standard output of executing cmd in a shell.
136
137    Accepts the same arguments as os.system().
138
139    Parameters
140    ----------
141    cmd : str or list
142      A command to be executed in the system shell.
143
144    Returns
145    -------
146    stdout : str
147    """
148
149    with AvoidUNCPath() as path:
150        if path is not None:
151            cmd = '"pushd %s &&"%s' % (path, cmd)
152        out = process_handler(cmd, lambda p: p.communicate()[0], STDOUT)
153
154    if out is None:
155        out = b''
156    return py3compat.bytes_to_str(out)
157
158try:
159    CommandLineToArgvW = ctypes.windll.shell32.CommandLineToArgvW
160    CommandLineToArgvW.arg_types = [LPCWSTR, POINTER(c_int)]
161    CommandLineToArgvW.restype = POINTER(LPCWSTR)
162    LocalFree = ctypes.windll.kernel32.LocalFree
163    LocalFree.res_type = HLOCAL
164    LocalFree.arg_types = [HLOCAL]
165
166    def arg_split(commandline, posix=False, strict=True):
167        """Split a command line's arguments in a shell-like manner.
168
169        This is a special version for windows that use a ctypes call to CommandLineToArgvW
170        to do the argv splitting. The posix paramter is ignored.
171
172        If strict=False, process_common.arg_split(...strict=False) is used instead.
173        """
174        #CommandLineToArgvW returns path to executable if called with empty string.
175        if commandline.strip() == "":
176            return []
177        if not strict:
178            # not really a cl-arg, fallback on _process_common
179            return py_arg_split(commandline, posix=posix, strict=strict)
180        argvn = c_int()
181        result_pointer = CommandLineToArgvW(py3compat.cast_unicode(commandline.lstrip()), ctypes.byref(argvn))
182        result_array_type = LPCWSTR * argvn.value
183        result = [arg for arg in result_array_type.from_address(ctypes.addressof(result_pointer.contents))]
184        retval = LocalFree(result_pointer)
185        return result
186except AttributeError:
187    arg_split = py_arg_split
188
189def check_pid(pid):
190    # OpenProcess returns 0 if no such process (of ours) exists
191    # positive int otherwise
192    return bool(ctypes.windll.kernel32.OpenProcess(1,0,pid))
193