1""" 2This file is based on the code from https://github.com/JustAMan/pyWinClobber/blob/master/win32elevate.py 3 4Copyright (c) 2013 by JustAMan at GitHub 5 6Permission is hereby granted, free of charge, to any person obtaining a copy of 7this software and associated documentation files (the "Software"), to deal in 8the Software without restriction, including without limitation the rights to 9use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 10the Software, and to permit persons to whom the Software is furnished to do so, 11subject to the following conditions: 12 13The above copyright notice and this permission notice shall be included in all 14copies or substantial portions of the Software. 15 16THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 18FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 19COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 20IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 21CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 22""" 23import os 24import ctypes 25import subprocess 26from ctypes import c_ulong, c_char_p, c_int, c_void_p, POINTER, byref 27from ctypes.wintypes import ( 28 HANDLE, 29 BOOL, 30 DWORD, 31 HWND, 32 HINSTANCE, 33 HKEY, 34 LPDWORD, 35 SHORT, 36 LPCWSTR, 37 WORD, 38 SMALL_RECT, 39 LPCSTR, 40) 41 42from xonsh.lazyasd import lazyobject 43from xonsh import lazyimps # we aren't amalgamated in this module. 44from xonsh import platform 45 46 47__all__ = ("sudo",) 48 49 50@lazyobject 51def CloseHandle(): 52 ch = ctypes.windll.kernel32.CloseHandle 53 ch.argtypes = (HANDLE,) 54 ch.restype = BOOL 55 return ch 56 57 58@lazyobject 59def GetActiveWindow(): 60 gaw = ctypes.windll.user32.GetActiveWindow 61 gaw.argtypes = () 62 gaw.restype = HANDLE 63 return gaw 64 65 66TOKEN_READ = 0x20008 67 68 69class ShellExecuteInfo(ctypes.Structure): 70 _fields_ = [ 71 ("cbSize", DWORD), 72 ("fMask", c_ulong), 73 ("hwnd", HWND), 74 ("lpVerb", c_char_p), 75 ("lpFile", c_char_p), 76 ("lpParameters", c_char_p), 77 ("lpDirectory", c_char_p), 78 ("nShow", c_int), 79 ("hInstApp", HINSTANCE), 80 ("lpIDList", c_void_p), 81 ("lpClass", c_char_p), 82 ("hKeyClass", HKEY), 83 ("dwHotKey", DWORD), 84 ("hIcon", HANDLE), 85 ("hProcess", HANDLE), 86 ] 87 88 def __init__(self, **kw): 89 ctypes.Structure.__init__(self) 90 self.cbSize = ctypes.sizeof(self) 91 for field_name, field_value in kw.items(): 92 setattr(self, field_name, field_value) 93 94 95@lazyobject 96def ShellExecuteEx(): 97 see = ctypes.windll.Shell32.ShellExecuteExA 98 PShellExecuteInfo = ctypes.POINTER(ShellExecuteInfo) 99 see.argtypes = (PShellExecuteInfo,) 100 see.restype = BOOL 101 return see 102 103 104@lazyobject 105def WaitForSingleObject(): 106 wfso = ctypes.windll.kernel32.WaitForSingleObject 107 wfso.argtypes = (HANDLE, DWORD) 108 wfso.restype = DWORD 109 return wfso 110 111 112# SW_HIDE = 0 113SW_SHOW = 5 114SEE_MASK_NOCLOSEPROCESS = 0x00000040 115SEE_MASK_NO_CONSOLE = 0x00008000 116INFINITE = -1 117 118 119def wait_and_close_handle(process_handle): 120 """ 121 Waits till spawned process finishes and closes the handle for it 122 123 Parameters 124 ---------- 125 process_handle : HANDLE 126 The Windows handle for the process 127 """ 128 WaitForSingleObject(process_handle, INFINITE) 129 CloseHandle(process_handle) 130 131 132def sudo(executable, args=None): 133 """ 134 This will re-run current Python script requesting to elevate administrative rights. 135 136 Parameters 137 ---------- 138 param executable : str 139 The path/name of the executable 140 args : list of str 141 The arguments to be passed to the executable 142 """ 143 if not args: 144 args = [] 145 146 execute_info = ShellExecuteInfo( 147 fMask=SEE_MASK_NOCLOSEPROCESS | SEE_MASK_NO_CONSOLE, 148 hwnd=GetActiveWindow(), 149 lpVerb=b"runas", 150 lpFile=executable.encode("utf-8"), 151 lpParameters=subprocess.list2cmdline(args).encode("utf-8"), 152 lpDirectory=None, 153 nShow=SW_SHOW, 154 ) 155 156 if not ShellExecuteEx(byref(execute_info)): 157 raise ctypes.WinError() 158 159 wait_and_close_handle(execute_info.hProcess) 160 161 162# 163# The following has been refactored from 164# http://stackoverflow.com/a/37505496/2312428 165# 166 167# input flags 168ENABLE_PROCESSED_INPUT = 0x0001 169ENABLE_LINE_INPUT = 0x0002 170ENABLE_ECHO_INPUT = 0x0004 171ENABLE_WINDOW_INPUT = 0x0008 172ENABLE_MOUSE_INPUT = 0x0010 173ENABLE_INSERT_MODE = 0x0020 174ENABLE_QUICK_EDIT_MODE = 0x0040 175 176# output flags 177ENABLE_PROCESSED_OUTPUT = 0x0001 178ENABLE_WRAP_AT_EOL_OUTPUT = 0x0002 179ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 # VT100 (Win 10) 180 181 182def check_zero(result, func, args): 183 if not result: 184 err = ctypes.get_last_error() 185 if err: 186 raise ctypes.WinError(err) 187 return args 188 189 190@lazyobject 191def GetStdHandle(): 192 return lazyimps._winapi.GetStdHandle 193 194 195@lazyobject 196def STDHANDLES(): 197 """Tuple of the Windows handles for (stdin, stdout, stderr).""" 198 hs = [ 199 lazyimps._winapi.STD_INPUT_HANDLE, 200 lazyimps._winapi.STD_OUTPUT_HANDLE, 201 lazyimps._winapi.STD_ERROR_HANDLE, 202 ] 203 hcons = [] 204 for h in hs: 205 hcon = GetStdHandle(int(h)) 206 hcons.append(hcon) 207 return tuple(hcons) 208 209 210@lazyobject 211def GetConsoleMode(): 212 gcm = ctypes.windll.kernel32.GetConsoleMode 213 gcm.errcheck = check_zero 214 gcm.argtypes = (HANDLE, LPDWORD) # _In_ hConsoleHandle # _Out_ lpMode 215 return gcm 216 217 218def get_console_mode(fd=1): 219 """Get the mode of the active console input, output, or error 220 buffer. Note that if the process isn't attached to a 221 console, this function raises an EBADF IOError. 222 223 Parameters 224 ---------- 225 fd : int 226 Standard buffer file descriptor, 0 for stdin, 1 for stdout (default), 227 and 2 for stderr 228 """ 229 mode = DWORD() 230 hcon = STDHANDLES[fd] 231 GetConsoleMode(hcon, byref(mode)) 232 return mode.value 233 234 235@lazyobject 236def SetConsoleMode(): 237 scm = ctypes.windll.kernel32.SetConsoleMode 238 scm.errcheck = check_zero 239 scm.argtypes = (HANDLE, DWORD) # _In_ hConsoleHandle # _Out_ lpMode 240 return scm 241 242 243def set_console_mode(mode, fd=1): 244 """Set the mode of the active console input, output, or 245 error buffer. Note that if the process isn't attached to a 246 console, this function raises an EBADF IOError. 247 248 Parameters 249 ---------- 250 mode : int 251 Mode flags to set on the handle. 252 fd : int, optional 253 Standard buffer file descriptor, 0 for stdin, 1 for stdout (default), 254 and 2 for stderr 255 """ 256 hcon = STDHANDLES[fd] 257 SetConsoleMode(hcon, mode) 258 259 260def enable_virtual_terminal_processing(): 261 """Enables virtual terminal processing on Windows. 262 This includes ANSI escape sequence interpretation. 263 See http://stackoverflow.com/a/36760881/2312428 264 """ 265 SetConsoleMode(GetStdHandle(-11), 7) 266 267 268@lazyobject 269def COORD(): 270 if platform.has_prompt_toolkit(): 271 # turns out that PTK has a separate ctype wrapper 272 # for this struct and also wraps similar function calls 273 # we need to use the same struct to prevent clashes. 274 import prompt_toolkit.win32_types 275 276 return prompt_toolkit.win32_types.COORD 277 278 class _COORD(ctypes.Structure): 279 """Struct from the winapi, representing coordinates in the console. 280 281 Attributes 282 ---------- 283 X : int 284 Column position 285 Y : int 286 Row position 287 """ 288 289 _fields_ = [("X", SHORT), ("Y", SHORT)] 290 291 return _COORD 292 293 294@lazyobject 295def ReadConsoleOutputCharacterA(): 296 rcoc = ctypes.windll.kernel32.ReadConsoleOutputCharacterA 297 rcoc.errcheck = check_zero 298 rcoc.argtypes = ( 299 HANDLE, # _In_ hConsoleOutput 300 LPCSTR, # _Out_ LPTSTR lpMode 301 DWORD, # _In_ nLength 302 COORD, # _In_ dwReadCoord, 303 LPDWORD, 304 ) # _Out_ lpNumberOfCharsRead 305 rcoc.restype = BOOL 306 return rcoc 307 308 309@lazyobject 310def ReadConsoleOutputCharacterW(): 311 rcoc = ctypes.windll.kernel32.ReadConsoleOutputCharacterW 312 rcoc.errcheck = check_zero 313 rcoc.argtypes = ( 314 HANDLE, # _In_ hConsoleOutput 315 LPCWSTR, # _Out_ LPTSTR lpMode 316 DWORD, # _In_ nLength 317 COORD, # _In_ dwReadCoord, 318 LPDWORD, 319 ) # _Out_ lpNumberOfCharsRead 320 rcoc.restype = BOOL 321 return rcoc 322 323 324def read_console_output_character(x=0, y=0, fd=1, buf=None, bufsize=1024, raw=False): 325 """Reads characters from the console buffer. 326 327 Parameters 328 ---------- 329 x : int, optional 330 Starting column. 331 y : int, optional 332 Starting row. 333 fd : int, optional 334 Standard buffer file descriptor, 0 for stdin, 1 for stdout (default), 335 and 2 for stderr. 336 buf : ctypes.c_wchar_p if raw else ctypes.c_wchar_p, optional 337 An existing buffer to (re-)use. 338 bufsize : int, optional 339 The maximum read size. 340 raw : bool, optional 341 Whether to read in and return as bytes (True) or as a 342 unicode string (False, default). 343 344 Returns 345 ------- 346 value : str 347 Result of what was read, may be shorter than bufsize. 348 """ 349 hcon = STDHANDLES[fd] 350 if buf is None: 351 if raw: 352 buf = ctypes.c_char_p(b" " * bufsize) 353 else: 354 buf = ctypes.c_wchar_p(" " * bufsize) 355 coord = COORD(x, y) 356 n = DWORD() 357 if raw: 358 ReadConsoleOutputCharacterA(hcon, buf, bufsize, coord, byref(n)) 359 else: 360 ReadConsoleOutputCharacterW(hcon, buf, bufsize, coord, byref(n)) 361 return buf.value[: n.value] 362 363 364def pread_console(fd, buffersize, offset, buf=None): 365 """This is a console-based implementation of os.pread() for windows. 366 that uses read_console_output_character(). 367 """ 368 cols, rows = os.get_terminal_size(fd=fd) 369 x = offset % cols 370 y = offset // cols 371 return read_console_output_character( 372 x=x, y=y, fd=fd, buf=buf, bufsize=buffersize, raw=True 373 ) 374 375 376# 377# The following piece has been forked from colorama.win32 378# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. 379# 380 381 382@lazyobject 383def CONSOLE_SCREEN_BUFFER_INFO(): 384 if platform.has_prompt_toolkit(): 385 # turns out that PTK has a separate ctype wrapper 386 # for this struct and also wraps kernel32.GetConsoleScreenBufferInfo 387 # we need to use the same struct to prevent clashes. 388 import prompt_toolkit.win32_types 389 390 return prompt_toolkit.win32_types.CONSOLE_SCREEN_BUFFER_INFO 391 392 # Otherwise we should wrap it ourselves 393 COORD() # force COORD to load 394 395 class _CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure): 396 """Struct from in wincon.h. See Windows API docs 397 for more details. 398 399 Attributes 400 ---------- 401 dwSize : COORD 402 Size of 403 dwCursorPosition : COORD 404 Current cursor location. 405 wAttributes : WORD 406 Flags for screen buffer. 407 srWindow : SMALL_RECT 408 Actual size of screen 409 dwMaximumWindowSize : COORD 410 Maximum window scrollback size. 411 """ 412 413 _fields_ = [ 414 ("dwSize", COORD), 415 ("dwCursorPosition", COORD), 416 ("wAttributes", WORD), 417 ("srWindow", SMALL_RECT), 418 ("dwMaximumWindowSize", COORD), 419 ] 420 421 return _CONSOLE_SCREEN_BUFFER_INFO 422 423 424@lazyobject 425def GetConsoleScreenBufferInfo(): 426 """Returns the windows version of the get screen buffer.""" 427 gcsbi = ctypes.windll.kernel32.GetConsoleScreenBufferInfo 428 gcsbi.errcheck = check_zero 429 gcsbi.argtypes = (HANDLE, POINTER(CONSOLE_SCREEN_BUFFER_INFO)) 430 gcsbi.restype = BOOL 431 return gcsbi 432 433 434def get_console_screen_buffer_info(fd=1): 435 """Returns an screen buffer info object for the relevant stdbuf. 436 437 Parameters 438 ---------- 439 fd : int, optional 440 Standard buffer file descriptor, 0 for stdin, 1 for stdout (default), 441 and 2 for stderr. 442 443 Returns 444 ------- 445 csbi : CONSOLE_SCREEN_BUFFER_INFO 446 Information about the console screen buffer. 447 """ 448 hcon = STDHANDLES[fd] 449 csbi = CONSOLE_SCREEN_BUFFER_INFO() 450 GetConsoleScreenBufferInfo(hcon, byref(csbi)) 451 return csbi 452 453 454# 455# end colorama forked section 456# 457 458 459def get_cursor_position(fd=1): 460 """Gets the current cursor position as an (x, y) tuple.""" 461 csbi = get_console_screen_buffer_info(fd=fd) 462 coord = csbi.dwCursorPosition 463 return (coord.X, coord.Y) 464 465 466def get_cursor_offset(fd=1): 467 """Gets the current cursor position as a total offset value.""" 468 csbi = get_console_screen_buffer_info(fd=fd) 469 pos = csbi.dwCursorPosition 470 size = csbi.dwSize 471 return (pos.Y * size.X) + pos.X 472 473 474def get_position_size(fd=1): 475 """Gets the current cursor position and screen size tuple: 476 (x, y, columns, lines). 477 """ 478 info = get_console_screen_buffer_info(fd) 479 return ( 480 info.dwCursorPosition.X, 481 info.dwCursorPosition.Y, 482 info.dwSize.X, 483 info.dwSize.Y, 484 ) 485 486 487@lazyobject 488def SetConsoleScreenBufferSize(): 489 """Set screen buffer dimensions.""" 490 scsbs = ctypes.windll.kernel32.SetConsoleScreenBufferSize 491 scsbs.errcheck = check_zero 492 scsbs.argtypes = (HANDLE, COORD) # _In_ HANDLE hConsoleOutput # _In_ COORD dwSize 493 scsbs.restype = BOOL 494 return scsbs 495 496 497def set_console_screen_buffer_size(x, y, fd=1): 498 """Sets the console size for a standard buffer. 499 500 Parameters 501 ---------- 502 x : int 503 Number of columns. 504 y : int 505 Number of rows. 506 fd : int, optional 507 Standard buffer file descriptor, 0 for stdin, 1 for stdout (default), 508 and 2 for stderr. 509 """ 510 coord = COORD() 511 coord.X = x 512 coord.Y = y 513 hcon = STDHANDLES[fd] 514 rtn = SetConsoleScreenBufferSize(hcon, coord) 515 return rtn 516 517 518@lazyobject 519def SetConsoleCursorPosition(): 520 """Set cursor position in console.""" 521 sccp = ctypes.windll.kernel32.SetConsoleCursorPosition 522 sccp.errcheck = check_zero 523 sccp.argtypes = ( 524 HANDLE, # _In_ HANDLE hConsoleOutput 525 COORD, # _In_ COORD dwCursorPosition 526 ) 527 sccp.restype = BOOL 528 return sccp 529 530 531def set_console_cursor_position(x, y, fd=1): 532 """Sets the console cursor position for a standard buffer. 533 534 Parameters 535 ---------- 536 x : int 537 Number of columns. 538 y : int 539 Number of rows. 540 fd : int, optional 541 Standard buffer file descriptor, 0 for stdin, 1 for stdout (default), 542 and 2 for stderr. 543 """ 544 coord = COORD() 545 coord.X = x 546 coord.Y = y 547 hcon = STDHANDLES[fd] 548 rtn = SetConsoleCursorPosition(hcon, coord) 549 return rtn 550