mirror of
https://github.com/Gallopsled/pwntools.git
synced 2026-10-12 02:02:27 +00:00
peace-maker
e88054fe49
* Start to gradually add type hints Use mypy-baseline to ignore all existing errors for now while avoiding adding new ones in new code. * Remove Python 2 fallback imports * Remove py2compat layer * Avoid redefining a tube helper function * Remove explicit new-style class (object) syntax old-style classes were removed in Python 3 so there is no distinction anymore. * Ignore type error in shellcraft helper script * Fix repeatedly changing `linux_dirent.d_name` * Avoid typechecking windows_termcap on Linux We run mypy on Linux in CI, ignore Windows related import errors * Run mypy-baseline in CI * Add type hint plan to CONTRIBUTING * Sort the mypy baseline Run with `--sort-baseline` for better diffs. `mypy | mypy-baseline sync --sort-baseline` * Switch dev dependencies to dependency-groups The dependency groups aren't baked into build metadata and won't appear on Pypi.
94 lines
2.6 KiB
Python
94 lines
2.6 KiB
Python
import os
|
|
import msvcrt
|
|
import ctypes
|
|
import sys
|
|
from ctypes import wintypes
|
|
|
|
if sys.platform != 'win32':
|
|
raise ImportError("This module is only for Windows")
|
|
|
|
__all__ = ['get']
|
|
|
|
cache = None
|
|
|
|
def get(cap, *args, **kwargs):
|
|
default = kwargs.pop('default', '')
|
|
|
|
if 'PWNLIB_NOTERM' in os.environ:
|
|
return default
|
|
|
|
if kwargs != {}:
|
|
raise TypeError("get(): No such argument %r" % kwargs.popitem()[0])
|
|
|
|
if cache is None:
|
|
init()
|
|
|
|
s = cache.get(cap)
|
|
if s:
|
|
if args:
|
|
return s(*args)
|
|
return s
|
|
return default
|
|
|
|
def init():
|
|
global cache
|
|
cache = {}
|
|
|
|
if 'PWNLIB_NOTERM' not in os.environ:
|
|
try:
|
|
enable_vt_mode()
|
|
except:
|
|
# If the terminal doesn't support ANSI escape codes, don't use them.
|
|
return
|
|
|
|
# Setup capabilities similar to curses on unix.
|
|
# https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
|
|
cache['colors'] = 256
|
|
cache['reset'] = '\x1b[0m'
|
|
cache['bold'] = '\x1b[1m'
|
|
cache['smul'] = '\x1b[4m'
|
|
cache['rev'] = '\x1b[7m'
|
|
cache['setaf'] = lambda c: f'\x1b[3{c}m' if c < 8 else f'\x1b[9{c-8}m'
|
|
cache['setab'] = lambda c: f'\x1b[4{c}m' if c < 8 else f'\x1b[10{c-8}m'
|
|
|
|
# Enable ANSI escape sequences on Windows 10.
|
|
# https://bugs.python.org/issue30075
|
|
kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
|
|
|
|
ERROR_INVALID_PARAMETER = 0x0057
|
|
ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
|
|
|
|
def _check_bool(result, func, args):
|
|
if not result:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
return args
|
|
|
|
LPDWORD = ctypes.POINTER(wintypes.DWORD)
|
|
kernel32.GetConsoleMode.errcheck = _check_bool
|
|
kernel32.GetConsoleMode.argtypes = (wintypes.HANDLE, LPDWORD)
|
|
kernel32.SetConsoleMode.errcheck = _check_bool
|
|
kernel32.SetConsoleMode.argtypes = (wintypes.HANDLE, wintypes.DWORD)
|
|
|
|
def set_conout_mode(new_mode, mask=0xffffffff):
|
|
# don't assume StandardOutput is a console.
|
|
# open CONOUT$ instead
|
|
fdout = os.open('CONOUT$', os.O_RDWR)
|
|
try:
|
|
hout = msvcrt.get_osfhandle(fdout)
|
|
old_mode = wintypes.DWORD()
|
|
kernel32.GetConsoleMode(hout, ctypes.byref(old_mode))
|
|
mode = (new_mode & mask) | (old_mode.value & ~mask)
|
|
kernel32.SetConsoleMode(hout, mode)
|
|
return old_mode.value
|
|
finally:
|
|
os.close(fdout)
|
|
|
|
def enable_vt_mode():
|
|
mode = mask = ENABLE_VIRTUAL_TERMINAL_PROCESSING
|
|
try:
|
|
return set_conout_mode(mode, mask)
|
|
except WindowsError as e:
|
|
if e.winerror == ERROR_INVALID_PARAMETER:
|
|
raise NotImplementedError
|
|
raise
|