"""
This file is based on the code from https://github.com/JustAMan/pyWinClobber/blob/master/win32elevate.py
Copyright (c) 2013 by JustAMan at GitHub
Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
import ctypes
import os
import subprocess
from ctypes import POINTER, byref, c_char_p, c_int, c_ulong, c_void_p
from ctypes.wintypes import (
BOOL,
DWORD,
HANDLE,
HINSTANCE,
HKEY,
HWND,
LPCSTR,
LPCWSTR,
LPDWORD,
SHORT,
SMALL_RECT,
WORD,
)
from xonsh import platform
from xonsh.lib import lazyimps
from xonsh.lib.lazyasd import lazyobject
__all__ = ("sudo",)
@lazyobject
def CloseHandle():
ch = ctypes.windll.kernel32.CloseHandle
ch.argtypes = (HANDLE,)
ch.restype = BOOL
return ch
@lazyobject
def GetActiveWindow():
gaw = ctypes.windll.user32.GetActiveWindow
gaw.argtypes = ()
gaw.restype = HANDLE
return gaw
TOKEN_READ = 0x20008
class ShellExecuteInfo(ctypes.Structure):
_fields_ = [
("cbSize", DWORD),
("fMask", c_ulong),
("hwnd", HWND),
("lpVerb", c_char_p),
("lpFile", c_char_p),
("lpParameters", c_char_p),
("lpDirectory", c_char_p),
("nShow", c_int),
("hInstApp", HINSTANCE),
("lpIDList", c_void_p),
("lpClass", c_char_p),
("hKeyClass", HKEY),
("dwHotKey", DWORD),
("hIcon", HANDLE),
("hProcess", HANDLE),
]
def __init__(self, **kw):
ctypes.Structure.__init__(self)
self.cbSize = ctypes.sizeof(self)
for field_name, field_value in kw.items():
setattr(self, field_name, field_value)
@lazyobject
def ShellExecuteEx():
see = ctypes.windll.Shell32.ShellExecuteExA
PShellExecuteInfo = ctypes.POINTER(ShellExecuteInfo)
see.argtypes = (PShellExecuteInfo,)
see.restype = BOOL
return see
@lazyobject
def WaitForSingleObject():
wfso = ctypes.windll.kernel32.WaitForSingleObject
wfso.argtypes = (HANDLE, DWORD)
wfso.restype = DWORD
return wfso
# SW_HIDE = 0
SW_SHOW = 5
SEE_MASK_NOCLOSEPROCESS = 0x00000040
SEE_MASK_NO_CONSOLE = 0x00008000
INFINITE = -1
def wait_and_close_handle(process_handle):
"""
Waits till spawned process finishes and closes the handle for it
Parameters
----------
process_handle : HANDLE
The Windows handle for the process
"""
WaitForSingleObject(process_handle, INFINITE)
CloseHandle(process_handle)
[docs]
def sudo(executable, args=None):
"""
This will re-run current Python script requesting to elevate administrative rights.
Parameters
----------
executable : str
The path/name of the executable
args : list of str
The arguments to be passed to the executable
"""
if not args:
args = []
execute_info = ShellExecuteInfo(
fMask=SEE_MASK_NOCLOSEPROCESS | SEE_MASK_NO_CONSOLE,
hwnd=GetActiveWindow(),
lpVerb=b"runas",
lpFile=executable.encode("utf-8"),
lpParameters=subprocess.list2cmdline(args).encode("utf-8"),
lpDirectory=None,
nShow=SW_SHOW,
)
if not ShellExecuteEx(byref(execute_info)):
raise ctypes.WinError()
wait_and_close_handle(execute_info.hProcess)
#
# The following has been refactored from
# http://stackoverflow.com/a/37505496/2312428
#
# input flags
ENABLE_PROCESSED_INPUT = 0x0001
ENABLE_LINE_INPUT = 0x0002
ENABLE_ECHO_INPUT = 0x0004
ENABLE_WINDOW_INPUT = 0x0008
ENABLE_MOUSE_INPUT = 0x0010
ENABLE_INSERT_MODE = 0x0020
ENABLE_QUICK_EDIT_MODE = 0x0040
# output flags
ENABLE_PROCESSED_OUTPUT = 0x0001
ENABLE_WRAP_AT_EOL_OUTPUT = 0x0002
ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 # VT100 (Win 10)
def check_zero(result, func, args):
if not result:
err = ctypes.get_last_error()
if err:
raise ctypes.WinError(err)
return args
@lazyobject
def GetStdHandle():
return lazyimps._winapi.GetStdHandle
@lazyobject
def STDHANDLES():
"""Tuple of the Windows handles for (stdin, stdout, stderr)."""
hs = [
lazyimps._winapi.STD_INPUT_HANDLE,
lazyimps._winapi.STD_OUTPUT_HANDLE,
lazyimps._winapi.STD_ERROR_HANDLE,
]
hcons = []
for h in hs:
hcon = GetStdHandle(int(h))
hcons.append(hcon)
return tuple(hcons)
@lazyobject
def GetConsoleMode():
gcm = ctypes.windll.kernel32.GetConsoleMode
gcm.errcheck = check_zero
gcm.argtypes = (HANDLE, LPDWORD) # _In_ hConsoleHandle # _Out_ lpMode
return gcm
def get_console_mode(fd=1):
"""Get the mode of the active console input, output, or error
buffer. Note that if the process isn't attached to a
console, this function raises an EBADF IOError.
Parameters
----------
fd : int
Standard buffer file descriptor, 0 for stdin, 1 for stdout (default),
and 2 for stderr
"""
mode = DWORD()
hcon = STDHANDLES[fd]
GetConsoleMode(hcon, byref(mode))
return mode.value
@lazyobject
def SetConsoleMode():
scm = ctypes.windll.kernel32.SetConsoleMode
scm.errcheck = check_zero
scm.argtypes = (HANDLE, DWORD) # _In_ hConsoleHandle # _Out_ lpMode
return scm
def set_console_mode(mode, fd=1):
"""Set the mode of the active console input, output, or
error buffer. Note that if the process isn't attached to a
console, this function raises an EBADF IOError.
Parameters
----------
mode : int
Mode flags to set on the handle.
fd : int, optional
Standard buffer file descriptor, 0 for stdin, 1 for stdout (default),
and 2 for stderr
"""
hcon = STDHANDLES[fd]
SetConsoleMode(hcon, mode)
def enable_virtual_terminal_processing():
"""Enables virtual terminal processing on Windows.
This includes ANSI escape sequence interpretation.
See http://stackoverflow.com/a/36760881/2312428
"""
SetConsoleMode(GetStdHandle(-11), 7)
@lazyobject
def COORD():
if platform.has_prompt_toolkit():
# turns out that PTK has a separate ctype wrapper
# for this struct and also wraps similar function calls
# we need to use the same struct to prevent clashes.
import prompt_toolkit.win32_types
return prompt_toolkit.win32_types.COORD
class _COORD(ctypes.Structure):
"""Struct from the winapi, representing coordinates in the console.
Attributes
----------
X : int
Column position
Y : int
Row position
"""
_fields_ = [("X", SHORT), ("Y", SHORT)]
return _COORD
@lazyobject
def ReadConsoleOutputCharacterA():
rcoc = ctypes.windll.kernel32.ReadConsoleOutputCharacterA
rcoc.errcheck = check_zero
rcoc.argtypes = (
HANDLE, # _In_ hConsoleOutput
LPCSTR, # _Out_ LPTSTR lpMode
DWORD, # _In_ nLength
COORD, # _In_ dwReadCoord,
LPDWORD,
) # _Out_ lpNumberOfCharsRead
rcoc.restype = BOOL
return rcoc
@lazyobject
def ReadConsoleOutputCharacterW():
rcoc = ctypes.windll.kernel32.ReadConsoleOutputCharacterW
rcoc.errcheck = check_zero
rcoc.argtypes = (
HANDLE, # _In_ hConsoleOutput
LPCWSTR, # _Out_ LPTSTR lpMode
DWORD, # _In_ nLength
COORD, # _In_ dwReadCoord,
LPDWORD,
) # _Out_ lpNumberOfCharsRead
rcoc.restype = BOOL
return rcoc
def read_console_output_character(x=0, y=0, fd=1, buf=None, bufsize=1024, raw=False):
"""Reads characters from the console buffer.
Parameters
----------
x : int, optional
Starting column.
y : int, optional
Starting row.
fd : int, optional
Standard buffer file descriptor, 0 for stdin, 1 for stdout (default),
and 2 for stderr.
buf : ctypes.c_wchar_p if raw else ctypes.c_wchar_p, optional
An existing buffer to (re-)use.
bufsize : int, optional
The maximum read size.
raw : bool, optional
Whether to read in and return as bytes (True) or as a
unicode string (False, default).
Returns
-------
value : str
Result of what was read, may be shorter than bufsize.
"""
hcon = STDHANDLES[fd]
if buf is None:
if raw:
buf = ctypes.c_char_p(b" " * bufsize)
else:
buf = ctypes.c_wchar_p(" " * bufsize)
coord = COORD(x, y)
n = DWORD()
if raw:
ReadConsoleOutputCharacterA(hcon, buf, bufsize, coord, byref(n))
else:
ReadConsoleOutputCharacterW(hcon, buf, bufsize, coord, byref(n))
return buf.value[: n.value]
def pread_console(fd, buffersize, offset, buf=None):
"""This is a console-based implementation of os.pread() for windows.
that uses read_console_output_character().
"""
cols, rows = os.get_terminal_size(fd=fd)
x = offset % cols
y = offset // cols
return read_console_output_character(
x=x, y=y, fd=fd, buf=buf, bufsize=buffersize, raw=True
)
#
# The following piece has been forked from colorama.win32
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
#
@lazyobject
def CONSOLE_SCREEN_BUFFER_INFO():
if platform.has_prompt_toolkit():
# turns out that PTK has a separate ctype wrapper
# for this struct and also wraps kernel32.GetConsoleScreenBufferInfo
# we need to use the same struct to prevent clashes.
import prompt_toolkit.win32_types
return prompt_toolkit.win32_types.CONSOLE_SCREEN_BUFFER_INFO
# Otherwise we should wrap it ourselves
COORD() # force COORD to load
class _CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
"""Struct from in wincon.h. See Windows API docs
for more details.
Attributes
----------
dwSize : COORD
Size of
dwCursorPosition : COORD
Current cursor location.
wAttributes : WORD
Flags for screen buffer.
srWindow : SMALL_RECT
Actual size of screen
dwMaximumWindowSize : COORD
Maximum window scrollback size.
"""
_fields_ = [
("dwSize", COORD),
("dwCursorPosition", COORD),
("wAttributes", WORD),
("srWindow", SMALL_RECT),
("dwMaximumWindowSize", COORD),
]
return _CONSOLE_SCREEN_BUFFER_INFO
@lazyobject
def GetConsoleScreenBufferInfo():
"""Returns the windows version of the get screen buffer."""
gcsbi = ctypes.windll.kernel32.GetConsoleScreenBufferInfo
gcsbi.errcheck = check_zero
gcsbi.argtypes = (HANDLE, POINTER(CONSOLE_SCREEN_BUFFER_INFO))
gcsbi.restype = BOOL
return gcsbi
def get_console_screen_buffer_info(fd=1):
"""Returns an screen buffer info object for the relevant stdbuf.
Parameters
----------
fd : int, optional
Standard buffer file descriptor, 0 for stdin, 1 for stdout (default),
and 2 for stderr.
Returns
-------
csbi : CONSOLE_SCREEN_BUFFER_INFO
Information about the console screen buffer.
"""
hcon = STDHANDLES[fd]
csbi = CONSOLE_SCREEN_BUFFER_INFO()
GetConsoleScreenBufferInfo(hcon, byref(csbi))
return csbi
#
# end colorama forked section
#
def get_cursor_position(fd=1):
"""Gets the current cursor position as an (x, y) tuple."""
csbi = get_console_screen_buffer_info(fd=fd)
coord = csbi.dwCursorPosition
return (coord.X, coord.Y)
def get_cursor_offset(fd=1):
"""Gets the current cursor position as a total offset value."""
csbi = get_console_screen_buffer_info(fd=fd)
pos = csbi.dwCursorPosition
size = csbi.dwSize
return (pos.Y * size.X) + pos.X
def get_position_size(fd=1):
"""Gets the current cursor position and screen size tuple:
(x, y, columns, lines).
"""
info = get_console_screen_buffer_info(fd)
return (
info.dwCursorPosition.X,
info.dwCursorPosition.Y,
info.dwSize.X,
info.dwSize.Y,
)
@lazyobject
def SetConsoleScreenBufferSize():
"""Set screen buffer dimensions."""
scsbs = ctypes.windll.kernel32.SetConsoleScreenBufferSize
scsbs.errcheck = check_zero
scsbs.argtypes = (HANDLE, COORD) # _In_ HANDLE hConsoleOutput # _In_ COORD dwSize
scsbs.restype = BOOL
return scsbs
def set_console_screen_buffer_size(x, y, fd=1):
"""Sets the console size for a standard buffer.
Parameters
----------
x : int
Number of columns.
y : int
Number of rows.
fd : int, optional
Standard buffer file descriptor, 0 for stdin, 1 for stdout (default),
and 2 for stderr.
"""
coord = COORD()
coord.X = x
coord.Y = y
hcon = STDHANDLES[fd]
rtn = SetConsoleScreenBufferSize(hcon, coord)
return rtn
@lazyobject
def SetConsoleCursorPosition():
"""Set cursor position in console."""
sccp = ctypes.windll.kernel32.SetConsoleCursorPosition
sccp.errcheck = check_zero
sccp.argtypes = (
HANDLE, # _In_ HANDLE hConsoleOutput
COORD, # _In_ COORD dwCursorPosition
)
sccp.restype = BOOL
return sccp
def set_console_cursor_position(x, y, fd=1):
"""Sets the console cursor position for a standard buffer.
Parameters
----------
x : int
Number of columns.
y : int
Number of rows.
fd : int, optional
Standard buffer file descriptor, 0 for stdin, 1 for stdout (default),
and 2 for stderr.
"""
coord = COORD()
coord.X = x
coord.Y = y
hcon = STDHANDLES[fd]
rtn = SetConsoleCursorPosition(hcon, coord)
return rtn