vendor all the things!

This commit is contained in:
2017-03-11 23:26:08 -05:00
parent 5da375b3ac
commit b6fb357873
221 changed files with 105904 additions and 113 deletions
+3
View File
@@ -1,3 +1,6 @@
3.5.2:
- Vendor all the things!
- get-pipenv.py.
3.5.1:
- Basic Windows support!
3.5.0
+3
View File
@@ -0,0 +1,3 @@
The contents of the vendor directory are subject to different licenses
than the rest of this project. Their respective licenses can be looked
up at pypi.python.org.
-12
View File
@@ -4,19 +4,7 @@ mock = "*"
Sphinx = "*"
[packages]
click = "*"
crayons = "*"
toml = "*"
"delegator.py" = ">=0.0.6"
requests = ">=2.4.0"
requirements-parser = "*"
parse = "*"
pipfile = "==0.0.2"
click-completion = "*"
psutil = "*"
pew = ">=0.1.26"
blindspin = "*"
"backports.shutil_get_terminal_size" = "*"
[requires]
python_version = "2.7"
Generated
+24 -88
View File
@@ -1,80 +1,20 @@
{
"default": {
"crayons": {
"version": "==0.1.2",
"hash": "sha256:6f51241d0c4faec1c04c1c0ac6a68f1d66a4655476ce1570b3f37e5166a599cc"
},
"requirements-parser": {
"version": "==0.1.0",
"hash": "sha256:fee2380a469ffe4067bc7f0096a6fcfb27539da7496fae12b74b8d5d0f33a4ee"
},
"requests": {
"version": "==2.13.0",
"hash": "sha256:1a720e8862a41aa22e339373b526f508ef0c8988baf48b84d3fc891a8e237efb"
},
"blindspin": {
"version": "==2.0.1",
"hash": "sha256:31c4e93d4ae2ef6765e3c42460456814c17addd5add8298dced21fe9dc50f496"
},
"Jinja2": {
"version": "==2.9.5",
"hash": "sha256:a7b7438120dbe76a8e735ef7eba6048eaf4e0b7dbc530e100812f8ec462a4d50"
},
"packaging": {
"version": "==16.8",
"hash": "sha256:99276dc6e3a7851f32027a68f1095cd3f77c148091b092ea867a351811cfe388"
},
"MarkupSafe": {
"version": "==1.0",
"hash": "sha256:a6be69091dac236ea9c6bc7d012beab42010fa914c459791d627dad4910eb665"
},
"colorama": {
"version": "==0.3.7",
"hash": "sha256:a4c0f5bc358a62849653471e309dcc991223cf86abafbec17cd8f41327279e89"
},
"pyparsing": {
"version": "==2.2.0",
"hash": "sha256:fee43f17a9c4087e7ed1605bd6df994c6173c1e977d7ade7b651292fab2bd010"
},
"click-completion": {
"version": "==0.2.1",
"hash": "sha256:079fb138887d4de12a0b7fbebf8d92d396b7c1a9c49f63475d9f3909d2588976"
},
"click": {
"version": "==6.7",
"hash": "sha256:29f99fc6125fbc931b758dc053b3114e55c77a6e4c6c3a2674a2dc986016381d"
},
"appdirs": {
"version": "==1.4.3",
"hash": "sha256:d8b24664561d0d34ddfaec54636d502d7cea6e29c3eaf68f3df6180863e2166e"
},
"virtualenv": {
"version": "==15.1.0",
"hash": "sha256:39d88b533b422825d644087a21e78c45cf5af0ef7a99a1fc9fbb7b481e5c85b0"
},
"parse": {
"version": "==1.8.0",
"hash": "sha256:8b4f28bbe7c0f24981669ea92b2ba704ee63b5346027e82be30118bb5788ff10"
},
"pexpect": {
"version": "==4.2.1",
"hash": "sha256:f853b52afaf3b064d29854771e2db509ef80392509bde2dd7a6ecf2dfc3f0018"
},
"delegator.py": {
"version": "==0.0.8",
"hash": "sha256:603c1c1c76b1340520d95a1e768cf2d3a8df7b3c15a1eb9df83aac29b8d025a4"
},
"pythonz-bd": {
"version": "==1.11.4",
"hash": "sha256:30fa48c5b542e1ebfca167f10699b149768dd18a90185d98b8a766636b6343b9"
"pew": {
"version": "==0.1.26",
"hash": "sha256:259ac7a4603fe41b1fa950f30b2e4ccf4a23f7c89be2c34e0a4cec176c3ec581"
},
"virtualenv-clone": {
"version": "==0.2.6",
"hash": "sha256:6b3be5cab59e455f08c9eda573d23006b7d6fb41fae974ddaa2b275c93cc4405"
},
"pew": {
"version": "==0.1.26",
"hash": "sha256:259ac7a4603fe41b1fa950f30b2e4ccf4a23f7c89be2c34e0a4cec176c3ec581"
"six": {
"version": "==1.10.0",
"hash": "sha256:0ff78c403d9bccf5a425a6d31a12aa6b47f1c21ca4dc2573a7e2f32a97335eb1"
},
"backports.shutil_get_terminal_size": {
"version": "==1.0.0",
@@ -84,33 +24,29 @@
"version": "==1.0.1",
"hash": "sha256:6940718dfc3eff4258203ad5021090933e5c04707d5ca8cc9e73c94a7894ea9f"
},
"ptyprocess": {
"version": "==0.5.1",
"hash": "sha256:464cb76f7a7122743dd25507650db89cd447c51f38e4671602b3eaa2e38e05ae"
"packaging": {
"version": "==16.8",
"hash": "sha256:99276dc6e3a7851f32027a68f1095cd3f77c148091b092ea867a351811cfe388"
},
"six": {
"version": "==1.10.0",
"hash": "sha256:0ff78c403d9bccf5a425a6d31a12aa6b47f1c21ca4dc2573a7e2f32a97335eb1"
},
"psutil": {
"version": "==5.2.0",
"hash": "sha256:2fc91d068faa5613c093335f0e758673ef8c722ad4bfa4aded64c13ae69089eb"
"pythonz-bd": {
"version": "==1.11.4",
"hash": "sha256:30fa48c5b542e1ebfca167f10699b149768dd18a90185d98b8a766636b6343b9"
},
"shutilwhich": {
"version": "==1.1.0",
"hash": "sha256:db1f39c6461e42f630fa617bb8c79090f7711c9ca493e615e43d0610ecb64dc6"
},
"toml": {
"version": "==0.9.2",
"hash": "sha256:b3953bffe848ad9a6d554114d82f2dcb3e23945e90b4d9addc9956f37f336594"
},
"setuptools": {
"version": "==34.3.1",
"hash": "sha256:9349b47bcdef2e695065791260f9416b7aae8739fd30c6849f807bce5281d046"
"version": "==34.3.2",
"hash": "sha256:6483f8412313ec787fa71379147a4605d3b1cc303c3648d02542a9160d3db72b"
},
"pipfile": {
"version": "==0.0.2",
"hash": "sha256:f7d9f15de8b660986557eb3cc5391aa1a16207ac41bc378d03f414762d36c984"
"pyparsing": {
"version": "==2.2.0",
"hash": "sha256:fee43f17a9c4087e7ed1605bd6df994c6173c1e977d7ade7b651292fab2bd010"
},
"appdirs": {
"version": "==1.4.3",
"hash": "sha256:d8b24664561d0d34ddfaec54636d502d7cea6e29c3eaf68f3df6180863e2166e"
}
},
"develop": {
@@ -187,8 +123,8 @@
"hash": "sha256:a7b7438120dbe76a8e735ef7eba6048eaf4e0b7dbc530e100812f8ec462a4d50"
},
"setuptools": {
"version": "==34.3.1",
"hash": "sha256:9349b47bcdef2e695065791260f9416b7aae8739fd30c6849f807bce5281d046"
"version": "==34.3.2",
"hash": "sha256:6483f8412313ec787fa71379147a4605d3b1cc303c3648d02542a9160d3db72b"
},
"requests": {
"version": "==2.13.0",
@@ -210,7 +146,7 @@
"python_version": "2.7"
},
"hash": {
"sha256": "4e215f83ada415726a064a01ad5bdd5daf169d1ef41623b1c05e6cd8fe1f0fc5"
"sha256": "a1ebdaccf29fc67e8b827cd7a04d0cf7a453baeae55508b930b4e4af2e4c47c5"
}
}
}
Executable
+20068
View File
File diff suppressed because it is too large Load Diff
+8
View File
@@ -3,6 +3,14 @@
# | ||__/|__| |\/
# |
import os
import sys
# Inject vendored directory into system path.
v_path = os.path.sep.join([os.path.dirname(os.path.realpath(__file__)), 'vendor'])
sys.path.append(v_path)
from .cli import cli
if __name__ == '__main__':
BIN
View File
Binary file not shown.
+2
View File
@@ -0,0 +1,2 @@
from pkgutil import extend_path
__path__ = extend_path(__path__, __name__)
@@ -0,0 +1,11 @@
"""A backport of the get_terminal_size function from Python 3.3's shutil."""
__title__ = "backports.shutil_get_terminal_size"
__version__ = "1.0.0"
__license__ = "MIT"
__author__ = "Christopher Rosell"
__copyright__ = "Copyright 2014 Christopher Rosell"
__all__ = ["get_terminal_size"]
from .get_terminal_size import get_terminal_size
@@ -0,0 +1,101 @@
"""This is a backport of shutil.get_terminal_size from Python 3.3.
The original implementation is in C, but here we use the ctypes and
fcntl modules to create a pure Python version of os.get_terminal_size.
"""
import os
import struct
import sys
from collections import namedtuple
__all__ = ["get_terminal_size"]
terminal_size = namedtuple("terminal_size", "columns lines")
try:
from ctypes import windll, create_string_buffer
_handles = {
0: windll.kernel32.GetStdHandle(-10),
1: windll.kernel32.GetStdHandle(-11),
2: windll.kernel32.GetStdHandle(-12),
}
def _get_terminal_size(fd):
columns = lines = 0
try:
handle = _handles[fd]
csbi = create_string_buffer(22)
res = windll.kernel32.GetConsoleScreenBufferInfo(handle, csbi)
if res:
res = struct.unpack("hhhhHhhhhhh", csbi.raw)
left, top, right, bottom = res[5:9]
columns = right - left + 1
lines = bottom - top + 1
except Exception:
pass
return terminal_size(columns, lines)
except ImportError:
import fcntl
import termios
def _get_terminal_size(fd):
try:
res = fcntl.ioctl(fd, termios.TIOCGWINSZ, b"\x00" * 4)
lines, columns = struct.unpack("hh", res)
except Exception:
columns = lines = 0
return terminal_size(columns, lines)
def get_terminal_size(fallback=(80, 24)):
"""Get the size of the terminal window.
For each of the two dimensions, the environment variable, COLUMNS
and LINES respectively, is checked. If the variable is defined and
the value is a positive integer, it is used.
When COLUMNS or LINES is not defined, which is the common case,
the terminal connected to sys.__stdout__ is queried
by invoking os.get_terminal_size.
If the terminal size cannot be successfully queried, either because
the system doesn't support querying, or because we are not
connected to a terminal, the value given in fallback parameter
is used. Fallback defaults to (80, 24) which is the default
size used by many terminal emulators.
The value returned is a named tuple of type os.terminal_size.
"""
# Try the environment first
try:
columns = int(os.environ["COLUMNS"])
except (KeyError, ValueError):
columns = 0
try:
lines = int(os.environ["LINES"])
except (KeyError, ValueError):
lines = 0
# Only query if necessary
if columns <= 0 or lines <= 0:
try:
size = _get_terminal_size(sys.__stdout__.fileno())
except (NameError, OSError):
size = terminal_size(*fallback)
if columns <= 0:
columns = size.columns
if lines <= 0:
lines = size.lines
return terminal_size(columns, lines)
+73
View File
@@ -0,0 +1,73 @@
# -*- coding: utf-8 -*-
import sys
import threading
import time
import itertools
class Spinner(object):
spinner_cycle = itertools.cycle(u'⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏')
def __init__(self, beep=False, force=False):
self.beep = beep
self.force = force
self.stop_running = None
self.spin_thread = None
def start(self):
if sys.stdout.isatty() or self.force:
self.stop_running = threading.Event()
self.spin_thread = threading.Thread(target=self.init_spin)
self.spin_thread.start()
def stop(self):
if self.spin_thread:
self.stop_running.set()
self.spin_thread.join()
def init_spin(self):
while not self.stop_running.is_set():
next_val = next(self.spinner_cycle)
if sys.version_info[0] == 2:
next_val = next_val.encode('utf-8')
sys.stdout.write(next_val)
sys.stdout.flush()
time.sleep(0.07)
sys.stdout.write('\b')
def __enter__(self):
self.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.stop()
if self.beep:
sys.stdout.write('\7')
sys.stdout.flush()
return False
def spinner(beep=False, force=False):
"""This function creates a context manager that is used to display a
spinner on stdout as long as the context has not exited.
The spinner is created only if stdout is not redirected, or if the spinner
is forced using the `force` parameter.
Parameters
----------
beep : bool
Beep when spinner finishes.
force : bool
Force creation of spinner even when stdout is redirected.
Example
-------
with spinner():
do_something()
do_something_else()
"""
return Spinner(beep, force)
+98
View File
@@ -0,0 +1,98 @@
# -*- coding: utf-8 -*-
"""
click
~~~~~
Click is a simple Python module that wraps the stdlib's optparse to make
writing command line scripts fun. Unlike other modules, it's based around
a simple API that does not come with too much magic and is composable.
In case optparse ever gets removed from the stdlib, it will be shipped by
this module.
:copyright: (c) 2014 by Armin Ronacher.
:license: BSD, see LICENSE for more details.
"""
# Core classes
from .core import Context, BaseCommand, Command, MultiCommand, Group, \
CommandCollection, Parameter, Option, Argument
# Globals
from .globals import get_current_context
# Decorators
from .decorators import pass_context, pass_obj, make_pass_decorator, \
command, group, argument, option, confirmation_option, \
password_option, version_option, help_option
# Types
from .types import ParamType, File, Path, Choice, IntRange, Tuple, \
STRING, INT, FLOAT, BOOL, UUID, UNPROCESSED
# Utilities
from .utils import echo, get_binary_stream, get_text_stream, open_file, \
format_filename, get_app_dir, get_os_args
# Terminal functions
from .termui import prompt, confirm, get_terminal_size, echo_via_pager, \
progressbar, clear, style, unstyle, secho, edit, launch, getchar, \
pause
# Exceptions
from .exceptions import ClickException, UsageError, BadParameter, \
FileError, Abort, NoSuchOption, BadOptionUsage, BadArgumentUsage, \
MissingParameter
# Formatting
from .formatting import HelpFormatter, wrap_text
# Parsing
from .parser import OptionParser
__all__ = [
# Core classes
'Context', 'BaseCommand', 'Command', 'MultiCommand', 'Group',
'CommandCollection', 'Parameter', 'Option', 'Argument',
# Globals
'get_current_context',
# Decorators
'pass_context', 'pass_obj', 'make_pass_decorator', 'command', 'group',
'argument', 'option', 'confirmation_option', 'password_option',
'version_option', 'help_option',
# Types
'ParamType', 'File', 'Path', 'Choice', 'IntRange', 'Tuple', 'STRING',
'INT', 'FLOAT', 'BOOL', 'UUID', 'UNPROCESSED',
# Utilities
'echo', 'get_binary_stream', 'get_text_stream', 'open_file',
'format_filename', 'get_app_dir', 'get_os_args',
# Terminal functions
'prompt', 'confirm', 'get_terminal_size', 'echo_via_pager',
'progressbar', 'clear', 'style', 'unstyle', 'secho', 'edit', 'launch',
'getchar', 'pause',
# Exceptions
'ClickException', 'UsageError', 'BadParameter', 'FileError',
'Abort', 'NoSuchOption', 'BadOptionUsage', 'BadArgumentUsage',
'MissingParameter',
# Formatting
'HelpFormatter', 'wrap_text',
# Parsing
'OptionParser',
]
# Controls if click should emit the warning about the use of unicode
# literals.
disable_unicode_literals_warning = False
__version__ = '6.7'
+83
View File
@@ -0,0 +1,83 @@
import os
import re
from .utils import echo
from .parser import split_arg_string
from .core import MultiCommand, Option
COMPLETION_SCRIPT = '''
%(complete_func)s() {
COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \\
COMP_CWORD=$COMP_CWORD \\
%(autocomplete_var)s=complete $1 ) )
return 0
}
complete -F %(complete_func)s -o default %(script_names)s
'''
_invalid_ident_char_re = re.compile(r'[^a-zA-Z0-9_]')
def get_completion_script(prog_name, complete_var):
cf_name = _invalid_ident_char_re.sub('', prog_name.replace('-', '_'))
return (COMPLETION_SCRIPT % {
'complete_func': '_%s_completion' % cf_name,
'script_names': prog_name,
'autocomplete_var': complete_var,
}).strip() + ';'
def resolve_ctx(cli, prog_name, args):
ctx = cli.make_context(prog_name, args, resilient_parsing=True)
while ctx.protected_args + ctx.args and isinstance(ctx.command, MultiCommand):
a = ctx.protected_args + ctx.args
cmd = ctx.command.get_command(ctx, a[0])
if cmd is None:
return None
ctx = cmd.make_context(a[0], a[1:], parent=ctx, resilient_parsing=True)
return ctx
def get_choices(cli, prog_name, args, incomplete):
ctx = resolve_ctx(cli, prog_name, args)
if ctx is None:
return
choices = []
if incomplete and not incomplete[:1].isalnum():
for param in ctx.command.params:
if not isinstance(param, Option):
continue
choices.extend(param.opts)
choices.extend(param.secondary_opts)
elif isinstance(ctx.command, MultiCommand):
choices.extend(ctx.command.list_commands(ctx))
for item in choices:
if item.startswith(incomplete):
yield item
def do_complete(cli, prog_name):
cwords = split_arg_string(os.environ['COMP_WORDS'])
cword = int(os.environ['COMP_CWORD'])
args = cwords[1:cword]
try:
incomplete = cwords[cword]
except IndexError:
incomplete = ''
for item in get_choices(cli, prog_name, args, incomplete):
echo(item)
return True
def bashcomplete(cli, prog_name, complete_var, complete_instr):
if complete_instr == 'source':
echo(get_completion_script(prog_name, complete_var))
return True
elif complete_instr == 'complete':
return do_complete(cli, prog_name)
return False
+648
View File
@@ -0,0 +1,648 @@
import re
import io
import os
import sys
import codecs
from weakref import WeakKeyDictionary
PY2 = sys.version_info[0] == 2
WIN = sys.platform.startswith('win')
DEFAULT_COLUMNS = 80
_ansi_re = re.compile('\033\[((?:\d|;)*)([a-zA-Z])')
def get_filesystem_encoding():
return sys.getfilesystemencoding() or sys.getdefaultencoding()
def _make_text_stream(stream, encoding, errors):
if encoding is None:
encoding = get_best_encoding(stream)
if errors is None:
errors = 'replace'
return _NonClosingTextIOWrapper(stream, encoding, errors,
line_buffering=True)
def is_ascii_encoding(encoding):
"""Checks if a given encoding is ascii."""
try:
return codecs.lookup(encoding).name == 'ascii'
except LookupError:
return False
def get_best_encoding(stream):
"""Returns the default stream encoding if not found."""
rv = getattr(stream, 'encoding', None) or sys.getdefaultencoding()
if is_ascii_encoding(rv):
return 'utf-8'
return rv
class _NonClosingTextIOWrapper(io.TextIOWrapper):
def __init__(self, stream, encoding, errors, **extra):
self._stream = stream = _FixupStream(stream)
io.TextIOWrapper.__init__(self, stream, encoding, errors, **extra)
# The io module is a place where the Python 3 text behavior
# was forced upon Python 2, so we need to unbreak
# it to look like Python 2.
if PY2:
def write(self, x):
if isinstance(x, str) or is_bytes(x):
try:
self.flush()
except Exception:
pass
return self.buffer.write(str(x))
return io.TextIOWrapper.write(self, x)
def writelines(self, lines):
for line in lines:
self.write(line)
def __del__(self):
try:
self.detach()
except Exception:
pass
def isatty(self):
# https://bitbucket.org/pypy/pypy/issue/1803
return self._stream.isatty()
class _FixupStream(object):
"""The new io interface needs more from streams than streams
traditionally implement. As such, this fix-up code is necessary in
some circumstances.
"""
def __init__(self, stream):
self._stream = stream
def __getattr__(self, name):
return getattr(self._stream, name)
def read1(self, size):
f = getattr(self._stream, 'read1', None)
if f is not None:
return f(size)
# We only dispatch to readline instead of read in Python 2 as we
# do not want cause problems with the different implementation
# of line buffering.
if PY2:
return self._stream.readline(size)
return self._stream.read(size)
def readable(self):
x = getattr(self._stream, 'readable', None)
if x is not None:
return x()
try:
self._stream.read(0)
except Exception:
return False
return True
def writable(self):
x = getattr(self._stream, 'writable', None)
if x is not None:
return x()
try:
self._stream.write('')
except Exception:
try:
self._stream.write(b'')
except Exception:
return False
return True
def seekable(self):
x = getattr(self._stream, 'seekable', None)
if x is not None:
return x()
try:
self._stream.seek(self._stream.tell())
except Exception:
return False
return True
if PY2:
text_type = unicode
bytes = str
raw_input = raw_input
string_types = (str, unicode)
iteritems = lambda x: x.iteritems()
range_type = xrange
def is_bytes(x):
return isinstance(x, (buffer, bytearray))
_identifier_re = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')
# For Windows, we need to force stdout/stdin/stderr to binary if it's
# fetched for that. This obviously is not the most correct way to do
# it as it changes global state. Unfortunately, there does not seem to
# be a clear better way to do it as just reopening the file in binary
# mode does not change anything.
#
# An option would be to do what Python 3 does and to open the file as
# binary only, patch it back to the system, and then use a wrapper
# stream that converts newlines. It's not quite clear what's the
# correct option here.
#
# This code also lives in _winconsole for the fallback to the console
# emulation stream.
#
# There are also Windows environments where the `msvcrt` module is not
# available (which is why we use try-catch instead of the WIN variable
# here), such as the Google App Engine development server on Windows. In
# those cases there is just nothing we can do.
try:
import msvcrt
except ImportError:
set_binary_mode = lambda x: x
else:
def set_binary_mode(f):
try:
fileno = f.fileno()
except Exception:
pass
else:
msvcrt.setmode(fileno, os.O_BINARY)
return f
def isidentifier(x):
return _identifier_re.search(x) is not None
def get_binary_stdin():
return set_binary_mode(sys.stdin)
def get_binary_stdout():
return set_binary_mode(sys.stdout)
def get_binary_stderr():
return set_binary_mode(sys.stderr)
def get_text_stdin(encoding=None, errors=None):
rv = _get_windows_console_stream(sys.stdin, encoding, errors)
if rv is not None:
return rv
return _make_text_stream(sys.stdin, encoding, errors)
def get_text_stdout(encoding=None, errors=None):
rv = _get_windows_console_stream(sys.stdout, encoding, errors)
if rv is not None:
return rv
return _make_text_stream(sys.stdout, encoding, errors)
def get_text_stderr(encoding=None, errors=None):
rv = _get_windows_console_stream(sys.stderr, encoding, errors)
if rv is not None:
return rv
return _make_text_stream(sys.stderr, encoding, errors)
def filename_to_ui(value):
if isinstance(value, bytes):
value = value.decode(get_filesystem_encoding(), 'replace')
return value
else:
import io
text_type = str
raw_input = input
string_types = (str,)
range_type = range
isidentifier = lambda x: x.isidentifier()
iteritems = lambda x: iter(x.items())
def is_bytes(x):
return isinstance(x, (bytes, memoryview, bytearray))
def _is_binary_reader(stream, default=False):
try:
return isinstance(stream.read(0), bytes)
except Exception:
return default
# This happens in some cases where the stream was already
# closed. In this case, we assume the default.
def _is_binary_writer(stream, default=False):
try:
stream.write(b'')
except Exception:
try:
stream.write('')
return False
except Exception:
pass
return default
return True
def _find_binary_reader(stream):
# We need to figure out if the given stream is already binary.
# This can happen because the official docs recommend detaching
# the streams to get binary streams. Some code might do this, so
# we need to deal with this case explicitly.
if _is_binary_reader(stream, False):
return stream
buf = getattr(stream, 'buffer', None)
# Same situation here; this time we assume that the buffer is
# actually binary in case it's closed.
if buf is not None and _is_binary_reader(buf, True):
return buf
def _find_binary_writer(stream):
# We need to figure out if the given stream is already binary.
# This can happen because the official docs recommend detatching
# the streams to get binary streams. Some code might do this, so
# we need to deal with this case explicitly.
if _is_binary_writer(stream, False):
return stream
buf = getattr(stream, 'buffer', None)
# Same situation here; this time we assume that the buffer is
# actually binary in case it's closed.
if buf is not None and _is_binary_writer(buf, True):
return buf
def _stream_is_misconfigured(stream):
"""A stream is misconfigured if its encoding is ASCII."""
# If the stream does not have an encoding set, we assume it's set
# to ASCII. This appears to happen in certain unittest
# environments. It's not quite clear what the correct behavior is
# but this at least will force Click to recover somehow.
return is_ascii_encoding(getattr(stream, 'encoding', None) or 'ascii')
def _is_compatible_text_stream(stream, encoding, errors):
stream_encoding = getattr(stream, 'encoding', None)
stream_errors = getattr(stream, 'errors', None)
# Perfect match.
if stream_encoding == encoding and stream_errors == errors:
return True
# Otherwise, it's only a compatible stream if we did not ask for
# an encoding.
if encoding is None:
return stream_encoding is not None
return False
def _force_correct_text_reader(text_reader, encoding, errors):
if _is_binary_reader(text_reader, False):
binary_reader = text_reader
else:
# If there is no target encoding set, we need to verify that the
# reader is not actually misconfigured.
if encoding is None and not _stream_is_misconfigured(text_reader):
return text_reader
if _is_compatible_text_stream(text_reader, encoding, errors):
return text_reader
# If the reader has no encoding, we try to find the underlying
# binary reader for it. If that fails because the environment is
# misconfigured, we silently go with the same reader because this
# is too common to happen. In that case, mojibake is better than
# exceptions.
binary_reader = _find_binary_reader(text_reader)
if binary_reader is None:
return text_reader
# At this point, we default the errors to replace instead of strict
# because nobody handles those errors anyways and at this point
# we're so fundamentally fucked that nothing can repair it.
if errors is None:
errors = 'replace'
return _make_text_stream(binary_reader, encoding, errors)
def _force_correct_text_writer(text_writer, encoding, errors):
if _is_binary_writer(text_writer, False):
binary_writer = text_writer
else:
# If there is no target encoding set, we need to verify that the
# writer is not actually misconfigured.
if encoding is None and not _stream_is_misconfigured(text_writer):
return text_writer
if _is_compatible_text_stream(text_writer, encoding, errors):
return text_writer
# If the writer has no encoding, we try to find the underlying
# binary writer for it. If that fails because the environment is
# misconfigured, we silently go with the same writer because this
# is too common to happen. In that case, mojibake is better than
# exceptions.
binary_writer = _find_binary_writer(text_writer)
if binary_writer is None:
return text_writer
# At this point, we default the errors to replace instead of strict
# because nobody handles those errors anyways and at this point
# we're so fundamentally fucked that nothing can repair it.
if errors is None:
errors = 'replace'
return _make_text_stream(binary_writer, encoding, errors)
def get_binary_stdin():
reader = _find_binary_reader(sys.stdin)
if reader is None:
raise RuntimeError('Was not able to determine binary '
'stream for sys.stdin.')
return reader
def get_binary_stdout():
writer = _find_binary_writer(sys.stdout)
if writer is None:
raise RuntimeError('Was not able to determine binary '
'stream for sys.stdout.')
return writer
def get_binary_stderr():
writer = _find_binary_writer(sys.stderr)
if writer is None:
raise RuntimeError('Was not able to determine binary '
'stream for sys.stderr.')
return writer
def get_text_stdin(encoding=None, errors=None):
rv = _get_windows_console_stream(sys.stdin, encoding, errors)
if rv is not None:
return rv
return _force_correct_text_reader(sys.stdin, encoding, errors)
def get_text_stdout(encoding=None, errors=None):
rv = _get_windows_console_stream(sys.stdout, encoding, errors)
if rv is not None:
return rv
return _force_correct_text_writer(sys.stdout, encoding, errors)
def get_text_stderr(encoding=None, errors=None):
rv = _get_windows_console_stream(sys.stderr, encoding, errors)
if rv is not None:
return rv
return _force_correct_text_writer(sys.stderr, encoding, errors)
def filename_to_ui(value):
if isinstance(value, bytes):
value = value.decode(get_filesystem_encoding(), 'replace')
else:
value = value.encode('utf-8', 'surrogateescape') \
.decode('utf-8', 'replace')
return value
def get_streerror(e, default=None):
if hasattr(e, 'strerror'):
msg = e.strerror
else:
if default is not None:
msg = default
else:
msg = str(e)
if isinstance(msg, bytes):
msg = msg.decode('utf-8', 'replace')
return msg
def open_stream(filename, mode='r', encoding=None, errors='strict',
atomic=False):
# Standard streams first. These are simple because they don't need
# special handling for the atomic flag. It's entirely ignored.
if filename == '-':
if 'w' in mode:
if 'b' in mode:
return get_binary_stdout(), False
return get_text_stdout(encoding=encoding, errors=errors), False
if 'b' in mode:
return get_binary_stdin(), False
return get_text_stdin(encoding=encoding, errors=errors), False
# Non-atomic writes directly go out through the regular open functions.
if not atomic:
if encoding is None:
return open(filename, mode), True
return io.open(filename, mode, encoding=encoding, errors=errors), True
# Some usability stuff for atomic writes
if 'a' in mode:
raise ValueError(
'Appending to an existing file is not supported, because that '
'would involve an expensive `copy`-operation to a temporary '
'file. Open the file in normal `w`-mode and copy explicitly '
'if that\'s what you\'re after.'
)
if 'x' in mode:
raise ValueError('Use the `overwrite`-parameter instead.')
if 'w' not in mode:
raise ValueError('Atomic writes only make sense with `w`-mode.')
# Atomic writes are more complicated. They work by opening a file
# as a proxy in the same folder and then using the fdopen
# functionality to wrap it in a Python file. Then we wrap it in an
# atomic file that moves the file over on close.
import tempfile
fd, tmp_filename = tempfile.mkstemp(dir=os.path.dirname(filename),
prefix='.__atomic-write')
if encoding is not None:
f = io.open(fd, mode, encoding=encoding, errors=errors)
else:
f = os.fdopen(fd, mode)
return _AtomicFile(f, tmp_filename, filename), True
# Used in a destructor call, needs extra protection from interpreter cleanup.
if hasattr(os, 'replace'):
_replace = os.replace
_can_replace = True
else:
_replace = os.rename
_can_replace = not WIN
class _AtomicFile(object):
def __init__(self, f, tmp_filename, real_filename):
self._f = f
self._tmp_filename = tmp_filename
self._real_filename = real_filename
self.closed = False
@property
def name(self):
return self._real_filename
def close(self, delete=False):
if self.closed:
return
self._f.close()
if not _can_replace:
try:
os.remove(self._real_filename)
except OSError:
pass
_replace(self._tmp_filename, self._real_filename)
self.closed = True
def __getattr__(self, name):
return getattr(self._f, name)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, tb):
self.close(delete=exc_type is not None)
def __repr__(self):
return repr(self._f)
auto_wrap_for_ansi = None
colorama = None
get_winterm_size = None
def strip_ansi(value):
return _ansi_re.sub('', value)
def should_strip_ansi(stream=None, color=None):
if color is None:
if stream is None:
stream = sys.stdin
return not isatty(stream)
return not color
# If we're on Windows, we provide transparent integration through
# colorama. This will make ANSI colors through the echo function
# work automatically.
if WIN:
# Windows has a smaller terminal
DEFAULT_COLUMNS = 79
from ._winconsole import _get_windows_console_stream
def _get_argv_encoding():
import locale
return locale.getpreferredencoding()
if PY2:
def raw_input(prompt=''):
sys.stderr.flush()
if prompt:
stdout = _default_text_stdout()
stdout.write(prompt)
stdin = _default_text_stdin()
return stdin.readline().rstrip('\r\n')
try:
import colorama
except ImportError:
pass
else:
_ansi_stream_wrappers = WeakKeyDictionary()
def auto_wrap_for_ansi(stream, color=None):
"""This function wraps a stream so that calls through colorama
are issued to the win32 console API to recolor on demand. It
also ensures to reset the colors if a write call is interrupted
to not destroy the console afterwards.
"""
try:
cached = _ansi_stream_wrappers.get(stream)
except Exception:
cached = None
if cached is not None:
return cached
strip = should_strip_ansi(stream, color)
ansi_wrapper = colorama.AnsiToWin32(stream, strip=strip)
rv = ansi_wrapper.stream
_write = rv.write
def _safe_write(s):
try:
return _write(s)
except:
ansi_wrapper.reset_all()
raise
rv.write = _safe_write
try:
_ansi_stream_wrappers[stream] = rv
except Exception:
pass
return rv
def get_winterm_size():
win = colorama.win32.GetConsoleScreenBufferInfo(
colorama.win32.STDOUT).srWindow
return win.Right - win.Left, win.Bottom - win.Top
else:
def _get_argv_encoding():
return getattr(sys.stdin, 'encoding', None) or get_filesystem_encoding()
_get_windows_console_stream = lambda *x: None
def term_len(x):
return len(strip_ansi(x))
def isatty(stream):
try:
return stream.isatty()
except Exception:
return False
def _make_cached_stream_func(src_func, wrapper_func):
cache = WeakKeyDictionary()
def func():
stream = src_func()
try:
rv = cache.get(stream)
except Exception:
rv = None
if rv is not None:
return rv
rv = wrapper_func()
try:
cache[stream] = rv
except Exception:
pass
return rv
return func
_default_text_stdin = _make_cached_stream_func(
lambda: sys.stdin, get_text_stdin)
_default_text_stdout = _make_cached_stream_func(
lambda: sys.stdout, get_text_stdout)
_default_text_stderr = _make_cached_stream_func(
lambda: sys.stderr, get_text_stderr)
binary_streams = {
'stdin': get_binary_stdin,
'stdout': get_binary_stdout,
'stderr': get_binary_stderr,
}
text_streams = {
'stdin': get_text_stdin,
'stdout': get_text_stdout,
'stderr': get_text_stderr,
}
+547
View File
@@ -0,0 +1,547 @@
"""
click._termui_impl
~~~~~~~~~~~~~~~~~~
This module contains implementations for the termui module. To keep the
import time of Click down, some infrequently used functionality is placed
in this module and only imported as needed.
:copyright: (c) 2014 by Armin Ronacher.
:license: BSD, see LICENSE for more details.
"""
import os
import sys
import time
import math
from ._compat import _default_text_stdout, range_type, PY2, isatty, \
open_stream, strip_ansi, term_len, get_best_encoding, WIN
from .utils import echo
from .exceptions import ClickException
if os.name == 'nt':
BEFORE_BAR = '\r'
AFTER_BAR = '\n'
else:
BEFORE_BAR = '\r\033[?25l'
AFTER_BAR = '\033[?25h\n'
def _length_hint(obj):
"""Returns the length hint of an object."""
try:
return len(obj)
except (AttributeError, TypeError):
try:
get_hint = type(obj).__length_hint__
except AttributeError:
return None
try:
hint = get_hint(obj)
except TypeError:
return None
if hint is NotImplemented or \
not isinstance(hint, (int, long)) or \
hint < 0:
return None
return hint
class ProgressBar(object):
def __init__(self, iterable, length=None, fill_char='#', empty_char=' ',
bar_template='%(bar)s', info_sep=' ', show_eta=True,
show_percent=None, show_pos=False, item_show_func=None,
label=None, file=None, color=None, width=30):
self.fill_char = fill_char
self.empty_char = empty_char
self.bar_template = bar_template
self.info_sep = info_sep
self.show_eta = show_eta
self.show_percent = show_percent
self.show_pos = show_pos
self.item_show_func = item_show_func
self.label = label or ''
if file is None:
file = _default_text_stdout()
self.file = file
self.color = color
self.width = width
self.autowidth = width == 0
if length is None:
length = _length_hint(iterable)
if iterable is None:
if length is None:
raise TypeError('iterable or length is required')
iterable = range_type(length)
self.iter = iter(iterable)
self.length = length
self.length_known = length is not None
self.pos = 0
self.avg = []
self.start = self.last_eta = time.time()
self.eta_known = False
self.finished = False
self.max_width = None
self.entered = False
self.current_item = None
self.is_hidden = not isatty(self.file)
self._last_line = None
def __enter__(self):
self.entered = True
self.render_progress()
return self
def __exit__(self, exc_type, exc_value, tb):
self.render_finish()
def __iter__(self):
if not self.entered:
raise RuntimeError('You need to use progress bars in a with block.')
self.render_progress()
return self
def render_finish(self):
if self.is_hidden:
return
self.file.write(AFTER_BAR)
self.file.flush()
@property
def pct(self):
if self.finished:
return 1.0
return min(self.pos / (float(self.length) or 1), 1.0)
@property
def time_per_iteration(self):
if not self.avg:
return 0.0
return sum(self.avg) / float(len(self.avg))
@property
def eta(self):
if self.length_known and not self.finished:
return self.time_per_iteration * (self.length - self.pos)
return 0.0
def format_eta(self):
if self.eta_known:
t = self.eta + 1
seconds = t % 60
t /= 60
minutes = t % 60
t /= 60
hours = t % 24
t /= 24
if t > 0:
days = t
return '%dd %02d:%02d:%02d' % (days, hours, minutes, seconds)
else:
return '%02d:%02d:%02d' % (hours, minutes, seconds)
return ''
def format_pos(self):
pos = str(self.pos)
if self.length_known:
pos += '/%s' % self.length
return pos
def format_pct(self):
return ('% 4d%%' % int(self.pct * 100))[1:]
def format_progress_line(self):
show_percent = self.show_percent
info_bits = []
if self.length_known:
bar_length = int(self.pct * self.width)
bar = self.fill_char * bar_length
bar += self.empty_char * (self.width - bar_length)
if show_percent is None:
show_percent = not self.show_pos
else:
if self.finished:
bar = self.fill_char * self.width
else:
bar = list(self.empty_char * (self.width or 1))
if self.time_per_iteration != 0:
bar[int((math.cos(self.pos * self.time_per_iteration)
/ 2.0 + 0.5) * self.width)] = self.fill_char
bar = ''.join(bar)
if self.show_pos:
info_bits.append(self.format_pos())
if show_percent:
info_bits.append(self.format_pct())
if self.show_eta and self.eta_known and not self.finished:
info_bits.append(self.format_eta())
if self.item_show_func is not None:
item_info = self.item_show_func(self.current_item)
if item_info is not None:
info_bits.append(item_info)
return (self.bar_template % {
'label': self.label,
'bar': bar,
'info': self.info_sep.join(info_bits)
}).rstrip()
def render_progress(self):
from .termui import get_terminal_size
nl = False
if self.is_hidden:
buf = [self.label]
nl = True
else:
buf = []
# Update width in case the terminal has been resized
if self.autowidth:
old_width = self.width
self.width = 0
clutter_length = term_len(self.format_progress_line())
new_width = max(0, get_terminal_size()[0] - clutter_length)
if new_width < old_width:
buf.append(BEFORE_BAR)
buf.append(' ' * self.max_width)
self.max_width = new_width
self.width = new_width
clear_width = self.width
if self.max_width is not None:
clear_width = self.max_width
buf.append(BEFORE_BAR)
line = self.format_progress_line()
line_len = term_len(line)
if self.max_width is None or self.max_width < line_len:
self.max_width = line_len
buf.append(line)
buf.append(' ' * (clear_width - line_len))
line = ''.join(buf)
# Render the line only if it changed.
if line != self._last_line:
self._last_line = line
echo(line, file=self.file, color=self.color, nl=nl)
self.file.flush()
def make_step(self, n_steps):
self.pos += n_steps
if self.length_known and self.pos >= self.length:
self.finished = True
if (time.time() - self.last_eta) < 1.0:
return
self.last_eta = time.time()
self.avg = self.avg[-6:] + [-(self.start - time.time()) / (self.pos)]
self.eta_known = self.length_known
def update(self, n_steps):
self.make_step(n_steps)
self.render_progress()
def finish(self):
self.eta_known = 0
self.current_item = None
self.finished = True
def next(self):
if self.is_hidden:
return next(self.iter)
try:
rv = next(self.iter)
self.current_item = rv
except StopIteration:
self.finish()
self.render_progress()
raise StopIteration()
else:
self.update(1)
return rv
if not PY2:
__next__ = next
del next
def pager(text, color=None):
"""Decide what method to use for paging through text."""
stdout = _default_text_stdout()
if not isatty(sys.stdin) or not isatty(stdout):
return _nullpager(stdout, text, color)
pager_cmd = (os.environ.get('PAGER', None) or '').strip()
if pager_cmd:
if WIN:
return _tempfilepager(text, pager_cmd, color)
return _pipepager(text, pager_cmd, color)
if os.environ.get('TERM') in ('dumb', 'emacs'):
return _nullpager(stdout, text, color)
if WIN or sys.platform.startswith('os2'):
return _tempfilepager(text, 'more <', color)
if hasattr(os, 'system') and os.system('(less) 2>/dev/null') == 0:
return _pipepager(text, 'less', color)
import tempfile
fd, filename = tempfile.mkstemp()
os.close(fd)
try:
if hasattr(os, 'system') and os.system('more "%s"' % filename) == 0:
return _pipepager(text, 'more', color)
return _nullpager(stdout, text, color)
finally:
os.unlink(filename)
def _pipepager(text, cmd, color):
"""Page through text by feeding it to another program. Invoking a
pager through this might support colors.
"""
import subprocess
env = dict(os.environ)
# If we're piping to less we might support colors under the
# condition that
cmd_detail = cmd.rsplit('/', 1)[-1].split()
if color is None and cmd_detail[0] == 'less':
less_flags = os.environ.get('LESS', '') + ' '.join(cmd_detail[1:])
if not less_flags:
env['LESS'] = '-R'
color = True
elif 'r' in less_flags or 'R' in less_flags:
color = True
if not color:
text = strip_ansi(text)
c = subprocess.Popen(cmd, shell=True, stdin=subprocess.PIPE,
env=env)
encoding = get_best_encoding(c.stdin)
try:
c.stdin.write(text.encode(encoding, 'replace'))
c.stdin.close()
except (IOError, KeyboardInterrupt):
pass
# Less doesn't respect ^C, but catches it for its own UI purposes (aborting
# search or other commands inside less).
#
# That means when the user hits ^C, the parent process (click) terminates,
# but less is still alive, paging the output and messing up the terminal.
#
# If the user wants to make the pager exit on ^C, they should set
# `LESS='-K'`. It's not our decision to make.
while True:
try:
c.wait()
except KeyboardInterrupt:
pass
else:
break
def _tempfilepager(text, cmd, color):
"""Page through text by invoking a program on a temporary file."""
import tempfile
filename = tempfile.mktemp()
if not color:
text = strip_ansi(text)
encoding = get_best_encoding(sys.stdout)
with open_stream(filename, 'wb')[0] as f:
f.write(text.encode(encoding))
try:
os.system(cmd + ' "' + filename + '"')
finally:
os.unlink(filename)
def _nullpager(stream, text, color):
"""Simply print unformatted text. This is the ultimate fallback."""
if not color:
text = strip_ansi(text)
stream.write(text)
class Editor(object):
def __init__(self, editor=None, env=None, require_save=True,
extension='.txt'):
self.editor = editor
self.env = env
self.require_save = require_save
self.extension = extension
def get_editor(self):
if self.editor is not None:
return self.editor
for key in 'VISUAL', 'EDITOR':
rv = os.environ.get(key)
if rv:
return rv
if WIN:
return 'notepad'
for editor in 'vim', 'nano':
if os.system('which %s >/dev/null 2>&1' % editor) == 0:
return editor
return 'vi'
def edit_file(self, filename):
import subprocess
editor = self.get_editor()
if self.env:
environ = os.environ.copy()
environ.update(self.env)
else:
environ = None
try:
c = subprocess.Popen('%s "%s"' % (editor, filename),
env=environ, shell=True)
exit_code = c.wait()
if exit_code != 0:
raise ClickException('%s: Editing failed!' % editor)
except OSError as e:
raise ClickException('%s: Editing failed: %s' % (editor, e))
def edit(self, text):
import tempfile
text = text or ''
if text and not text.endswith('\n'):
text += '\n'
fd, name = tempfile.mkstemp(prefix='editor-', suffix=self.extension)
try:
if WIN:
encoding = 'utf-8-sig'
text = text.replace('\n', '\r\n')
else:
encoding = 'utf-8'
text = text.encode(encoding)
f = os.fdopen(fd, 'wb')
f.write(text)
f.close()
timestamp = os.path.getmtime(name)
self.edit_file(name)
if self.require_save \
and os.path.getmtime(name) == timestamp:
return None
f = open(name, 'rb')
try:
rv = f.read()
finally:
f.close()
return rv.decode('utf-8-sig').replace('\r\n', '\n')
finally:
os.unlink(name)
def open_url(url, wait=False, locate=False):
import subprocess
def _unquote_file(url):
try:
import urllib
except ImportError:
import urllib
if url.startswith('file://'):
url = urllib.unquote(url[7:])
return url
if sys.platform == 'darwin':
args = ['open']
if wait:
args.append('-W')
if locate:
args.append('-R')
args.append(_unquote_file(url))
null = open('/dev/null', 'w')
try:
return subprocess.Popen(args, stderr=null).wait()
finally:
null.close()
elif WIN:
if locate:
url = _unquote_file(url)
args = 'explorer /select,"%s"' % _unquote_file(
url.replace('"', ''))
else:
args = 'start %s "" "%s"' % (
wait and '/WAIT' or '', url.replace('"', ''))
return os.system(args)
try:
if locate:
url = os.path.dirname(_unquote_file(url)) or '.'
else:
url = _unquote_file(url)
c = subprocess.Popen(['xdg-open', url])
if wait:
return c.wait()
return 0
except OSError:
if url.startswith(('http://', 'https://')) and not locate and not wait:
import webbrowser
webbrowser.open(url)
return 0
return 1
def _translate_ch_to_exc(ch):
if ch == '\x03':
raise KeyboardInterrupt()
if ch == '\x04':
raise EOFError()
if WIN:
import msvcrt
def getchar(echo):
rv = msvcrt.getch()
if echo:
msvcrt.putchar(rv)
_translate_ch_to_exc(rv)
if PY2:
enc = getattr(sys.stdin, 'encoding', None)
if enc is not None:
rv = rv.decode(enc, 'replace')
else:
rv = rv.decode('cp1252', 'replace')
return rv
else:
import tty
import termios
def getchar(echo):
if not isatty(sys.stdin):
f = open('/dev/tty')
fd = f.fileno()
else:
fd = sys.stdin.fileno()
f = None
try:
old_settings = termios.tcgetattr(fd)
try:
tty.setraw(fd)
ch = os.read(fd, 32)
if echo and isatty(sys.stdout):
sys.stdout.write(ch)
finally:
termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
sys.stdout.flush()
if f is not None:
f.close()
except termios.error:
pass
_translate_ch_to_exc(ch)
return ch.decode(get_best_encoding(sys.stdin), 'replace')
+38
View File
@@ -0,0 +1,38 @@
import textwrap
from contextlib import contextmanager
class TextWrapper(textwrap.TextWrapper):
def _handle_long_word(self, reversed_chunks, cur_line, cur_len, width):
space_left = max(width - cur_len, 1)
if self.break_long_words:
last = reversed_chunks[-1]
cut = last[:space_left]
res = last[space_left:]
cur_line.append(cut)
reversed_chunks[-1] = res
elif not cur_line:
cur_line.append(reversed_chunks.pop())
@contextmanager
def extra_indent(self, indent):
old_initial_indent = self.initial_indent
old_subsequent_indent = self.subsequent_indent
self.initial_indent += indent
self.subsequent_indent += indent
try:
yield
finally:
self.initial_indent = old_initial_indent
self.subsequent_indent = old_subsequent_indent
def indent_only(self, text):
rv = []
for idx, line in enumerate(text.splitlines()):
indent = self.initial_indent
if idx > 0:
indent = self.subsequent_indent
rv.append(indent + line)
return '\n'.join(rv)
+118
View File
@@ -0,0 +1,118 @@
import os
import sys
import codecs
from ._compat import PY2
# If someone wants to vendor click, we want to ensure the
# correct package is discovered. Ideally we could use a
# relative import here but unfortunately Python does not
# support that.
click = sys.modules[__name__.rsplit('.', 1)[0]]
def _find_unicode_literals_frame():
import __future__
frm = sys._getframe(1)
idx = 1
while frm is not None:
if frm.f_globals.get('__name__', '').startswith('click.'):
frm = frm.f_back
idx += 1
elif frm.f_code.co_flags & __future__.unicode_literals.compiler_flag:
return idx
else:
break
return 0
def _check_for_unicode_literals():
if not __debug__:
return
if not PY2 or click.disable_unicode_literals_warning:
return
bad_frame = _find_unicode_literals_frame()
if bad_frame <= 0:
return
from warnings import warn
warn(Warning('Click detected the use of the unicode_literals '
'__future__ import. This is heavily discouraged '
'because it can introduce subtle bugs in your '
'code. You should instead use explicit u"" literals '
'for your unicode strings. For more information see '
'http://click.pocoo.org/python3/'),
stacklevel=bad_frame)
def _verify_python3_env():
"""Ensures that the environment is good for unicode on Python 3."""
if PY2:
return
try:
import locale
fs_enc = codecs.lookup(locale.getpreferredencoding()).name
except Exception:
fs_enc = 'ascii'
if fs_enc != 'ascii':
return
extra = ''
if os.name == 'posix':
import subprocess
rv = subprocess.Popen(['locale', '-a'], stdout=subprocess.PIPE,
stderr=subprocess.PIPE).communicate()[0]
good_locales = set()
has_c_utf8 = False
# Make sure we're operating on text here.
if isinstance(rv, bytes):
rv = rv.decode('ascii', 'replace')
for line in rv.splitlines():
locale = line.strip()
if locale.lower().endswith(('.utf-8', '.utf8')):
good_locales.add(locale)
if locale.lower() in ('c.utf8', 'c.utf-8'):
has_c_utf8 = True
extra += '\n\n'
if not good_locales:
extra += (
'Additional information: on this system no suitable UTF-8\n'
'locales were discovered. This most likely requires resolving\n'
'by reconfiguring the locale system.'
)
elif has_c_utf8:
extra += (
'This system supports the C.UTF-8 locale which is recommended.\n'
'You might be able to resolve your issue by exporting the\n'
'following environment variables:\n\n'
' export LC_ALL=C.UTF-8\n'
' export LANG=C.UTF-8'
)
else:
extra += (
'This system lists a couple of UTF-8 supporting locales that\n'
'you can pick from. The following suitable locales where\n'
'discovered: %s'
) % ', '.join(sorted(good_locales))
bad_locale = None
for locale in os.environ.get('LC_ALL'), os.environ.get('LANG'):
if locale and locale.lower().endswith(('.utf-8', '.utf8')):
bad_locale = locale
if locale is not None:
break
if bad_locale is not None:
extra += (
'\n\nClick discovered that you exported a UTF-8 locale\n'
'but the locale system could not pick up from it because\n'
'it does not exist. The exported locale is "%s" but it\n'
'is not supported'
) % bad_locale
raise RuntimeError('Click will abort further execution because Python 3 '
'was configured to use ASCII as encoding for the '
'environment. Consult http://click.pocoo.org/python3/'
'for mitigation steps.' + extra)
+273
View File
@@ -0,0 +1,273 @@
# -*- coding: utf-8 -*-
# This module is based on the excellent work by Adam Bartoš who
# provided a lot of what went into the implementation here in
# the discussion to issue1602 in the Python bug tracker.
#
# There are some general differences in regards to how this works
# compared to the original patches as we do not need to patch
# the entire interpreter but just work in our little world of
# echo and prmopt.
import io
import os
import sys
import zlib
import time
import ctypes
import msvcrt
from click._compat import _NonClosingTextIOWrapper, text_type, PY2
from ctypes import byref, POINTER, c_int, c_char, c_char_p, \
c_void_p, py_object, c_ssize_t, c_ulong, windll, WINFUNCTYPE
try:
from ctypes import pythonapi
PyObject_GetBuffer = pythonapi.PyObject_GetBuffer
PyBuffer_Release = pythonapi.PyBuffer_Release
except ImportError:
pythonapi = None
from ctypes.wintypes import LPWSTR, LPCWSTR
c_ssize_p = POINTER(c_ssize_t)
kernel32 = windll.kernel32
GetStdHandle = kernel32.GetStdHandle
ReadConsoleW = kernel32.ReadConsoleW
WriteConsoleW = kernel32.WriteConsoleW
GetLastError = kernel32.GetLastError
GetCommandLineW = WINFUNCTYPE(LPWSTR)(
('GetCommandLineW', windll.kernel32))
CommandLineToArgvW = WINFUNCTYPE(
POINTER(LPWSTR), LPCWSTR, POINTER(c_int))(
('CommandLineToArgvW', windll.shell32))
STDIN_HANDLE = GetStdHandle(-10)
STDOUT_HANDLE = GetStdHandle(-11)
STDERR_HANDLE = GetStdHandle(-12)
PyBUF_SIMPLE = 0
PyBUF_WRITABLE = 1
ERROR_SUCCESS = 0
ERROR_NOT_ENOUGH_MEMORY = 8
ERROR_OPERATION_ABORTED = 995
STDIN_FILENO = 0
STDOUT_FILENO = 1
STDERR_FILENO = 2
EOF = b'\x1a'
MAX_BYTES_WRITTEN = 32767
class Py_buffer(ctypes.Structure):
_fields_ = [
('buf', c_void_p),
('obj', py_object),
('len', c_ssize_t),
('itemsize', c_ssize_t),
('readonly', c_int),
('ndim', c_int),
('format', c_char_p),
('shape', c_ssize_p),
('strides', c_ssize_p),
('suboffsets', c_ssize_p),
('internal', c_void_p)
]
if PY2:
_fields_.insert(-1, ('smalltable', c_ssize_t * 2))
# On PyPy we cannot get buffers so our ability to operate here is
# serverly limited.
if pythonapi is None:
get_buffer = None
else:
def get_buffer(obj, writable=False):
buf = Py_buffer()
flags = PyBUF_WRITABLE if writable else PyBUF_SIMPLE
PyObject_GetBuffer(py_object(obj), byref(buf), flags)
try:
buffer_type = c_char * buf.len
return buffer_type.from_address(buf.buf)
finally:
PyBuffer_Release(byref(buf))
class _WindowsConsoleRawIOBase(io.RawIOBase):
def __init__(self, handle):
self.handle = handle
def isatty(self):
io.RawIOBase.isatty(self)
return True
class _WindowsConsoleReader(_WindowsConsoleRawIOBase):
def readable(self):
return True
def readinto(self, b):
bytes_to_be_read = len(b)
if not bytes_to_be_read:
return 0
elif bytes_to_be_read % 2:
raise ValueError('cannot read odd number of bytes from '
'UTF-16-LE encoded console')
buffer = get_buffer(b, writable=True)
code_units_to_be_read = bytes_to_be_read // 2
code_units_read = c_ulong()
rv = ReadConsoleW(self.handle, buffer, code_units_to_be_read,
byref(code_units_read), None)
if GetLastError() == ERROR_OPERATION_ABORTED:
# wait for KeyboardInterrupt
time.sleep(0.1)
if not rv:
raise OSError('Windows error: %s' % GetLastError())
if buffer[0] == EOF:
return 0
return 2 * code_units_read.value
class _WindowsConsoleWriter(_WindowsConsoleRawIOBase):
def writable(self):
return True
@staticmethod
def _get_error_message(errno):
if errno == ERROR_SUCCESS:
return 'ERROR_SUCCESS'
elif errno == ERROR_NOT_ENOUGH_MEMORY:
return 'ERROR_NOT_ENOUGH_MEMORY'
return 'Windows error %s' % errno
def write(self, b):
bytes_to_be_written = len(b)
buf = get_buffer(b)
code_units_to_be_written = min(bytes_to_be_written,
MAX_BYTES_WRITTEN) // 2
code_units_written = c_ulong()
WriteConsoleW(self.handle, buf, code_units_to_be_written,
byref(code_units_written), None)
bytes_written = 2 * code_units_written.value
if bytes_written == 0 and bytes_to_be_written > 0:
raise OSError(self._get_error_message(GetLastError()))
return bytes_written
class ConsoleStream(object):
def __init__(self, text_stream, byte_stream):
self._text_stream = text_stream
self.buffer = byte_stream
@property
def name(self):
return self.buffer.name
def write(self, x):
if isinstance(x, text_type):
return self._text_stream.write(x)
try:
self.flush()
except Exception:
pass
return self.buffer.write(x)
def writelines(self, lines):
for line in lines:
self.write(line)
def __getattr__(self, name):
return getattr(self._text_stream, name)
def isatty(self):
return self.buffer.isatty()
def __repr__(self):
return '<ConsoleStream name=%r encoding=%r>' % (
self.name,
self.encoding,
)
def _get_text_stdin(buffer_stream):
text_stream = _NonClosingTextIOWrapper(
io.BufferedReader(_WindowsConsoleReader(STDIN_HANDLE)),
'utf-16-le', 'strict', line_buffering=True)
return ConsoleStream(text_stream, buffer_stream)
def _get_text_stdout(buffer_stream):
text_stream = _NonClosingTextIOWrapper(
_WindowsConsoleWriter(STDOUT_HANDLE),
'utf-16-le', 'strict', line_buffering=True)
return ConsoleStream(text_stream, buffer_stream)
def _get_text_stderr(buffer_stream):
text_stream = _NonClosingTextIOWrapper(
_WindowsConsoleWriter(STDERR_HANDLE),
'utf-16-le', 'strict', line_buffering=True)
return ConsoleStream(text_stream, buffer_stream)
if PY2:
def _hash_py_argv():
return zlib.crc32('\x00'.join(sys.argv[1:]))
_initial_argv_hash = _hash_py_argv()
def _get_windows_argv():
argc = c_int(0)
argv_unicode = CommandLineToArgvW(GetCommandLineW(), byref(argc))
argv = [argv_unicode[i] for i in range(0, argc.value)]
if not hasattr(sys, 'frozen'):
argv = argv[1:]
while len(argv) > 0:
arg = argv[0]
if not arg.startswith('-') or arg == '-':
break
argv = argv[1:]
if arg.startswith(('-c', '-m')):
break
return argv[1:]
_stream_factories = {
0: _get_text_stdin,
1: _get_text_stdout,
2: _get_text_stderr,
}
def _get_windows_console_stream(f, encoding, errors):
if get_buffer is not None and \
encoding in ('utf-16-le', None) \
and errors in ('strict', None) and \
hasattr(f, 'isatty') and f.isatty():
func = _stream_factories.get(f.fileno())
if func is not None:
if not PY2:
f = getattr(f, 'buffer')
if f is None:
return None
else:
# If we are on Python 2 we need to set the stream that we
# deal with to binary mode as otherwise the exercise if a
# bit moot. The same problems apply as for
# get_binary_stdin and friends from _compat.
msvcrt.setmode(f.fileno(), os.O_BINARY)
return func(f)
+1744
View File
File diff suppressed because it is too large Load Diff
+304
View File
@@ -0,0 +1,304 @@
import sys
import inspect
from functools import update_wrapper
from ._compat import iteritems
from ._unicodefun import _check_for_unicode_literals
from .utils import echo
from .globals import get_current_context
def pass_context(f):
"""Marks a callback as wanting to receive the current context
object as first argument.
"""
def new_func(*args, **kwargs):
return f(get_current_context(), *args, **kwargs)
return update_wrapper(new_func, f)
def pass_obj(f):
"""Similar to :func:`pass_context`, but only pass the object on the
context onwards (:attr:`Context.obj`). This is useful if that object
represents the state of a nested system.
"""
def new_func(*args, **kwargs):
return f(get_current_context().obj, *args, **kwargs)
return update_wrapper(new_func, f)
def make_pass_decorator(object_type, ensure=False):
"""Given an object type this creates a decorator that will work
similar to :func:`pass_obj` but instead of passing the object of the
current context, it will find the innermost context of type
:func:`object_type`.
This generates a decorator that works roughly like this::
from functools import update_wrapper
def decorator(f):
@pass_context
def new_func(ctx, *args, **kwargs):
obj = ctx.find_object(object_type)
return ctx.invoke(f, obj, *args, **kwargs)
return update_wrapper(new_func, f)
return decorator
:param object_type: the type of the object to pass.
:param ensure: if set to `True`, a new object will be created and
remembered on the context if it's not there yet.
"""
def decorator(f):
def new_func(*args, **kwargs):
ctx = get_current_context()
if ensure:
obj = ctx.ensure_object(object_type)
else:
obj = ctx.find_object(object_type)
if obj is None:
raise RuntimeError('Managed to invoke callback without a '
'context object of type %r existing'
% object_type.__name__)
return ctx.invoke(f, obj, *args[1:], **kwargs)
return update_wrapper(new_func, f)
return decorator
def _make_command(f, name, attrs, cls):
if isinstance(f, Command):
raise TypeError('Attempted to convert a callback into a '
'command twice.')
try:
params = f.__click_params__
params.reverse()
del f.__click_params__
except AttributeError:
params = []
help = attrs.get('help')
if help is None:
help = inspect.getdoc(f)
if isinstance(help, bytes):
help = help.decode('utf-8')
else:
help = inspect.cleandoc(help)
attrs['help'] = help
_check_for_unicode_literals()
return cls(name=name or f.__name__.lower(),
callback=f, params=params, **attrs)
def command(name=None, cls=None, **attrs):
"""Creates a new :class:`Command` and uses the decorated function as
callback. This will also automatically attach all decorated
:func:`option`\s and :func:`argument`\s as parameters to the command.
The name of the command defaults to the name of the function. If you
want to change that, you can pass the intended name as the first
argument.
All keyword arguments are forwarded to the underlying command class.
Once decorated the function turns into a :class:`Command` instance
that can be invoked as a command line utility or be attached to a
command :class:`Group`.
:param name: the name of the command. This defaults to the function
name.
:param cls: the command class to instantiate. This defaults to
:class:`Command`.
"""
if cls is None:
cls = Command
def decorator(f):
cmd = _make_command(f, name, attrs, cls)
cmd.__doc__ = f.__doc__
return cmd
return decorator
def group(name=None, **attrs):
"""Creates a new :class:`Group` with a function as callback. This
works otherwise the same as :func:`command` just that the `cls`
parameter is set to :class:`Group`.
"""
attrs.setdefault('cls', Group)
return command(name, **attrs)
def _param_memo(f, param):
if isinstance(f, Command):
f.params.append(param)
else:
if not hasattr(f, '__click_params__'):
f.__click_params__ = []
f.__click_params__.append(param)
def argument(*param_decls, **attrs):
"""Attaches an argument to the command. All positional arguments are
passed as parameter declarations to :class:`Argument`; all keyword
arguments are forwarded unchanged (except ``cls``).
This is equivalent to creating an :class:`Argument` instance manually
and attaching it to the :attr:`Command.params` list.
:param cls: the argument class to instantiate. This defaults to
:class:`Argument`.
"""
def decorator(f):
ArgumentClass = attrs.pop('cls', Argument)
_param_memo(f, ArgumentClass(param_decls, **attrs))
return f
return decorator
def option(*param_decls, **attrs):
"""Attaches an option to the command. All positional arguments are
passed as parameter declarations to :class:`Option`; all keyword
arguments are forwarded unchanged (except ``cls``).
This is equivalent to creating an :class:`Option` instance manually
and attaching it to the :attr:`Command.params` list.
:param cls: the option class to instantiate. This defaults to
:class:`Option`.
"""
def decorator(f):
if 'help' in attrs:
attrs['help'] = inspect.cleandoc(attrs['help'])
OptionClass = attrs.pop('cls', Option)
_param_memo(f, OptionClass(param_decls, **attrs))
return f
return decorator
def confirmation_option(*param_decls, **attrs):
"""Shortcut for confirmation prompts that can be ignored by passing
``--yes`` as parameter.
This is equivalent to decorating a function with :func:`option` with
the following parameters::
def callback(ctx, param, value):
if not value:
ctx.abort()
@click.command()
@click.option('--yes', is_flag=True, callback=callback,
expose_value=False, prompt='Do you want to continue?')
def dropdb():
pass
"""
def decorator(f):
def callback(ctx, param, value):
if not value:
ctx.abort()
attrs.setdefault('is_flag', True)
attrs.setdefault('callback', callback)
attrs.setdefault('expose_value', False)
attrs.setdefault('prompt', 'Do you want to continue?')
attrs.setdefault('help', 'Confirm the action without prompting.')
return option(*(param_decls or ('--yes',)), **attrs)(f)
return decorator
def password_option(*param_decls, **attrs):
"""Shortcut for password prompts.
This is equivalent to decorating a function with :func:`option` with
the following parameters::
@click.command()
@click.option('--password', prompt=True, confirmation_prompt=True,
hide_input=True)
def changeadmin(password):
pass
"""
def decorator(f):
attrs.setdefault('prompt', True)
attrs.setdefault('confirmation_prompt', True)
attrs.setdefault('hide_input', True)
return option(*(param_decls or ('--password',)), **attrs)(f)
return decorator
def version_option(version=None, *param_decls, **attrs):
"""Adds a ``--version`` option which immediately ends the program
printing out the version number. This is implemented as an eager
option that prints the version and exits the program in the callback.
:param version: the version number to show. If not provided Click
attempts an auto discovery via setuptools.
:param prog_name: the name of the program (defaults to autodetection)
:param message: custom message to show instead of the default
(``'%(prog)s, version %(version)s'``)
:param others: everything else is forwarded to :func:`option`.
"""
if version is None:
module = sys._getframe(1).f_globals.get('__name__')
def decorator(f):
prog_name = attrs.pop('prog_name', None)
message = attrs.pop('message', '%(prog)s, version %(version)s')
def callback(ctx, param, value):
if not value or ctx.resilient_parsing:
return
prog = prog_name
if prog is None:
prog = ctx.find_root().info_name
ver = version
if ver is None:
try:
import pkg_resources
except ImportError:
pass
else:
for dist in pkg_resources.working_set:
scripts = dist.get_entry_map().get('console_scripts') or {}
for script_name, entry_point in iteritems(scripts):
if entry_point.module_name == module:
ver = dist.version
break
if ver is None:
raise RuntimeError('Could not determine version')
echo(message % {
'prog': prog,
'version': ver,
}, color=ctx.color)
ctx.exit()
attrs.setdefault('is_flag', True)
attrs.setdefault('expose_value', False)
attrs.setdefault('is_eager', True)
attrs.setdefault('help', 'Show the version and exit.')
attrs['callback'] = callback
return option(*(param_decls or ('--version',)), **attrs)(f)
return decorator
def help_option(*param_decls, **attrs):
"""Adds a ``--help`` option which immediately ends the program
printing out the help page. This is usually unnecessary to add as
this is added by default to all commands unless suppressed.
Like :func:`version_option`, this is implemented as eager option that
prints in the callback and exits.
All arguments are forwarded to :func:`option`.
"""
def decorator(f):
def callback(ctx, param, value):
if value and not ctx.resilient_parsing:
echo(ctx.get_help(), color=ctx.color)
ctx.exit()
attrs.setdefault('is_flag', True)
attrs.setdefault('expose_value', False)
attrs.setdefault('help', 'Show this message and exit.')
attrs.setdefault('is_eager', True)
attrs['callback'] = callback
return option(*(param_decls or ('--help',)), **attrs)(f)
return decorator
# Circular dependencies between core and decorators
from .core import Command, Group, Argument, Option
+201
View File
@@ -0,0 +1,201 @@
from ._compat import PY2, filename_to_ui, get_text_stderr
from .utils import echo
class ClickException(Exception):
"""An exception that Click can handle and show to the user."""
#: The exit code for this exception
exit_code = 1
def __init__(self, message):
if PY2:
if message is not None:
message = message.encode('utf-8')
Exception.__init__(self, message)
self.message = message
def format_message(self):
return self.message
def show(self, file=None):
if file is None:
file = get_text_stderr()
echo('Error: %s' % self.format_message(), file=file)
class UsageError(ClickException):
"""An internal exception that signals a usage error. This typically
aborts any further handling.
:param message: the error message to display.
:param ctx: optionally the context that caused this error. Click will
fill in the context automatically in some situations.
"""
exit_code = 2
def __init__(self, message, ctx=None):
ClickException.__init__(self, message)
self.ctx = ctx
def show(self, file=None):
if file is None:
file = get_text_stderr()
color = None
if self.ctx is not None:
color = self.ctx.color
echo(self.ctx.get_usage() + '\n', file=file, color=color)
echo('Error: %s' % self.format_message(), file=file, color=color)
class BadParameter(UsageError):
"""An exception that formats out a standardized error message for a
bad parameter. This is useful when thrown from a callback or type as
Click will attach contextual information to it (for instance, which
parameter it is).
.. versionadded:: 2.0
:param param: the parameter object that caused this error. This can
be left out, and Click will attach this info itself
if possible.
:param param_hint: a string that shows up as parameter name. This
can be used as alternative to `param` in cases
where custom validation should happen. If it is
a string it's used as such, if it's a list then
each item is quoted and separated.
"""
def __init__(self, message, ctx=None, param=None,
param_hint=None):
UsageError.__init__(self, message, ctx)
self.param = param
self.param_hint = param_hint
def format_message(self):
if self.param_hint is not None:
param_hint = self.param_hint
elif self.param is not None:
param_hint = self.param.opts or [self.param.human_readable_name]
else:
return 'Invalid value: %s' % self.message
if isinstance(param_hint, (tuple, list)):
param_hint = ' / '.join('"%s"' % x for x in param_hint)
return 'Invalid value for %s: %s' % (param_hint, self.message)
class MissingParameter(BadParameter):
"""Raised if click required an option or argument but it was not
provided when invoking the script.
.. versionadded:: 4.0
:param param_type: a string that indicates the type of the parameter.
The default is to inherit the parameter type from
the given `param`. Valid values are ``'parameter'``,
``'option'`` or ``'argument'``.
"""
def __init__(self, message=None, ctx=None, param=None,
param_hint=None, param_type=None):
BadParameter.__init__(self, message, ctx, param, param_hint)
self.param_type = param_type
def format_message(self):
if self.param_hint is not None:
param_hint = self.param_hint
elif self.param is not None:
param_hint = self.param.opts or [self.param.human_readable_name]
else:
param_hint = None
if isinstance(param_hint, (tuple, list)):
param_hint = ' / '.join('"%s"' % x for x in param_hint)
param_type = self.param_type
if param_type is None and self.param is not None:
param_type = self.param.param_type_name
msg = self.message
if self.param is not None:
msg_extra = self.param.type.get_missing_message(self.param)
if msg_extra:
if msg:
msg += '. ' + msg_extra
else:
msg = msg_extra
return 'Missing %s%s%s%s' % (
param_type,
param_hint and ' %s' % param_hint or '',
msg and '. ' or '.',
msg or '',
)
class NoSuchOption(UsageError):
"""Raised if click attempted to handle an option that does not
exist.
.. versionadded:: 4.0
"""
def __init__(self, option_name, message=None, possibilities=None,
ctx=None):
if message is None:
message = 'no such option: %s' % option_name
UsageError.__init__(self, message, ctx)
self.option_name = option_name
self.possibilities = possibilities
def format_message(self):
bits = [self.message]
if self.possibilities:
if len(self.possibilities) == 1:
bits.append('Did you mean %s?' % self.possibilities[0])
else:
possibilities = sorted(self.possibilities)
bits.append('(Possible options: %s)' % ', '.join(possibilities))
return ' '.join(bits)
class BadOptionUsage(UsageError):
"""Raised if an option is generally supplied but the use of the option
was incorrect. This is for instance raised if the number of arguments
for an option is not correct.
.. versionadded:: 4.0
"""
def __init__(self, message, ctx=None):
UsageError.__init__(self, message, ctx)
class BadArgumentUsage(UsageError):
"""Raised if an argument is generally supplied but the use of the argument
was incorrect. This is for instance raised if the number of values
for an argument is not correct.
.. versionadded:: 6.0
"""
def __init__(self, message, ctx=None):
UsageError.__init__(self, message, ctx)
class FileError(ClickException):
"""Raised if a file cannot be opened."""
def __init__(self, filename, hint=None):
ui_filename = filename_to_ui(filename)
if hint is None:
hint = 'unknown error'
ClickException.__init__(self, hint)
self.ui_filename = ui_filename
self.filename = filename
def format_message(self):
return 'Could not open file %s: %s' % (self.ui_filename, self.message)
class Abort(RuntimeError):
"""An internal signalling exception that signals Click to abort."""
+256
View File
@@ -0,0 +1,256 @@
from contextlib import contextmanager
from .termui import get_terminal_size
from .parser import split_opt
from ._compat import term_len
# Can force a width. This is used by the test system
FORCED_WIDTH = None
def measure_table(rows):
widths = {}
for row in rows:
for idx, col in enumerate(row):
widths[idx] = max(widths.get(idx, 0), term_len(col))
return tuple(y for x, y in sorted(widths.items()))
def iter_rows(rows, col_count):
for row in rows:
row = tuple(row)
yield row + ('',) * (col_count - len(row))
def wrap_text(text, width=78, initial_indent='', subsequent_indent='',
preserve_paragraphs=False):
"""A helper function that intelligently wraps text. By default, it
assumes that it operates on a single paragraph of text but if the
`preserve_paragraphs` parameter is provided it will intelligently
handle paragraphs (defined by two empty lines).
If paragraphs are handled, a paragraph can be prefixed with an empty
line containing the ``\\b`` character (``\\x08``) to indicate that
no rewrapping should happen in that block.
:param text: the text that should be rewrapped.
:param width: the maximum width for the text.
:param initial_indent: the initial indent that should be placed on the
first line as a string.
:param subsequent_indent: the indent string that should be placed on
each consecutive line.
:param preserve_paragraphs: if this flag is set then the wrapping will
intelligently handle paragraphs.
"""
from ._textwrap import TextWrapper
text = text.expandtabs()
wrapper = TextWrapper(width, initial_indent=initial_indent,
subsequent_indent=subsequent_indent,
replace_whitespace=False)
if not preserve_paragraphs:
return wrapper.fill(text)
p = []
buf = []
indent = None
def _flush_par():
if not buf:
return
if buf[0].strip() == '\b':
p.append((indent or 0, True, '\n'.join(buf[1:])))
else:
p.append((indent or 0, False, ' '.join(buf)))
del buf[:]
for line in text.splitlines():
if not line:
_flush_par()
indent = None
else:
if indent is None:
orig_len = term_len(line)
line = line.lstrip()
indent = orig_len - term_len(line)
buf.append(line)
_flush_par()
rv = []
for indent, raw, text in p:
with wrapper.extra_indent(' ' * indent):
if raw:
rv.append(wrapper.indent_only(text))
else:
rv.append(wrapper.fill(text))
return '\n\n'.join(rv)
class HelpFormatter(object):
"""This class helps with formatting text-based help pages. It's
usually just needed for very special internal cases, but it's also
exposed so that developers can write their own fancy outputs.
At present, it always writes into memory.
:param indent_increment: the additional increment for each level.
:param width: the width for the text. This defaults to the terminal
width clamped to a maximum of 78.
"""
def __init__(self, indent_increment=2, width=None, max_width=None):
self.indent_increment = indent_increment
if max_width is None:
max_width = 80
if width is None:
width = FORCED_WIDTH
if width is None:
width = max(min(get_terminal_size()[0], max_width) - 2, 50)
self.width = width
self.current_indent = 0
self.buffer = []
def write(self, string):
"""Writes a unicode string into the internal buffer."""
self.buffer.append(string)
def indent(self):
"""Increases the indentation."""
self.current_indent += self.indent_increment
def dedent(self):
"""Decreases the indentation."""
self.current_indent -= self.indent_increment
def write_usage(self, prog, args='', prefix='Usage: '):
"""Writes a usage line into the buffer.
:param prog: the program name.
:param args: whitespace separated list of arguments.
:param prefix: the prefix for the first line.
"""
usage_prefix = '%*s%s ' % (self.current_indent, prefix, prog)
text_width = self.width - self.current_indent
if text_width >= (term_len(usage_prefix) + 20):
# The arguments will fit to the right of the prefix.
indent = ' ' * term_len(usage_prefix)
self.write(wrap_text(args, text_width,
initial_indent=usage_prefix,
subsequent_indent=indent))
else:
# The prefix is too long, put the arguments on the next line.
self.write(usage_prefix)
self.write('\n')
indent = ' ' * (max(self.current_indent, term_len(prefix)) + 4)
self.write(wrap_text(args, text_width,
initial_indent=indent,
subsequent_indent=indent))
self.write('\n')
def write_heading(self, heading):
"""Writes a heading into the buffer."""
self.write('%*s%s:\n' % (self.current_indent, '', heading))
def write_paragraph(self):
"""Writes a paragraph into the buffer."""
if self.buffer:
self.write('\n')
def write_text(self, text):
"""Writes re-indented text into the buffer. This rewraps and
preserves paragraphs.
"""
text_width = max(self.width - self.current_indent, 11)
indent = ' ' * self.current_indent
self.write(wrap_text(text, text_width,
initial_indent=indent,
subsequent_indent=indent,
preserve_paragraphs=True))
self.write('\n')
def write_dl(self, rows, col_max=30, col_spacing=2):
"""Writes a definition list into the buffer. This is how options
and commands are usually formatted.
:param rows: a list of two item tuples for the terms and values.
:param col_max: the maximum width of the first column.
:param col_spacing: the number of spaces between the first and
second column.
"""
rows = list(rows)
widths = measure_table(rows)
if len(widths) != 2:
raise TypeError('Expected two columns for definition list')
first_col = min(widths[0], col_max) + col_spacing
for first, second in iter_rows(rows, len(widths)):
self.write('%*s%s' % (self.current_indent, '', first))
if not second:
self.write('\n')
continue
if term_len(first) <= first_col - col_spacing:
self.write(' ' * (first_col - term_len(first)))
else:
self.write('\n')
self.write(' ' * (first_col + self.current_indent))
text_width = max(self.width - first_col - 2, 10)
lines = iter(wrap_text(second, text_width).splitlines())
if lines:
self.write(next(lines) + '\n')
for line in lines:
self.write('%*s%s\n' % (
first_col + self.current_indent, '', line))
else:
self.write('\n')
@contextmanager
def section(self, name):
"""Helpful context manager that writes a paragraph, a heading,
and the indents.
:param name: the section name that is written as heading.
"""
self.write_paragraph()
self.write_heading(name)
self.indent()
try:
yield
finally:
self.dedent()
@contextmanager
def indentation(self):
"""A context manager that increases the indentation."""
self.indent()
try:
yield
finally:
self.dedent()
def getvalue(self):
"""Returns the buffer contents."""
return ''.join(self.buffer)
def join_options(options):
"""Given a list of option strings this joins them in the most appropriate
way and returns them in the form ``(formatted_string,
any_prefix_is_slash)`` where the second item in the tuple is a flag that
indicates if any of the option prefixes was a slash.
"""
rv = []
any_prefix_is_slash = False
for opt in options:
prefix = split_opt(opt)[0]
if prefix == '/':
any_prefix_is_slash = True
rv.append((len(prefix), opt))
rv.sort(key=lambda x: x[0])
rv = ', '.join(x[1] for x in rv)
return rv, any_prefix_is_slash
+48
View File
@@ -0,0 +1,48 @@
from threading import local
_local = local()
def get_current_context(silent=False):
"""Returns the current click context. This can be used as a way to
access the current context object from anywhere. This is a more implicit
alternative to the :func:`pass_context` decorator. This function is
primarily useful for helpers such as :func:`echo` which might be
interested in changing it's behavior based on the current context.
To push the current context, :meth:`Context.scope` can be used.
.. versionadded:: 5.0
:param silent: is set to `True` the return value is `None` if no context
is available. The default behavior is to raise a
:exc:`RuntimeError`.
"""
try:
return getattr(_local, 'stack')[-1]
except (AttributeError, IndexError):
if not silent:
raise RuntimeError('There is no active click context.')
def push_context(ctx):
"""Pushes a new context to the current stack."""
_local.__dict__.setdefault('stack', []).append(ctx)
def pop_context():
"""Removes the top level from the stack."""
_local.stack.pop()
def resolve_color_default(color=None):
""""Internal helper to get the default value of the color flag. If a
value is passed it's returned unchanged, otherwise it's looked up from
the current context.
"""
if color is not None:
return color
ctx = get_current_context(silent=True)
if ctx is not None:
return ctx.color
+426
View File
@@ -0,0 +1,426 @@
# -*- coding: utf-8 -*-
"""
click.parser
~~~~~~~~~~~~
This module started out as largely a copy paste from the stdlib's
optparse module with the features removed that we do not need from
optparse because we implement them in Click on a higher level (for
instance type handling, help formatting and a lot more).
The plan is to remove more and more from here over time.
The reason this is a different module and not optparse from the stdlib
is that there are differences in 2.x and 3.x about the error messages
generated and optparse in the stdlib uses gettext for no good reason
and might cause us issues.
"""
import re
from collections import deque
from .exceptions import UsageError, NoSuchOption, BadOptionUsage, \
BadArgumentUsage
def _unpack_args(args, nargs_spec):
"""Given an iterable of arguments and an iterable of nargs specifications,
it returns a tuple with all the unpacked arguments at the first index
and all remaining arguments as the second.
The nargs specification is the number of arguments that should be consumed
or `-1` to indicate that this position should eat up all the remainders.
Missing items are filled with `None`.
"""
args = deque(args)
nargs_spec = deque(nargs_spec)
rv = []
spos = None
def _fetch(c):
try:
if spos is None:
return c.popleft()
else:
return c.pop()
except IndexError:
return None
while nargs_spec:
nargs = _fetch(nargs_spec)
if nargs == 1:
rv.append(_fetch(args))
elif nargs > 1:
x = [_fetch(args) for _ in range(nargs)]
# If we're reversed, we're pulling in the arguments in reverse,
# so we need to turn them around.
if spos is not None:
x.reverse()
rv.append(tuple(x))
elif nargs < 0:
if spos is not None:
raise TypeError('Cannot have two nargs < 0')
spos = len(rv)
rv.append(None)
# spos is the position of the wildcard (star). If it's not `None`,
# we fill it with the remainder.
if spos is not None:
rv[spos] = tuple(args)
args = []
rv[spos + 1:] = reversed(rv[spos + 1:])
return tuple(rv), list(args)
def _error_opt_args(nargs, opt):
if nargs == 1:
raise BadOptionUsage('%s option requires an argument' % opt)
raise BadOptionUsage('%s option requires %d arguments' % (opt, nargs))
def split_opt(opt):
first = opt[:1]
if first.isalnum():
return '', opt
if opt[1:2] == first:
return opt[:2], opt[2:]
return first, opt[1:]
def normalize_opt(opt, ctx):
if ctx is None or ctx.token_normalize_func is None:
return opt
prefix, opt = split_opt(opt)
return prefix + ctx.token_normalize_func(opt)
def split_arg_string(string):
"""Given an argument string this attempts to split it into small parts."""
rv = []
for match in re.finditer(r"('([^'\\]*(?:\\.[^'\\]*)*)'"
r'|"([^"\\]*(?:\\.[^"\\]*)*)"'
r'|\S+)\s*', string, re.S):
arg = match.group().strip()
if arg[:1] == arg[-1:] and arg[:1] in '"\'':
arg = arg[1:-1].encode('ascii', 'backslashreplace') \
.decode('unicode-escape')
try:
arg = type(string)(arg)
except UnicodeError:
pass
rv.append(arg)
return rv
class Option(object):
def __init__(self, opts, dest, action=None, nargs=1, const=None, obj=None):
self._short_opts = []
self._long_opts = []
self.prefixes = set()
for opt in opts:
prefix, value = split_opt(opt)
if not prefix:
raise ValueError('Invalid start character for option (%s)'
% opt)
self.prefixes.add(prefix[0])
if len(prefix) == 1 and len(value) == 1:
self._short_opts.append(opt)
else:
self._long_opts.append(opt)
self.prefixes.add(prefix)
if action is None:
action = 'store'
self.dest = dest
self.action = action
self.nargs = nargs
self.const = const
self.obj = obj
@property
def takes_value(self):
return self.action in ('store', 'append')
def process(self, value, state):
if self.action == 'store':
state.opts[self.dest] = value
elif self.action == 'store_const':
state.opts[self.dest] = self.const
elif self.action == 'append':
state.opts.setdefault(self.dest, []).append(value)
elif self.action == 'append_const':
state.opts.setdefault(self.dest, []).append(self.const)
elif self.action == 'count':
state.opts[self.dest] = state.opts.get(self.dest, 0) + 1
else:
raise ValueError('unknown action %r' % self.action)
state.order.append(self.obj)
class Argument(object):
def __init__(self, dest, nargs=1, obj=None):
self.dest = dest
self.nargs = nargs
self.obj = obj
def process(self, value, state):
if self.nargs > 1:
holes = sum(1 for x in value if x is None)
if holes == len(value):
value = None
elif holes != 0:
raise BadArgumentUsage('argument %s takes %d values'
% (self.dest, self.nargs))
state.opts[self.dest] = value
state.order.append(self.obj)
class ParsingState(object):
def __init__(self, rargs):
self.opts = {}
self.largs = []
self.rargs = rargs
self.order = []
class OptionParser(object):
"""The option parser is an internal class that is ultimately used to
parse options and arguments. It's modelled after optparse and brings
a similar but vastly simplified API. It should generally not be used
directly as the high level Click classes wrap it for you.
It's not nearly as extensible as optparse or argparse as it does not
implement features that are implemented on a higher level (such as
types or defaults).
:param ctx: optionally the :class:`~click.Context` where this parser
should go with.
"""
def __init__(self, ctx=None):
#: The :class:`~click.Context` for this parser. This might be
#: `None` for some advanced use cases.
self.ctx = ctx
#: This controls how the parser deals with interspersed arguments.
#: If this is set to `False`, the parser will stop on the first
#: non-option. Click uses this to implement nested subcommands
#: safely.
self.allow_interspersed_args = True
#: This tells the parser how to deal with unknown options. By
#: default it will error out (which is sensible), but there is a
#: second mode where it will ignore it and continue processing
#: after shifting all the unknown options into the resulting args.
self.ignore_unknown_options = False
if ctx is not None:
self.allow_interspersed_args = ctx.allow_interspersed_args
self.ignore_unknown_options = ctx.ignore_unknown_options
self._short_opt = {}
self._long_opt = {}
self._opt_prefixes = set(['-', '--'])
self._args = []
def add_option(self, opts, dest, action=None, nargs=1, const=None,
obj=None):
"""Adds a new option named `dest` to the parser. The destination
is not inferred (unlike with optparse) and needs to be explicitly
provided. Action can be any of ``store``, ``store_const``,
``append``, ``appnd_const`` or ``count``.
The `obj` can be used to identify the option in the order list
that is returned from the parser.
"""
if obj is None:
obj = dest
opts = [normalize_opt(opt, self.ctx) for opt in opts]
option = Option(opts, dest, action=action, nargs=nargs,
const=const, obj=obj)
self._opt_prefixes.update(option.prefixes)
for opt in option._short_opts:
self._short_opt[opt] = option
for opt in option._long_opts:
self._long_opt[opt] = option
def add_argument(self, dest, nargs=1, obj=None):
"""Adds a positional argument named `dest` to the parser.
The `obj` can be used to identify the option in the order list
that is returned from the parser.
"""
if obj is None:
obj = dest
self._args.append(Argument(dest=dest, nargs=nargs, obj=obj))
def parse_args(self, args):
"""Parses positional arguments and returns ``(values, args, order)``
for the parsed options and arguments as well as the leftover
arguments if there are any. The order is a list of objects as they
appear on the command line. If arguments appear multiple times they
will be memorized multiple times as well.
"""
state = ParsingState(args)
try:
self._process_args_for_options(state)
self._process_args_for_args(state)
except UsageError:
if self.ctx is None or not self.ctx.resilient_parsing:
raise
return state.opts, state.largs, state.order
def _process_args_for_args(self, state):
pargs, args = _unpack_args(state.largs + state.rargs,
[x.nargs for x in self._args])
for idx, arg in enumerate(self._args):
arg.process(pargs[idx], state)
state.largs = args
state.rargs = []
def _process_args_for_options(self, state):
while state.rargs:
arg = state.rargs.pop(0)
arglen = len(arg)
# Double dashes always handled explicitly regardless of what
# prefixes are valid.
if arg == '--':
return
elif arg[:1] in self._opt_prefixes and arglen > 1:
self._process_opts(arg, state)
elif self.allow_interspersed_args:
state.largs.append(arg)
else:
state.rargs.insert(0, arg)
return
# Say this is the original argument list:
# [arg0, arg1, ..., arg(i-1), arg(i), arg(i+1), ..., arg(N-1)]
# ^
# (we are about to process arg(i)).
#
# Then rargs is [arg(i), ..., arg(N-1)] and largs is a *subset* of
# [arg0, ..., arg(i-1)] (any options and their arguments will have
# been removed from largs).
#
# The while loop will usually consume 1 or more arguments per pass.
# If it consumes 1 (eg. arg is an option that takes no arguments),
# then after _process_arg() is done the situation is:
#
# largs = subset of [arg0, ..., arg(i)]
# rargs = [arg(i+1), ..., arg(N-1)]
#
# If allow_interspersed_args is false, largs will always be
# *empty* -- still a subset of [arg0, ..., arg(i-1)], but
# not a very interesting subset!
def _match_long_opt(self, opt, explicit_value, state):
if opt not in self._long_opt:
possibilities = [word for word in self._long_opt
if word.startswith(opt)]
raise NoSuchOption(opt, possibilities=possibilities)
option = self._long_opt[opt]
if option.takes_value:
# At this point it's safe to modify rargs by injecting the
# explicit value, because no exception is raised in this
# branch. This means that the inserted value will be fully
# consumed.
if explicit_value is not None:
state.rargs.insert(0, explicit_value)
nargs = option.nargs
if len(state.rargs) < nargs:
_error_opt_args(nargs, opt)
elif nargs == 1:
value = state.rargs.pop(0)
else:
value = tuple(state.rargs[:nargs])
del state.rargs[:nargs]
elif explicit_value is not None:
raise BadOptionUsage('%s option does not take a value' % opt)
else:
value = None
option.process(value, state)
def _match_short_opt(self, arg, state):
stop = False
i = 1
prefix = arg[0]
unknown_options = []
for ch in arg[1:]:
opt = normalize_opt(prefix + ch, self.ctx)
option = self._short_opt.get(opt)
i += 1
if not option:
if self.ignore_unknown_options:
unknown_options.append(ch)
continue
raise NoSuchOption(opt)
if option.takes_value:
# Any characters left in arg? Pretend they're the
# next arg, and stop consuming characters of arg.
if i < len(arg):
state.rargs.insert(0, arg[i:])
stop = True
nargs = option.nargs
if len(state.rargs) < nargs:
_error_opt_args(nargs, opt)
elif nargs == 1:
value = state.rargs.pop(0)
else:
value = tuple(state.rargs[:nargs])
del state.rargs[:nargs]
else:
value = None
option.process(value, state)
if stop:
break
# If we got any unknown options we re-combinate the string of the
# remaining options and re-attach the prefix, then report that
# to the state as new larg. This way there is basic combinatorics
# that can be achieved while still ignoring unknown arguments.
if self.ignore_unknown_options and unknown_options:
state.largs.append(prefix + ''.join(unknown_options))
def _process_opts(self, arg, state):
explicit_value = None
# Long option handling happens in two parts. The first part is
# supporting explicitly attached values. In any case, we will try
# to long match the option first.
if '=' in arg:
long_opt, explicit_value = arg.split('=', 1)
else:
long_opt = arg
norm_long_opt = normalize_opt(long_opt, self.ctx)
# At this point we will match the (assumed) long option through
# the long option matching code. Note that this allows options
# like "-foo" to be matched as long options.
try:
self._match_long_opt(norm_long_opt, explicit_value, state)
except NoSuchOption:
# At this point the long option matching failed, and we need
# to try with short options. However there is a special rule
# which says, that if we have a two character options prefix
# (applies to "--foo" for instance), we do not dispatch to the
# short option code and will instead raise the no option
# error.
if arg[:2] not in self._opt_prefixes:
return self._match_short_opt(arg, state)
if not self.ignore_unknown_options:
raise
state.largs.append(arg)
+539
View File
@@ -0,0 +1,539 @@
import os
import sys
import struct
from ._compat import raw_input, text_type, string_types, \
isatty, strip_ansi, get_winterm_size, DEFAULT_COLUMNS, WIN
from .utils import echo
from .exceptions import Abort, UsageError
from .types import convert_type
from .globals import resolve_color_default
# The prompt functions to use. The doc tools currently override these
# functions to customize how they work.
visible_prompt_func = raw_input
_ansi_colors = ('black', 'red', 'green', 'yellow', 'blue', 'magenta',
'cyan', 'white', 'reset')
_ansi_reset_all = '\033[0m'
def hidden_prompt_func(prompt):
import getpass
return getpass.getpass(prompt)
def _build_prompt(text, suffix, show_default=False, default=None):
prompt = text
if default is not None and show_default:
prompt = '%s [%s]' % (prompt, default)
return prompt + suffix
def prompt(text, default=None, hide_input=False,
confirmation_prompt=False, type=None,
value_proc=None, prompt_suffix=': ',
show_default=True, err=False):
"""Prompts a user for input. This is a convenience function that can
be used to prompt a user for input later.
If the user aborts the input by sending a interrupt signal, this
function will catch it and raise a :exc:`Abort` exception.
.. versionadded:: 6.0
Added unicode support for cmd.exe on Windows.
.. versionadded:: 4.0
Added the `err` parameter.
:param text: the text to show for the prompt.
:param default: the default value to use if no input happens. If this
is not given it will prompt until it's aborted.
:param hide_input: if this is set to true then the input value will
be hidden.
:param confirmation_prompt: asks for confirmation for the value.
:param type: the type to use to check the value against.
:param value_proc: if this parameter is provided it's a function that
is invoked instead of the type conversion to
convert a value.
:param prompt_suffix: a suffix that should be added to the prompt.
:param show_default: shows or hides the default value in the prompt.
:param err: if set to true the file defaults to ``stderr`` instead of
``stdout``, the same as with echo.
"""
result = None
def prompt_func(text):
f = hide_input and hidden_prompt_func or visible_prompt_func
try:
# Write the prompt separately so that we get nice
# coloring through colorama on Windows
echo(text, nl=False, err=err)
return f('')
except (KeyboardInterrupt, EOFError):
# getpass doesn't print a newline if the user aborts input with ^C.
# Allegedly this behavior is inherited from getpass(3).
# A doc bug has been filed at https://bugs.python.org/issue24711
if hide_input:
echo(None, err=err)
raise Abort()
if value_proc is None:
value_proc = convert_type(type, default)
prompt = _build_prompt(text, prompt_suffix, show_default, default)
while 1:
while 1:
value = prompt_func(prompt)
if value:
break
# If a default is set and used, then the confirmation
# prompt is always skipped because that's the only thing
# that really makes sense.
elif default is not None:
return default
try:
result = value_proc(value)
except UsageError as e:
echo('Error: %s' % e.message, err=err)
continue
if not confirmation_prompt:
return result
while 1:
value2 = prompt_func('Repeat for confirmation: ')
if value2:
break
if value == value2:
return result
echo('Error: the two entered values do not match', err=err)
def confirm(text, default=False, abort=False, prompt_suffix=': ',
show_default=True, err=False):
"""Prompts for confirmation (yes/no question).
If the user aborts the input by sending a interrupt signal this
function will catch it and raise a :exc:`Abort` exception.
.. versionadded:: 4.0
Added the `err` parameter.
:param text: the question to ask.
:param default: the default for the prompt.
:param abort: if this is set to `True` a negative answer aborts the
exception by raising :exc:`Abort`.
:param prompt_suffix: a suffix that should be added to the prompt.
:param show_default: shows or hides the default value in the prompt.
:param err: if set to true the file defaults to ``stderr`` instead of
``stdout``, the same as with echo.
"""
prompt = _build_prompt(text, prompt_suffix, show_default,
default and 'Y/n' or 'y/N')
while 1:
try:
# Write the prompt separately so that we get nice
# coloring through colorama on Windows
echo(prompt, nl=False, err=err)
value = visible_prompt_func('').lower().strip()
except (KeyboardInterrupt, EOFError):
raise Abort()
if value in ('y', 'yes'):
rv = True
elif value in ('n', 'no'):
rv = False
elif value == '':
rv = default
else:
echo('Error: invalid input', err=err)
continue
break
if abort and not rv:
raise Abort()
return rv
def get_terminal_size():
"""Returns the current size of the terminal as tuple in the form
``(width, height)`` in columns and rows.
"""
# If shutil has get_terminal_size() (Python 3.3 and later) use that
if sys.version_info >= (3, 3):
import shutil
shutil_get_terminal_size = getattr(shutil, 'get_terminal_size', None)
if shutil_get_terminal_size:
sz = shutil_get_terminal_size()
return sz.columns, sz.lines
if get_winterm_size is not None:
return get_winterm_size()
def ioctl_gwinsz(fd):
try:
import fcntl
import termios
cr = struct.unpack(
'hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234'))
except Exception:
return
return cr
cr = ioctl_gwinsz(0) or ioctl_gwinsz(1) or ioctl_gwinsz(2)
if not cr:
try:
fd = os.open(os.ctermid(), os.O_RDONLY)
try:
cr = ioctl_gwinsz(fd)
finally:
os.close(fd)
except Exception:
pass
if not cr or not cr[0] or not cr[1]:
cr = (os.environ.get('LINES', 25),
os.environ.get('COLUMNS', DEFAULT_COLUMNS))
return int(cr[1]), int(cr[0])
def echo_via_pager(text, color=None):
"""This function takes a text and shows it via an environment specific
pager on stdout.
.. versionchanged:: 3.0
Added the `color` flag.
:param text: the text to page.
:param color: controls if the pager supports ANSI colors or not. The
default is autodetection.
"""
color = resolve_color_default(color)
if not isinstance(text, string_types):
text = text_type(text)
from ._termui_impl import pager
return pager(text + '\n', color)
def progressbar(iterable=None, length=None, label=None, show_eta=True,
show_percent=None, show_pos=False,
item_show_func=None, fill_char='#', empty_char='-',
bar_template='%(label)s [%(bar)s] %(info)s',
info_sep=' ', width=36, file=None, color=None):
"""This function creates an iterable context manager that can be used
to iterate over something while showing a progress bar. It will
either iterate over the `iterable` or `length` items (that are counted
up). While iteration happens, this function will print a rendered
progress bar to the given `file` (defaults to stdout) and will attempt
to calculate remaining time and more. By default, this progress bar
will not be rendered if the file is not a terminal.
The context manager creates the progress bar. When the context
manager is entered the progress bar is already displayed. With every
iteration over the progress bar, the iterable passed to the bar is
advanced and the bar is updated. When the context manager exits,
a newline is printed and the progress bar is finalized on screen.
No printing must happen or the progress bar will be unintentionally
destroyed.
Example usage::
with progressbar(items) as bar:
for item in bar:
do_something_with(item)
Alternatively, if no iterable is specified, one can manually update the
progress bar through the `update()` method instead of directly
iterating over the progress bar. The update method accepts the number
of steps to increment the bar with::
with progressbar(length=chunks.total_bytes) as bar:
for chunk in chunks:
process_chunk(chunk)
bar.update(chunks.bytes)
.. versionadded:: 2.0
.. versionadded:: 4.0
Added the `color` parameter. Added a `update` method to the
progressbar object.
:param iterable: an iterable to iterate over. If not provided the length
is required.
:param length: the number of items to iterate over. By default the
progressbar will attempt to ask the iterator about its
length, which might or might not work. If an iterable is
also provided this parameter can be used to override the
length. If an iterable is not provided the progress bar
will iterate over a range of that length.
:param label: the label to show next to the progress bar.
:param show_eta: enables or disables the estimated time display. This is
automatically disabled if the length cannot be
determined.
:param show_percent: enables or disables the percentage display. The
default is `True` if the iterable has a length or
`False` if not.
:param show_pos: enables or disables the absolute position display. The
default is `False`.
:param item_show_func: a function called with the current item which
can return a string to show the current item
next to the progress bar. Note that the current
item can be `None`!
:param fill_char: the character to use to show the filled part of the
progress bar.
:param empty_char: the character to use to show the non-filled part of
the progress bar.
:param bar_template: the format string to use as template for the bar.
The parameters in it are ``label`` for the label,
``bar`` for the progress bar and ``info`` for the
info section.
:param info_sep: the separator between multiple info items (eta etc.)
:param width: the width of the progress bar in characters, 0 means full
terminal width
:param file: the file to write to. If this is not a terminal then
only the label is printed.
:param color: controls if the terminal supports ANSI colors or not. The
default is autodetection. This is only needed if ANSI
codes are included anywhere in the progress bar output
which is not the case by default.
"""
from ._termui_impl import ProgressBar
color = resolve_color_default(color)
return ProgressBar(iterable=iterable, length=length, show_eta=show_eta,
show_percent=show_percent, show_pos=show_pos,
item_show_func=item_show_func, fill_char=fill_char,
empty_char=empty_char, bar_template=bar_template,
info_sep=info_sep, file=file, label=label,
width=width, color=color)
def clear():
"""Clears the terminal screen. This will have the effect of clearing
the whole visible space of the terminal and moving the cursor to the
top left. This does not do anything if not connected to a terminal.
.. versionadded:: 2.0
"""
if not isatty(sys.stdout):
return
# If we're on Windows and we don't have colorama available, then we
# clear the screen by shelling out. Otherwise we can use an escape
# sequence.
if WIN:
os.system('cls')
else:
sys.stdout.write('\033[2J\033[1;1H')
def style(text, fg=None, bg=None, bold=None, dim=None, underline=None,
blink=None, reverse=None, reset=True):
"""Styles a text with ANSI styles and returns the new string. By
default the styling is self contained which means that at the end
of the string a reset code is issued. This can be prevented by
passing ``reset=False``.
Examples::
click.echo(click.style('Hello World!', fg='green'))
click.echo(click.style('ATTENTION!', blink=True))
click.echo(click.style('Some things', reverse=True, fg='cyan'))
Supported color names:
* ``black`` (might be a gray)
* ``red``
* ``green``
* ``yellow`` (might be an orange)
* ``blue``
* ``magenta``
* ``cyan``
* ``white`` (might be light gray)
* ``reset`` (reset the color code only)
.. versionadded:: 2.0
:param text: the string to style with ansi codes.
:param fg: if provided this will become the foreground color.
:param bg: if provided this will become the background color.
:param bold: if provided this will enable or disable bold mode.
:param dim: if provided this will enable or disable dim mode. This is
badly supported.
:param underline: if provided this will enable or disable underline.
:param blink: if provided this will enable or disable blinking.
:param reverse: if provided this will enable or disable inverse
rendering (foreground becomes background and the
other way round).
:param reset: by default a reset-all code is added at the end of the
string which means that styles do not carry over. This
can be disabled to compose styles.
"""
bits = []
if fg:
try:
bits.append('\033[%dm' % (_ansi_colors.index(fg) + 30))
except ValueError:
raise TypeError('Unknown color %r' % fg)
if bg:
try:
bits.append('\033[%dm' % (_ansi_colors.index(bg) + 40))
except ValueError:
raise TypeError('Unknown color %r' % bg)
if bold is not None:
bits.append('\033[%dm' % (1 if bold else 22))
if dim is not None:
bits.append('\033[%dm' % (2 if dim else 22))
if underline is not None:
bits.append('\033[%dm' % (4 if underline else 24))
if blink is not None:
bits.append('\033[%dm' % (5 if blink else 25))
if reverse is not None:
bits.append('\033[%dm' % (7 if reverse else 27))
bits.append(text)
if reset:
bits.append(_ansi_reset_all)
return ''.join(bits)
def unstyle(text):
"""Removes ANSI styling information from a string. Usually it's not
necessary to use this function as Click's echo function will
automatically remove styling if necessary.
.. versionadded:: 2.0
:param text: the text to remove style information from.
"""
return strip_ansi(text)
def secho(text, file=None, nl=True, err=False, color=None, **styles):
"""This function combines :func:`echo` and :func:`style` into one
call. As such the following two calls are the same::
click.secho('Hello World!', fg='green')
click.echo(click.style('Hello World!', fg='green'))
All keyword arguments are forwarded to the underlying functions
depending on which one they go with.
.. versionadded:: 2.0
"""
return echo(style(text, **styles), file=file, nl=nl, err=err, color=color)
def edit(text=None, editor=None, env=None, require_save=True,
extension='.txt', filename=None):
r"""Edits the given text in the defined editor. If an editor is given
(should be the full path to the executable but the regular operating
system search path is used for finding the executable) it overrides
the detected editor. Optionally, some environment variables can be
used. If the editor is closed without changes, `None` is returned. In
case a file is edited directly the return value is always `None` and
`require_save` and `extension` are ignored.
If the editor cannot be opened a :exc:`UsageError` is raised.
Note for Windows: to simplify cross-platform usage, the newlines are
automatically converted from POSIX to Windows and vice versa. As such,
the message here will have ``\n`` as newline markers.
:param text: the text to edit.
:param editor: optionally the editor to use. Defaults to automatic
detection.
:param env: environment variables to forward to the editor.
:param require_save: if this is true, then not saving in the editor
will make the return value become `None`.
:param extension: the extension to tell the editor about. This defaults
to `.txt` but changing this might change syntax
highlighting.
:param filename: if provided it will edit this file instead of the
provided text contents. It will not use a temporary
file as an indirection in that case.
"""
from ._termui_impl import Editor
editor = Editor(editor=editor, env=env, require_save=require_save,
extension=extension)
if filename is None:
return editor.edit(text)
editor.edit_file(filename)
def launch(url, wait=False, locate=False):
"""This function launches the given URL (or filename) in the default
viewer application for this file type. If this is an executable, it
might launch the executable in a new session. The return value is
the exit code of the launched application. Usually, ``0`` indicates
success.
Examples::
click.launch('http://click.pocoo.org/')
click.launch('/my/downloaded/file', locate=True)
.. versionadded:: 2.0
:param url: URL or filename of the thing to launch.
:param wait: waits for the program to stop.
:param locate: if this is set to `True` then instead of launching the
application associated with the URL it will attempt to
launch a file manager with the file located. This
might have weird effects if the URL does not point to
the filesystem.
"""
from ._termui_impl import open_url
return open_url(url, wait=wait, locate=locate)
# If this is provided, getchar() calls into this instead. This is used
# for unittesting purposes.
_getchar = None
def getchar(echo=False):
"""Fetches a single character from the terminal and returns it. This
will always return a unicode character and under certain rare
circumstances this might return more than one character. The
situations which more than one character is returned is when for
whatever reason multiple characters end up in the terminal buffer or
standard input was not actually a terminal.
Note that this will always read from the terminal, even if something
is piped into the standard input.
.. versionadded:: 2.0
:param echo: if set to `True`, the character read will also show up on
the terminal. The default is to not show it.
"""
f = _getchar
if f is None:
from ._termui_impl import getchar as f
return f(echo)
def pause(info='Press any key to continue ...', err=False):
"""This command stops execution and waits for the user to press any
key to continue. This is similar to the Windows batch "pause"
command. If the program is not run through a terminal, this command
will instead do nothing.
.. versionadded:: 2.0
.. versionadded:: 4.0
Added the `err` parameter.
:param info: the info string to print before pausing.
:param err: if set to message goes to ``stderr`` instead of
``stdout``, the same as with echo.
"""
if not isatty(sys.stdin) or not isatty(sys.stdout):
return
try:
if info:
echo(info, nl=False, err=err)
try:
getchar()
except (KeyboardInterrupt, EOFError):
pass
finally:
if info:
echo(err=err)
+322
View File
@@ -0,0 +1,322 @@
import os
import sys
import shutil
import tempfile
import contextlib
from ._compat import iteritems, PY2
# If someone wants to vendor click, we want to ensure the
# correct package is discovered. Ideally we could use a
# relative import here but unfortunately Python does not
# support that.
clickpkg = sys.modules[__name__.rsplit('.', 1)[0]]
if PY2:
from cStringIO import StringIO
else:
import io
from ._compat import _find_binary_reader
class EchoingStdin(object):
def __init__(self, input, output):
self._input = input
self._output = output
def __getattr__(self, x):
return getattr(self._input, x)
def _echo(self, rv):
self._output.write(rv)
return rv
def read(self, n=-1):
return self._echo(self._input.read(n))
def readline(self, n=-1):
return self._echo(self._input.readline(n))
def readlines(self):
return [self._echo(x) for x in self._input.readlines()]
def __iter__(self):
return iter(self._echo(x) for x in self._input)
def __repr__(self):
return repr(self._input)
def make_input_stream(input, charset):
# Is already an input stream.
if hasattr(input, 'read'):
if PY2:
return input
rv = _find_binary_reader(input)
if rv is not None:
return rv
raise TypeError('Could not find binary reader for input stream.')
if input is None:
input = b''
elif not isinstance(input, bytes):
input = input.encode(charset)
if PY2:
return StringIO(input)
return io.BytesIO(input)
class Result(object):
"""Holds the captured result of an invoked CLI script."""
def __init__(self, runner, output_bytes, exit_code, exception,
exc_info=None):
#: The runner that created the result
self.runner = runner
#: The output as bytes.
self.output_bytes = output_bytes
#: The exit code as integer.
self.exit_code = exit_code
#: The exception that happend if one did.
self.exception = exception
#: The traceback
self.exc_info = exc_info
@property
def output(self):
"""The output as unicode string."""
return self.output_bytes.decode(self.runner.charset, 'replace') \
.replace('\r\n', '\n')
def __repr__(self):
return '<Result %s>' % (
self.exception and repr(self.exception) or 'okay',
)
class CliRunner(object):
"""The CLI runner provides functionality to invoke a Click command line
script for unittesting purposes in a isolated environment. This only
works in single-threaded systems without any concurrency as it changes the
global interpreter state.
:param charset: the character set for the input and output data. This is
UTF-8 by default and should not be changed currently as
the reporting to Click only works in Python 2 properly.
:param env: a dictionary with environment variables for overriding.
:param echo_stdin: if this is set to `True`, then reading from stdin writes
to stdout. This is useful for showing examples in
some circumstances. Note that regular prompts
will automatically echo the input.
"""
def __init__(self, charset=None, env=None, echo_stdin=False):
if charset is None:
charset = 'utf-8'
self.charset = charset
self.env = env or {}
self.echo_stdin = echo_stdin
def get_default_prog_name(self, cli):
"""Given a command object it will return the default program name
for it. The default is the `name` attribute or ``"root"`` if not
set.
"""
return cli.name or 'root'
def make_env(self, overrides=None):
"""Returns the environment overrides for invoking a script."""
rv = dict(self.env)
if overrides:
rv.update(overrides)
return rv
@contextlib.contextmanager
def isolation(self, input=None, env=None, color=False):
"""A context manager that sets up the isolation for invoking of a
command line tool. This sets up stdin with the given input data
and `os.environ` with the overrides from the given dictionary.
This also rebinds some internals in Click to be mocked (like the
prompt functionality).
This is automatically done in the :meth:`invoke` method.
.. versionadded:: 4.0
The ``color`` parameter was added.
:param input: the input stream to put into sys.stdin.
:param env: the environment overrides as dictionary.
:param color: whether the output should contain color codes. The
application can still override this explicitly.
"""
input = make_input_stream(input, self.charset)
old_stdin = sys.stdin
old_stdout = sys.stdout
old_stderr = sys.stderr
old_forced_width = clickpkg.formatting.FORCED_WIDTH
clickpkg.formatting.FORCED_WIDTH = 80
env = self.make_env(env)
if PY2:
sys.stdout = sys.stderr = bytes_output = StringIO()
if self.echo_stdin:
input = EchoingStdin(input, bytes_output)
else:
bytes_output = io.BytesIO()
if self.echo_stdin:
input = EchoingStdin(input, bytes_output)
input = io.TextIOWrapper(input, encoding=self.charset)
sys.stdout = sys.stderr = io.TextIOWrapper(
bytes_output, encoding=self.charset)
sys.stdin = input
def visible_input(prompt=None):
sys.stdout.write(prompt or '')
val = input.readline().rstrip('\r\n')
sys.stdout.write(val + '\n')
sys.stdout.flush()
return val
def hidden_input(prompt=None):
sys.stdout.write((prompt or '') + '\n')
sys.stdout.flush()
return input.readline().rstrip('\r\n')
def _getchar(echo):
char = sys.stdin.read(1)
if echo:
sys.stdout.write(char)
sys.stdout.flush()
return char
default_color = color
def should_strip_ansi(stream=None, color=None):
if color is None:
return not default_color
return not color
old_visible_prompt_func = clickpkg.termui.visible_prompt_func
old_hidden_prompt_func = clickpkg.termui.hidden_prompt_func
old__getchar_func = clickpkg.termui._getchar
old_should_strip_ansi = clickpkg.utils.should_strip_ansi
clickpkg.termui.visible_prompt_func = visible_input
clickpkg.termui.hidden_prompt_func = hidden_input
clickpkg.termui._getchar = _getchar
clickpkg.utils.should_strip_ansi = should_strip_ansi
old_env = {}
try:
for key, value in iteritems(env):
old_env[key] = os.environ.get(key)
if value is None:
try:
del os.environ[key]
except Exception:
pass
else:
os.environ[key] = value
yield bytes_output
finally:
for key, value in iteritems(old_env):
if value is None:
try:
del os.environ[key]
except Exception:
pass
else:
os.environ[key] = value
sys.stdout = old_stdout
sys.stderr = old_stderr
sys.stdin = old_stdin
clickpkg.termui.visible_prompt_func = old_visible_prompt_func
clickpkg.termui.hidden_prompt_func = old_hidden_prompt_func
clickpkg.termui._getchar = old__getchar_func
clickpkg.utils.should_strip_ansi = old_should_strip_ansi
clickpkg.formatting.FORCED_WIDTH = old_forced_width
def invoke(self, cli, args=None, input=None, env=None,
catch_exceptions=True, color=False, **extra):
"""Invokes a command in an isolated environment. The arguments are
forwarded directly to the command line script, the `extra` keyword
arguments are passed to the :meth:`~clickpkg.Command.main` function of
the command.
This returns a :class:`Result` object.
.. versionadded:: 3.0
The ``catch_exceptions`` parameter was added.
.. versionchanged:: 3.0
The result object now has an `exc_info` attribute with the
traceback if available.
.. versionadded:: 4.0
The ``color`` parameter was added.
:param cli: the command to invoke
:param args: the arguments to invoke
:param input: the input data for `sys.stdin`.
:param env: the environment overrides.
:param catch_exceptions: Whether to catch any other exceptions than
``SystemExit``.
:param extra: the keyword arguments to pass to :meth:`main`.
:param color: whether the output should contain color codes. The
application can still override this explicitly.
"""
exc_info = None
with self.isolation(input=input, env=env, color=color) as out:
exception = None
exit_code = 0
try:
cli.main(args=args or (),
prog_name=self.get_default_prog_name(cli), **extra)
except SystemExit as e:
if e.code != 0:
exception = e
exc_info = sys.exc_info()
exit_code = e.code
if not isinstance(exit_code, int):
sys.stdout.write(str(exit_code))
sys.stdout.write('\n')
exit_code = 1
except Exception as e:
if not catch_exceptions:
raise
exception = e
exit_code = -1
exc_info = sys.exc_info()
finally:
sys.stdout.flush()
output = out.getvalue()
return Result(runner=self,
output_bytes=output,
exit_code=exit_code,
exception=exception,
exc_info=exc_info)
@contextlib.contextmanager
def isolated_filesystem(self):
"""A context manager that creates a temporary folder and changes
the current working directory to it for isolated filesystem tests.
"""
cwd = os.getcwd()
t = tempfile.mkdtemp()
os.chdir(t)
try:
yield t
finally:
os.chdir(cwd)
try:
shutil.rmtree(t)
except (OSError, IOError):
pass
+550
View File
@@ -0,0 +1,550 @@
import os
import stat
from ._compat import open_stream, text_type, filename_to_ui, \
get_filesystem_encoding, get_streerror, _get_argv_encoding, PY2
from .exceptions import BadParameter
from .utils import safecall, LazyFile
class ParamType(object):
"""Helper for converting values through types. The following is
necessary for a valid type:
* it needs a name
* it needs to pass through None unchanged
* it needs to convert from a string
* it needs to convert its result type through unchanged
(eg: needs to be idempotent)
* it needs to be able to deal with param and context being `None`.
This can be the case when the object is used with prompt
inputs.
"""
is_composite = False
#: the descriptive name of this type
name = None
#: if a list of this type is expected and the value is pulled from a
#: string environment variable, this is what splits it up. `None`
#: means any whitespace. For all parameters the general rule is that
#: whitespace splits them up. The exception are paths and files which
#: are split by ``os.path.pathsep`` by default (":" on Unix and ";" on
#: Windows).
envvar_list_splitter = None
def __call__(self, value, param=None, ctx=None):
if value is not None:
return self.convert(value, param, ctx)
def get_metavar(self, param):
"""Returns the metavar default for this param if it provides one."""
def get_missing_message(self, param):
"""Optionally might return extra information about a missing
parameter.
.. versionadded:: 2.0
"""
def convert(self, value, param, ctx):
"""Converts the value. This is not invoked for values that are
`None` (the missing value).
"""
return value
def split_envvar_value(self, rv):
"""Given a value from an environment variable this splits it up
into small chunks depending on the defined envvar list splitter.
If the splitter is set to `None`, which means that whitespace splits,
then leading and trailing whitespace is ignored. Otherwise, leading
and trailing splitters usually lead to empty items being included.
"""
return (rv or '').split(self.envvar_list_splitter)
def fail(self, message, param=None, ctx=None):
"""Helper method to fail with an invalid value message."""
raise BadParameter(message, ctx=ctx, param=param)
class CompositeParamType(ParamType):
is_composite = True
@property
def arity(self):
raise NotImplementedError()
class FuncParamType(ParamType):
def __init__(self, func):
self.name = func.__name__
self.func = func
def convert(self, value, param, ctx):
try:
return self.func(value)
except ValueError:
try:
value = text_type(value)
except UnicodeError:
value = str(value).decode('utf-8', 'replace')
self.fail(value, param, ctx)
class UnprocessedParamType(ParamType):
name = 'text'
def convert(self, value, param, ctx):
return value
def __repr__(self):
return 'UNPROCESSED'
class StringParamType(ParamType):
name = 'text'
def convert(self, value, param, ctx):
if isinstance(value, bytes):
enc = _get_argv_encoding()
try:
value = value.decode(enc)
except UnicodeError:
fs_enc = get_filesystem_encoding()
if fs_enc != enc:
try:
value = value.decode(fs_enc)
except UnicodeError:
value = value.decode('utf-8', 'replace')
return value
return value
def __repr__(self):
return 'STRING'
class Choice(ParamType):
"""The choice type allows a value to be checked against a fixed set of
supported values. All of these values have to be strings.
See :ref:`choice-opts` for an example.
"""
name = 'choice'
def __init__(self, choices):
self.choices = choices
def get_metavar(self, param):
return '[%s]' % '|'.join(self.choices)
def get_missing_message(self, param):
return 'Choose from %s.' % ', '.join(self.choices)
def convert(self, value, param, ctx):
# Exact match
if value in self.choices:
return value
# Match through normalization
if ctx is not None and \
ctx.token_normalize_func is not None:
value = ctx.token_normalize_func(value)
for choice in self.choices:
if ctx.token_normalize_func(choice) == value:
return choice
self.fail('invalid choice: %s. (choose from %s)' %
(value, ', '.join(self.choices)), param, ctx)
def __repr__(self):
return 'Choice(%r)' % list(self.choices)
class IntParamType(ParamType):
name = 'integer'
def convert(self, value, param, ctx):
try:
return int(value)
except (ValueError, UnicodeError):
self.fail('%s is not a valid integer' % value, param, ctx)
def __repr__(self):
return 'INT'
class IntRange(IntParamType):
"""A parameter that works similar to :data:`click.INT` but restricts
the value to fit into a range. The default behavior is to fail if the
value falls outside the range, but it can also be silently clamped
between the two edges.
See :ref:`ranges` for an example.
"""
name = 'integer range'
def __init__(self, min=None, max=None, clamp=False):
self.min = min
self.max = max
self.clamp = clamp
def convert(self, value, param, ctx):
rv = IntParamType.convert(self, value, param, ctx)
if self.clamp:
if self.min is not None and rv < self.min:
return self.min
if self.max is not None and rv > self.max:
return self.max
if self.min is not None and rv < self.min or \
self.max is not None and rv > self.max:
if self.min is None:
self.fail('%s is bigger than the maximum valid value '
'%s.' % (rv, self.max), param, ctx)
elif self.max is None:
self.fail('%s is smaller than the minimum valid value '
'%s.' % (rv, self.min), param, ctx)
else:
self.fail('%s is not in the valid range of %s to %s.'
% (rv, self.min, self.max), param, ctx)
return rv
def __repr__(self):
return 'IntRange(%r, %r)' % (self.min, self.max)
class BoolParamType(ParamType):
name = 'boolean'
def convert(self, value, param, ctx):
if isinstance(value, bool):
return bool(value)
value = value.lower()
if value in ('true', '1', 'yes', 'y'):
return True
elif value in ('false', '0', 'no', 'n'):
return False
self.fail('%s is not a valid boolean' % value, param, ctx)
def __repr__(self):
return 'BOOL'
class FloatParamType(ParamType):
name = 'float'
def convert(self, value, param, ctx):
try:
return float(value)
except (UnicodeError, ValueError):
self.fail('%s is not a valid floating point value' %
value, param, ctx)
def __repr__(self):
return 'FLOAT'
class UUIDParameterType(ParamType):
name = 'uuid'
def convert(self, value, param, ctx):
import uuid
try:
if PY2 and isinstance(value, text_type):
value = value.encode('ascii')
return uuid.UUID(value)
except (UnicodeError, ValueError):
self.fail('%s is not a valid UUID value' % value, param, ctx)
def __repr__(self):
return 'UUID'
class File(ParamType):
"""Declares a parameter to be a file for reading or writing. The file
is automatically closed once the context tears down (after the command
finished working).
Files can be opened for reading or writing. The special value ``-``
indicates stdin or stdout depending on the mode.
By default, the file is opened for reading text data, but it can also be
opened in binary mode or for writing. The encoding parameter can be used
to force a specific encoding.
The `lazy` flag controls if the file should be opened immediately or
upon first IO. The default is to be non lazy for standard input and
output streams as well as files opened for reading, lazy otherwise.
Starting with Click 2.0, files can also be opened atomically in which
case all writes go into a separate file in the same folder and upon
completion the file will be moved over to the original location. This
is useful if a file regularly read by other users is modified.
See :ref:`file-args` for more information.
"""
name = 'filename'
envvar_list_splitter = os.path.pathsep
def __init__(self, mode='r', encoding=None, errors='strict', lazy=None,
atomic=False):
self.mode = mode
self.encoding = encoding
self.errors = errors
self.lazy = lazy
self.atomic = atomic
def resolve_lazy_flag(self, value):
if self.lazy is not None:
return self.lazy
if value == '-':
return False
elif 'w' in self.mode:
return True
return False
def convert(self, value, param, ctx):
try:
if hasattr(value, 'read') or hasattr(value, 'write'):
return value
lazy = self.resolve_lazy_flag(value)
if lazy:
f = LazyFile(value, self.mode, self.encoding, self.errors,
atomic=self.atomic)
if ctx is not None:
ctx.call_on_close(f.close_intelligently)
return f
f, should_close = open_stream(value, self.mode,
self.encoding, self.errors,
atomic=self.atomic)
# If a context is provided, we automatically close the file
# at the end of the context execution (or flush out). If a
# context does not exist, it's the caller's responsibility to
# properly close the file. This for instance happens when the
# type is used with prompts.
if ctx is not None:
if should_close:
ctx.call_on_close(safecall(f.close))
else:
ctx.call_on_close(safecall(f.flush))
return f
except (IOError, OSError) as e:
self.fail('Could not open file: %s: %s' % (
filename_to_ui(value),
get_streerror(e),
), param, ctx)
class Path(ParamType):
"""The path type is similar to the :class:`File` type but it performs
different checks. First of all, instead of returning an open file
handle it returns just the filename. Secondly, it can perform various
basic checks about what the file or directory should be.
.. versionchanged:: 6.0
`allow_dash` was added.
:param exists: if set to true, the file or directory needs to exist for
this value to be valid. If this is not required and a
file does indeed not exist, then all further checks are
silently skipped.
:param file_okay: controls if a file is a possible value.
:param dir_okay: controls if a directory is a possible value.
:param writable: if true, a writable check is performed.
:param readable: if true, a readable check is performed.
:param resolve_path: if this is true, then the path is fully resolved
before the value is passed onwards. This means
that it's absolute and symlinks are resolved.
:param allow_dash: If this is set to `True`, a single dash to indicate
standard streams is permitted.
:param type: optionally a string type that should be used to
represent the path. The default is `None` which
means the return value will be either bytes or
unicode depending on what makes most sense given the
input data Click deals with.
"""
envvar_list_splitter = os.path.pathsep
def __init__(self, exists=False, file_okay=True, dir_okay=True,
writable=False, readable=True, resolve_path=False,
allow_dash=False, path_type=None):
self.exists = exists
self.file_okay = file_okay
self.dir_okay = dir_okay
self.writable = writable
self.readable = readable
self.resolve_path = resolve_path
self.allow_dash = allow_dash
self.type = path_type
if self.file_okay and not self.dir_okay:
self.name = 'file'
self.path_type = 'File'
if self.dir_okay and not self.file_okay:
self.name = 'directory'
self.path_type = 'Directory'
else:
self.name = 'path'
self.path_type = 'Path'
def coerce_path_result(self, rv):
if self.type is not None and not isinstance(rv, self.type):
if self.type is text_type:
rv = rv.decode(get_filesystem_encoding())
else:
rv = rv.encode(get_filesystem_encoding())
return rv
def convert(self, value, param, ctx):
rv = value
is_dash = self.file_okay and self.allow_dash and rv in (b'-', '-')
if not is_dash:
if self.resolve_path:
rv = os.path.realpath(rv)
try:
st = os.stat(rv)
except OSError:
if not self.exists:
return self.coerce_path_result(rv)
self.fail('%s "%s" does not exist.' % (
self.path_type,
filename_to_ui(value)
), param, ctx)
if not self.file_okay and stat.S_ISREG(st.st_mode):
self.fail('%s "%s" is a file.' % (
self.path_type,
filename_to_ui(value)
), param, ctx)
if not self.dir_okay and stat.S_ISDIR(st.st_mode):
self.fail('%s "%s" is a directory.' % (
self.path_type,
filename_to_ui(value)
), param, ctx)
if self.writable and not os.access(value, os.W_OK):
self.fail('%s "%s" is not writable.' % (
self.path_type,
filename_to_ui(value)
), param, ctx)
if self.readable and not os.access(value, os.R_OK):
self.fail('%s "%s" is not readable.' % (
self.path_type,
filename_to_ui(value)
), param, ctx)
return self.coerce_path_result(rv)
class Tuple(CompositeParamType):
"""The default behavior of Click is to apply a type on a value directly.
This works well in most cases, except for when `nargs` is set to a fixed
count and different types should be used for different items. In this
case the :class:`Tuple` type can be used. This type can only be used
if `nargs` is set to a fixed number.
For more information see :ref:`tuple-type`.
This can be selected by using a Python tuple literal as a type.
:param types: a list of types that should be used for the tuple items.
"""
def __init__(self, types):
self.types = [convert_type(ty) for ty in types]
@property
def name(self):
return "<" + " ".join(ty.name for ty in self.types) + ">"
@property
def arity(self):
return len(self.types)
def convert(self, value, param, ctx):
if len(value) != len(self.types):
raise TypeError('It would appear that nargs is set to conflict '
'with the composite type arity.')
return tuple(ty(x, param, ctx) for ty, x in zip(self.types, value))
def convert_type(ty, default=None):
"""Converts a callable or python ty into the most appropriate param
ty.
"""
guessed_type = False
if ty is None and default is not None:
if isinstance(default, tuple):
ty = tuple(map(type, default))
else:
ty = type(default)
guessed_type = True
if isinstance(ty, tuple):
return Tuple(ty)
if isinstance(ty, ParamType):
return ty
if ty is text_type or ty is str or ty is None:
return STRING
if ty is int:
return INT
# Booleans are only okay if not guessed. This is done because for
# flags the default value is actually a bit of a lie in that it
# indicates which of the flags is the one we want. See get_default()
# for more information.
if ty is bool and not guessed_type:
return BOOL
if ty is float:
return FLOAT
if guessed_type:
return STRING
# Catch a common mistake
if __debug__:
try:
if issubclass(ty, ParamType):
raise AssertionError('Attempted to use an uninstantiated '
'parameter type (%s).' % ty)
except TypeError:
pass
return FuncParamType(ty)
#: A dummy parameter type that just does nothing. From a user's
#: perspective this appears to just be the same as `STRING` but internally
#: no string conversion takes place. This is necessary to achieve the
#: same bytes/unicode behavior on Python 2/3 in situations where you want
#: to not convert argument types. This is usually useful when working
#: with file paths as they can appear in bytes and unicode.
#:
#: For path related uses the :class:`Path` type is a better choice but
#: there are situations where an unprocessed type is useful which is why
#: it is is provided.
#:
#: .. versionadded:: 4.0
UNPROCESSED = UnprocessedParamType()
#: A unicode string parameter type which is the implicit default. This
#: can also be selected by using ``str`` as type.
STRING = StringParamType()
#: An integer parameter. This can also be selected by using ``int`` as
#: type.
INT = IntParamType()
#: A floating point value parameter. This can also be selected by using
#: ``float`` as type.
FLOAT = FloatParamType()
#: A boolean parameter. This is the default for boolean flags. This can
#: also be selected by using ``bool`` as a type.
BOOL = BoolParamType()
#: A UUID parameter.
UUID = UUIDParameterType()
+415
View File
@@ -0,0 +1,415 @@
import os
import sys
from .globals import resolve_color_default
from ._compat import text_type, open_stream, get_filesystem_encoding, \
get_streerror, string_types, PY2, binary_streams, text_streams, \
filename_to_ui, auto_wrap_for_ansi, strip_ansi, should_strip_ansi, \
_default_text_stdout, _default_text_stderr, is_bytes, WIN
if not PY2:
from ._compat import _find_binary_writer
elif WIN:
from ._winconsole import _get_windows_argv, \
_hash_py_argv, _initial_argv_hash
echo_native_types = string_types + (bytes, bytearray)
def _posixify(name):
return '-'.join(name.split()).lower()
def safecall(func):
"""Wraps a function so that it swallows exceptions."""
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except Exception:
pass
return wrapper
def make_str(value):
"""Converts a value into a valid string."""
if isinstance(value, bytes):
try:
return value.decode(get_filesystem_encoding())
except UnicodeError:
return value.decode('utf-8', 'replace')
return text_type(value)
def make_default_short_help(help, max_length=45):
words = help.split()
total_length = 0
result = []
done = False
for word in words:
if word[-1:] == '.':
done = True
new_length = result and 1 + len(word) or len(word)
if total_length + new_length > max_length:
result.append('...')
done = True
else:
if result:
result.append(' ')
result.append(word)
if done:
break
total_length += new_length
return ''.join(result)
class LazyFile(object):
"""A lazy file works like a regular file but it does not fully open
the file but it does perform some basic checks early to see if the
filename parameter does make sense. This is useful for safely opening
files for writing.
"""
def __init__(self, filename, mode='r', encoding=None, errors='strict',
atomic=False):
self.name = filename
self.mode = mode
self.encoding = encoding
self.errors = errors
self.atomic = atomic
if filename == '-':
self._f, self.should_close = open_stream(filename, mode,
encoding, errors)
else:
if 'r' in mode:
# Open and close the file in case we're opening it for
# reading so that we can catch at least some errors in
# some cases early.
open(filename, mode).close()
self._f = None
self.should_close = True
def __getattr__(self, name):
return getattr(self.open(), name)
def __repr__(self):
if self._f is not None:
return repr(self._f)
return '<unopened file %r %s>' % (self.name, self.mode)
def open(self):
"""Opens the file if it's not yet open. This call might fail with
a :exc:`FileError`. Not handling this error will produce an error
that Click shows.
"""
if self._f is not None:
return self._f
try:
rv, self.should_close = open_stream(self.name, self.mode,
self.encoding,
self.errors,
atomic=self.atomic)
except (IOError, OSError) as e:
from .exceptions import FileError
raise FileError(self.name, hint=get_streerror(e))
self._f = rv
return rv
def close(self):
"""Closes the underlying file, no matter what."""
if self._f is not None:
self._f.close()
def close_intelligently(self):
"""This function only closes the file if it was opened by the lazy
file wrapper. For instance this will never close stdin.
"""
if self.should_close:
self.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, tb):
self.close_intelligently()
def __iter__(self):
self.open()
return iter(self._f)
class KeepOpenFile(object):
def __init__(self, file):
self._file = file
def __getattr__(self, name):
return getattr(self._file, name)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, tb):
pass
def __repr__(self):
return repr(self._file)
def __iter__(self):
return iter(self._file)
def echo(message=None, file=None, nl=True, err=False, color=None):
"""Prints a message plus a newline to the given file or stdout. On
first sight, this looks like the print function, but it has improved
support for handling Unicode and binary data that does not fail no
matter how badly configured the system is.
Primarily it means that you can print binary data as well as Unicode
data on both 2.x and 3.x to the given file in the most appropriate way
possible. This is a very carefree function as in that it will try its
best to not fail. As of Click 6.0 this includes support for unicode
output on the Windows console.
In addition to that, if `colorama`_ is installed, the echo function will
also support clever handling of ANSI codes. Essentially it will then
do the following:
- add transparent handling of ANSI color codes on Windows.
- hide ANSI codes automatically if the destination file is not a
terminal.
.. _colorama: http://pypi.python.org/pypi/colorama
.. versionchanged:: 6.0
As of Click 6.0 the echo function will properly support unicode
output on the windows console. Not that click does not modify
the interpreter in any way which means that `sys.stdout` or the
print statement or function will still not provide unicode support.
.. versionchanged:: 2.0
Starting with version 2.0 of Click, the echo function will work
with colorama if it's installed.
.. versionadded:: 3.0
The `err` parameter was added.
.. versionchanged:: 4.0
Added the `color` flag.
:param message: the message to print
:param file: the file to write to (defaults to ``stdout``)
:param err: if set to true the file defaults to ``stderr`` instead of
``stdout``. This is faster and easier than calling
:func:`get_text_stderr` yourself.
:param nl: if set to `True` (the default) a newline is printed afterwards.
:param color: controls if the terminal supports ANSI colors or not. The
default is autodetection.
"""
if file is None:
if err:
file = _default_text_stderr()
else:
file = _default_text_stdout()
# Convert non bytes/text into the native string type.
if message is not None and not isinstance(message, echo_native_types):
message = text_type(message)
if nl:
message = message or u''
if isinstance(message, text_type):
message += u'\n'
else:
message += b'\n'
# If there is a message, and we're in Python 3, and the value looks
# like bytes, we manually need to find the binary stream and write the
# message in there. This is done separately so that most stream
# types will work as you would expect. Eg: you can write to StringIO
# for other cases.
if message and not PY2 and is_bytes(message):
binary_file = _find_binary_writer(file)
if binary_file is not None:
file.flush()
binary_file.write(message)
binary_file.flush()
return
# ANSI-style support. If there is no message or we are dealing with
# bytes nothing is happening. If we are connected to a file we want
# to strip colors. If we are on windows we either wrap the stream
# to strip the color or we use the colorama support to translate the
# ansi codes to API calls.
if message and not is_bytes(message):
color = resolve_color_default(color)
if should_strip_ansi(file, color):
message = strip_ansi(message)
elif WIN:
if auto_wrap_for_ansi is not None:
file = auto_wrap_for_ansi(file)
elif not color:
message = strip_ansi(message)
if message:
file.write(message)
file.flush()
def get_binary_stream(name):
"""Returns a system stream for byte processing. This essentially
returns the stream from the sys module with the given name but it
solves some compatibility issues between different Python versions.
Primarily this function is necessary for getting binary streams on
Python 3.
:param name: the name of the stream to open. Valid names are ``'stdin'``,
``'stdout'`` and ``'stderr'``
"""
opener = binary_streams.get(name)
if opener is None:
raise TypeError('Unknown standard stream %r' % name)
return opener()
def get_text_stream(name, encoding=None, errors='strict'):
"""Returns a system stream for text processing. This usually returns
a wrapped stream around a binary stream returned from
:func:`get_binary_stream` but it also can take shortcuts on Python 3
for already correctly configured streams.
:param name: the name of the stream to open. Valid names are ``'stdin'``,
``'stdout'`` and ``'stderr'``
:param encoding: overrides the detected default encoding.
:param errors: overrides the default error mode.
"""
opener = text_streams.get(name)
if opener is None:
raise TypeError('Unknown standard stream %r' % name)
return opener(encoding, errors)
def open_file(filename, mode='r', encoding=None, errors='strict',
lazy=False, atomic=False):
"""This is similar to how the :class:`File` works but for manual
usage. Files are opened non lazy by default. This can open regular
files as well as stdin/stdout if ``'-'`` is passed.
If stdin/stdout is returned the stream is wrapped so that the context
manager will not close the stream accidentally. This makes it possible
to always use the function like this without having to worry to
accidentally close a standard stream::
with open_file(filename) as f:
...
.. versionadded:: 3.0
:param filename: the name of the file to open (or ``'-'`` for stdin/stdout).
:param mode: the mode in which to open the file.
:param encoding: the encoding to use.
:param errors: the error handling for this file.
:param lazy: can be flipped to true to open the file lazily.
:param atomic: in atomic mode writes go into a temporary file and it's
moved on close.
"""
if lazy:
return LazyFile(filename, mode, encoding, errors, atomic=atomic)
f, should_close = open_stream(filename, mode, encoding, errors,
atomic=atomic)
if not should_close:
f = KeepOpenFile(f)
return f
def get_os_args():
"""This returns the argument part of sys.argv in the most appropriate
form for processing. What this means is that this return value is in
a format that works for Click to process but does not necessarily
correspond well to what's actually standard for the interpreter.
On most environments the return value is ``sys.argv[:1]`` unchanged.
However if you are on Windows and running Python 2 the return value
will actually be a list of unicode strings instead because the
default behavior on that platform otherwise will not be able to
carry all possible values that sys.argv can have.
.. versionadded:: 6.0
"""
# We can only extract the unicode argv if sys.argv has not been
# changed since the startup of the application.
if PY2 and WIN and _initial_argv_hash == _hash_py_argv():
return _get_windows_argv()
return sys.argv[1:]
def format_filename(filename, shorten=False):
"""Formats a filename for user display. The main purpose of this
function is to ensure that the filename can be displayed at all. This
will decode the filename to unicode if necessary in a way that it will
not fail. Optionally, it can shorten the filename to not include the
full path to the filename.
:param filename: formats a filename for UI display. This will also convert
the filename into unicode without failing.
:param shorten: this optionally shortens the filename to strip of the
path that leads up to it.
"""
if shorten:
filename = os.path.basename(filename)
return filename_to_ui(filename)
def get_app_dir(app_name, roaming=True, force_posix=False):
r"""Returns the config folder for the application. The default behavior
is to return whatever is most appropriate for the operating system.
To give you an idea, for an app called ``"Foo Bar"``, something like
the following folders could be returned:
Mac OS X:
``~/Library/Application Support/Foo Bar``
Mac OS X (POSIX):
``~/.foo-bar``
Unix:
``~/.config/foo-bar``
Unix (POSIX):
``~/.foo-bar``
Win XP (roaming):
``C:\Documents and Settings\<user>\Local Settings\Application Data\Foo Bar``
Win XP (not roaming):
``C:\Documents and Settings\<user>\Application Data\Foo Bar``
Win 7 (roaming):
``C:\Users\<user>\AppData\Roaming\Foo Bar``
Win 7 (not roaming):
``C:\Users\<user>\AppData\Local\Foo Bar``
.. versionadded:: 2.0
:param app_name: the application name. This should be properly capitalized
and can contain whitespace.
:param roaming: controls if the folder should be roaming or not on Windows.
Has no affect otherwise.
:param force_posix: if this is set to `True` then on any POSIX system the
folder will be stored in the home folder with a leading
dot instead of the XDG config home or darwin's
application support folder.
"""
if WIN:
key = roaming and 'APPDATA' or 'LOCALAPPDATA'
folder = os.environ.get(key)
if folder is None:
folder = os.path.expanduser('~')
return os.path.join(folder, app_name)
if force_posix:
return os.path.join(os.path.expanduser('~/.' + _posixify(app_name)))
if sys.platform == 'darwin':
return os.path.join(os.path.expanduser(
'~/Library/Application Support'), app_name)
return os.path.join(
os.environ.get('XDG_CONFIG_HOME', os.path.expanduser('~/.config')),
_posixify(app_name))
+470
View File
@@ -0,0 +1,470 @@
#!/usr/bin/env python
# -*- coding:utf-8 -*-
from __future__ import print_function, absolute_import
import os
import platform
import re
import sys
import shlex
import subprocess
import click
import six
from click import echo, MultiCommand, Option, Argument, ParamType
__version__ = '0.2.1'
FISH_TEMPLATE = 'complete --command {{prog_name}} --arguments "(env {{complete_var}}=complete-fish COMMANDLINE=(commandline -cp){% for k, v in extra_env.items() %} {{k}}={{v}}{% endfor %} {{prog_name}})" -f'
ZSH_TEMPLATE = '''
#compdef {{prog_name}}
_{{prog_name}}() {
eval $(env COMMANDLINE="${words[1,$CURRENT]}" {{complete_var}}=complete-zsh {% for k, v in extra_env.items() %} {{k}}={{v}}{% endfor %} {{prog_name}})
}
if [[ "$(basename ${(%):-%x})" != "_{{prog_name}}" ]]; then
autoload -U compinit && compinit
compdef _{{prog_name}} {{prog_name}}
fi
'''
POWERSHELL_COMPLETION_SCRIPT = '''
if ((Test-Path Function:\TabExpansion) -and -not (Test-Path Function:\{{prog_name}}TabExpansionBackup)) {
Rename-Item Function:\TabExpansion {{prog_name}}TabExpansionBackup
}
function TabExpansion($line, $lastWord) {
$lastBlock = [regex]::Split($line, '[|;]')[-1].TrimStart()
$aliases = @("{{prog_name}}") + @(Get-Alias | where { $_.Definition -eq "{{prog_name}}" } | select -Exp Name)
$aliasPattern = "($($aliases -join '|'))"
if($lastBlock -match "^$aliasPattern ") {
$Env:{{complete_var}} = "complete-powershell"
$Env:COMMANDLINE = "$lastBlock"
{%- for k, v in extra_env.items() %}
$Env:{{k}} = "{{v}}"
{%- endfor %}
({{prog_name}}) | ? {$_.trim() -ne "" }
Remove-Item Env:{{complete_var}}
Remove-Item Env:COMMANDLINE
{%- for k in extra_env.keys() %}
Remove-Item $Env:{{k}}
{%- endfor %}
}
elseif (Test-Path Function:\{{prog_name}}TabExpansionBackup) {
# Fall back on existing tab expansion
{{prog_name}}TabExpansionBackup $line $lastWord
}
}
'''
BASH_COMPLETION_SCRIPT = '''
_{{prog_name}}_completion() {
local IFS=$'\\t'
COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \\
COMP_CWORD=$COMP_CWORD \\
{%- for k, v in extra_env.items() %}
{{k}}={{v}} \\
{%- endfor %}
{{complete_var}}=complete-bash $1 ) )
return 0
}
complete -F _{{prog_name}}_completion -o default {{prog_name}}
'''
_invalid_ident_char_re = re.compile(r'[^a-zA-Z0-9_]')
def resolve_ctx(cli, prog_name, args):
ctx = cli.make_context(prog_name, list(args), resilient_parsing=True)
while ctx.args + ctx.protected_args and isinstance(ctx.command, MultiCommand):
a = ctx.protected_args + ctx.args
cmd = ctx.command.get_command(ctx, a[0])
if cmd is None:
return None
ctx = cmd.make_context(a[0], a[1:], parent=ctx, resilient_parsing=True)
return ctx
def startswith(string, incomplete):
"""Returns True when string starts with incomplete
It might be overridden with a fuzzier version - for example a case insensitive version"""
return string.startswith(incomplete)
def get_choices(cli, prog_name, args, incomplete):
ctx = resolve_ctx(cli, prog_name, args)
if ctx is None:
return
optctx = None
if args:
for param in ctx.command.get_params(ctx):
if isinstance(param, Option) and not param.is_flag and args[-1] in param.opts + param.secondary_opts:
optctx = param
choices = []
if optctx:
choices += [c if isinstance(c, tuple) else (c, None) for c in optctx.type.complete(ctx, incomplete)]
elif incomplete and not incomplete[:1].isalnum():
for param in ctx.command.get_params(ctx):
if not isinstance(param, Option):
continue
for opt in param.opts:
if startswith(opt, incomplete):
choices.append((opt, param.help))
for opt in param.secondary_opts:
if startswith(opt, incomplete):
# don't put the doc so fish won't group the primary and
# and secondary options
choices.append((opt, None))
elif isinstance(ctx.command, MultiCommand):
for name in ctx.command.list_commands(ctx):
if startswith(name, incomplete):
choices.append((name, ctx.command.get_command_short_help(ctx, name)))
else:
for param in ctx.command.get_params(ctx):
if isinstance(param, Argument):
choices += [c if isinstance(c, tuple) else (c, None) for c in param.type.complete(ctx, incomplete)]
for item, help in choices:
yield (item, help)
def split_args(line):
"""Version of shlex.split that silently accept incomplete strings."""
lex = shlex.shlex(line, posix=True)
lex.whitespace_split = True
lex.commenters = ''
res = []
try:
while True:
res.append(next(lex))
except ValueError: # No closing quotation
pass
except StopIteration: # End of loop
pass
if lex.token:
res.append(lex.token)
return res
def decode_args(strings):
res = []
for s in strings:
s = split_args(s)
s = s[0] if s else ''
res.append(s)
return res
def do_bash_complete(cli, prog_name):
comp_words = os.environ['COMP_WORDS']
try:
cwords = shlex.split(comp_words)
quoted = False
except ValueError: # No closing quotation
cwords = split_args(comp_words)
quoted = True
cword = int(os.environ['COMP_CWORD'])
args = cwords[1:cword]
try:
incomplete = cwords[cword]
except IndexError:
incomplete = ''
choices = get_choices(cli, prog_name, args, incomplete)
if quoted:
echo('\t'.join(opt for opt, _ in choices), nl=False)
else:
echo('\t'.join(re.sub(r"""([\s\\"'])""", r'\\\1', opt) for opt, _ in choices), nl=False)
return True
def do_fish_complete(cli, prog_name):
commandline = os.environ['COMMANDLINE']
args = split_args(commandline)[1:]
if args and not commandline.endswith(' '):
incomplete = args[-1]
args = args[:-1]
else:
incomplete = ''
for item, help in get_choices(cli, prog_name, args, incomplete):
if help:
echo("%s\t%s" % (item, help))
else:
echo(item)
return True
def do_zsh_complete(cli, prog_name):
commandline = os.environ['COMMANDLINE']
args = split_args(commandline)[1:]
if args and not commandline.endswith(' '):
incomplete = args[-1]
args = args[:-1]
else:
incomplete = ''
def escape(s):
return s.replace('"', '""').replace("'", "''").replace('$', '\\$')
res = []
for item, help in get_choices(cli, prog_name, args, incomplete):
if help:
res.append('"%s"\:"%s"' % (escape(item), escape(help)))
else:
res.append('"%s"' % escape(item))
echo("_arguments '*: :((%s))'" % '\n'.join(res))
return True
def do_powershell_complete(cli, prog_name):
commandline = os.environ['COMMANDLINE']
args = split_args(commandline)[1:]
quote = single_quote
incomplete = ''
if args and not commandline.endswith(' '):
incomplete = args[-1]
args = args[:-1]
quote_pos = commandline.rfind(incomplete) - 1
if quote_pos >= 0 and commandline[quote_pos] == '"':
quote = double_quote
for item, help in get_choices(cli, prog_name, args, incomplete):
echo(quote(item))
return True
find_unsafe = re.compile(r'[^\w@%+=:,./-]').search
def single_quote(s):
"""Return a shell-escaped version of the string *s*."""
if not s:
return "''"
if find_unsafe(s) is None:
return s
# use single quotes, and put single quotes into double quotes
# the string $'b is then quoted as '$'"'"'b'
return "'" + s.replace("'", "'\"'\"'") + "'"
def double_quote(s):
'''Return a shell-escaped version of the string *s*.'''
if not s:
return '""'
if find_unsafe(s) is None:
return s
# use double quotes, and put double quotes into single quotes
# the string $"b is then quoted as "$"'"'"b"
return '"' + s.replace('"', '"\'"\'"') + '"'
# extend click completion features
def param_type_complete(self, ctx, incomplete):
return []
def choice_complete(self, ctx, incomplete):
return [c for c in self.choices if c.startswith(incomplete)]
def multicommand_get_command_short_help(self, ctx, cmd_name):
return self.get_command(ctx, cmd_name).short_help
def _shellcomplete(cli, prog_name, complete_var=None):
"""Internal handler for the bash completion support."""
if complete_var is None:
complete_var = '_%s_COMPLETE' % (prog_name.replace('-', '_')).upper()
complete_instr = os.environ.get(complete_var)
if not complete_instr:
return
if complete_instr == 'source':
echo(get_code(prog_name=prog_name, env_name=complete_var))
elif complete_instr == 'source-bash':
echo(get_code('bash', prog_name, complete_var))
elif complete_instr == 'source-fish':
echo(get_code('fish', prog_name, complete_var))
elif complete_instr == 'source-powershell':
echo(get_code('powershell', prog_name, complete_var))
elif complete_instr == 'source-zsh':
echo(get_code('zsh', prog_name, complete_var))
elif complete_instr in ['complete', 'complete-bash']:
# keep 'complete' for bash for backward compatibility
do_bash_complete(cli, prog_name)
elif complete_instr == 'complete-fish':
do_fish_complete(cli, prog_name)
elif complete_instr == 'complete-powershell':
do_powershell_complete(cli, prog_name)
elif complete_instr == 'complete-zsh':
do_zsh_complete(cli, prog_name)
elif complete_instr == 'install':
shell, path = install(prog_name=prog_name, env_name=complete_var)
click.echo('%s completion installed in %s' % (shell, path))
elif complete_instr == 'install-bash':
shell, path = install(shell='bash', prog_name=prog_name, env_name=complete_var)
click.echo('%s completion installed in %s' % (shell, path))
elif complete_instr == 'install-fish':
shell, path = install(shell='fish', prog_name=prog_name, env_name=complete_var)
click.echo('%s completion installed in %s' % (shell, path))
elif complete_instr == 'install-zsh':
shell, path = install(shell='zsh', prog_name=prog_name, env_name=complete_var)
click.echo('%s completion installed in %s' % (shell, path))
elif complete_instr == 'install-powershell':
shell, path = install(shell='powershell', prog_name=prog_name, env_name=complete_var)
click.echo('%s completion installed in %s' % (shell, path))
sys.exit()
def init():
"""patch click to support enhanced completion"""
import click
click.types.ParamType.complete = param_type_complete
click.types.Choice.complete = choice_complete
click.core.MultiCommand.get_command_short_help = multicommand_get_command_short_help
click.core._bashcomplete = _shellcomplete
class DocumentedChoice(ParamType):
"""The choice type allows a value to be checked against a fixed set of
supported values. All of these values have to be strings. Each value may
be associated to a help message that will be display in the error message
and during the completion.
"""
name = 'choice'
def __init__(self, choices):
self.choices = dict(choices)
def get_metavar(self, param):
return '[%s]' % '|'.join(self.choices.keys())
def get_missing_message(self, param):
formated_choices = ['{:<12} {}'.format(k, self.choices[k] or '') for k in sorted(self.choices.keys())]
return 'Choose from\n ' + '\n '.join(formated_choices)
def convert(self, value, param, ctx):
# Exact match
if value in self.choices:
return value
# Match through normalization
if ctx is not None and \
ctx.token_normalize_func is not None:
value = ctx.token_normalize_func(value)
for choice in self.choices:
if ctx.token_normalize_func(choice) == value:
return choice
self.fail('invalid choice: %s. %s' %
(value, self.get_missing_message(param)), param, ctx)
def __repr__(self):
return 'DocumentedChoice(%r)' % list(self.choices.keys())
def complete(self, ctx, incomplete):
return [(c, v) for c, v in six.iteritems(self.choices) if startswith(c, incomplete)]
def get_code(shell=None, prog_name=None, env_name=None, extra_env=None):
"""Return the specified completion code"""
from jinja2 import Template
if shell in [None, 'auto']:
shell = get_auto_shell()
prog_name = prog_name or click.get_current_context().find_root().info_name
env_name = env_name or '_%s_COMPLETE' % prog_name.upper().replace('-', '_')
extra_env = extra_env if extra_env else {}
if shell == 'fish':
return Template(FISH_TEMPLATE).render(prog_name=prog_name, complete_var=env_name, extra_env=extra_env)
elif shell == 'bash':
return Template(BASH_COMPLETION_SCRIPT).render(prog_name=prog_name, complete_var=env_name, extra_env=extra_env)
elif shell == 'zsh':
return Template(ZSH_TEMPLATE).render(prog_name=prog_name, complete_var=env_name, extra_env=extra_env)
elif shell == 'powershell':
return Template(POWERSHELL_COMPLETION_SCRIPT).render(prog_name=prog_name, complete_var=env_name, extra_env=extra_env)
else:
raise click.ClickException('%s is not supported.' % shell)
def get_auto_shell():
"""Return the shell that is calling this process"""
try:
import psutil
parent = psutil.Process(os.getpid()).parent()
if platform.system() == 'Windows':
parent = parent.parent() or parent
return parent.name().replace('.exe', '')
except ImportError:
raise click.UsageError("Please explicitly give the shell type or install the psutil package to activate the"
" automatic shell detection.")
def install(shell=None, prog_name=None, env_name=None, path=None, append=None, extra_env=None):
"""Install the completion"""
prog_name = prog_name or click.get_current_context().find_root().info_name
shell = shell or get_auto_shell()
if append is None and path is not None:
append = True
if append is not None:
mode = 'a' if append else 'w'
else:
mode = None
if shell == 'fish':
path = path or os.path.expanduser('~') + '/.config/fish/completions/%s.fish' % prog_name
mode = mode or 'w'
elif shell == 'bash':
path = path or os.path.expanduser('~') + '/.bash_completion'
mode = mode or 'a'
elif shell == 'zsh':
ohmyzsh = os.path.expanduser('~') + '/.oh-my-zsh'
if os.path.exists(ohmyzsh):
path = path or ohmyzsh + '/completions/_%s' % prog_name
mode = mode or 'w'
else:
path = path or os.path.expanduser('~') + '/.zshrc'
mode = mode or 'a'
elif shell == 'powershell':
subprocess.check_call(['powershell', 'Set-ExecutionPolicy Unrestricted -Scope CurrentUser'])
path = path or subprocess.check_output(['powershell', '-NoProfile', 'echo $profile']).strip() if install else ''
mode = mode or 'a'
else:
raise click.ClickException('%s is not supported.' % shell)
if append is not None:
mode = 'a' if append else 'w'
else:
mode = mode
d = os.path.dirname(path)
if not os.path.exists(d):
os.makedirs(d)
f = open(path, mode)
f.write(get_code(shell, prog_name, env_name, extra_env))
f.write("\n")
f.close()
return shell, path
shells = {
'bash': 'Bourne again shell',
'fish': 'Friendly interactive shell',
'zsh': 'Z shell',
'powershell': 'Windows PowerShell'
}
+7
View File
@@ -0,0 +1,7 @@
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
from .initialise import init, deinit, reinit, colorama_text
from .ansi import Fore, Back, Style, Cursor
from .ansitowin32 import AnsiToWin32
__version__ = '0.3.7'
+102
View File
@@ -0,0 +1,102 @@
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
'''
This module generates ANSI character codes to printing colors to terminals.
See: http://en.wikipedia.org/wiki/ANSI_escape_code
'''
CSI = '\033['
OSC = '\033]'
BEL = '\007'
def code_to_chars(code):
return CSI + str(code) + 'm'
def set_title(title):
return OSC + '2;' + title + BEL
def clear_screen(mode=2):
return CSI + str(mode) + 'J'
def clear_line(mode=2):
return CSI + str(mode) + 'K'
class AnsiCodes(object):
def __init__(self):
# the subclasses declare class attributes which are numbers.
# Upon instantiation we define instance attributes, which are the same
# as the class attributes but wrapped with the ANSI escape sequence
for name in dir(self):
if not name.startswith('_'):
value = getattr(self, name)
setattr(self, name, code_to_chars(value))
class AnsiCursor(object):
def UP(self, n=1):
return CSI + str(n) + 'A'
def DOWN(self, n=1):
return CSI + str(n) + 'B'
def FORWARD(self, n=1):
return CSI + str(n) + 'C'
def BACK(self, n=1):
return CSI + str(n) + 'D'
def POS(self, x=1, y=1):
return CSI + str(y) + ';' + str(x) + 'H'
class AnsiFore(AnsiCodes):
BLACK = 30
RED = 31
GREEN = 32
YELLOW = 33
BLUE = 34
MAGENTA = 35
CYAN = 36
WHITE = 37
RESET = 39
# These are fairly well supported, but not part of the standard.
LIGHTBLACK_EX = 90
LIGHTRED_EX = 91
LIGHTGREEN_EX = 92
LIGHTYELLOW_EX = 93
LIGHTBLUE_EX = 94
LIGHTMAGENTA_EX = 95
LIGHTCYAN_EX = 96
LIGHTWHITE_EX = 97
class AnsiBack(AnsiCodes):
BLACK = 40
RED = 41
GREEN = 42
YELLOW = 43
BLUE = 44
MAGENTA = 45
CYAN = 46
WHITE = 47
RESET = 49
# These are fairly well supported, but not part of the standard.
LIGHTBLACK_EX = 100
LIGHTRED_EX = 101
LIGHTGREEN_EX = 102
LIGHTYELLOW_EX = 103
LIGHTBLUE_EX = 104
LIGHTMAGENTA_EX = 105
LIGHTCYAN_EX = 106
LIGHTWHITE_EX = 107
class AnsiStyle(AnsiCodes):
BRIGHT = 1
DIM = 2
NORMAL = 22
RESET_ALL = 0
Fore = AnsiFore()
Back = AnsiBack()
Style = AnsiStyle()
Cursor = AnsiCursor()
+236
View File
@@ -0,0 +1,236 @@
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
import re
import sys
import os
from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style
from .winterm import WinTerm, WinColor, WinStyle
from .win32 import windll, winapi_test
winterm = None
if windll is not None:
winterm = WinTerm()
def is_stream_closed(stream):
return not hasattr(stream, 'closed') or stream.closed
def is_a_tty(stream):
return hasattr(stream, 'isatty') and stream.isatty()
class StreamWrapper(object):
'''
Wraps a stream (such as stdout), acting as a transparent proxy for all
attribute access apart from method 'write()', which is delegated to our
Converter instance.
'''
def __init__(self, wrapped, converter):
# double-underscore everything to prevent clashes with names of
# attributes on the wrapped stream object.
self.__wrapped = wrapped
self.__convertor = converter
def __getattr__(self, name):
return getattr(self.__wrapped, name)
def write(self, text):
self.__convertor.write(text)
class AnsiToWin32(object):
'''
Implements a 'write()' method which, on Windows, will strip ANSI character
sequences from the text, and if outputting to a tty, will convert them into
win32 function calls.
'''
ANSI_CSI_RE = re.compile('\001?\033\[((?:\d|;)*)([a-zA-Z])\002?') # Control Sequence Introducer
ANSI_OSC_RE = re.compile('\001?\033\]((?:.|;)*?)(\x07)\002?') # Operating System Command
def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
# The wrapped stream (normally sys.stdout or sys.stderr)
self.wrapped = wrapped
# should we reset colors to defaults after every .write()
self.autoreset = autoreset
# create the proxy wrapping our output stream
self.stream = StreamWrapper(wrapped, self)
on_windows = os.name == 'nt'
# We test if the WinAPI works, because even if we are on Windows
# we may be using a terminal that doesn't support the WinAPI
# (e.g. Cygwin Terminal). In this case it's up to the terminal
# to support the ANSI codes.
conversion_supported = on_windows and winapi_test()
# should we strip ANSI sequences from our output?
if strip is None:
strip = conversion_supported or (not is_stream_closed(wrapped) and not is_a_tty(wrapped))
self.strip = strip
# should we should convert ANSI sequences into win32 calls?
if convert is None:
convert = conversion_supported and not is_stream_closed(wrapped) and is_a_tty(wrapped)
self.convert = convert
# dict of ansi codes to win32 functions and parameters
self.win32_calls = self.get_win32_calls()
# are we wrapping stderr?
self.on_stderr = self.wrapped is sys.stderr
def should_wrap(self):
'''
True if this class is actually needed. If false, then the output
stream will not be affected, nor will win32 calls be issued, so
wrapping stdout is not actually required. This will generally be
False on non-Windows platforms, unless optional functionality like
autoreset has been requested using kwargs to init()
'''
return self.convert or self.strip or self.autoreset
def get_win32_calls(self):
if self.convert and winterm:
return {
AnsiStyle.RESET_ALL: (winterm.reset_all, ),
AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT),
AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL),
AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL),
AnsiFore.BLACK: (winterm.fore, WinColor.BLACK),
AnsiFore.RED: (winterm.fore, WinColor.RED),
AnsiFore.GREEN: (winterm.fore, WinColor.GREEN),
AnsiFore.YELLOW: (winterm.fore, WinColor.YELLOW),
AnsiFore.BLUE: (winterm.fore, WinColor.BLUE),
AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA),
AnsiFore.CYAN: (winterm.fore, WinColor.CYAN),
AnsiFore.WHITE: (winterm.fore, WinColor.GREY),
AnsiFore.RESET: (winterm.fore, ),
AnsiFore.LIGHTBLACK_EX: (winterm.fore, WinColor.BLACK, True),
AnsiFore.LIGHTRED_EX: (winterm.fore, WinColor.RED, True),
AnsiFore.LIGHTGREEN_EX: (winterm.fore, WinColor.GREEN, True),
AnsiFore.LIGHTYELLOW_EX: (winterm.fore, WinColor.YELLOW, True),
AnsiFore.LIGHTBLUE_EX: (winterm.fore, WinColor.BLUE, True),
AnsiFore.LIGHTMAGENTA_EX: (winterm.fore, WinColor.MAGENTA, True),
AnsiFore.LIGHTCYAN_EX: (winterm.fore, WinColor.CYAN, True),
AnsiFore.LIGHTWHITE_EX: (winterm.fore, WinColor.GREY, True),
AnsiBack.BLACK: (winterm.back, WinColor.BLACK),
AnsiBack.RED: (winterm.back, WinColor.RED),
AnsiBack.GREEN: (winterm.back, WinColor.GREEN),
AnsiBack.YELLOW: (winterm.back, WinColor.YELLOW),
AnsiBack.BLUE: (winterm.back, WinColor.BLUE),
AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA),
AnsiBack.CYAN: (winterm.back, WinColor.CYAN),
AnsiBack.WHITE: (winterm.back, WinColor.GREY),
AnsiBack.RESET: (winterm.back, ),
AnsiBack.LIGHTBLACK_EX: (winterm.back, WinColor.BLACK, True),
AnsiBack.LIGHTRED_EX: (winterm.back, WinColor.RED, True),
AnsiBack.LIGHTGREEN_EX: (winterm.back, WinColor.GREEN, True),
AnsiBack.LIGHTYELLOW_EX: (winterm.back, WinColor.YELLOW, True),
AnsiBack.LIGHTBLUE_EX: (winterm.back, WinColor.BLUE, True),
AnsiBack.LIGHTMAGENTA_EX: (winterm.back, WinColor.MAGENTA, True),
AnsiBack.LIGHTCYAN_EX: (winterm.back, WinColor.CYAN, True),
AnsiBack.LIGHTWHITE_EX: (winterm.back, WinColor.GREY, True),
}
return dict()
def write(self, text):
if self.strip or self.convert:
self.write_and_convert(text)
else:
self.wrapped.write(text)
self.wrapped.flush()
if self.autoreset:
self.reset_all()
def reset_all(self):
if self.convert:
self.call_win32('m', (0,))
elif not self.strip and not is_stream_closed(self.wrapped):
self.wrapped.write(Style.RESET_ALL)
def write_and_convert(self, text):
'''
Write the given text to our wrapped stream, stripping any ANSI
sequences from the text, and optionally converting them into win32
calls.
'''
cursor = 0
text = self.convert_osc(text)
for match in self.ANSI_CSI_RE.finditer(text):
start, end = match.span()
self.write_plain_text(text, cursor, start)
self.convert_ansi(*match.groups())
cursor = end
self.write_plain_text(text, cursor, len(text))
def write_plain_text(self, text, start, end):
if start < end:
self.wrapped.write(text[start:end])
self.wrapped.flush()
def convert_ansi(self, paramstring, command):
if self.convert:
params = self.extract_params(command, paramstring)
self.call_win32(command, params)
def extract_params(self, command, paramstring):
if command in 'Hf':
params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(';'))
while len(params) < 2:
# defaults:
params = params + (1,)
else:
params = tuple(int(p) for p in paramstring.split(';') if len(p) != 0)
if len(params) == 0:
# defaults:
if command in 'JKm':
params = (0,)
elif command in 'ABCD':
params = (1,)
return params
def call_win32(self, command, params):
if command == 'm':
for param in params:
if param in self.win32_calls:
func_args = self.win32_calls[param]
func = func_args[0]
args = func_args[1:]
kwargs = dict(on_stderr=self.on_stderr)
func(*args, **kwargs)
elif command in 'J':
winterm.erase_screen(params[0], on_stderr=self.on_stderr)
elif command in 'K':
winterm.erase_line(params[0], on_stderr=self.on_stderr)
elif command in 'Hf': # cursor position - absolute
winterm.set_cursor_position(params, on_stderr=self.on_stderr)
elif command in 'ABCD': # cursor position - relative
n = params[0]
# A - up, B - down, C - forward, D - back
x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command]
winterm.cursor_adjust(x, y, on_stderr=self.on_stderr)
def convert_osc(self, text):
for match in self.ANSI_OSC_RE.finditer(text):
start, end = match.span()
text = text[:start] + text[end:]
paramstring, command = match.groups()
if command in '\x07': # \x07 = BEL
params = paramstring.split(";")
# 0 - change title and icon (we will only change title)
# 1 - change icon (we don't support this)
# 2 - change title
if params[0] in '02':
winterm.set_title(params[1])
return text
+82
View File
@@ -0,0 +1,82 @@
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
import atexit
import contextlib
import sys
from .ansitowin32 import AnsiToWin32
orig_stdout = None
orig_stderr = None
wrapped_stdout = None
wrapped_stderr = None
atexit_done = False
def reset_all():
if AnsiToWin32 is not None: # Issue #74: objects might become None at exit
AnsiToWin32(orig_stdout).reset_all()
def init(autoreset=False, convert=None, strip=None, wrap=True):
if not wrap and any([autoreset, convert, strip]):
raise ValueError('wrap=False conflicts with any other arg=True')
global wrapped_stdout, wrapped_stderr
global orig_stdout, orig_stderr
orig_stdout = sys.stdout
orig_stderr = sys.stderr
if sys.stdout is None:
wrapped_stdout = None
else:
sys.stdout = wrapped_stdout = \
wrap_stream(orig_stdout, convert, strip, autoreset, wrap)
if sys.stderr is None:
wrapped_stderr = None
else:
sys.stderr = wrapped_stderr = \
wrap_stream(orig_stderr, convert, strip, autoreset, wrap)
global atexit_done
if not atexit_done:
atexit.register(reset_all)
atexit_done = True
def deinit():
if orig_stdout is not None:
sys.stdout = orig_stdout
if orig_stderr is not None:
sys.stderr = orig_stderr
@contextlib.contextmanager
def colorama_text(*args, **kwargs):
init(*args, **kwargs)
try:
yield
finally:
deinit()
def reinit():
if wrapped_stdout is not None:
sys.stdout = wrapped_stdout
if wrapped_stderr is not None:
sys.stderr = wrapped_stderr
def wrap_stream(stream, convert, strip, autoreset, wrap):
if wrap:
wrapper = AnsiToWin32(stream,
convert=convert, strip=strip, autoreset=autoreset)
if wrapper.should_wrap():
stream = wrapper.stream
return stream
+154
View File
@@ -0,0 +1,154 @@
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
# from winbase.h
STDOUT = -11
STDERR = -12
try:
import ctypes
from ctypes import LibraryLoader
windll = LibraryLoader(ctypes.WinDLL)
from ctypes import wintypes
except (AttributeError, ImportError):
windll = None
SetConsoleTextAttribute = lambda *_: None
winapi_test = lambda *_: None
else:
from ctypes import byref, Structure, c_char, POINTER
COORD = wintypes._COORD
class CONSOLE_SCREEN_BUFFER_INFO(Structure):
"""struct in wincon.h."""
_fields_ = [
("dwSize", COORD),
("dwCursorPosition", COORD),
("wAttributes", wintypes.WORD),
("srWindow", wintypes.SMALL_RECT),
("dwMaximumWindowSize", COORD),
]
def __str__(self):
return '(%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d)' % (
self.dwSize.Y, self.dwSize.X
, self.dwCursorPosition.Y, self.dwCursorPosition.X
, self.wAttributes
, self.srWindow.Top, self.srWindow.Left, self.srWindow.Bottom, self.srWindow.Right
, self.dwMaximumWindowSize.Y, self.dwMaximumWindowSize.X
)
_GetStdHandle = windll.kernel32.GetStdHandle
_GetStdHandle.argtypes = [
wintypes.DWORD,
]
_GetStdHandle.restype = wintypes.HANDLE
_GetConsoleScreenBufferInfo = windll.kernel32.GetConsoleScreenBufferInfo
_GetConsoleScreenBufferInfo.argtypes = [
wintypes.HANDLE,
POINTER(CONSOLE_SCREEN_BUFFER_INFO),
]
_GetConsoleScreenBufferInfo.restype = wintypes.BOOL
_SetConsoleTextAttribute = windll.kernel32.SetConsoleTextAttribute
_SetConsoleTextAttribute.argtypes = [
wintypes.HANDLE,
wintypes.WORD,
]
_SetConsoleTextAttribute.restype = wintypes.BOOL
_SetConsoleCursorPosition = windll.kernel32.SetConsoleCursorPosition
_SetConsoleCursorPosition.argtypes = [
wintypes.HANDLE,
COORD,
]
_SetConsoleCursorPosition.restype = wintypes.BOOL
_FillConsoleOutputCharacterA = windll.kernel32.FillConsoleOutputCharacterA
_FillConsoleOutputCharacterA.argtypes = [
wintypes.HANDLE,
c_char,
wintypes.DWORD,
COORD,
POINTER(wintypes.DWORD),
]
_FillConsoleOutputCharacterA.restype = wintypes.BOOL
_FillConsoleOutputAttribute = windll.kernel32.FillConsoleOutputAttribute
_FillConsoleOutputAttribute.argtypes = [
wintypes.HANDLE,
wintypes.WORD,
wintypes.DWORD,
COORD,
POINTER(wintypes.DWORD),
]
_FillConsoleOutputAttribute.restype = wintypes.BOOL
_SetConsoleTitleW = windll.kernel32.SetConsoleTitleA
_SetConsoleTitleW.argtypes = [
wintypes.LPCSTR
]
_SetConsoleTitleW.restype = wintypes.BOOL
handles = {
STDOUT: _GetStdHandle(STDOUT),
STDERR: _GetStdHandle(STDERR),
}
def winapi_test():
handle = handles[STDOUT]
csbi = CONSOLE_SCREEN_BUFFER_INFO()
success = _GetConsoleScreenBufferInfo(
handle, byref(csbi))
return bool(success)
def GetConsoleScreenBufferInfo(stream_id=STDOUT):
handle = handles[stream_id]
csbi = CONSOLE_SCREEN_BUFFER_INFO()
success = _GetConsoleScreenBufferInfo(
handle, byref(csbi))
return csbi
def SetConsoleTextAttribute(stream_id, attrs):
handle = handles[stream_id]
return _SetConsoleTextAttribute(handle, attrs)
def SetConsoleCursorPosition(stream_id, position, adjust=True):
position = COORD(*position)
# If the position is out of range, do nothing.
if position.Y <= 0 or position.X <= 0:
return
# Adjust for Windows' SetConsoleCursorPosition:
# 1. being 0-based, while ANSI is 1-based.
# 2. expecting (x,y), while ANSI uses (y,x).
adjusted_position = COORD(position.Y - 1, position.X - 1)
if adjust:
# Adjust for viewport's scroll position
sr = GetConsoleScreenBufferInfo(STDOUT).srWindow
adjusted_position.Y += sr.Top
adjusted_position.X += sr.Left
# Resume normal processing
handle = handles[stream_id]
return _SetConsoleCursorPosition(handle, adjusted_position)
def FillConsoleOutputCharacter(stream_id, char, length, start):
handle = handles[stream_id]
char = c_char(char.encode())
length = wintypes.DWORD(length)
num_written = wintypes.DWORD(0)
# Note that this is hard-coded for ANSI (vs wide) bytes.
success = _FillConsoleOutputCharacterA(
handle, char, length, start, byref(num_written))
return num_written.value
def FillConsoleOutputAttribute(stream_id, attr, length, start):
''' FillConsoleOutputAttribute( hConsole, csbi.wAttributes, dwConSize, coordScreen, &cCharsWritten )'''
handle = handles[stream_id]
attribute = wintypes.WORD(attr)
length = wintypes.DWORD(length)
num_written = wintypes.DWORD(0)
# Note that this is hard-coded for ANSI (vs wide) bytes.
return _FillConsoleOutputAttribute(
handle, attribute, length, start, byref(num_written))
def SetConsoleTitle(title):
return _SetConsoleTitleW(title)
+162
View File
@@ -0,0 +1,162 @@
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
from . import win32
# from wincon.h
class WinColor(object):
BLACK = 0
BLUE = 1
GREEN = 2
CYAN = 3
RED = 4
MAGENTA = 5
YELLOW = 6
GREY = 7
# from wincon.h
class WinStyle(object):
NORMAL = 0x00 # dim text, dim background
BRIGHT = 0x08 # bright text, dim background
BRIGHT_BACKGROUND = 0x80 # dim text, bright background
class WinTerm(object):
def __init__(self):
self._default = win32.GetConsoleScreenBufferInfo(win32.STDOUT).wAttributes
self.set_attrs(self._default)
self._default_fore = self._fore
self._default_back = self._back
self._default_style = self._style
# In order to emulate LIGHT_EX in windows, we borrow the BRIGHT style.
# So that LIGHT_EX colors and BRIGHT style do not clobber each other,
# we track them separately, since LIGHT_EX is overwritten by Fore/Back
# and BRIGHT is overwritten by Style codes.
self._light = 0
def get_attrs(self):
return self._fore + self._back * 16 + (self._style | self._light)
def set_attrs(self, value):
self._fore = value & 7
self._back = (value >> 4) & 7
self._style = value & (WinStyle.BRIGHT | WinStyle.BRIGHT_BACKGROUND)
def reset_all(self, on_stderr=None):
self.set_attrs(self._default)
self.set_console(attrs=self._default)
def fore(self, fore=None, light=False, on_stderr=False):
if fore is None:
fore = self._default_fore
self._fore = fore
# Emulate LIGHT_EX with BRIGHT Style
if light:
self._light |= WinStyle.BRIGHT
else:
self._light &= ~WinStyle.BRIGHT
self.set_console(on_stderr=on_stderr)
def back(self, back=None, light=False, on_stderr=False):
if back is None:
back = self._default_back
self._back = back
# Emulate LIGHT_EX with BRIGHT_BACKGROUND Style
if light:
self._light |= WinStyle.BRIGHT_BACKGROUND
else:
self._light &= ~WinStyle.BRIGHT_BACKGROUND
self.set_console(on_stderr=on_stderr)
def style(self, style=None, on_stderr=False):
if style is None:
style = self._default_style
self._style = style
self.set_console(on_stderr=on_stderr)
def set_console(self, attrs=None, on_stderr=False):
if attrs is None:
attrs = self.get_attrs()
handle = win32.STDOUT
if on_stderr:
handle = win32.STDERR
win32.SetConsoleTextAttribute(handle, attrs)
def get_position(self, handle):
position = win32.GetConsoleScreenBufferInfo(handle).dwCursorPosition
# Because Windows coordinates are 0-based,
# and win32.SetConsoleCursorPosition expects 1-based.
position.X += 1
position.Y += 1
return position
def set_cursor_position(self, position=None, on_stderr=False):
if position is None:
# I'm not currently tracking the position, so there is no default.
# position = self.get_position()
return
handle = win32.STDOUT
if on_stderr:
handle = win32.STDERR
win32.SetConsoleCursorPosition(handle, position)
def cursor_adjust(self, x, y, on_stderr=False):
handle = win32.STDOUT
if on_stderr:
handle = win32.STDERR
position = self.get_position(handle)
adjusted_position = (position.Y + y, position.X + x)
win32.SetConsoleCursorPosition(handle, adjusted_position, adjust=False)
def erase_screen(self, mode=0, on_stderr=False):
# 0 should clear from the cursor to the end of the screen.
# 1 should clear from the cursor to the beginning of the screen.
# 2 should clear the entire screen, and move cursor to (1,1)
handle = win32.STDOUT
if on_stderr:
handle = win32.STDERR
csbi = win32.GetConsoleScreenBufferInfo(handle)
# get the number of character cells in the current buffer
cells_in_screen = csbi.dwSize.X * csbi.dwSize.Y
# get number of character cells before current cursor position
cells_before_cursor = csbi.dwSize.X * csbi.dwCursorPosition.Y + csbi.dwCursorPosition.X
if mode == 0:
from_coord = csbi.dwCursorPosition
cells_to_erase = cells_in_screen - cells_before_cursor
if mode == 1:
from_coord = win32.COORD(0, 0)
cells_to_erase = cells_before_cursor
elif mode == 2:
from_coord = win32.COORD(0, 0)
cells_to_erase = cells_in_screen
# fill the entire screen with blanks
win32.FillConsoleOutputCharacter(handle, ' ', cells_to_erase, from_coord)
# now set the buffer's attributes accordingly
win32.FillConsoleOutputAttribute(handle, self.get_attrs(), cells_to_erase, from_coord)
if mode == 2:
# put the cursor where needed
win32.SetConsoleCursorPosition(handle, (1, 1))
def erase_line(self, mode=0, on_stderr=False):
# 0 should clear from the cursor to the end of the line.
# 1 should clear from the cursor to the beginning of the line.
# 2 should clear the entire line.
handle = win32.STDOUT
if on_stderr:
handle = win32.STDERR
csbi = win32.GetConsoleScreenBufferInfo(handle)
if mode == 0:
from_coord = csbi.dwCursorPosition
cells_to_erase = csbi.dwSize.X - csbi.dwCursorPosition.X
if mode == 1:
from_coord = win32.COORD(0, csbi.dwCursorPosition.Y)
cells_to_erase = csbi.dwCursorPosition.X
elif mode == 2:
from_coord = win32.COORD(0, csbi.dwCursorPosition.Y)
cells_to_erase = csbi.dwSize.X
# fill the entire screen with blanks
win32.FillConsoleOutputCharacter(handle, ' ', cells_to_erase, from_coord)
# now set the buffer's attributes accordingly
win32.FillConsoleOutputAttribute(handle, self.get_attrs(), cells_to_erase, from_coord)
def set_title(self, title):
win32.SetConsoleTitle(title)
+152
View File
@@ -0,0 +1,152 @@
# -*- coding: utf-8 -*-
"""
clint.colored
~~~~~~~~~~~~~
This module provides a simple and elegant wrapper for colorama.
"""
import os
import re
import sys
PY3 = sys.version_info[0] >= 3
import colorama
__all__ = (
'red', 'green', 'yellow', 'blue',
'black', 'magenta', 'cyan', 'white',
'clean', 'disable'
)
COLORS = __all__[:-2]
if 'get_ipython' in dir():
"""
when ipython is fired lot of variables like _oh, etc are used.
There are so many ways to find current python interpreter is ipython.
get_ipython is easiest is most appealing for readers to understand.
"""
DISABLE_COLOR = True
else:
DISABLE_COLOR = False
class ColoredString(object):
"""Enhanced string for __len__ operations on Colored output."""
def __init__(self, color, s, always_color=False, bold=False):
super(ColoredString, self).__init__()
if not PY3 and isinstance(s, unicode):
self.s = s.encode('utf-8')
else:
self.s = s
self.color = color
self.always_color = always_color
self.bold = bold
if os.environ.get('CLINT_FORCE_COLOR'):
self.always_color = True
def __getattr__(self, att):
def func_help(*args, **kwargs):
result = getattr(self.s, att)(*args, **kwargs)
try:
is_result_string = isinstance(result, basestring)
except NameError:
is_result_string = isinstance(result, str)
if is_result_string:
return self._new(result)
elif isinstance(result, list):
return [self._new(x) for x in result]
else:
return result
return func_help
@property
def color_str(self):
style = 'BRIGHT' if self.bold else 'NORMAL'
c = '%s%s%s%s%s' % (getattr(colorama.Fore, self.color), getattr(colorama.Style, style), self.s, colorama.Fore.RESET, getattr(colorama.Style, 'NORMAL'))
if self.always_color:
return c
elif sys.stdout.isatty() and not DISABLE_COLOR:
return c
else:
return self.s
def __len__(self):
return len(self.s)
def __repr__(self):
return "<%s-string: '%s'>" % (self.color, self.s)
def __unicode__(self):
value = self.color_str
if isinstance(value, bytes):
return value.decode('utf8')
return value
if PY3:
__str__ = __unicode__
else:
def __str__(self):
return self.color_str
def __iter__(self):
return iter(self.color_str)
def __add__(self, other):
return str(self.color_str) + str(other)
def __radd__(self, other):
return str(other) + str(self.color_str)
def __mul__(self, other):
return (self.color_str * other)
def _new(self, s):
return ColoredString(self.color, s)
def clean(s):
strip = re.compile("([^-_a-zA-Z0-9!@#%&=,/'\";:~`\$\^\*\(\)\+\[\]\.\{\}\|\?\<\>\\]+|[^\s]+)")
txt = strip.sub('', str(s))
strip = re.compile(r'\[\d+m')
txt = strip.sub('', txt)
return txt
def black(string, always=False, bold=False):
return ColoredString('BLACK', string, always_color=always, bold=bold)
def red(string, always=False, bold=False):
return ColoredString('RED', string, always_color=always, bold=bold)
def green(string, always=False, bold=False):
return ColoredString('GREEN', string, always_color=always, bold=bold)
def yellow(string, always=False, bold=False):
return ColoredString('YELLOW', string, always_color=always, bold=bold)
def blue(string, always=False, bold=False):
return ColoredString('BLUE', string, always_color=always, bold=bold)
def magenta(string, always=False, bold=False):
return ColoredString('MAGENTA', string, always_color=always, bold=bold)
def cyan(string, always=False, bold=False):
return ColoredString('CYAN', string, always_color=always, bold=bold)
def white(string, always=False, bold=False):
return ColoredString('WHITE', string, always_color=always, bold=bold)
def disable():
"""Disables colors."""
global DISABLE_COLOR
DISABLE_COLOR = True
+220
View File
@@ -0,0 +1,220 @@
import os
import subprocess
import shlex
from pexpect.popen_spawn import PopenSpawn
# Enable Python subprocesses to work with expect functionality.
os.environ['PYTHONUNBUFFERED'] = '1'
class Command(object):
def __init__(self, cmd):
super(Command, self).__init__()
self.cmd = cmd
self.subprocess = None
self.blocking = None
self.was_run = False
self.__out = None
def __repr__(self):
return '<Commmand {!r}>'.format(self.cmd)
@property
def _popen_args(self):
return self.cmd
@property
def _default_popen_kwargs(self):
return {
'env': os.environ.copy(),
'stdin': subprocess.PIPE,
'stdout': subprocess.PIPE,
'stderr': subprocess.PIPE,
'shell': True,
'universal_newlines': True,
'bufsize': 0,
}
@property
def _default_pexpect_kwargs(self):
return {
'env': os.environ.copy(),
}
@property
def _uses_subprocess(self):
return isinstance(self.subprocess, subprocess.Popen)
@property
def _uses_pexpect(self):
return isinstance(self.subprocess, PopenSpawn)
@property
def std_out(self):
return self.subprocess.stdout
@property
def _pexpect_out(self):
result = ''
if self.subprocess.before:
result += self.subprocess.before
if isinstance(self.subprocess.after, str):
result += self.subprocess.after
result += self.subprocess.read().decode('utf-8')
return result
@property
def out(self):
"""Std/out output (cached), as well as stderr for non-blocking runs."""
if self.__out:
return self.__out
if self._uses_subprocess:
self.__out = self.std_out.read()
else:
self.__out = self._pexpect_out
return self.__out
@property
def std_err(self):
return self.subprocess.stderr
@property
def err(self):
if self._uses_subprocess:
return self.std_err.read()
else:
return self._pexpect_out
@property
def pid(self):
"""The process' PID."""
# Support for pexpect's functionality.
if hasattr(self.subprocess, 'proc'):
return self.subprocess.proc.pid
# Standard subprocess method.
return self.subprocess.pid
@property
def return_code(self):
return self.subprocess.returncode
@property
def std_in(self):
return self.subprocess.stdin
def run(self, block=True):
"""Runs the given command, with or without pexpect functionality enabled."""
self.blocking = block
# Use subprocess.
if self.blocking:
s = subprocess.Popen(self._popen_args, **self._default_popen_kwargs)
# Otherwise, use pexpect.
else:
s = PopenSpawn(self._popen_args, **self._default_pexpect_kwargs)
self.subprocess = s
self.was_run = True
def expect(self, pattern, timeout=-1):
"""Waits on the given pattern to appear in std_out"""
if self.blocking:
raise RuntimeError('expect can only be used on non-blocking commands.')
self.subprocess.expect(pattern=pattern, timeout=timeout)
def send(self, s, end=os.linesep, signal=False):
"""Sends the given string or signal to std_in."""
if self.blocking:
raise RuntimeError('send can only be used on non-blocking commands.')
if not signal:
if self._uses_subprocess:
return self.subprocess.communicate(s + end)
else:
return self.subprocess.send(s + end)
else:
self.subprocess.send_signal(s)
def terminate(self):
self.subprocess.terminate()
def kill(self):
self.subprocess.kill()
def block(self):
"""Blocks until process is complete."""
self.subprocess.wait()
def pipe(self, command):
"""Runs the current command and passes its output to the next
given process.
"""
if not self.was_run:
self.run(block=False)
data = self.out
c = Command(command)
c.run(block=False)
if data:
c.send(data)
c.subprocess.sendeof()
c.block()
return c
def _expand_args(command):
"""Parses command strings and returns a Popen-ready list."""
# Prepare arguments.
if isinstance(command, (str, unicode)):
splitter = shlex.shlex(command.encode('utf-8'))
splitter.whitespace = '|'
splitter.whitespace_split = True
command = []
while True:
token = splitter.get_token()
if token:
command.append(token)
else:
break
command = list(map(shlex.split, command))
return command
def chain(command):
commands = _expand_args(command)
data = None
for command in commands:
c = run(command, block=False)
if data:
c.send(data)
c.subprocess.sendeof()
data = c.out
return c
def run(command, block=True):
c = Command(command)
c.run(block=block)
if block:
c.block()
return c
Vendored Executable
+1209
View File
File diff suppressed because it is too large Load Diff
+1280
View File
File diff suppressed because it is too large Load Diff
+351
View File
@@ -0,0 +1,351 @@
'''This implements an ANSI (VT100) terminal emulator as a subclass of screen.
PEXPECT LICENSE
This license is approved by the OSI and FSF as GPL-compatible.
http://opensource.org/licenses/isc-license.txt
Copyright (c) 2012, Noah Spurrier <noah@noah.org>
PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''
# references:
# http://en.wikipedia.org/wiki/ANSI_escape_code
# http://www.retards.org/terminals/vt102.html
# http://vt100.net/docs/vt102-ug/contents.html
# http://vt100.net/docs/vt220-rm/
# http://www.termsys.demon.co.uk/vtansi.htm
from . import screen
from . import FSM
import string
#
# The 'Do.*' functions are helper functions for the ANSI class.
#
def DoEmit (fsm):
screen = fsm.memory[0]
screen.write_ch(fsm.input_symbol)
def DoStartNumber (fsm):
fsm.memory.append (fsm.input_symbol)
def DoBuildNumber (fsm):
ns = fsm.memory.pop()
ns = ns + fsm.input_symbol
fsm.memory.append (ns)
def DoBackOne (fsm):
screen = fsm.memory[0]
screen.cursor_back ()
def DoBack (fsm):
count = int(fsm.memory.pop())
screen = fsm.memory[0]
screen.cursor_back (count)
def DoDownOne (fsm):
screen = fsm.memory[0]
screen.cursor_down ()
def DoDown (fsm):
count = int(fsm.memory.pop())
screen = fsm.memory[0]
screen.cursor_down (count)
def DoForwardOne (fsm):
screen = fsm.memory[0]
screen.cursor_forward ()
def DoForward (fsm):
count = int(fsm.memory.pop())
screen = fsm.memory[0]
screen.cursor_forward (count)
def DoUpReverse (fsm):
screen = fsm.memory[0]
screen.cursor_up_reverse()
def DoUpOne (fsm):
screen = fsm.memory[0]
screen.cursor_up ()
def DoUp (fsm):
count = int(fsm.memory.pop())
screen = fsm.memory[0]
screen.cursor_up (count)
def DoHome (fsm):
c = int(fsm.memory.pop())
r = int(fsm.memory.pop())
screen = fsm.memory[0]
screen.cursor_home (r,c)
def DoHomeOrigin (fsm):
c = 1
r = 1
screen = fsm.memory[0]
screen.cursor_home (r,c)
def DoEraseDown (fsm):
screen = fsm.memory[0]
screen.erase_down()
def DoErase (fsm):
arg = int(fsm.memory.pop())
screen = fsm.memory[0]
if arg == 0:
screen.erase_down()
elif arg == 1:
screen.erase_up()
elif arg == 2:
screen.erase_screen()
def DoEraseEndOfLine (fsm):
screen = fsm.memory[0]
screen.erase_end_of_line()
def DoEraseLine (fsm):
arg = int(fsm.memory.pop())
screen = fsm.memory[0]
if arg == 0:
screen.erase_end_of_line()
elif arg == 1:
screen.erase_start_of_line()
elif arg == 2:
screen.erase_line()
def DoEnableScroll (fsm):
screen = fsm.memory[0]
screen.scroll_screen()
def DoCursorSave (fsm):
screen = fsm.memory[0]
screen.cursor_save_attrs()
def DoCursorRestore (fsm):
screen = fsm.memory[0]
screen.cursor_restore_attrs()
def DoScrollRegion (fsm):
screen = fsm.memory[0]
r2 = int(fsm.memory.pop())
r1 = int(fsm.memory.pop())
screen.scroll_screen_rows (r1,r2)
def DoMode (fsm):
screen = fsm.memory[0]
mode = fsm.memory.pop() # Should be 4
# screen.setReplaceMode ()
def DoLog (fsm):
screen = fsm.memory[0]
fsm.memory = [screen]
fout = open ('log', 'a')
fout.write (fsm.input_symbol + ',' + fsm.current_state + '\n')
fout.close()
class term (screen.screen):
'''This class is an abstract, generic terminal.
This does nothing. This is a placeholder that
provides a common base class for other terminals
such as an ANSI terminal. '''
def __init__ (self, r=24, c=80, *args, **kwargs):
screen.screen.__init__(self, r,c,*args,**kwargs)
class ANSI (term):
'''This class implements an ANSI (VT100) terminal.
It is a stream filter that recognizes ANSI terminal
escape sequences and maintains the state of a screen object. '''
def __init__ (self, r=24,c=80,*args,**kwargs):
term.__init__(self,r,c,*args,**kwargs)
#self.screen = screen (24,80)
self.state = FSM.FSM ('INIT',[self])
self.state.set_default_transition (DoLog, 'INIT')
self.state.add_transition_any ('INIT', DoEmit, 'INIT')
self.state.add_transition ('\x1b', 'INIT', None, 'ESC')
self.state.add_transition_any ('ESC', DoLog, 'INIT')
self.state.add_transition ('(', 'ESC', None, 'G0SCS')
self.state.add_transition (')', 'ESC', None, 'G1SCS')
self.state.add_transition_list ('AB012', 'G0SCS', None, 'INIT')
self.state.add_transition_list ('AB012', 'G1SCS', None, 'INIT')
self.state.add_transition ('7', 'ESC', DoCursorSave, 'INIT')
self.state.add_transition ('8', 'ESC', DoCursorRestore, 'INIT')
self.state.add_transition ('M', 'ESC', DoUpReverse, 'INIT')
self.state.add_transition ('>', 'ESC', DoUpReverse, 'INIT')
self.state.add_transition ('<', 'ESC', DoUpReverse, 'INIT')
self.state.add_transition ('=', 'ESC', None, 'INIT') # Selects application keypad.
self.state.add_transition ('#', 'ESC', None, 'GRAPHICS_POUND')
self.state.add_transition_any ('GRAPHICS_POUND', None, 'INIT')
self.state.add_transition ('[', 'ESC', None, 'ELB')
# ELB means Escape Left Bracket. That is ^[[
self.state.add_transition ('H', 'ELB', DoHomeOrigin, 'INIT')
self.state.add_transition ('D', 'ELB', DoBackOne, 'INIT')
self.state.add_transition ('B', 'ELB', DoDownOne, 'INIT')
self.state.add_transition ('C', 'ELB', DoForwardOne, 'INIT')
self.state.add_transition ('A', 'ELB', DoUpOne, 'INIT')
self.state.add_transition ('J', 'ELB', DoEraseDown, 'INIT')
self.state.add_transition ('K', 'ELB', DoEraseEndOfLine, 'INIT')
self.state.add_transition ('r', 'ELB', DoEnableScroll, 'INIT')
self.state.add_transition ('m', 'ELB', self.do_sgr, 'INIT')
self.state.add_transition ('?', 'ELB', None, 'MODECRAP')
self.state.add_transition_list (string.digits, 'ELB', DoStartNumber, 'NUMBER_1')
self.state.add_transition_list (string.digits, 'NUMBER_1', DoBuildNumber, 'NUMBER_1')
self.state.add_transition ('D', 'NUMBER_1', DoBack, 'INIT')
self.state.add_transition ('B', 'NUMBER_1', DoDown, 'INIT')
self.state.add_transition ('C', 'NUMBER_1', DoForward, 'INIT')
self.state.add_transition ('A', 'NUMBER_1', DoUp, 'INIT')
self.state.add_transition ('J', 'NUMBER_1', DoErase, 'INIT')
self.state.add_transition ('K', 'NUMBER_1', DoEraseLine, 'INIT')
self.state.add_transition ('l', 'NUMBER_1', DoMode, 'INIT')
### It gets worse... the 'm' code can have infinite number of
### number;number;number before it. I've never seen more than two,
### but the specs say it's allowed. crap!
self.state.add_transition ('m', 'NUMBER_1', self.do_sgr, 'INIT')
### LED control. Same implementation problem as 'm' code.
self.state.add_transition ('q', 'NUMBER_1', self.do_decsca, 'INIT')
# \E[?47h switch to alternate screen
# \E[?47l restores to normal screen from alternate screen.
self.state.add_transition_list (string.digits, 'MODECRAP', DoStartNumber, 'MODECRAP_NUM')
self.state.add_transition_list (string.digits, 'MODECRAP_NUM', DoBuildNumber, 'MODECRAP_NUM')
self.state.add_transition ('l', 'MODECRAP_NUM', self.do_modecrap, 'INIT')
self.state.add_transition ('h', 'MODECRAP_NUM', self.do_modecrap, 'INIT')
#RM Reset Mode Esc [ Ps l none
self.state.add_transition (';', 'NUMBER_1', None, 'SEMICOLON')
self.state.add_transition_any ('SEMICOLON', DoLog, 'INIT')
self.state.add_transition_list (string.digits, 'SEMICOLON', DoStartNumber, 'NUMBER_2')
self.state.add_transition_list (string.digits, 'NUMBER_2', DoBuildNumber, 'NUMBER_2')
self.state.add_transition_any ('NUMBER_2', DoLog, 'INIT')
self.state.add_transition ('H', 'NUMBER_2', DoHome, 'INIT')
self.state.add_transition ('f', 'NUMBER_2', DoHome, 'INIT')
self.state.add_transition ('r', 'NUMBER_2', DoScrollRegion, 'INIT')
### It gets worse... the 'm' code can have infinite number of
### number;number;number before it. I've never seen more than two,
### but the specs say it's allowed. crap!
self.state.add_transition ('m', 'NUMBER_2', self.do_sgr, 'INIT')
### LED control. Same problem as 'm' code.
self.state.add_transition ('q', 'NUMBER_2', self.do_decsca, 'INIT')
self.state.add_transition (';', 'NUMBER_2', None, 'SEMICOLON_X')
# Create a state for 'q' and 'm' which allows an infinite number of ignored numbers
self.state.add_transition_any ('SEMICOLON_X', DoLog, 'INIT')
self.state.add_transition_list (string.digits, 'SEMICOLON_X', DoStartNumber, 'NUMBER_X')
self.state.add_transition_list (string.digits, 'NUMBER_X', DoBuildNumber, 'NUMBER_X')
self.state.add_transition_any ('NUMBER_X', DoLog, 'INIT')
self.state.add_transition ('m', 'NUMBER_X', self.do_sgr, 'INIT')
self.state.add_transition ('q', 'NUMBER_X', self.do_decsca, 'INIT')
self.state.add_transition (';', 'NUMBER_X', None, 'SEMICOLON_X')
def process (self, c):
"""Process a single character. Called by :meth:`write`."""
if isinstance(c, bytes):
c = self._decode(c)
self.state.process(c)
def process_list (self, l):
self.write(l)
def write (self, s):
"""Process text, writing it to the virtual screen while handling
ANSI escape codes.
"""
if isinstance(s, bytes):
s = self._decode(s)
for c in s:
self.process(c)
def flush (self):
pass
def write_ch (self, ch):
'''This puts a character at the current cursor position. The cursor
position is moved forward with wrap-around, but no scrolling is done if
the cursor hits the lower-right corner of the screen. '''
if isinstance(ch, bytes):
ch = self._decode(ch)
#\r and \n both produce a call to cr() and lf(), respectively.
ch = ch[0]
if ch == u'\r':
self.cr()
return
if ch == u'\n':
self.crlf()
return
if ch == chr(screen.BS):
self.cursor_back()
return
self.put_abs(self.cur_r, self.cur_c, ch)
old_r = self.cur_r
old_c = self.cur_c
self.cursor_forward()
if old_c == self.cur_c:
self.cursor_down()
if old_r != self.cur_r:
self.cursor_home (self.cur_r, 1)
else:
self.scroll_up ()
self.cursor_home (self.cur_r, 1)
self.erase_line()
def do_sgr (self, fsm):
'''Select Graphic Rendition, e.g. color. '''
screen = fsm.memory[0]
fsm.memory = [screen]
def do_decsca (self, fsm):
'''Select character protection attribute. '''
screen = fsm.memory[0]
fsm.memory = [screen]
def do_modecrap (self, fsm):
'''Handler for \x1b[?<number>h and \x1b[?<number>l. If anyone
wanted to actually use these, they'd need to add more states to the
FSM rather than just improve or override this method. '''
screen = fsm.memory[0]
fsm.memory = [screen]
+334
View File
@@ -0,0 +1,334 @@
#!/usr/bin/env python
'''This module implements a Finite State Machine (FSM). In addition to state
this FSM also maintains a user defined "memory". So this FSM can be used as a
Push-down Automata (PDA) since a PDA is a FSM + memory.
The following describes how the FSM works, but you will probably also need to
see the example function to understand how the FSM is used in practice.
You define an FSM by building tables of transitions. For a given input symbol
the process() method uses these tables to decide what action to call and what
the next state will be. The FSM has a table of transitions that associate:
(input_symbol, current_state) --> (action, next_state)
Where "action" is a function you define. The symbols and states can be any
objects. You use the add_transition() and add_transition_list() methods to add
to the transition table. The FSM also has a table of transitions that
associate:
(current_state) --> (action, next_state)
You use the add_transition_any() method to add to this transition table. The
FSM also has one default transition that is not associated with any specific
input_symbol or state. You use the set_default_transition() method to set the
default transition.
When an action function is called it is passed a reference to the FSM. The
action function may then access attributes of the FSM such as input_symbol,
current_state, or "memory". The "memory" attribute can be any object that you
want to pass along to the action functions. It is not used by the FSM itself.
For parsing you would typically pass a list to be used as a stack.
The processing sequence is as follows. The process() method is given an
input_symbol to process. The FSM will search the table of transitions that
associate:
(input_symbol, current_state) --> (action, next_state)
If the pair (input_symbol, current_state) is found then process() will call the
associated action function and then set the current state to the next_state.
If the FSM cannot find a match for (input_symbol, current_state) it will then
search the table of transitions that associate:
(current_state) --> (action, next_state)
If the current_state is found then the process() method will call the
associated action function and then set the current state to the next_state.
Notice that this table lacks an input_symbol. It lets you define transitions
for a current_state and ANY input_symbol. Hence, it is called the "any" table.
Remember, it is always checked after first searching the table for a specific
(input_symbol, current_state).
For the case where the FSM did not match either of the previous two cases the
FSM will try to use the default transition. If the default transition is
defined then the process() method will call the associated action function and
then set the current state to the next_state. This lets you define a default
transition as a catch-all case. You can think of it as an exception handler.
There can be only one default transition.
Finally, if none of the previous cases are defined for an input_symbol and
current_state then the FSM will raise an exception. This may be desirable, but
you can always prevent this just by defining a default transition.
Noah Spurrier 20020822
PEXPECT LICENSE
This license is approved by the OSI and FSF as GPL-compatible.
http://opensource.org/licenses/isc-license.txt
Copyright (c) 2012, Noah Spurrier <noah@noah.org>
PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''
class ExceptionFSM(Exception):
'''This is the FSM Exception class.'''
def __init__(self, value):
self.value = value
def __str__(self):
return 'ExceptionFSM: ' + str(self.value)
class FSM:
'''This is a Finite State Machine (FSM).
'''
def __init__(self, initial_state, memory=None):
'''This creates the FSM. You set the initial state here. The "memory"
attribute is any object that you want to pass along to the action
functions. It is not used by the FSM. For parsing you would typically
pass a list to be used as a stack. '''
# Map (input_symbol, current_state) --> (action, next_state).
self.state_transitions = {}
# Map (current_state) --> (action, next_state).
self.state_transitions_any = {}
self.default_transition = None
self.input_symbol = None
self.initial_state = initial_state
self.current_state = self.initial_state
self.next_state = None
self.action = None
self.memory = memory
def reset (self):
'''This sets the current_state to the initial_state and sets
input_symbol to None. The initial state was set by the constructor
__init__(). '''
self.current_state = self.initial_state
self.input_symbol = None
def add_transition (self, input_symbol, state, action=None, next_state=None):
'''This adds a transition that associates:
(input_symbol, current_state) --> (action, next_state)
The action may be set to None in which case the process() method will
ignore the action and only set the next_state. The next_state may be
set to None in which case the current state will be unchanged.
You can also set transitions for a list of symbols by using
add_transition_list(). '''
if next_state is None:
next_state = state
self.state_transitions[(input_symbol, state)] = (action, next_state)
def add_transition_list (self, list_input_symbols, state, action=None, next_state=None):
'''This adds the same transition for a list of input symbols.
You can pass a list or a string. Note that it is handy to use
string.digits, string.whitespace, string.letters, etc. to add
transitions that match character classes.
The action may be set to None in which case the process() method will
ignore the action and only set the next_state. The next_state may be
set to None in which case the current state will be unchanged. '''
if next_state is None:
next_state = state
for input_symbol in list_input_symbols:
self.add_transition (input_symbol, state, action, next_state)
def add_transition_any (self, state, action=None, next_state=None):
'''This adds a transition that associates:
(current_state) --> (action, next_state)
That is, any input symbol will match the current state.
The process() method checks the "any" state associations after it first
checks for an exact match of (input_symbol, current_state).
The action may be set to None in which case the process() method will
ignore the action and only set the next_state. The next_state may be
set to None in which case the current state will be unchanged. '''
if next_state is None:
next_state = state
self.state_transitions_any [state] = (action, next_state)
def set_default_transition (self, action, next_state):
'''This sets the default transition. This defines an action and
next_state if the FSM cannot find the input symbol and the current
state in the transition list and if the FSM cannot find the
current_state in the transition_any list. This is useful as a final
fall-through state for catching errors and undefined states.
The default transition can be removed by setting the attribute
default_transition to None. '''
self.default_transition = (action, next_state)
def get_transition (self, input_symbol, state):
'''This returns (action, next state) given an input_symbol and state.
This does not modify the FSM state, so calling this method has no side
effects. Normally you do not call this method directly. It is called by
process().
The sequence of steps to check for a defined transition goes from the
most specific to the least specific.
1. Check state_transitions[] that match exactly the tuple,
(input_symbol, state)
2. Check state_transitions_any[] that match (state)
In other words, match a specific state and ANY input_symbol.
3. Check if the default_transition is defined.
This catches any input_symbol and any state.
This is a handler for errors, undefined states, or defaults.
4. No transition was defined. If we get here then raise an exception.
'''
if (input_symbol, state) in self.state_transitions:
return self.state_transitions[(input_symbol, state)]
elif state in self.state_transitions_any:
return self.state_transitions_any[state]
elif self.default_transition is not None:
return self.default_transition
else:
raise ExceptionFSM ('Transition is undefined: (%s, %s).' %
(str(input_symbol), str(state)) )
def process (self, input_symbol):
'''This is the main method that you call to process input. This may
cause the FSM to change state and call an action. This method calls
get_transition() to find the action and next_state associated with the
input_symbol and current_state. If the action is None then the action
is not called and only the current state is changed. This method
processes one complete input symbol. You can process a list of symbols
(or a string) by calling process_list(). '''
self.input_symbol = input_symbol
(self.action, self.next_state) = self.get_transition (self.input_symbol, self.current_state)
if self.action is not None:
self.action (self)
self.current_state = self.next_state
self.next_state = None
def process_list (self, input_symbols):
'''This takes a list and sends each element to process(). The list may
be a string or any iterable object. '''
for s in input_symbols:
self.process (s)
##############################################################################
# The following is an example that demonstrates the use of the FSM class to
# process an RPN expression. Run this module from the command line. You will
# get a prompt > for input. Enter an RPN Expression. Numbers may be integers.
# Operators are * / + - Use the = sign to evaluate and print the expression.
# For example:
#
# 167 3 2 2 * * * 1 - =
#
# will print:
#
# 2003
##############################################################################
import sys
import string
PY3 = (sys.version_info[0] >= 3)
#
# These define the actions.
# Note that "memory" is a list being used as a stack.
#
def BeginBuildNumber (fsm):
fsm.memory.append (fsm.input_symbol)
def BuildNumber (fsm):
s = fsm.memory.pop ()
s = s + fsm.input_symbol
fsm.memory.append (s)
def EndBuildNumber (fsm):
s = fsm.memory.pop ()
fsm.memory.append (int(s))
def DoOperator (fsm):
ar = fsm.memory.pop()
al = fsm.memory.pop()
if fsm.input_symbol == '+':
fsm.memory.append (al + ar)
elif fsm.input_symbol == '-':
fsm.memory.append (al - ar)
elif fsm.input_symbol == '*':
fsm.memory.append (al * ar)
elif fsm.input_symbol == '/':
fsm.memory.append (al / ar)
def DoEqual (fsm):
print(str(fsm.memory.pop()))
def Error (fsm):
print('That does not compute.')
print(str(fsm.input_symbol))
def main():
'''This is where the example starts and the FSM state transitions are
defined. Note that states are strings (such as 'INIT'). This is not
necessary, but it makes the example easier to read. '''
f = FSM ('INIT', [])
f.set_default_transition (Error, 'INIT')
f.add_transition_any ('INIT', None, 'INIT')
f.add_transition ('=', 'INIT', DoEqual, 'INIT')
f.add_transition_list (string.digits, 'INIT', BeginBuildNumber, 'BUILDING_NUMBER')
f.add_transition_list (string.digits, 'BUILDING_NUMBER', BuildNumber, 'BUILDING_NUMBER')
f.add_transition_list (string.whitespace, 'BUILDING_NUMBER', EndBuildNumber, 'INIT')
f.add_transition_list ('+-*/', 'INIT', DoOperator, 'INIT')
print()
print('Enter an RPN Expression.')
print('Numbers may be integers. Operators are * / + -')
print('Use the = sign to evaluate and print the expression.')
print('For example: ')
print(' 167 3 2 2 * * * 1 - =')
inputstr = (input if PY3 else raw_input)('> ') # analysis:ignore
f.process_list(inputstr)
if __name__ == '__main__':
main()
+85
View File
@@ -0,0 +1,85 @@
'''Pexpect is a Python module for spawning child applications and controlling
them automatically. Pexpect can be used for automating interactive applications
such as ssh, ftp, passwd, telnet, etc. It can be used to a automate setup
scripts for duplicating software package installations on different servers. It
can be used for automated software testing. Pexpect is in the spirit of Don
Libes' Expect, but Pexpect is pure Python. Other Expect-like modules for Python
require TCL and Expect or require C extensions to be compiled. Pexpect does not
use C, Expect, or TCL extensions. It should work on any platform that supports
the standard Python pty module. The Pexpect interface focuses on ease of use so
that simple tasks are easy.
There are two main interfaces to the Pexpect system; these are the function,
run() and the class, spawn. The spawn class is more powerful. The run()
function is simpler than spawn, and is good for quickly calling program. When
you call the run() function it executes a given program and then returns the
output. This is a handy replacement for os.system().
For example::
pexpect.run('ls -la')
The spawn class is the more powerful interface to the Pexpect system. You can
use this to spawn a child program then interact with it by sending input and
expecting responses (waiting for patterns in the child's output).
For example::
child = pexpect.spawn('scp foo user@example.com:.')
child.expect('Password:')
child.sendline(mypassword)
This works even for commands that ask for passwords or other input outside of
the normal stdio streams. For example, ssh reads input directly from the TTY
device which bypasses stdin.
Credits: Noah Spurrier, Richard Holden, Marco Molteni, Kimberley Burchett,
Robert Stone, Hartmut Goebel, Chad Schroeder, Erick Tryzelaar, Dave Kirby, Ids
vander Molen, George Todd, Noel Taylor, Nicolas D. Cesar, Alexander Gattin,
Jacques-Etienne Baudoux, Geoffrey Marshall, Francisco Lourenco, Glen Mabey,
Karthik Gurusamy, Fernando Perez, Corey Minyard, Jon Cohen, Guillaume
Chazarain, Andrew Ryan, Nick Craig-Wood, Andrew Stone, Jorgen Grahn, John
Spiegel, Jan Grant, and Shane Kerr. Let me know if I forgot anyone.
Pexpect is free, open source, and all that good stuff.
http://pexpect.sourceforge.net/
PEXPECT LICENSE
This license is approved by the OSI and FSF as GPL-compatible.
http://opensource.org/licenses/isc-license.txt
Copyright (c) 2012, Noah Spurrier <noah@noah.org>
PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''
import sys
PY3 = (sys.version_info[0] >= 3)
from .exceptions import ExceptionPexpect, EOF, TIMEOUT
from .utils import split_command_line, which, is_executable_file
from .expect import Expecter, searcher_re, searcher_string
if sys.platform != 'win32':
# On Unix, these are available at the top level for backwards compatibility
from .pty_spawn import spawn, spawnu
from .run import run, runu
__version__ = '4.2.1'
__revision__ = ''
__all__ = ['ExceptionPexpect', 'EOF', 'TIMEOUT', 'spawn', 'spawnu', 'run', 'runu',
'which', 'split_command_line', '__version__', '__revision__']
# vim: set shiftround expandtab tabstop=4 shiftwidth=4 ft=python autoindent :
+78
View File
@@ -0,0 +1,78 @@
import asyncio
import errno
from pexpect import EOF
@asyncio.coroutine
def expect_async(expecter, timeout=None):
# First process data that was previously read - if it maches, we don't need
# async stuff.
previously_read = expecter.spawn.buffer
expecter.spawn.buffer = expecter.spawn.string_type()
idx = expecter.new_data(previously_read)
if idx is not None:
return idx
transport, pw = yield from asyncio.get_event_loop()\
.connect_read_pipe(lambda: PatternWaiter(expecter), expecter.spawn)
try:
return (yield from asyncio.wait_for(pw.fut, timeout))
except asyncio.TimeoutError as e:
transport.pause_reading()
return expecter.timeout(e)
class PatternWaiter(asyncio.Protocol):
transport = None
def __init__(self, expecter):
self.expecter = expecter
self.fut = asyncio.Future()
def found(self, result):
if not self.fut.done():
self.fut.set_result(result)
self.transport.pause_reading()
def error(self, exc):
if not self.fut.done():
self.fut.set_exception(exc)
self.transport.pause_reading()
def connection_made(self, transport):
self.transport = transport
def data_received(self, data):
spawn = self.expecter.spawn
s = spawn._decoder.decode(data)
spawn._log(s, 'read')
if self.fut.done():
spawn.buffer += s
return
try:
index = self.expecter.new_data(s)
if index is not None:
# Found a match
self.found(index)
except Exception as e:
self.expecter.errored()
self.error(e)
def eof_received(self):
# N.B. If this gets called, async will close the pipe (the spawn object)
# for us
try:
self.expecter.spawn.flag_eof = True
index = self.expecter.eof()
except EOF as e:
self.error(e)
else:
self.found(index)
def connection_lost(self, exc):
if isinstance(exc, OSError) and exc.errno == errno.EIO:
# We may get here without eof_received being called, e.g on Linux
self.eof_received()
elif exc is not None:
self.error(exc)
+5
View File
@@ -0,0 +1,5 @@
source /etc/bash.bashrc
source ~/.bashrc
# Reset PS1 so pexpect can find it
PS1="$"
+35
View File
@@ -0,0 +1,35 @@
"""Exception classes used by Pexpect"""
import traceback
import sys
class ExceptionPexpect(Exception):
'''Base class for all exceptions raised by this module.
'''
def __init__(self, value):
super(ExceptionPexpect, self).__init__(value)
self.value = value
def __str__(self):
return str(self.value)
def get_trace(self):
'''This returns an abbreviated stack trace with lines that only concern
the caller. In other words, the stack trace inside the Pexpect module
is not included. '''
tblist = traceback.extract_tb(sys.exc_info()[2])
tblist = [item for item in tblist if ('pexpect/__init__' not in item[0])
and ('pexpect/expect' not in item[0])]
tblist = traceback.format_list(tblist)
return ''.join(tblist)
class EOF(ExceptionPexpect):
'''Raised when EOF is read from a child.
This usually means the child has exited.'''
class TIMEOUT(ExceptionPexpect):
'''Raised when a read time exceeds the timeout. '''
+300
View File
@@ -0,0 +1,300 @@
import time
from .exceptions import EOF, TIMEOUT
class Expecter(object):
def __init__(self, spawn, searcher, searchwindowsize=-1):
self.spawn = spawn
self.searcher = searcher
if searchwindowsize == -1:
searchwindowsize = spawn.searchwindowsize
self.searchwindowsize = searchwindowsize
def new_data(self, data):
spawn = self.spawn
searcher = self.searcher
incoming = spawn.buffer + data
freshlen = len(data)
index = searcher.search(incoming, freshlen, self.searchwindowsize)
if index >= 0:
spawn.buffer = incoming[searcher.end:]
spawn.before = incoming[: searcher.start]
spawn.after = incoming[searcher.start: searcher.end]
spawn.match = searcher.match
spawn.match_index = index
# Found a match
return index
spawn.buffer = incoming
def eof(self, err=None):
spawn = self.spawn
from . import EOF
spawn.before = spawn.buffer
spawn.buffer = spawn.string_type()
spawn.after = EOF
index = self.searcher.eof_index
if index >= 0:
spawn.match = EOF
spawn.match_index = index
return index
else:
spawn.match = None
spawn.match_index = None
msg = str(spawn)
msg += '\nsearcher: %s' % self.searcher
if err is not None:
msg = str(err) + '\n' + msg
raise EOF(msg)
def timeout(self, err=None):
spawn = self.spawn
from . import TIMEOUT
spawn.before = spawn.buffer
spawn.after = TIMEOUT
index = self.searcher.timeout_index
if index >= 0:
spawn.match = TIMEOUT
spawn.match_index = index
return index
else:
spawn.match = None
spawn.match_index = None
msg = str(spawn)
msg += '\nsearcher: %s' % self.searcher
if err is not None:
msg = str(err) + '\n' + msg
raise TIMEOUT(msg)
def errored(self):
spawn = self.spawn
spawn.before = spawn.buffer
spawn.after = None
spawn.match = None
spawn.match_index = None
def expect_loop(self, timeout=-1):
"""Blocking expect"""
spawn = self.spawn
from . import EOF, TIMEOUT
if timeout is not None:
end_time = time.time() + timeout
try:
incoming = spawn.buffer
spawn.buffer = spawn.string_type() # Treat buffer as new data
while True:
idx = self.new_data(incoming)
# Keep reading until exception or return.
if idx is not None:
return idx
# No match at this point
if (timeout is not None) and (timeout < 0):
return self.timeout()
# Still have time left, so read more data
incoming = spawn.read_nonblocking(spawn.maxread, timeout)
if self.spawn.delayafterread is not None:
time.sleep(self.spawn.delayafterread)
if timeout is not None:
timeout = end_time - time.time()
except EOF as e:
return self.eof(e)
except TIMEOUT as e:
return self.timeout(e)
except:
self.errored()
raise
class searcher_string(object):
'''This is a plain string search helper for the spawn.expect_any() method.
This helper class is for speed. For more powerful regex patterns
see the helper class, searcher_re.
Attributes:
eof_index - index of EOF, or -1
timeout_index - index of TIMEOUT, or -1
After a successful match by the search() method the following attributes
are available:
start - index into the buffer, first byte of match
end - index into the buffer, first byte after match
match - the matching string itself
'''
def __init__(self, strings):
'''This creates an instance of searcher_string. This argument 'strings'
may be a list; a sequence of strings; or the EOF or TIMEOUT types. '''
self.eof_index = -1
self.timeout_index = -1
self._strings = []
for n, s in enumerate(strings):
if s is EOF:
self.eof_index = n
continue
if s is TIMEOUT:
self.timeout_index = n
continue
self._strings.append((n, s))
def __str__(self):
'''This returns a human-readable string that represents the state of
the object.'''
ss = [(ns[0], ' %d: "%s"' % ns) for ns in self._strings]
ss.append((-1, 'searcher_string:'))
if self.eof_index >= 0:
ss.append((self.eof_index, ' %d: EOF' % self.eof_index))
if self.timeout_index >= 0:
ss.append((self.timeout_index,
' %d: TIMEOUT' % self.timeout_index))
ss.sort()
ss = list(zip(*ss))[1]
return '\n'.join(ss)
def search(self, buffer, freshlen, searchwindowsize=None):
'''This searches 'buffer' for the first occurence of one of the search
strings. 'freshlen' must indicate the number of bytes at the end of
'buffer' which have not been searched before. It helps to avoid
searching the same, possibly big, buffer over and over again.
See class spawn for the 'searchwindowsize' argument.
If there is a match this returns the index of that string, and sets
'start', 'end' and 'match'. Otherwise, this returns -1. '''
first_match = None
# 'freshlen' helps a lot here. Further optimizations could
# possibly include:
#
# using something like the Boyer-Moore Fast String Searching
# Algorithm; pre-compiling the search through a list of
# strings into something that can scan the input once to
# search for all N strings; realize that if we search for
# ['bar', 'baz'] and the input is '...foo' we need not bother
# rescanning until we've read three more bytes.
#
# Sadly, I don't know enough about this interesting topic. /grahn
for index, s in self._strings:
if searchwindowsize is None:
# the match, if any, can only be in the fresh data,
# or at the very end of the old data
offset = -(freshlen + len(s))
else:
# better obey searchwindowsize
offset = -searchwindowsize
n = buffer.find(s, offset)
if n >= 0 and (first_match is None or n < first_match):
first_match = n
best_index, best_match = index, s
if first_match is None:
return -1
self.match = best_match
self.start = first_match
self.end = self.start + len(self.match)
return best_index
class searcher_re(object):
'''This is regular expression string search helper for the
spawn.expect_any() method. This helper class is for powerful
pattern matching. For speed, see the helper class, searcher_string.
Attributes:
eof_index - index of EOF, or -1
timeout_index - index of TIMEOUT, or -1
After a successful match by the search() method the following attributes
are available:
start - index into the buffer, first byte of match
end - index into the buffer, first byte after match
match - the re.match object returned by a succesful re.search
'''
def __init__(self, patterns):
'''This creates an instance that searches for 'patterns' Where
'patterns' may be a list or other sequence of compiled regular
expressions, or the EOF or TIMEOUT types.'''
self.eof_index = -1
self.timeout_index = -1
self._searches = []
for n, s in zip(list(range(len(patterns))), patterns):
if s is EOF:
self.eof_index = n
continue
if s is TIMEOUT:
self.timeout_index = n
continue
self._searches.append((n, s))
def __str__(self):
'''This returns a human-readable string that represents the state of
the object.'''
#ss = [(n, ' %d: re.compile("%s")' %
# (n, repr(s.pattern))) for n, s in self._searches]
ss = list()
for n, s in self._searches:
try:
ss.append((n, ' %d: re.compile("%s")' % (n, s.pattern)))
except UnicodeEncodeError:
# for test cases that display __str__ of searches, dont throw
# another exception just because stdout is ascii-only, using
# repr()
ss.append((n, ' %d: re.compile(%r)' % (n, s.pattern)))
ss.append((-1, 'searcher_re:'))
if self.eof_index >= 0:
ss.append((self.eof_index, ' %d: EOF' % self.eof_index))
if self.timeout_index >= 0:
ss.append((self.timeout_index, ' %d: TIMEOUT' %
self.timeout_index))
ss.sort()
ss = list(zip(*ss))[1]
return '\n'.join(ss)
def search(self, buffer, freshlen, searchwindowsize=None):
'''This searches 'buffer' for the first occurence of one of the regular
expressions. 'freshlen' must indicate the number of bytes at the end of
'buffer' which have not been searched before.
See class spawn for the 'searchwindowsize' argument.
If there is a match this returns the index of that string, and sets
'start', 'end' and 'match'. Otherwise, returns -1.'''
first_match = None
# 'freshlen' doesn't help here -- we cannot predict the
# length of a match, and the re module provides no help.
if searchwindowsize is None:
searchstart = 0
else:
searchstart = max(0, len(buffer) - searchwindowsize)
for index, s in self._searches:
match = s.search(buffer, searchstart)
if match is None:
continue
n = match.start()
if first_match is None or n < first_match:
first_match = n
the_match = match
best_index = index
if first_match is None:
return -1
self.start = first_match
self.match = the_match
self.end = self.match.end()
return best_index
+142
View File
@@ -0,0 +1,142 @@
'''This is like pexpect, but it will work with any file descriptor that you
pass it. You are reponsible for opening and close the file descriptor.
This allows you to use Pexpect with sockets and named pipes (FIFOs).
PEXPECT LICENSE
This license is approved by the OSI and FSF as GPL-compatible.
http://opensource.org/licenses/isc-license.txt
Copyright (c) 2012, Noah Spurrier <noah@noah.org>
PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''
from .spawnbase import SpawnBase
from .exceptions import ExceptionPexpect, TIMEOUT
from .utils import select_ignore_interrupts
import os
__all__ = ['fdspawn']
class fdspawn(SpawnBase):
'''This is like pexpect.spawn but allows you to supply your own open file
descriptor. For example, you could use it to read through a file looking
for patterns, or to control a modem or serial device. '''
def __init__ (self, fd, args=None, timeout=30, maxread=2000, searchwindowsize=None,
logfile=None, encoding=None, codec_errors='strict'):
'''This takes a file descriptor (an int) or an object that support the
fileno() method (returning an int). All Python file-like objects
support fileno(). '''
if type(fd) != type(0) and hasattr(fd, 'fileno'):
fd = fd.fileno()
if type(fd) != type(0):
raise ExceptionPexpect('The fd argument is not an int. If this is a command string then maybe you want to use pexpect.spawn.')
try: # make sure fd is a valid file descriptor
os.fstat(fd)
except OSError:
raise ExceptionPexpect('The fd argument is not a valid file descriptor.')
self.args = None
self.command = None
SpawnBase.__init__(self, timeout, maxread, searchwindowsize, logfile,
encoding=encoding, codec_errors=codec_errors)
self.child_fd = fd
self.own_fd = False
self.closed = False
self.name = '<file descriptor %d>' % fd
def close (self):
"""Close the file descriptor.
Calling this method a second time does nothing, but if the file
descriptor was closed elsewhere, :class:`OSError` will be raised.
"""
if self.child_fd == -1:
return
self.flush()
os.close(self.child_fd)
self.child_fd = -1
self.closed = True
def isalive (self):
'''This checks if the file descriptor is still valid. If :func:`os.fstat`
does not raise an exception then we assume it is alive. '''
if self.child_fd == -1:
return False
try:
os.fstat(self.child_fd)
return True
except:
return False
def terminate (self, force=False): # pragma: no cover
'''Deprecated and invalid. Just raises an exception.'''
raise ExceptionPexpect('This method is not valid for file descriptors.')
# These four methods are left around for backwards compatibility, but not
# documented as part of fdpexpect. You're encouraged to use os.write
# directly.
def send(self, s):
"Write to fd, return number of bytes written"
s = self._coerce_send_string(s)
self._log(s, 'send')
b = self._encoder.encode(s, final=False)
return os.write(self.child_fd, b)
def sendline(self, s):
"Write to fd with trailing newline, return number of bytes written"
s = self._coerce_send_string(s)
return self.send(s + self.linesep)
def write(self, s):
"Write to fd, return None"
self.send(s)
def writelines(self, sequence):
"Call self.write() for each item in sequence"
for s in sequence:
self.write(s)
def read_nonblocking(self, size=1, timeout=-1):
"""
Read from the file descriptor and return the result as a string.
The read_nonblocking method of :class:`SpawnBase` assumes that a call
to os.read will not block (timeout parameter is ignored). This is not
the case for POSIX file-like objects such as sockets and serial ports.
Use :func:`select.select`, timeout is implemented conditionally for
POSIX systems.
:param int size: Read at most *size* bytes.
:param int timeout: Wait timeout seconds for file descriptor to be
ready to read. When -1 (default), use self.timeout. When 0, poll.
:return: String containing the bytes read
"""
if os.name == 'posix':
if timeout == -1:
timeout = self.timeout
rlist = [self.child_fd]
wlist = []
xlist = []
rlist, wlist, xlist = select_ignore_interrupts(rlist, wlist, xlist, timeout)
if self.child_fd not in rlist:
raise TIMEOUT('Timeout exceeded.')
return super(fdspawn, self).read_nonblocking(size)
+179
View File
@@ -0,0 +1,179 @@
"""Provides an interface like pexpect.spawn interface using subprocess.Popen
"""
import os
import threading
import subprocess
import sys
import time
import signal
import shlex
try:
from queue import Queue, Empty # Python 3
except ImportError:
from Queue import Queue, Empty # Python 2
from .spawnbase import SpawnBase, PY3
from .exceptions import EOF
class PopenSpawn(SpawnBase):
if PY3:
crlf = '\n'.encode('ascii')
else:
crlf = '\n'
def __init__(self, cmd, timeout=30, maxread=2000, searchwindowsize=None,
logfile=None, cwd=None, env=None, encoding=None,
codec_errors='strict'):
super(PopenSpawn, self).__init__(timeout=timeout, maxread=maxread,
searchwindowsize=searchwindowsize, logfile=logfile,
encoding=encoding, codec_errors=codec_errors)
kwargs = dict(bufsize=0, stdin=subprocess.PIPE,
stderr=subprocess.STDOUT, stdout=subprocess.PIPE,
cwd=cwd, env=env)
if sys.platform == 'win32':
startupinfo = subprocess.STARTUPINFO()
startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
kwargs['startupinfo'] = startupinfo
kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
if not isinstance(cmd, (list, tuple)):
cmd = shlex.split(cmd)
self.proc = subprocess.Popen(cmd, **kwargs)
self.closed = False
self._buf = self.string_type()
self._read_queue = Queue()
self._read_thread = threading.Thread(target=self._read_incoming)
self._read_thread.setDaemon(True)
self._read_thread.start()
_read_reached_eof = False
def read_nonblocking(self, size, timeout):
buf = self._buf
if self._read_reached_eof:
# We have already finished reading. Use up any buffered data,
# then raise EOF
if buf:
self._buf = buf[size:]
return buf[:size]
else:
self.flag_eof = True
raise EOF('End Of File (EOF).')
if timeout == -1:
timeout = self.timeout
elif timeout is None:
timeout = 1e6
t0 = time.time()
while (time.time() - t0) < timeout and size and len(buf) < size:
try:
incoming = self._read_queue.get_nowait()
except Empty:
break
else:
if incoming is None:
self._read_reached_eof = True
break
buf += self._decoder.decode(incoming, final=False)
r, self._buf = buf[:size], buf[size:]
self._log(r, 'read')
return r
def _read_incoming(self):
"""Run in a thread to move output from a pipe to a queue."""
fileno = self.proc.stdout.fileno()
while 1:
buf = b''
try:
buf = os.read(fileno, 1024)
except OSError as e:
self._log(e, 'read')
if not buf:
# This indicates we have reached EOF
self._read_queue.put(None)
return
self._read_queue.put(buf)
def write(self, s):
'''This is similar to send() except that there is no return value.
'''
self.send(s)
def writelines(self, sequence):
'''This calls write() for each element in the sequence.
The sequence can be any iterable object producing strings, typically a
list of strings. This does not add line separators. There is no return
value.
'''
for s in sequence:
self.send(s)
def send(self, s):
'''Send data to the subprocess' stdin.
Returns the number of bytes written.
'''
s = self._coerce_send_string(s)
self._log(s, 'send')
b = self._encoder.encode(s, final=False)
if PY3:
return self.proc.stdin.write(b)
else:
# On Python 2, .write() returns None, so we return the length of
# bytes written ourselves. This assumes they all got written.
self.proc.stdin.write(b)
return len(b)
def sendline(self, s=''):
'''Wraps send(), sending string ``s`` to child process, with os.linesep
automatically appended. Returns number of bytes written. '''
n = self.send(s)
return n + self.send(self.linesep)
def wait(self):
'''Wait for the subprocess to finish.
Returns the exit code.
'''
status = self.proc.wait()
if status >= 0:
self.exitstatus = status
self.signalstatus = None
else:
self.exitstatus = None
self.signalstatus = -status
self.terminated = True
return status
def kill(self, sig):
'''Sends a Unix signal to the subprocess.
Use constants from the :mod:`signal` module to specify which signal.
'''
if sys.platform == 'win32':
if sig in [signal.SIGINT, signal.CTRL_C_EVENT]:
sig = signal.CTRL_C_EVENT
elif sig in [signal.SIGBREAK, signal.CTRL_BREAK_EVENT]:
sig = signal.CTRL_BREAK_EVENT
else:
sig = signal.SIGTERM
os.kill(self.proc.pid, sig)
def sendeof(self):
'''Closes the stdin pipe from the writing end.'''
self.proc.stdin.close()
+806
View File
@@ -0,0 +1,806 @@
import os
import sys
import time
import pty
import tty
import errno
import signal
from contextlib import contextmanager
import ptyprocess
from ptyprocess.ptyprocess import use_native_pty_fork
from .exceptions import ExceptionPexpect, EOF, TIMEOUT
from .spawnbase import SpawnBase
from .utils import which, split_command_line, select_ignore_interrupts
@contextmanager
def _wrap_ptyprocess_err():
"""Turn ptyprocess errors into our own ExceptionPexpect errors"""
try:
yield
except ptyprocess.PtyProcessError as e:
raise ExceptionPexpect(*e.args)
PY3 = (sys.version_info[0] >= 3)
class spawn(SpawnBase):
'''This is the main class interface for Pexpect. Use this class to start
and control child applications. '''
# This is purely informational now - changing it has no effect
use_native_pty_fork = use_native_pty_fork
def __init__(self, command, args=[], timeout=30, maxread=2000,
searchwindowsize=None, logfile=None, cwd=None, env=None,
ignore_sighup=False, echo=True, preexec_fn=None,
encoding=None, codec_errors='strict', dimensions=None):
'''This is the constructor. The command parameter may be a string that
includes a command and any arguments to the command. For example::
child = pexpect.spawn('/usr/bin/ftp')
child = pexpect.spawn('/usr/bin/ssh user@example.com')
child = pexpect.spawn('ls -latr /tmp')
You may also construct it with a list of arguments like so::
child = pexpect.spawn('/usr/bin/ftp', [])
child = pexpect.spawn('/usr/bin/ssh', ['user@example.com'])
child = pexpect.spawn('ls', ['-latr', '/tmp'])
After this the child application will be created and will be ready to
talk to. For normal use, see expect() and send() and sendline().
Remember that Pexpect does NOT interpret shell meta characters such as
redirect, pipe, or wild cards (``>``, ``|``, or ``*``). This is a
common mistake. If you want to run a command and pipe it through
another command then you must also start a shell. For example::
child = pexpect.spawn('/bin/bash -c "ls -l | grep LOG > logs.txt"')
child.expect(pexpect.EOF)
The second form of spawn (where you pass a list of arguments) is useful
in situations where you wish to spawn a command and pass it its own
argument list. This can make syntax more clear. For example, the
following is equivalent to the previous example::
shell_cmd = 'ls -l | grep LOG > logs.txt'
child = pexpect.spawn('/bin/bash', ['-c', shell_cmd])
child.expect(pexpect.EOF)
The maxread attribute sets the read buffer size. This is maximum number
of bytes that Pexpect will try to read from a TTY at one time. Setting
the maxread size to 1 will turn off buffering. Setting the maxread
value higher may help performance in cases where large amounts of
output are read back from the child. This feature is useful in
conjunction with searchwindowsize.
When the keyword argument *searchwindowsize* is None (default), the
full buffer is searched at each iteration of receiving incoming data.
The default number of bytes scanned at each iteration is very large
and may be reduced to collaterally reduce search cost. After
:meth:`~.expect` returns, the full buffer attribute remains up to
size *maxread* irrespective of *searchwindowsize* value.
When the keyword argument ``timeout`` is specified as a number,
(default: *30*), then :class:`TIMEOUT` will be raised after the value
specified has elapsed, in seconds, for any of the :meth:`~.expect`
family of method calls. When None, TIMEOUT will not be raised, and
:meth:`~.expect` may block indefinitely until match.
The logfile member turns on or off logging. All input and output will
be copied to the given file object. Set logfile to None to stop
logging. This is the default. Set logfile to sys.stdout to echo
everything to standard output. The logfile is flushed after each write.
Example log input and output to a file::
child = pexpect.spawn('some_command')
fout = open('mylog.txt','wb')
child.logfile = fout
Example log to stdout::
# In Python 2:
child = pexpect.spawn('some_command')
child.logfile = sys.stdout
# In Python 3, spawnu should be used to give str to stdout:
child = pexpect.spawnu('some_command')
child.logfile = sys.stdout
The logfile_read and logfile_send members can be used to separately log
the input from the child and output sent to the child. Sometimes you
don't want to see everything you write to the child. You only want to
log what the child sends back. For example::
child = pexpect.spawn('some_command')
child.logfile_read = sys.stdout
You will need to pass an encoding to spawn in the above code if you are
using Python 3.
To separately log output sent to the child use logfile_send::
child.logfile_send = fout
If ``ignore_sighup`` is True, the child process will ignore SIGHUP
signals. The default is False from Pexpect 4.0, meaning that SIGHUP
will be handled normally by the child.
The delaybeforesend helps overcome a weird behavior that many users
were experiencing. The typical problem was that a user would expect() a
"Password:" prompt and then immediately call sendline() to send the
password. The user would then see that their password was echoed back
to them. Passwords don't normally echo. The problem is caused by the
fact that most applications print out the "Password" prompt and then
turn off stdin echo, but if you send your password before the
application turned off echo, then you get your password echoed.
Normally this wouldn't be a problem when interacting with a human at a
real keyboard. If you introduce a slight delay just before writing then
this seems to clear up the problem. This was such a common problem for
many users that I decided that the default pexpect behavior should be
to sleep just before writing to the child application. 1/20th of a
second (50 ms) seems to be enough to clear up the problem. You can set
delaybeforesend to None to return to the old behavior.
Note that spawn is clever about finding commands on your path.
It uses the same logic that "which" uses to find executables.
If you wish to get the exit status of the child you must call the
close() method. The exit or signal status of the child will be stored
in self.exitstatus or self.signalstatus. If the child exited normally
then exitstatus will store the exit return code and signalstatus will
be None. If the child was terminated abnormally with a signal then
signalstatus will store the signal value and exitstatus will be None::
child = pexpect.spawn('some_command')
child.close()
print(child.exitstatus, child.signalstatus)
If you need more detail you can also read the self.status member which
stores the status returned by os.waitpid. You can interpret this using
os.WIFEXITED/os.WEXITSTATUS or os.WIFSIGNALED/os.TERMSIG.
The echo attribute may be set to False to disable echoing of input.
As a pseudo-terminal, all input echoed by the "keyboard" (send()
or sendline()) will be repeated to output. For many cases, it is
not desirable to have echo enabled, and it may be later disabled
using setecho(False) followed by waitnoecho(). However, for some
platforms such as Solaris, this is not possible, and should be
disabled immediately on spawn.
If preexec_fn is given, it will be called in the child process before
launching the given command. This is useful to e.g. reset inherited
signal handlers.
The dimensions attribute specifies the size of the pseudo-terminal as
seen by the subprocess, and is specified as a two-entry tuple (rows,
columns). If this is unspecified, the defaults in ptyprocess will apply.
'''
super(spawn, self).__init__(timeout=timeout, maxread=maxread, searchwindowsize=searchwindowsize,
logfile=logfile, encoding=encoding, codec_errors=codec_errors)
self.STDIN_FILENO = pty.STDIN_FILENO
self.STDOUT_FILENO = pty.STDOUT_FILENO
self.STDERR_FILENO = pty.STDERR_FILENO
self.cwd = cwd
self.env = env
self.echo = echo
self.ignore_sighup = ignore_sighup
self.__irix_hack = sys.platform.lower().startswith('irix')
if command is None:
self.command = None
self.args = None
self.name = '<pexpect factory incomplete>'
else:
self._spawn(command, args, preexec_fn, dimensions)
def __str__(self):
'''This returns a human-readable string that represents the state of
the object. '''
s = []
s.append(repr(self))
s.append('command: ' + str(self.command))
s.append('args: %r' % (self.args,))
s.append('buffer (last 100 chars): %r' % (
self.buffer[-100:] if self.buffer else self.buffer,))
s.append('before (last 100 chars): %r' % (
self.before[-100:] if self.before else self.before,))
s.append('after: %r' % (self.after,))
s.append('match: %r' % (self.match,))
s.append('match_index: ' + str(self.match_index))
s.append('exitstatus: ' + str(self.exitstatus))
if hasattr(self, 'ptyproc'):
s.append('flag_eof: ' + str(self.flag_eof))
s.append('pid: ' + str(self.pid))
s.append('child_fd: ' + str(self.child_fd))
s.append('closed: ' + str(self.closed))
s.append('timeout: ' + str(self.timeout))
s.append('delimiter: ' + str(self.delimiter))
s.append('logfile: ' + str(self.logfile))
s.append('logfile_read: ' + str(self.logfile_read))
s.append('logfile_send: ' + str(self.logfile_send))
s.append('maxread: ' + str(self.maxread))
s.append('ignorecase: ' + str(self.ignorecase))
s.append('searchwindowsize: ' + str(self.searchwindowsize))
s.append('delaybeforesend: ' + str(self.delaybeforesend))
s.append('delayafterclose: ' + str(self.delayafterclose))
s.append('delayafterterminate: ' + str(self.delayafterterminate))
return '\n'.join(s)
def _spawn(self, command, args=[], preexec_fn=None, dimensions=None):
'''This starts the given command in a child process. This does all the
fork/exec type of stuff for a pty. This is called by __init__. If args
is empty then command will be parsed (split on spaces) and args will be
set to parsed arguments. '''
# The pid and child_fd of this object get set by this method.
# Note that it is difficult for this method to fail.
# You cannot detect if the child process cannot start.
# So the only way you can tell if the child process started
# or not is to try to read from the file descriptor. If you get
# EOF immediately then it means that the child is already dead.
# That may not necessarily be bad because you may have spawned a child
# that performs some task; creates no stdout output; and then dies.
# If command is an int type then it may represent a file descriptor.
if isinstance(command, type(0)):
raise ExceptionPexpect('Command is an int type. ' +
'If this is a file descriptor then maybe you want to ' +
'use fdpexpect.fdspawn which takes an existing ' +
'file descriptor instead of a command string.')
if not isinstance(args, type([])):
raise TypeError('The argument, args, must be a list.')
if args == []:
self.args = split_command_line(command)
self.command = self.args[0]
else:
# Make a shallow copy of the args list.
self.args = args[:]
self.args.insert(0, command)
self.command = command
command_with_path = which(self.command, env=self.env)
if command_with_path is None:
raise ExceptionPexpect('The command was not found or was not ' +
'executable: %s.' % self.command)
self.command = command_with_path
self.args[0] = self.command
self.name = '<' + ' '.join(self.args) + '>'
assert self.pid is None, 'The pid member must be None.'
assert self.command is not None, 'The command member must not be None.'
kwargs = {'echo': self.echo, 'preexec_fn': preexec_fn}
if self.ignore_sighup:
def preexec_wrapper():
"Set SIGHUP to be ignored, then call the real preexec_fn"
signal.signal(signal.SIGHUP, signal.SIG_IGN)
if preexec_fn is not None:
preexec_fn()
kwargs['preexec_fn'] = preexec_wrapper
if dimensions is not None:
kwargs['dimensions'] = dimensions
if self.encoding is not None:
# Encode command line using the specified encoding
self.args = [a if isinstance(a, bytes) else a.encode(self.encoding)
for a in self.args]
self.ptyproc = self._spawnpty(self.args, env=self.env,
cwd=self.cwd, **kwargs)
self.pid = self.ptyproc.pid
self.child_fd = self.ptyproc.fd
self.terminated = False
self.closed = False
def _spawnpty(self, args, **kwargs):
'''Spawn a pty and return an instance of PtyProcess.'''
return ptyprocess.PtyProcess.spawn(args, **kwargs)
def close(self, force=True):
'''This closes the connection with the child application. Note that
calling close() more than once is valid. This emulates standard Python
behavior with files. Set force to True if you want to make sure that
the child is terminated (SIGKILL is sent if the child ignores SIGHUP
and SIGINT). '''
self.flush()
self.ptyproc.close(force=force)
self.isalive() # Update exit status from ptyproc
self.child_fd = -1
def isatty(self):
'''This returns True if the file descriptor is open and connected to a
tty(-like) device, else False.
On SVR4-style platforms implementing streams, such as SunOS and HP-UX,
the child pty may not appear as a terminal device. This means
methods such as setecho(), setwinsize(), getwinsize() may raise an
IOError. '''
return os.isatty(self.child_fd)
def waitnoecho(self, timeout=-1):
'''This waits until the terminal ECHO flag is set False. This returns
True if the echo mode is off. This returns False if the ECHO flag was
not set False before the timeout. This can be used to detect when the
child is waiting for a password. Usually a child application will turn
off echo mode when it is waiting for the user to enter a password. For
example, instead of expecting the "password:" prompt you can wait for
the child to set ECHO off::
p = pexpect.spawn('ssh user@example.com')
p.waitnoecho()
p.sendline(mypassword)
If timeout==-1 then this method will use the value in self.timeout.
If timeout==None then this method to block until ECHO flag is False.
'''
if timeout == -1:
timeout = self.timeout
if timeout is not None:
end_time = time.time() + timeout
while True:
if not self.getecho():
return True
if timeout < 0 and timeout is not None:
return False
if timeout is not None:
timeout = end_time - time.time()
time.sleep(0.1)
def getecho(self):
'''This returns the terminal echo mode. This returns True if echo is
on or False if echo is off. Child applications that are expecting you
to enter a password often set ECHO False. See waitnoecho().
Not supported on platforms where ``isatty()`` returns False. '''
return self.ptyproc.getecho()
def setecho(self, state):
'''This sets the terminal echo mode on or off. Note that anything the
child sent before the echo will be lost, so you should be sure that
your input buffer is empty before you call setecho(). For example, the
following will work as expected::
p = pexpect.spawn('cat') # Echo is on by default.
p.sendline('1234') # We expect see this twice from the child...
p.expect(['1234']) # ... once from the tty echo...
p.expect(['1234']) # ... and again from cat itself.
p.setecho(False) # Turn off tty echo
p.sendline('abcd') # We will set this only once (echoed by cat).
p.sendline('wxyz') # We will set this only once (echoed by cat)
p.expect(['abcd'])
p.expect(['wxyz'])
The following WILL NOT WORK because the lines sent before the setecho
will be lost::
p = pexpect.spawn('cat')
p.sendline('1234')
p.setecho(False) # Turn off tty echo
p.sendline('abcd') # We will set this only once (echoed by cat).
p.sendline('wxyz') # We will set this only once (echoed by cat)
p.expect(['1234'])
p.expect(['1234'])
p.expect(['abcd'])
p.expect(['wxyz'])
Not supported on platforms where ``isatty()`` returns False.
'''
return self.ptyproc.setecho(state)
def read_nonblocking(self, size=1, timeout=-1):
'''This reads at most size characters from the child application. It
includes a timeout. If the read does not complete within the timeout
period then a TIMEOUT exception is raised. If the end of file is read
then an EOF exception will be raised. If a logfile is specified, a
copy is written to that log.
If timeout is None then the read may block indefinitely.
If timeout is -1 then the self.timeout value is used. If timeout is 0
then the child is polled and if there is no data immediately ready
then this will raise a TIMEOUT exception.
The timeout refers only to the amount of time to read at least one
character. This is not affected by the 'size' parameter, so if you call
read_nonblocking(size=100, timeout=30) and only one character is
available right away then one character will be returned immediately.
It will not wait for 30 seconds for another 99 characters to come in.
This is a wrapper around os.read(). It uses select.select() to
implement the timeout. '''
if self.closed:
raise ValueError('I/O operation on closed file.')
if timeout == -1:
timeout = self.timeout
# Note that some systems such as Solaris do not give an EOF when
# the child dies. In fact, you can still try to read
# from the child_fd -- it will block forever or until TIMEOUT.
# For this case, I test isalive() before doing any reading.
# If isalive() is false, then I pretend that this is the same as EOF.
if not self.isalive():
# timeout of 0 means "poll"
r, w, e = select_ignore_interrupts([self.child_fd], [], [], 0)
if not r:
self.flag_eof = True
raise EOF('End Of File (EOF). Braindead platform.')
elif self.__irix_hack:
# Irix takes a long time before it realizes a child was terminated.
# FIXME So does this mean Irix systems are forced to always have
# FIXME a 2 second delay when calling read_nonblocking? That sucks.
r, w, e = select_ignore_interrupts([self.child_fd], [], [], 2)
if not r and not self.isalive():
self.flag_eof = True
raise EOF('End Of File (EOF). Slow platform.')
r, w, e = select_ignore_interrupts([self.child_fd], [], [], timeout)
if not r:
if not self.isalive():
# Some platforms, such as Irix, will claim that their
# processes are alive; timeout on the select; and
# then finally admit that they are not alive.
self.flag_eof = True
raise EOF('End of File (EOF). Very slow platform.')
else:
raise TIMEOUT('Timeout exceeded.')
if self.child_fd in r:
return super(spawn, self).read_nonblocking(size)
raise ExceptionPexpect('Reached an unexpected state.') # pragma: no cover
def write(self, s):
'''This is similar to send() except that there is no return value.
'''
self.send(s)
def writelines(self, sequence):
'''This calls write() for each element in the sequence. The sequence
can be any iterable object producing strings, typically a list of
strings. This does not add line separators. There is no return value.
'''
for s in sequence:
self.write(s)
def send(self, s):
'''Sends string ``s`` to the child process, returning the number of
bytes written. If a logfile is specified, a copy is written to that
log.
The default terminal input mode is canonical processing unless set
otherwise by the child process. This allows backspace and other line
processing to be performed prior to transmitting to the receiving
program. As this is buffered, there is a limited size of such buffer.
On Linux systems, this is 4096 (defined by N_TTY_BUF_SIZE). All
other systems honor the POSIX.1 definition PC_MAX_CANON -- 1024
on OSX, 256 on OpenSolaris, and 1920 on FreeBSD.
This value may be discovered using fpathconf(3)::
>>> from os import fpathconf
>>> print(fpathconf(0, 'PC_MAX_CANON'))
256
On such a system, only 256 bytes may be received per line. Any
subsequent bytes received will be discarded. BEL (``'\a'``) is then
sent to output if IMAXBEL (termios.h) is set by the tty driver.
This is usually enabled by default. Linux does not honor this as
an option -- it behaves as though it is always set on.
Canonical input processing may be disabled altogether by executing
a shell, then stty(1), before executing the final program::
>>> bash = pexpect.spawn('/bin/bash', echo=False)
>>> bash.sendline('stty -icanon')
>>> bash.sendline('base64')
>>> bash.sendline('x' * 5000)
'''
if self.delaybeforesend is not None:
time.sleep(self.delaybeforesend)
s = self._coerce_send_string(s)
self._log(s, 'send')
b = self._encoder.encode(s, final=False)
return os.write(self.child_fd, b)
def sendline(self, s=''):
'''Wraps send(), sending string ``s`` to child process, with
``os.linesep`` automatically appended. Returns number of bytes
written. Only a limited number of bytes may be sent for each
line in the default terminal mode, see docstring of :meth:`send`.
'''
s = self._coerce_send_string(s)
return self.send(s + self.linesep)
def _log_control(self, s):
"""Write control characters to the appropriate log files"""
if self.encoding is not None:
s = s.decode(self.encoding, 'replace')
self._log(s, 'send')
def sendcontrol(self, char):
'''Helper method that wraps send() with mnemonic access for sending control
character to the child (such as Ctrl-C or Ctrl-D). For example, to send
Ctrl-G (ASCII 7, bell, '\a')::
child.sendcontrol('g')
See also, sendintr() and sendeof().
'''
n, byte = self.ptyproc.sendcontrol(char)
self._log_control(byte)
return n
def sendeof(self):
'''This sends an EOF to the child. This sends a character which causes
the pending parent output buffer to be sent to the waiting child
program without waiting for end-of-line. If it is the first character
of the line, the read() in the user program returns 0, which signifies
end-of-file. This means to work as expected a sendeof() has to be
called at the beginning of a line. This method does not send a newline.
It is the responsibility of the caller to ensure the eof is sent at the
beginning of a line. '''
n, byte = self.ptyproc.sendeof()
self._log_control(byte)
def sendintr(self):
'''This sends a SIGINT to the child. It does not require
the SIGINT to be the first character on a line. '''
n, byte = self.ptyproc.sendintr()
self._log_control(byte)
@property
def flag_eof(self):
return self.ptyproc.flag_eof
@flag_eof.setter
def flag_eof(self, value):
self.ptyproc.flag_eof = value
def eof(self):
'''This returns True if the EOF exception was ever raised.
'''
return self.flag_eof
def terminate(self, force=False):
'''This forces a child process to terminate. It starts nicely with
SIGHUP and SIGINT. If "force" is True then moves onto SIGKILL. This
returns True if the child was terminated. This returns False if the
child could not be terminated. '''
if not self.isalive():
return True
try:
self.kill(signal.SIGHUP)
time.sleep(self.delayafterterminate)
if not self.isalive():
return True
self.kill(signal.SIGCONT)
time.sleep(self.delayafterterminate)
if not self.isalive():
return True
self.kill(signal.SIGINT)
time.sleep(self.delayafterterminate)
if not self.isalive():
return True
if force:
self.kill(signal.SIGKILL)
time.sleep(self.delayafterterminate)
if not self.isalive():
return True
else:
return False
return False
except OSError:
# I think there are kernel timing issues that sometimes cause
# this to happen. I think isalive() reports True, but the
# process is dead to the kernel.
# Make one last attempt to see if the kernel is up to date.
time.sleep(self.delayafterterminate)
if not self.isalive():
return True
else:
return False
def wait(self):
'''This waits until the child exits. This is a blocking call. This will
not read any data from the child, so this will block forever if the
child has unread output and has terminated. In other words, the child
may have printed output then called exit(), but, the child is
technically still alive until its output is read by the parent.
This method is non-blocking if :meth:`wait` has already been called
previously or :meth:`isalive` method returns False. It simply returns
the previously determined exit status.
'''
ptyproc = self.ptyproc
with _wrap_ptyprocess_err():
# exception may occur if "Is some other process attempting
# "job control with our child pid?"
exitstatus = ptyproc.wait()
self.status = ptyproc.status
self.exitstatus = ptyproc.exitstatus
self.signalstatus = ptyproc.signalstatus
self.terminated = True
return exitstatus
def isalive(self):
'''This tests if the child process is running or not. This is
non-blocking. If the child was terminated then this will read the
exitstatus or signalstatus of the child. This returns True if the child
process appears to be running or False if not. It can take literally
SECONDS for Solaris to return the right status. '''
ptyproc = self.ptyproc
with _wrap_ptyprocess_err():
alive = ptyproc.isalive()
if not alive:
self.status = ptyproc.status
self.exitstatus = ptyproc.exitstatus
self.signalstatus = ptyproc.signalstatus
self.terminated = True
return alive
def kill(self, sig):
'''This sends the given signal to the child application. In keeping
with UNIX tradition it has a misleading name. It does not necessarily
kill the child unless you send the right signal. '''
# Same as os.kill, but the pid is given for you.
if self.isalive():
os.kill(self.pid, sig)
def getwinsize(self):
'''This returns the terminal window size of the child tty. The return
value is a tuple of (rows, cols). '''
return self.ptyproc.getwinsize()
def setwinsize(self, rows, cols):
'''This sets the terminal window size of the child tty. This will cause
a SIGWINCH signal to be sent to the child. This does not change the
physical window size. It changes the size reported to TTY-aware
applications like vi or curses -- applications that respond to the
SIGWINCH signal. '''
return self.ptyproc.setwinsize(rows, cols)
def interact(self, escape_character=chr(29),
input_filter=None, output_filter=None):
'''This gives control of the child process to the interactive user (the
human at the keyboard). Keystrokes are sent to the child process, and
the stdout and stderr output of the child process is printed. This
simply echos the child stdout and child stderr to the real stdout and
it echos the real stdin to the child stdin. When the user types the
escape_character this method will return None. The escape_character
will not be transmitted. The default for escape_character is
entered as ``Ctrl - ]``, the very same as BSD telnet. To prevent
escaping, escape_character may be set to None.
If a logfile is specified, then the data sent and received from the
child process in interact mode is duplicated to the given log.
You may pass in optional input and output filter functions. These
functions should take a string and return a string. The output_filter
will be passed all the output from the child process. The input_filter
will be passed all the keyboard input from the user. The input_filter
is run BEFORE the check for the escape_character.
Note that if you change the window size of the parent the SIGWINCH
signal will not be passed through to the child. If you want the child
window size to change when the parent's window size changes then do
something like the following example::
import pexpect, struct, fcntl, termios, signal, sys
def sigwinch_passthrough (sig, data):
s = struct.pack("HHHH", 0, 0, 0, 0)
a = struct.unpack('hhhh', fcntl.ioctl(sys.stdout.fileno(),
termios.TIOCGWINSZ , s))
global p
p.setwinsize(a[0],a[1])
# Note this 'p' global and used in sigwinch_passthrough.
p = pexpect.spawn('/bin/bash')
signal.signal(signal.SIGWINCH, sigwinch_passthrough)
p.interact()
'''
# Flush the buffer.
self.write_to_stdout(self.buffer)
self.stdout.flush()
self.buffer = self.string_type()
mode = tty.tcgetattr(self.STDIN_FILENO)
tty.setraw(self.STDIN_FILENO)
if escape_character is not None and PY3:
escape_character = escape_character.encode('latin-1')
try:
self.__interact_copy(escape_character, input_filter, output_filter)
finally:
tty.tcsetattr(self.STDIN_FILENO, tty.TCSAFLUSH, mode)
def __interact_writen(self, fd, data):
'''This is used by the interact() method.
'''
while data != b'' and self.isalive():
n = os.write(fd, data)
data = data[n:]
def __interact_read(self, fd):
'''This is used by the interact() method.
'''
return os.read(fd, 1000)
def __interact_copy(self, escape_character=None,
input_filter=None, output_filter=None):
'''This is used by the interact() method.
'''
while self.isalive():
r, w, e = select_ignore_interrupts([self.child_fd, self.STDIN_FILENO], [], [])
if self.child_fd in r:
try:
data = self.__interact_read(self.child_fd)
except OSError as err:
if err.args[0] == errno.EIO:
# Linux-style EOF
break
raise
if data == b'':
# BSD-style EOF
break
if output_filter:
data = output_filter(data)
self._log(data, 'read')
os.write(self.STDOUT_FILENO, data)
if self.STDIN_FILENO in r:
data = self.__interact_read(self.STDIN_FILENO)
if input_filter:
data = input_filter(data)
i = -1
if escape_character is not None:
i = data.rfind(escape_character)
if i != -1:
data = data[:i]
if data:
self._log(data, 'send')
self.__interact_writen(self.child_fd, data)
break
self._log(data, 'send')
self.__interact_writen(self.child_fd, data)
def spawnu(*args, **kwargs):
"""Deprecated: pass encoding to spawn() instead."""
kwargs.setdefault('encoding', 'utf-8')
return spawn(*args, **kwargs)
+409
View File
@@ -0,0 +1,409 @@
'''This class extends pexpect.spawn to specialize setting up SSH connections.
This adds methods for login, logout, and expecting the shell prompt.
PEXPECT LICENSE
This license is approved by the OSI and FSF as GPL-compatible.
http://opensource.org/licenses/isc-license.txt
Copyright (c) 2012, Noah Spurrier <noah@noah.org>
PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''
from pexpect import ExceptionPexpect, TIMEOUT, EOF, spawn
import time
import os
__all__ = ['ExceptionPxssh', 'pxssh']
# Exception classes used by this module.
class ExceptionPxssh(ExceptionPexpect):
'''Raised for pxssh exceptions.
'''
class pxssh (spawn):
'''This class extends pexpect.spawn to specialize setting up SSH
connections. This adds methods for login, logout, and expecting the shell
prompt. It does various tricky things to handle many situations in the SSH
login process. For example, if the session is your first login, then pxssh
automatically accepts the remote certificate; or if you have public key
authentication setup then pxssh won't wait for the password prompt.
pxssh uses the shell prompt to synchronize output from the remote host. In
order to make this more robust it sets the shell prompt to something more
unique than just $ or #. This should work on most Borne/Bash or Csh style
shells.
Example that runs a few commands on a remote server and prints the result::
from pexpect import pxssh
import getpass
try:
s = pxssh.pxssh()
hostname = raw_input('hostname: ')
username = raw_input('username: ')
password = getpass.getpass('password: ')
s.login(hostname, username, password)
s.sendline('uptime') # run a command
s.prompt() # match the prompt
print(s.before) # print everything before the prompt.
s.sendline('ls -l')
s.prompt()
print(s.before)
s.sendline('df')
s.prompt()
print(s.before)
s.logout()
except pxssh.ExceptionPxssh as e:
print("pxssh failed on login.")
print(e)
Example showing how to specify SSH options::
from pexpect import pxssh
s = pxssh.pxssh(options={
"StrictHostKeyChecking": "no",
"UserKnownHostsFile": "/dev/null"})
...
Note that if you have ssh-agent running while doing development with pxssh
then this can lead to a lot of confusion. Many X display managers (xdm,
gdm, kdm, etc.) will automatically start a GUI agent. You may see a GUI
dialog box popup asking for a password during development. You should turn
off any key agents during testing. The 'force_password' attribute will turn
off public key authentication. This will only work if the remote SSH server
is configured to allow password logins. Example of using 'force_password'
attribute::
s = pxssh.pxssh()
s.force_password = True
hostname = raw_input('hostname: ')
username = raw_input('username: ')
password = getpass.getpass('password: ')
s.login (hostname, username, password)
'''
def __init__ (self, timeout=30, maxread=2000, searchwindowsize=None,
logfile=None, cwd=None, env=None, ignore_sighup=True, echo=True,
options={}, encoding=None, codec_errors='strict'):
spawn.__init__(self, None, timeout=timeout, maxread=maxread,
searchwindowsize=searchwindowsize, logfile=logfile,
cwd=cwd, env=env, ignore_sighup=ignore_sighup, echo=echo,
encoding=encoding, codec_errors=codec_errors)
self.name = '<pxssh>'
#SUBTLE HACK ALERT! Note that the command that SETS the prompt uses a
#slightly different string than the regular expression to match it. This
#is because when you set the prompt the command will echo back, but we
#don't want to match the echoed command. So if we make the set command
#slightly different than the regex we eliminate the problem. To make the
#set command different we add a backslash in front of $. The $ doesn't
#need to be escaped, but it doesn't hurt and serves to make the set
#prompt command different than the regex.
# used to match the command-line prompt
self.UNIQUE_PROMPT = "\[PEXPECT\][\$\#] "
self.PROMPT = self.UNIQUE_PROMPT
# used to set shell command-line prompt to UNIQUE_PROMPT.
self.PROMPT_SET_SH = "PS1='[PEXPECT]\$ '"
self.PROMPT_SET_CSH = "set prompt='[PEXPECT]\$ '"
self.SSH_OPTS = ("-o'RSAAuthentication=no'"
+ " -o 'PubkeyAuthentication=no'")
# Disabling host key checking, makes you vulnerable to MITM attacks.
# + " -o 'StrictHostKeyChecking=no'"
# + " -o 'UserKnownHostsFile /dev/null' ")
# Disabling X11 forwarding gets rid of the annoying SSH_ASKPASS from
# displaying a GUI password dialog. I have not figured out how to
# disable only SSH_ASKPASS without also disabling X11 forwarding.
# Unsetting SSH_ASKPASS on the remote side doesn't disable it! Annoying!
#self.SSH_OPTS = "-x -o'RSAAuthentication=no' -o 'PubkeyAuthentication=no'"
self.force_password = False
# User defined SSH options, eg,
# ssh.otions = dict(StrictHostKeyChecking="no",UserKnownHostsFile="/dev/null")
self.options = options
def levenshtein_distance(self, a, b):
'''This calculates the Levenshtein distance between a and b.
'''
n, m = len(a), len(b)
if n > m:
a,b = b,a
n,m = m,n
current = range(n+1)
for i in range(1,m+1):
previous, current = current, [i]+[0]*n
for j in range(1,n+1):
add, delete = previous[j]+1, current[j-1]+1
change = previous[j-1]
if a[j-1] != b[i-1]:
change = change + 1
current[j] = min(add, delete, change)
return current[n]
def try_read_prompt(self, timeout_multiplier):
'''This facilitates using communication timeouts to perform
synchronization as quickly as possible, while supporting high latency
connections with a tunable worst case performance. Fast connections
should be read almost immediately. Worst case performance for this
method is timeout_multiplier * 3 seconds.
'''
# maximum time allowed to read the first response
first_char_timeout = timeout_multiplier * 0.5
# maximum time allowed between subsequent characters
inter_char_timeout = timeout_multiplier * 0.1
# maximum time for reading the entire prompt
total_timeout = timeout_multiplier * 3.0
prompt = self.string_type()
begin = time.time()
expired = 0.0
timeout = first_char_timeout
while expired < total_timeout:
try:
prompt += self.read_nonblocking(size=1, timeout=timeout)
expired = time.time() - begin # updated total time expired
timeout = inter_char_timeout
except TIMEOUT:
break
return prompt
def sync_original_prompt (self, sync_multiplier=1.0):
'''This attempts to find the prompt. Basically, press enter and record
the response; press enter again and record the response; if the two
responses are similar then assume we are at the original prompt.
This can be a slow function. Worst case with the default sync_multiplier
can take 12 seconds. Low latency connections are more likely to fail
with a low sync_multiplier. Best case sync time gets worse with a
high sync multiplier (500 ms with default). '''
# All of these timing pace values are magic.
# I came up with these based on what seemed reliable for
# connecting to a heavily loaded machine I have.
self.sendline()
time.sleep(0.1)
try:
# Clear the buffer before getting the prompt.
self.try_read_prompt(sync_multiplier)
except TIMEOUT:
pass
self.sendline()
x = self.try_read_prompt(sync_multiplier)
self.sendline()
a = self.try_read_prompt(sync_multiplier)
self.sendline()
b = self.try_read_prompt(sync_multiplier)
ld = self.levenshtein_distance(a,b)
len_a = len(a)
if len_a == 0:
return False
if float(ld)/len_a < 0.4:
return True
return False
### TODO: This is getting messy and I'm pretty sure this isn't perfect.
### TODO: I need to draw a flow chart for this.
def login (self, server, username, password='', terminal_type='ansi',
original_prompt=r"[#$]", login_timeout=10, port=None,
auto_prompt_reset=True, ssh_key=None, quiet=True,
sync_multiplier=1, check_local_ip=True):
'''This logs the user into the given server.
It uses
'original_prompt' to try to find the prompt right after login. When it
finds the prompt it immediately tries to reset the prompt to something
more easily matched. The default 'original_prompt' is very optimistic
and is easily fooled. It's more reliable to try to match the original
prompt as exactly as possible to prevent false matches by server
strings such as the "Message Of The Day". On many systems you can
disable the MOTD on the remote server by creating a zero-length file
called :file:`~/.hushlogin` on the remote server. If a prompt cannot be found
then this will not necessarily cause the login to fail. In the case of
a timeout when looking for the prompt we assume that the original
prompt was so weird that we could not match it, so we use a few tricks
to guess when we have reached the prompt. Then we hope for the best and
blindly try to reset the prompt to something more unique. If that fails
then login() raises an :class:`ExceptionPxssh` exception.
In some situations it is not possible or desirable to reset the
original prompt. In this case, pass ``auto_prompt_reset=False`` to
inhibit setting the prompt to the UNIQUE_PROMPT. Remember that pxssh
uses a unique prompt in the :meth:`prompt` method. If the original prompt is
not reset then this will disable the :meth:`prompt` method unless you
manually set the :attr:`PROMPT` attribute.
'''
ssh_options = ''.join([" -o '%s=%s'" % (o, v) for (o, v) in self.options.items()])
if quiet:
ssh_options = ssh_options + ' -q'
if not check_local_ip:
ssh_options = ssh_options + " -o'NoHostAuthenticationForLocalhost=yes'"
if self.force_password:
ssh_options = ssh_options + ' ' + self.SSH_OPTS
if port is not None:
ssh_options = ssh_options + ' -p %s'%(str(port))
if ssh_key is not None:
try:
os.path.isfile(ssh_key)
except:
raise ExceptionPxssh('private ssh key does not exist')
ssh_options = ssh_options + ' -i %s' % (ssh_key)
cmd = "ssh %s -l %s %s" % (ssh_options, username, server)
# This does not distinguish between a remote server 'password' prompt
# and a local ssh 'passphrase' prompt (for unlocking a private key).
spawn._spawn(self, cmd)
i = self.expect(["(?i)are you sure you want to continue connecting", original_prompt, "(?i)(?:password)|(?:passphrase for key)", "(?i)permission denied", "(?i)terminal type", TIMEOUT, "(?i)connection closed by remote host", EOF], timeout=login_timeout)
# First phase
if i==0:
# New certificate -- always accept it.
# This is what you get if SSH does not have the remote host's
# public key stored in the 'known_hosts' cache.
self.sendline("yes")
i = self.expect(["(?i)are you sure you want to continue connecting", original_prompt, "(?i)(?:password)|(?:passphrase for key)", "(?i)permission denied", "(?i)terminal type", TIMEOUT])
if i==2: # password or passphrase
self.sendline(password)
i = self.expect(["(?i)are you sure you want to continue connecting", original_prompt, "(?i)(?:password)|(?:passphrase for key)", "(?i)permission denied", "(?i)terminal type", TIMEOUT])
if i==4:
self.sendline(terminal_type)
i = self.expect(["(?i)are you sure you want to continue connecting", original_prompt, "(?i)(?:password)|(?:passphrase for key)", "(?i)permission denied", "(?i)terminal type", TIMEOUT])
if i==7:
self.close()
raise ExceptionPxssh('Could not establish connection to host')
# Second phase
if i==0:
# This is weird. This should not happen twice in a row.
self.close()
raise ExceptionPxssh('Weird error. Got "are you sure" prompt twice.')
elif i==1: # can occur if you have a public key pair set to authenticate.
### TODO: May NOT be OK if expect() got tricked and matched a false prompt.
pass
elif i==2: # password prompt again
# For incorrect passwords, some ssh servers will
# ask for the password again, others return 'denied' right away.
# If we get the password prompt again then this means
# we didn't get the password right the first time.
self.close()
raise ExceptionPxssh('password refused')
elif i==3: # permission denied -- password was bad.
self.close()
raise ExceptionPxssh('permission denied')
elif i==4: # terminal type again? WTF?
self.close()
raise ExceptionPxssh('Weird error. Got "terminal type" prompt twice.')
elif i==5: # Timeout
#This is tricky... I presume that we are at the command-line prompt.
#It may be that the shell prompt was so weird that we couldn't match
#it. Or it may be that we couldn't log in for some other reason. I
#can't be sure, but it's safe to guess that we did login because if
#I presume wrong and we are not logged in then this should be caught
#later when I try to set the shell prompt.
pass
elif i==6: # Connection closed by remote host
self.close()
raise ExceptionPxssh('connection closed')
else: # Unexpected
self.close()
raise ExceptionPxssh('unexpected login response')
if not self.sync_original_prompt(sync_multiplier):
self.close()
raise ExceptionPxssh('could not synchronize with original prompt')
# We appear to be in.
# set shell prompt to something unique.
if auto_prompt_reset:
if not self.set_unique_prompt():
self.close()
raise ExceptionPxssh('could not set shell prompt '
'(received: %r, expected: %r).' % (
self.before, self.PROMPT,))
return True
def logout (self):
'''Sends exit to the remote shell.
If there are stopped jobs then this automatically sends exit twice.
'''
self.sendline("exit")
index = self.expect([EOF, "(?i)there are stopped jobs"])
if index==1:
self.sendline("exit")
self.expect(EOF)
self.close()
def prompt(self, timeout=-1):
'''Match the next shell prompt.
This is little more than a short-cut to the :meth:`~pexpect.spawn.expect`
method. Note that if you called :meth:`login` with
``auto_prompt_reset=False``, then before calling :meth:`prompt` you must
set the :attr:`PROMPT` attribute to a regex that it will use for
matching the prompt.
Calling :meth:`prompt` will erase the contents of the :attr:`before`
attribute even if no prompt is ever matched. If timeout is not given or
it is set to -1 then self.timeout is used.
:return: True if the shell prompt was matched, False if the timeout was
reached.
'''
if timeout == -1:
timeout = self.timeout
i = self.expect([self.PROMPT, TIMEOUT], timeout=timeout)
if i==1:
return False
return True
def set_unique_prompt(self):
'''This sets the remote prompt to something more unique than ``#`` or ``$``.
This makes it easier for the :meth:`prompt` method to match the shell prompt
unambiguously. This method is called automatically by the :meth:`login`
method, but you may want to call it manually if you somehow reset the
shell prompt. For example, if you 'su' to a different user then you
will need to manually reset the prompt. This sends shell commands to
the remote host to set the prompt, so this assumes the remote host is
ready to receive commands.
Alternatively, you may use your own prompt pattern. In this case you
should call :meth:`login` with ``auto_prompt_reset=False``; then set the
:attr:`PROMPT` attribute to a regular expression. After that, the
:meth:`prompt` method will try to match your prompt pattern.
'''
self.sendline("unset PROMPT_COMMAND")
self.sendline(self.PROMPT_SET_SH) # sh-style
i = self.expect ([TIMEOUT, self.PROMPT], timeout=10)
if i == 0: # csh-style
self.sendline(self.PROMPT_SET_CSH)
i = self.expect([TIMEOUT, self.PROMPT], timeout=10)
if i == 0:
return False
return True
# vi:ts=4:sw=4:expandtab:ft=python:
+122
View File
@@ -0,0 +1,122 @@
"""Generic wrapper for read-eval-print-loops, a.k.a. interactive shells
"""
import os.path
import signal
import sys
import pexpect
PY3 = (sys.version_info[0] >= 3)
if PY3:
basestring = str
PEXPECT_PROMPT = u'[PEXPECT_PROMPT>'
PEXPECT_CONTINUATION_PROMPT = u'[PEXPECT_PROMPT+'
class REPLWrapper(object):
"""Wrapper for a REPL.
:param cmd_or_spawn: This can either be an instance of :class:`pexpect.spawn`
in which a REPL has already been started, or a str command to start a new
REPL process.
:param str orig_prompt: The prompt to expect at first.
:param str prompt_change: A command to change the prompt to something more
unique. If this is ``None``, the prompt will not be changed. This will
be formatted with the new and continuation prompts as positional
parameters, so you can use ``{}`` style formatting to insert them into
the command.
:param str new_prompt: The more unique prompt to expect after the change.
:param str extra_init_cmd: Commands to do extra initialisation, such as
disabling pagers.
"""
def __init__(self, cmd_or_spawn, orig_prompt, prompt_change,
new_prompt=PEXPECT_PROMPT,
continuation_prompt=PEXPECT_CONTINUATION_PROMPT,
extra_init_cmd=None):
if isinstance(cmd_or_spawn, basestring):
self.child = pexpect.spawn(cmd_or_spawn, echo=False, encoding='utf-8')
else:
self.child = cmd_or_spawn
if self.child.echo:
# Existing spawn instance has echo enabled, disable it
# to prevent our input from being repeated to output.
self.child.setecho(False)
self.child.waitnoecho()
if prompt_change is None:
self.prompt = orig_prompt
else:
self.set_prompt(orig_prompt,
prompt_change.format(new_prompt, continuation_prompt))
self.prompt = new_prompt
self.continuation_prompt = continuation_prompt
self._expect_prompt()
if extra_init_cmd is not None:
self.run_command(extra_init_cmd)
def set_prompt(self, orig_prompt, prompt_change):
self.child.expect(orig_prompt)
self.child.sendline(prompt_change)
def _expect_prompt(self, timeout=-1):
return self.child.expect_exact([self.prompt, self.continuation_prompt],
timeout=timeout)
def run_command(self, command, timeout=-1):
"""Send a command to the REPL, wait for and return output.
:param str command: The command to send. Trailing newlines are not needed.
This should be a complete block of input that will trigger execution;
if a continuation prompt is found after sending input, :exc:`ValueError`
will be raised.
:param int timeout: How long to wait for the next prompt. -1 means the
default from the :class:`pexpect.spawn` object (default 30 seconds).
None means to wait indefinitely.
"""
# Split up multiline commands and feed them in bit-by-bit
cmdlines = command.splitlines()
# splitlines ignores trailing newlines - add it back in manually
if command.endswith('\n'):
cmdlines.append('')
if not cmdlines:
raise ValueError("No command was given")
res = []
self.child.sendline(cmdlines[0])
for line in cmdlines[1:]:
self._expect_prompt(timeout=timeout)
res.append(self.child.before)
self.child.sendline(line)
# Command was fully submitted, now wait for the next prompt
if self._expect_prompt(timeout=timeout) == 1:
# We got the continuation prompt - command was incomplete
self.child.kill(signal.SIGINT)
self._expect_prompt(timeout=1)
raise ValueError("Continuation prompt found - input was incomplete:\n"
+ command)
return u''.join(res + [self.child.before])
def python(command="python"):
"""Start a Python shell and return a :class:`REPLWrapper` object."""
return REPLWrapper(command, u">>> ", u"import sys; sys.ps1={0!r}; sys.ps2={1!r}")
def bash(command="bash"):
"""Start a bash shell and return a :class:`REPLWrapper` object."""
bashrc = os.path.join(os.path.dirname(__file__), 'bashrc.sh')
child = pexpect.spawn(command, ['--rcfile', bashrc], echo=False,
encoding='utf-8')
# If the user runs 'env', the value of PS1 will be in the output. To avoid
# replwrap seeing that as the next prompt, we'll embed the marker characters
# for invisible characters in the prompt; these show up when inspecting the
# environment variable, but not when bash displays the prompt.
ps1 = PEXPECT_PROMPT[:5] + u'\[\]' + PEXPECT_PROMPT[5:]
ps2 = PEXPECT_CONTINUATION_PROMPT[:5] + u'\[\]' + PEXPECT_CONTINUATION_PROMPT[5:]
prompt_change = u"PS1='{0}' PS2='{1}' PROMPT_COMMAND=''".format(ps1, ps2)
return REPLWrapper(child, u'\$', prompt_change,
extra_init_cmd="export PAGER=cat")
+157
View File
@@ -0,0 +1,157 @@
import sys
import types
from .exceptions import EOF, TIMEOUT
from .pty_spawn import spawn
def run(command, timeout=30, withexitstatus=False, events=None,
extra_args=None, logfile=None, cwd=None, env=None, **kwargs):
'''
This function runs the given command; waits for it to finish; then
returns all output as a string. STDERR is included in output. If the full
path to the command is not given then the path is searched.
Note that lines are terminated by CR/LF (\\r\\n) combination even on
UNIX-like systems because this is the standard for pseudottys. If you set
'withexitstatus' to true, then run will return a tuple of (command_output,
exitstatus). If 'withexitstatus' is false then this returns just
command_output.
The run() function can often be used instead of creating a spawn instance.
For example, the following code uses spawn::
from pexpect import *
child = spawn('scp foo user@example.com:.')
child.expect('(?i)password')
child.sendline(mypassword)
The previous code can be replace with the following::
from pexpect import *
run('scp foo user@example.com:.', events={'(?i)password': mypassword})
**Examples**
Start the apache daemon on the local machine::
from pexpect import *
run("/usr/local/apache/bin/apachectl start")
Check in a file using SVN::
from pexpect import *
run("svn ci -m 'automatic commit' my_file.py")
Run a command and capture exit status::
from pexpect import *
(command_output, exitstatus) = run('ls -l /bin', withexitstatus=1)
The following will run SSH and execute 'ls -l' on the remote machine. The
password 'secret' will be sent if the '(?i)password' pattern is ever seen::
run("ssh username@machine.example.com 'ls -l'",
events={'(?i)password':'secret\\n'})
This will start mencoder to rip a video from DVD. This will also display
progress ticks every 5 seconds as it runs. For example::
from pexpect import *
def print_ticks(d):
print d['event_count'],
run("mencoder dvd://1 -o video.avi -oac copy -ovc copy",
events={TIMEOUT:print_ticks}, timeout=5)
The 'events' argument should be either a dictionary or a tuple list that
contains patterns and responses. Whenever one of the patterns is seen
in the command output, run() will send the associated response string.
So, run() in the above example can be also written as:
run("mencoder dvd://1 -o video.avi -oac copy -ovc copy",
events=[(TIMEOUT,print_ticks)], timeout=5)
Use a tuple list for events if the command output requires a delicate
control over what pattern should be matched, since the tuple list is passed
to pexpect() as its pattern list, with the order of patterns preserved.
Note that you should put newlines in your string if Enter is necessary.
Like the example above, the responses may also contain a callback, either
a function or method. It should accept a dictionary value as an argument.
The dictionary contains all the locals from the run() function, so you can
access the child spawn object or any other variable defined in run()
(event_count, child, and extra_args are the most useful). A callback may
return True to stop the current run process. Otherwise run() continues
until the next event. A callback may also return a string which will be
sent to the child. 'extra_args' is not used by directly run(). It provides
a way to pass data to a callback function through run() through the locals
dictionary passed to a callback.
Like :class:`spawn`, passing *encoding* will make it work with unicode
instead of bytes. You can pass *codec_errors* to control how errors in
encoding and decoding are handled.
'''
if timeout == -1:
child = spawn(command, maxread=2000, logfile=logfile, cwd=cwd, env=env,
**kwargs)
else:
child = spawn(command, timeout=timeout, maxread=2000, logfile=logfile,
cwd=cwd, env=env, **kwargs)
if isinstance(events, list):
patterns= [x for x,y in events]
responses = [y for x,y in events]
elif isinstance(events, dict):
patterns = list(events.keys())
responses = list(events.values())
else:
# This assumes EOF or TIMEOUT will eventually cause run to terminate.
patterns = None
responses = None
child_result_list = []
event_count = 0
while True:
try:
index = child.expect(patterns)
if isinstance(child.after, child.allowed_string_types):
child_result_list.append(child.before + child.after)
else:
# child.after may have been a TIMEOUT or EOF,
# which we don't want appended to the list.
child_result_list.append(child.before)
if isinstance(responses[index], child.allowed_string_types):
child.send(responses[index])
elif (isinstance(responses[index], types.FunctionType) or
isinstance(responses[index], types.MethodType)):
callback_result = responses[index](locals())
sys.stdout.flush()
if isinstance(callback_result, child.allowed_string_types):
child.send(callback_result)
elif callback_result:
break
else:
raise TypeError("parameter `event' at index {index} must be "
"a string, method, or function: {value!r}"
.format(index=index, value=responses[index]))
event_count = event_count + 1
except TIMEOUT:
child_result_list.append(child.before)
break
except EOF:
child_result_list.append(child.before)
break
child_result = child.string_type().join(child_result_list)
if withexitstatus:
child.close()
return (child_result, child.exitstatus)
else:
return child_result
def runu(command, timeout=30, withexitstatus=False, events=None,
extra_args=None, logfile=None, cwd=None, env=None, **kwargs):
"""Deprecated: pass encoding to run() instead.
"""
kwargs.setdefault('encoding', 'utf-8')
return run(command, timeout=timeout, withexitstatus=withexitstatus,
events=events, extra_args=extra_args, logfile=logfile, cwd=cwd,
env=env, **kwargs)
+431
View File
@@ -0,0 +1,431 @@
'''This implements a virtual screen. This is used to support ANSI terminal
emulation. The screen representation and state is implemented in this class.
Most of the methods are inspired by ANSI screen control codes. The
:class:`~pexpect.ANSI.ANSI` class extends this class to add parsing of ANSI
escape codes.
PEXPECT LICENSE
This license is approved by the OSI and FSF as GPL-compatible.
http://opensource.org/licenses/isc-license.txt
Copyright (c) 2012, Noah Spurrier <noah@noah.org>
PERMISSION TO USE, COPY, MODIFY, AND/OR DISTRIBUTE THIS SOFTWARE FOR ANY
PURPOSE WITH OR WITHOUT FEE IS HEREBY GRANTED, PROVIDED THAT THE ABOVE
COPYRIGHT NOTICE AND THIS PERMISSION NOTICE APPEAR IN ALL COPIES.
THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
'''
import codecs
import copy
import sys
import warnings
warnings.warn(("pexpect.screen and pexpect.ANSI are deprecated. "
"We recommend using pyte to emulate a terminal screen: "
"https://pypi.python.org/pypi/pyte"),
stacklevel=2)
NUL = 0 # Fill character; ignored on input.
ENQ = 5 # Transmit answerback message.
BEL = 7 # Ring the bell.
BS = 8 # Move cursor left.
HT = 9 # Move cursor to next tab stop.
LF = 10 # Line feed.
VT = 11 # Same as LF.
FF = 12 # Same as LF.
CR = 13 # Move cursor to left margin or newline.
SO = 14 # Invoke G1 character set.
SI = 15 # Invoke G0 character set.
XON = 17 # Resume transmission.
XOFF = 19 # Halt transmission.
CAN = 24 # Cancel escape sequence.
SUB = 26 # Same as CAN.
ESC = 27 # Introduce a control sequence.
DEL = 127 # Fill character; ignored on input.
SPACE = u' ' # Space or blank character.
PY3 = (sys.version_info[0] >= 3)
if PY3:
unicode = str
def constrain (n, min, max):
'''This returns a number, n constrained to the min and max bounds. '''
if n < min:
return min
if n > max:
return max
return n
class screen:
'''This object maintains the state of a virtual text screen as a
rectangluar array. This maintains a virtual cursor position and handles
scrolling as characters are added. This supports most of the methods needed
by an ANSI text screen. Row and column indexes are 1-based (not zero-based,
like arrays).
Characters are represented internally using unicode. Methods that accept
input characters, when passed 'bytes' (which in Python 2 is equivalent to
'str'), convert them from the encoding specified in the 'encoding'
parameter to the constructor. Methods that return screen contents return
unicode strings, with the exception of __str__() under Python 2. Passing
``encoding=None`` limits the API to only accept unicode input, so passing
bytes in will raise :exc:`TypeError`.
'''
def __init__(self, r=24, c=80, encoding='latin-1', encoding_errors='replace'):
'''This initializes a blank screen of the given dimensions.'''
self.rows = r
self.cols = c
self.encoding = encoding
self.encoding_errors = encoding_errors
if encoding is not None:
self.decoder = codecs.getincrementaldecoder(encoding)(encoding_errors)
else:
self.decoder = None
self.cur_r = 1
self.cur_c = 1
self.cur_saved_r = 1
self.cur_saved_c = 1
self.scroll_row_start = 1
self.scroll_row_end = self.rows
self.w = [ [SPACE] * self.cols for _ in range(self.rows)]
def _decode(self, s):
'''This converts from the external coding system (as passed to
the constructor) to the internal one (unicode). '''
if self.decoder is not None:
return self.decoder.decode(s)
else:
raise TypeError("This screen was constructed with encoding=None, "
"so it does not handle bytes.")
def _unicode(self):
'''This returns a printable representation of the screen as a unicode
string (which, under Python 3.x, is the same as 'str'). The end of each
screen line is terminated by a newline.'''
return u'\n'.join ([ u''.join(c) for c in self.w ])
if PY3:
__str__ = _unicode
else:
__unicode__ = _unicode
def __str__(self):
'''This returns a printable representation of the screen. The end of
each screen line is terminated by a newline. '''
encoding = self.encoding or 'ascii'
return self._unicode().encode(encoding, 'replace')
def dump (self):
'''This returns a copy of the screen as a unicode string. This is similar to
__str__/__unicode__ except that lines are not terminated with line
feeds.'''
return u''.join ([ u''.join(c) for c in self.w ])
def pretty (self):
'''This returns a copy of the screen as a unicode string with an ASCII
text box around the screen border. This is similar to
__str__/__unicode__ except that it adds a box.'''
top_bot = u'+' + u'-'*self.cols + u'+\n'
return top_bot + u'\n'.join([u'|'+line+u'|' for line in unicode(self).split(u'\n')]) + u'\n' + top_bot
def fill (self, ch=SPACE):
if isinstance(ch, bytes):
ch = self._decode(ch)
self.fill_region (1,1,self.rows,self.cols, ch)
def fill_region (self, rs,cs, re,ce, ch=SPACE):
if isinstance(ch, bytes):
ch = self._decode(ch)
rs = constrain (rs, 1, self.rows)
re = constrain (re, 1, self.rows)
cs = constrain (cs, 1, self.cols)
ce = constrain (ce, 1, self.cols)
if rs > re:
rs, re = re, rs
if cs > ce:
cs, ce = ce, cs
for r in range (rs, re+1):
for c in range (cs, ce + 1):
self.put_abs (r,c,ch)
def cr (self):
'''This moves the cursor to the beginning (col 1) of the current row.
'''
self.cursor_home (self.cur_r, 1)
def lf (self):
'''This moves the cursor down with scrolling.
'''
old_r = self.cur_r
self.cursor_down()
if old_r == self.cur_r:
self.scroll_up ()
self.erase_line()
def crlf (self):
'''This advances the cursor with CRLF properties.
The cursor will line wrap and the screen may scroll.
'''
self.cr ()
self.lf ()
def newline (self):
'''This is an alias for crlf().
'''
self.crlf()
def put_abs (self, r, c, ch):
'''Screen array starts at 1 index.'''
r = constrain (r, 1, self.rows)
c = constrain (c, 1, self.cols)
if isinstance(ch, bytes):
ch = self._decode(ch)[0]
else:
ch = ch[0]
self.w[r-1][c-1] = ch
def put (self, ch):
'''This puts a characters at the current cursor position.
'''
if isinstance(ch, bytes):
ch = self._decode(ch)
self.put_abs (self.cur_r, self.cur_c, ch)
def insert_abs (self, r, c, ch):
'''This inserts a character at (r,c). Everything under
and to the right is shifted right one character.
The last character of the line is lost.
'''
if isinstance(ch, bytes):
ch = self._decode(ch)
r = constrain (r, 1, self.rows)
c = constrain (c, 1, self.cols)
for ci in range (self.cols, c, -1):
self.put_abs (r,ci, self.get_abs(r,ci-1))
self.put_abs (r,c,ch)
def insert (self, ch):
if isinstance(ch, bytes):
ch = self._decode(ch)
self.insert_abs (self.cur_r, self.cur_c, ch)
def get_abs (self, r, c):
r = constrain (r, 1, self.rows)
c = constrain (c, 1, self.cols)
return self.w[r-1][c-1]
def get (self):
self.get_abs (self.cur_r, self.cur_c)
def get_region (self, rs,cs, re,ce):
'''This returns a list of lines representing the region.
'''
rs = constrain (rs, 1, self.rows)
re = constrain (re, 1, self.rows)
cs = constrain (cs, 1, self.cols)
ce = constrain (ce, 1, self.cols)
if rs > re:
rs, re = re, rs
if cs > ce:
cs, ce = ce, cs
sc = []
for r in range (rs, re+1):
line = u''
for c in range (cs, ce + 1):
ch = self.get_abs (r,c)
line = line + ch
sc.append (line)
return sc
def cursor_constrain (self):
'''This keeps the cursor within the screen area.
'''
self.cur_r = constrain (self.cur_r, 1, self.rows)
self.cur_c = constrain (self.cur_c, 1, self.cols)
def cursor_home (self, r=1, c=1): # <ESC>[{ROW};{COLUMN}H
self.cur_r = r
self.cur_c = c
self.cursor_constrain ()
def cursor_back (self,count=1): # <ESC>[{COUNT}D (not confused with down)
self.cur_c = self.cur_c - count
self.cursor_constrain ()
def cursor_down (self,count=1): # <ESC>[{COUNT}B (not confused with back)
self.cur_r = self.cur_r + count
self.cursor_constrain ()
def cursor_forward (self,count=1): # <ESC>[{COUNT}C
self.cur_c = self.cur_c + count
self.cursor_constrain ()
def cursor_up (self,count=1): # <ESC>[{COUNT}A
self.cur_r = self.cur_r - count
self.cursor_constrain ()
def cursor_up_reverse (self): # <ESC> M (called RI -- Reverse Index)
old_r = self.cur_r
self.cursor_up()
if old_r == self.cur_r:
self.scroll_up()
def cursor_force_position (self, r, c): # <ESC>[{ROW};{COLUMN}f
'''Identical to Cursor Home.'''
self.cursor_home (r, c)
def cursor_save (self): # <ESC>[s
'''Save current cursor position.'''
self.cursor_save_attrs()
def cursor_unsave (self): # <ESC>[u
'''Restores cursor position after a Save Cursor.'''
self.cursor_restore_attrs()
def cursor_save_attrs (self): # <ESC>7
'''Save current cursor position.'''
self.cur_saved_r = self.cur_r
self.cur_saved_c = self.cur_c
def cursor_restore_attrs (self): # <ESC>8
'''Restores cursor position after a Save Cursor.'''
self.cursor_home (self.cur_saved_r, self.cur_saved_c)
def scroll_constrain (self):
'''This keeps the scroll region within the screen region.'''
if self.scroll_row_start <= 0:
self.scroll_row_start = 1
if self.scroll_row_end > self.rows:
self.scroll_row_end = self.rows
def scroll_screen (self): # <ESC>[r
'''Enable scrolling for entire display.'''
self.scroll_row_start = 1
self.scroll_row_end = self.rows
def scroll_screen_rows (self, rs, re): # <ESC>[{start};{end}r
'''Enable scrolling from row {start} to row {end}.'''
self.scroll_row_start = rs
self.scroll_row_end = re
self.scroll_constrain()
def scroll_down (self): # <ESC>D
'''Scroll display down one line.'''
# Screen is indexed from 1, but arrays are indexed from 0.
s = self.scroll_row_start - 1
e = self.scroll_row_end - 1
self.w[s+1:e+1] = copy.deepcopy(self.w[s:e])
def scroll_up (self): # <ESC>M
'''Scroll display up one line.'''
# Screen is indexed from 1, but arrays are indexed from 0.
s = self.scroll_row_start - 1
e = self.scroll_row_end - 1
self.w[s:e] = copy.deepcopy(self.w[s+1:e+1])
def erase_end_of_line (self): # <ESC>[0K -or- <ESC>[K
'''Erases from the current cursor position to the end of the current
line.'''
self.fill_region (self.cur_r, self.cur_c, self.cur_r, self.cols)
def erase_start_of_line (self): # <ESC>[1K
'''Erases from the current cursor position to the start of the current
line.'''
self.fill_region (self.cur_r, 1, self.cur_r, self.cur_c)
def erase_line (self): # <ESC>[2K
'''Erases the entire current line.'''
self.fill_region (self.cur_r, 1, self.cur_r, self.cols)
def erase_down (self): # <ESC>[0J -or- <ESC>[J
'''Erases the screen from the current line down to the bottom of the
screen.'''
self.erase_end_of_line ()
self.fill_region (self.cur_r + 1, 1, self.rows, self.cols)
def erase_up (self): # <ESC>[1J
'''Erases the screen from the current line up to the top of the
screen.'''
self.erase_start_of_line ()
self.fill_region (self.cur_r-1, 1, 1, self.cols)
def erase_screen (self): # <ESC>[2J
'''Erases the screen with the background color.'''
self.fill ()
def set_tab (self): # <ESC>H
'''Sets a tab at the current position.'''
pass
def clear_tab (self): # <ESC>[g
'''Clears tab at the current position.'''
pass
def clear_all_tabs (self): # <ESC>[3g
'''Clears all tabs.'''
pass
# Insert line Esc [ Pn L
# Delete line Esc [ Pn M
# Delete character Esc [ Pn P
# Scrolling region Esc [ Pn(top);Pn(bot) r
+494
View File
@@ -0,0 +1,494 @@
import codecs
import os
import sys
import re
import errno
from .exceptions import ExceptionPexpect, EOF, TIMEOUT
from .expect import Expecter, searcher_string, searcher_re
PY3 = (sys.version_info[0] >= 3)
text_type = str if PY3 else unicode
class _NullCoder(object):
"""Pass bytes through unchanged."""
@staticmethod
def encode(b, final=False):
return b
@staticmethod
def decode(b, final=False):
return b
class SpawnBase(object):
"""A base class providing the backwards-compatible spawn API for Pexpect.
This should not be instantiated directly: use :class:`pexpect.spawn` or
:class:`pexpect.fdpexpect.fdspawn`.
"""
encoding = None
pid = None
flag_eof = False
def __init__(self, timeout=30, maxread=2000, searchwindowsize=None,
logfile=None, encoding=None, codec_errors='strict'):
self.stdin = sys.stdin
self.stdout = sys.stdout
self.stderr = sys.stderr
self.searcher = None
self.ignorecase = False
self.before = None
self.after = None
self.match = None
self.match_index = None
self.terminated = True
self.exitstatus = None
self.signalstatus = None
# status returned by os.waitpid
self.status = None
# the child file descriptor is initially closed
self.child_fd = -1
self.timeout = timeout
self.delimiter = EOF
self.logfile = logfile
# input from child (read_nonblocking)
self.logfile_read = None
# output to send (send, sendline)
self.logfile_send = None
# max bytes to read at one time into buffer
self.maxread = maxread
# This is the read buffer. See maxread.
self.buffer = bytes() if (encoding is None) else text_type()
# Data before searchwindowsize point is preserved, but not searched.
self.searchwindowsize = searchwindowsize
# Delay used before sending data to child. Time in seconds.
# Set this to None to skip the time.sleep() call completely.
self.delaybeforesend = 0.05
# Used by close() to give kernel time to update process status.
# Time in seconds.
self.delayafterclose = 0.1
# Used by terminate() to give kernel time to update process status.
# Time in seconds.
self.delayafterterminate = 0.1
# Delay in seconds to sleep after each call to read_nonblocking().
# Set this to None to skip the time.sleep() call completely: that
# would restore the behavior from pexpect-2.0 (for performance
# reasons or because you don't want to release Python's global
# interpreter lock).
self.delayafterread = 0.0001
self.softspace = False
self.name = '<' + repr(self) + '>'
self.closed = True
# Unicode interface
self.encoding = encoding
self.codec_errors = codec_errors
if encoding is None:
# bytes mode (accepts some unicode for backwards compatibility)
self._encoder = self._decoder = _NullCoder()
self.string_type = bytes
self.crlf = b'\r\n'
if PY3:
self.allowed_string_types = (bytes, str)
self.linesep = os.linesep.encode('ascii')
def write_to_stdout(b):
try:
return sys.stdout.buffer.write(b)
except AttributeError:
# If stdout has been replaced, it may not have .buffer
return sys.stdout.write(b.decode('ascii', 'replace'))
self.write_to_stdout = write_to_stdout
else:
self.allowed_string_types = (basestring,) # analysis:ignore
self.linesep = os.linesep
self.write_to_stdout = sys.stdout.write
else:
# unicode mode
self._encoder = codecs.getincrementalencoder(encoding)(codec_errors)
self._decoder = codecs.getincrementaldecoder(encoding)(codec_errors)
self.string_type = text_type
self.crlf = u'\r\n'
self.allowed_string_types = (text_type, )
if PY3:
self.linesep = os.linesep
else:
self.linesep = os.linesep.decode('ascii')
# This can handle unicode in both Python 2 and 3
self.write_to_stdout = sys.stdout.write
def _log(self, s, direction):
if self.logfile is not None:
self.logfile.write(s)
self.logfile.flush()
second_log = self.logfile_send if (direction=='send') else self.logfile_read
if second_log is not None:
second_log.write(s)
second_log.flush()
# For backwards compatibility, in bytes mode (when encoding is None)
# unicode is accepted for send and expect. Unicode mode is strictly unicode
# only.
def _coerce_expect_string(self, s):
if self.encoding is None and not isinstance(s, bytes):
return s.encode('ascii')
return s
def _coerce_send_string(self, s):
if self.encoding is None and not isinstance(s, bytes):
return s.encode('utf-8')
return s
def read_nonblocking(self, size=1, timeout=None):
"""This reads data from the file descriptor.
This is a simple implementation suitable for a regular file. Subclasses using ptys or pipes should override it.
The timeout parameter is ignored.
"""
try:
s = os.read(self.child_fd, size)
except OSError as err:
if err.args[0] == errno.EIO:
# Linux-style EOF
self.flag_eof = True
raise EOF('End Of File (EOF). Exception style platform.')
raise
if s == b'':
# BSD-style EOF
self.flag_eof = True
raise EOF('End Of File (EOF). Empty string style platform.')
s = self._decoder.decode(s, final=False)
self._log(s, 'read')
return s
def _pattern_type_err(self, pattern):
raise TypeError('got {badtype} ({badobj!r}) as pattern, must be one'
' of: {goodtypes}, pexpect.EOF, pexpect.TIMEOUT'\
.format(badtype=type(pattern),
badobj=pattern,
goodtypes=', '.join([str(ast)\
for ast in self.allowed_string_types])
)
)
def compile_pattern_list(self, patterns):
'''This compiles a pattern-string or a list of pattern-strings.
Patterns must be a StringType, EOF, TIMEOUT, SRE_Pattern, or a list of
those. Patterns may also be None which results in an empty list (you
might do this if waiting for an EOF or TIMEOUT condition without
expecting any pattern).
This is used by expect() when calling expect_list(). Thus expect() is
nothing more than::
cpl = self.compile_pattern_list(pl)
return self.expect_list(cpl, timeout)
If you are using expect() within a loop it may be more
efficient to compile the patterns first and then call expect_list().
This avoid calls in a loop to compile_pattern_list()::
cpl = self.compile_pattern_list(my_pattern)
while some_condition:
...
i = self.expect_list(cpl, timeout)
...
'''
if patterns is None:
return []
if not isinstance(patterns, list):
patterns = [patterns]
# Allow dot to match \n
compile_flags = re.DOTALL
if self.ignorecase:
compile_flags = compile_flags | re.IGNORECASE
compiled_pattern_list = []
for idx, p in enumerate(patterns):
if isinstance(p, self.allowed_string_types):
p = self._coerce_expect_string(p)
compiled_pattern_list.append(re.compile(p, compile_flags))
elif p is EOF:
compiled_pattern_list.append(EOF)
elif p is TIMEOUT:
compiled_pattern_list.append(TIMEOUT)
elif isinstance(p, type(re.compile(''))):
compiled_pattern_list.append(p)
else:
self._pattern_type_err(p)
return compiled_pattern_list
def expect(self, pattern, timeout=-1, searchwindowsize=-1, async=False):
'''This seeks through the stream until a pattern is matched. The
pattern is overloaded and may take several types. The pattern can be a
StringType, EOF, a compiled re, or a list of any of those types.
Strings will be compiled to re types. This returns the index into the
pattern list. If the pattern was not a list this returns index 0 on a
successful match. This may raise exceptions for EOF or TIMEOUT. To
avoid the EOF or TIMEOUT exceptions add EOF or TIMEOUT to the pattern
list. That will cause expect to match an EOF or TIMEOUT condition
instead of raising an exception.
If you pass a list of patterns and more than one matches, the first
match in the stream is chosen. If more than one pattern matches at that
point, the leftmost in the pattern list is chosen. For example::
# the input is 'foobar'
index = p.expect(['bar', 'foo', 'foobar'])
# returns 1('foo') even though 'foobar' is a "better" match
Please note, however, that buffering can affect this behavior, since
input arrives in unpredictable chunks. For example::
# the input is 'foobar'
index = p.expect(['foobar', 'foo'])
# returns 0('foobar') if all input is available at once,
# but returs 1('foo') if parts of the final 'bar' arrive late
When a match is found for the given pattern, the class instance
attribute *match* becomes an re.MatchObject result. Should an EOF
or TIMEOUT pattern match, then the match attribute will be an instance
of that exception class. The pairing before and after class
instance attributes are views of the data preceding and following
the matching pattern. On general exception, class attribute
*before* is all data received up to the exception, while *match* and
*after* attributes are value None.
When the keyword argument timeout is -1 (default), then TIMEOUT will
raise after the default value specified by the class timeout
attribute. When None, TIMEOUT will not be raised and may block
indefinitely until match.
When the keyword argument searchwindowsize is -1 (default), then the
value specified by the class maxread attribute is used.
A list entry may be EOF or TIMEOUT instead of a string. This will
catch these exceptions and return the index of the list entry instead
of raising the exception. The attribute 'after' will be set to the
exception type. The attribute 'match' will be None. This allows you to
write code like this::
index = p.expect(['good', 'bad', pexpect.EOF, pexpect.TIMEOUT])
if index == 0:
do_something()
elif index == 1:
do_something_else()
elif index == 2:
do_some_other_thing()
elif index == 3:
do_something_completely_different()
instead of code like this::
try:
index = p.expect(['good', 'bad'])
if index == 0:
do_something()
elif index == 1:
do_something_else()
except EOF:
do_some_other_thing()
except TIMEOUT:
do_something_completely_different()
These two forms are equivalent. It all depends on what you want. You
can also just expect the EOF if you are waiting for all output of a
child to finish. For example::
p = pexpect.spawn('/bin/ls')
p.expect(pexpect.EOF)
print p.before
If you are trying to optimize for speed then see expect_list().
On Python 3.4, or Python 3.3 with asyncio installed, passing
``async=True`` will make this return an :mod:`asyncio` coroutine,
which you can yield from to get the same result that this method would
normally give directly. So, inside a coroutine, you can replace this code::
index = p.expect(patterns)
With this non-blocking form::
index = yield from p.expect(patterns, async=True)
'''
compiled_pattern_list = self.compile_pattern_list(pattern)
return self.expect_list(compiled_pattern_list,
timeout, searchwindowsize, async)
def expect_list(self, pattern_list, timeout=-1, searchwindowsize=-1,
async=False):
'''This takes a list of compiled regular expressions and returns the
index into the pattern_list that matched the child output. The list may
also contain EOF or TIMEOUT(which are not compiled regular
expressions). This method is similar to the expect() method except that
expect_list() does not recompile the pattern list on every call. This
may help if you are trying to optimize for speed, otherwise just use
the expect() method. This is called by expect().
Like :meth:`expect`, passing ``async=True`` will make this return an
asyncio coroutine.
'''
if timeout == -1:
timeout = self.timeout
exp = Expecter(self, searcher_re(pattern_list), searchwindowsize)
if async:
from .async import expect_async
return expect_async(exp, timeout)
else:
return exp.expect_loop(timeout)
def expect_exact(self, pattern_list, timeout=-1, searchwindowsize=-1,
async=False):
'''This is similar to expect(), but uses plain string matching instead
of compiled regular expressions in 'pattern_list'. The 'pattern_list'
may be a string; a list or other sequence of strings; or TIMEOUT and
EOF.
This call might be faster than expect() for two reasons: string
searching is faster than RE matching and it is possible to limit the
search to just the end of the input buffer.
This method is also useful when you don't want to have to worry about
escaping regular expression characters that you want to match.
Like :meth:`expect`, passing ``async=True`` will make this return an
asyncio coroutine.
'''
if timeout == -1:
timeout = self.timeout
if (isinstance(pattern_list, self.allowed_string_types) or
pattern_list in (TIMEOUT, EOF)):
pattern_list = [pattern_list]
def prepare_pattern(pattern):
if pattern in (TIMEOUT, EOF):
return pattern
if isinstance(pattern, self.allowed_string_types):
return self._coerce_expect_string(pattern)
self._pattern_type_err(pattern)
try:
pattern_list = iter(pattern_list)
except TypeError:
self._pattern_type_err(pattern_list)
pattern_list = [prepare_pattern(p) for p in pattern_list]
exp = Expecter(self, searcher_string(pattern_list), searchwindowsize)
if async:
from .async import expect_async
return expect_async(exp, timeout)
else:
return exp.expect_loop(timeout)
def expect_loop(self, searcher, timeout=-1, searchwindowsize=-1):
'''This is the common loop used inside expect. The 'searcher' should be
an instance of searcher_re or searcher_string, which describes how and
what to search for in the input.
See expect() for other arguments, return value and exceptions. '''
exp = Expecter(self, searcher, searchwindowsize)
return exp.expect_loop(timeout)
def read(self, size=-1):
'''This reads at most "size" bytes from the file (less if the read hits
EOF before obtaining size bytes). If the size argument is negative or
omitted, read all data until EOF is reached. The bytes are returned as
a string object. An empty string is returned when EOF is encountered
immediately. '''
if size == 0:
return self.string_type()
if size < 0:
# delimiter default is EOF
self.expect(self.delimiter)
return self.before
# I could have done this more directly by not using expect(), but
# I deliberately decided to couple read() to expect() so that
# I would catch any bugs early and ensure consistant behavior.
# It's a little less efficient, but there is less for me to
# worry about if I have to later modify read() or expect().
# Note, it's OK if size==-1 in the regex. That just means it
# will never match anything in which case we stop only on EOF.
cre = re.compile(self._coerce_expect_string('.{%d}' % size), re.DOTALL)
# delimiter default is EOF
index = self.expect([cre, self.delimiter])
if index == 0:
### FIXME self.before should be ''. Should I assert this?
return self.after
return self.before
def readline(self, size=-1):
'''This reads and returns one entire line. The newline at the end of
line is returned as part of the string, unless the file ends without a
newline. An empty string is returned if EOF is encountered immediately.
This looks for a newline as a CR/LF pair (\\r\\n) even on UNIX because
this is what the pseudotty device returns. So contrary to what you may
expect you will receive newlines as \\r\\n.
If the size argument is 0 then an empty string is returned. In all
other cases the size argument is ignored, which is not standard
behavior for a file-like object. '''
if size == 0:
return self.string_type()
# delimiter default is EOF
index = self.expect([self.crlf, self.delimiter])
if index == 0:
return self.before + self.crlf
else:
return self.before
def __iter__(self):
'''This is to support iterators over a file-like object.
'''
return iter(self.readline, self.string_type())
def readlines(self, sizehint=-1):
'''This reads until EOF using readline() and returns a list containing
the lines thus read. The optional 'sizehint' argument is ignored.
Remember, because this reads until EOF that means the child
process should have closed its stdout. If you run this method on
a child that is still running with its stdout open then this
method will block until it timesout.'''
lines = []
while True:
line = self.readline()
if not line:
break
lines.append(line)
return lines
def fileno(self):
'''Expose file descriptor for a file-like interface
'''
return self.child_fd
def flush(self):
'''This does nothing. It is here to support the interface for a
File-like object. '''
pass
def isatty(self):
"""Overridden in subclass using tty"""
return False
# For 'with spawn(...) as child:'
def __enter__(self):
return self
def __exit__(self, etype, evalue, tb):
# We rely on subclasses to implement close(). If they don't, it's not
# clear what a context manager should do.
self.close()
+151
View File
@@ -0,0 +1,151 @@
import os
import sys
import stat
import select
import time
import errno
try:
InterruptedError
except NameError:
# Alias Python2 exception to Python3
InterruptedError = select.error
def is_executable_file(path):
"""Checks that path is an executable regular file, or a symlink towards one.
This is roughly ``os.path isfile(path) and os.access(path, os.X_OK)``.
"""
# follow symlinks,
fpath = os.path.realpath(path)
if not os.path.isfile(fpath):
# non-files (directories, fifo, etc.)
return False
mode = os.stat(fpath).st_mode
if (sys.platform.startswith('sunos')
and os.getuid() == 0):
# When root on Solaris, os.X_OK is True for *all* files, irregardless
# of their executability -- instead, any permission bit of any user,
# group, or other is fine enough.
#
# (This may be true for other "Unix98" OS's such as HP-UX and AIX)
return bool(mode & (stat.S_IXUSR |
stat.S_IXGRP |
stat.S_IXOTH))
return os.access(fpath, os.X_OK)
def which(filename, env=None):
'''This takes a given filename; tries to find it in the environment path;
then checks if it is executable. This returns the full path to the filename
if found and executable. Otherwise this returns None.'''
# Special case where filename contains an explicit path.
if os.path.dirname(filename) != '' and is_executable_file(filename):
return filename
if env is None:
env = os.environ
p = env.get('PATH')
if not p:
p = os.defpath
pathlist = p.split(os.pathsep)
for path in pathlist:
ff = os.path.join(path, filename)
if is_executable_file(ff):
return ff
return None
def split_command_line(command_line):
'''This splits a command line into a list of arguments. It splits arguments
on spaces, but handles embedded quotes, doublequotes, and escaped
characters. It's impossible to do this with a regular expression, so I
wrote a little state machine to parse the command line. '''
arg_list = []
arg = ''
# Constants to name the states we can be in.
state_basic = 0
state_esc = 1
state_singlequote = 2
state_doublequote = 3
# The state when consuming whitespace between commands.
state_whitespace = 4
state = state_basic
for c in command_line:
if state == state_basic or state == state_whitespace:
if c == '\\':
# Escape the next character
state = state_esc
elif c == r"'":
# Handle single quote
state = state_singlequote
elif c == r'"':
# Handle double quote
state = state_doublequote
elif c.isspace():
# Add arg to arg_list if we aren't in the middle of whitespace.
if state == state_whitespace:
# Do nothing.
None
else:
arg_list.append(arg)
arg = ''
state = state_whitespace
else:
arg = arg + c
state = state_basic
elif state == state_esc:
arg = arg + c
state = state_basic
elif state == state_singlequote:
if c == r"'":
state = state_basic
else:
arg = arg + c
elif state == state_doublequote:
if c == r'"':
state = state_basic
else:
arg = arg + c
if arg != '':
arg_list.append(arg)
return arg_list
def select_ignore_interrupts(iwtd, owtd, ewtd, timeout=None):
'''This is a wrapper around select.select() that ignores signals. If
select.select raises a select.error exception and errno is an EINTR
error then it is ignored. Mainly this is used to ignore sigwinch
(terminal resize). '''
# if select() is interrupted by a signal (errno==EINTR) then
# we loop back and enter the select() again.
if timeout is not None:
end_time = time.time() + timeout
while True:
try:
return select.select(iwtd, owtd, ewtd, timeout)
except InterruptedError:
err = sys.exc_info()[1]
if err.args[0] == errno.EINTR:
# if we loop back we have to subtract the
# amount of time we already waited.
if timeout is not None:
timeout = end_time - time.time()
if timeout < 0:
return([], [], [])
else:
# something else caused the select.error, so
# this actually is an exception.
raise
+21
View File
@@ -0,0 +1,21 @@
# This file is dual licensed under the terms of the Apache License, Version
# 2.0, and the BSD License. See the LICENSE file in the root of this repository
# for complete details.
from __future__ import absolute_import, division, print_function
__all__ = [
"__title__", "__summary__", "__uri__", "__version__", "__author__",
"__email__", "__license__", "__copyright__",
]
__title__ = "pipfile"
__summary__ = ""
__uri__ = "https://github.com/pypa/pipfile"
__version__ = "0.0.2"
__author__ = "Kenneth Reitz and individual contributors"
__email__ = "me@kennethreitz.org"
__license__ = "BSD or Apache License, Version 2.0"
__copyright__ = "Copyright 2017 %s" % __author__
+11
View File
@@ -0,0 +1,11 @@
# This file is dual licensed under the terms of the Apache License, Version
# 2.0, and the BSD License. See the LICENSE file in the root of this repository
# for complete details.
from __future__ import absolute_import, division, print_function
from .__about__ import (
__author__, __copyright__, __email__, __license__, __summary__, __title__,
__uri__, __version__
)
from .api import load, Pipfile
+195
View File
@@ -0,0 +1,195 @@
import toml
import codecs
import json
import hashlib
import platform
import sys
import os
def format_full_version(info):
version = '{0.major}.{0.minor}.{0.micro}'.format(info)
kind = info.releaselevel
if kind != 'final':
version += kind[0] + str(info.serial)
return version
def walk_up(bottom):
"""mimic os.walk, but walk 'up' instead of down the directory tree.
From: https://gist.github.com/zdavkeos/1098474
"""
bottom = os.path.realpath(bottom)
# get files in current dir
try:
names = os.listdir(bottom)
except Exception as e:
return
dirs, nondirs = [], []
for name in names:
if os.path.isdir(os.path.join(bottom, name)):
dirs.append(name)
else:
nondirs.append(name)
yield bottom, dirs, nondirs
new_path = os.path.realpath(os.path.join(bottom, '..'))
# see if we are at the top
if new_path == bottom:
return
for x in walk_up(new_path):
yield x
class PipfileParser(object):
def __init__(self, filename='Pipfile'):
self.filename = filename
self.sources = []
self.groups = {
'default': [],
'develop': []
}
self.group_stack = ['default']
self.requirements = []
def __repr__(self):
return '<PipfileParser path={0!r}'.format(self.filename)
def parse(self):
# Open the Pipfile.
with open(self.filename) as f:
content = f.read()
# Load the default configuration.
default_config = {
u'source': [{u'url': u'https://pypi.python.org/simple', u'verify_ssl': True}],
u'packages': {},
u'requires': {},
u'dev-packages': {}
}
config = {}
config.update(default_config)
# Load the Pipfile's configuration.
config.update(toml.loads(content))
# Structure the data for output.
data = {
'_meta': {
'sources': config['source'],
'requires': config['requires']
},
}
# TODO: Validate given data here.
self.groups['default'] = config['packages']
self.groups['develop'] = config['dev-packages']
# Update the data structure with group information.
data.update(self.groups)
return data
class Pipfile(object):
def __init__(self, filename):
super(Pipfile, self).__init__()
self.filename = filename
self.data = None
@staticmethod
def find(max_depth=3):
"""Returns the path of a Pipfile in parent directories."""
i = 0
for c, d, f in walk_up(os.getcwd()):
i += 1
if i < max_depth:
if 'Pipfile':
p = os.path.join(c, 'Pipfile')
if os.path.isfile(p):
return p
raise RuntimeError('No Pipfile found!')
@classmethod
def load(klass, filename):
"""Load a Pipfile from a given filename."""
p = PipfileParser(filename=filename)
pipfile = klass(filename=filename)
pipfile.data = p.parse()
return pipfile
@property
def hash(self):
"""Returns the SHA256 of the pipfile's data."""
content = json.dumps(self.data, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(content.encode("utf8")).hexdigest()
@property
def contents(self):
"""Returns the contents of the pipfile."""
with codecs.open(self.filename, 'r', 'utf-8') as f:
return f.read()
def lock(self):
"""Returns a JSON representation of the Pipfile."""
data = self.data
data['_meta']['hash'] = {"sha256": self.hash}
# return _json.dumps(data)
return json.dumps(data, indent=4, separators=(',', ': '))
def assert_requirements(self):
""""Asserts PEP 508 specifiers."""
# Support for 508's implementation_version.
if hasattr(sys, 'implementation'):
implementation_version = format_full_version(sys.implementation.version)
else:
implementation_version = "0"
# Default to cpython for 2.7.
if hasattr(sys, 'implementation'):
implementation_name = sys.implementation.name
else:
implementation_name = 'cpython'
lookup = {
'os_name': os.name,
'sys_platform': sys.platform,
'platform_machine': platform.machine(),
'platform_python_implementation': platform.python_implementation(),
'platform_release': platform.release(),
'platform_system': platform.system(),
'platform_version': platform.version(),
'python_version': platform.python_version()[:3],
'python_full_version': platform.python_version(),
'implementation_name': implementation_name,
'implementation_version': implementation_version
}
# Assert each specified requirement.
for marker, specifier in self.data['_meta']['requires'].items():
if marker in lookup:
try:
assert lookup[marker] == specifier
except AssertionError:
raise AssertionError('Specifier {!r} does not match {!r}.'.format(marker, specifier))
def load(pipfile_path=None):
"""Loads a pipfile from a given path.
If none is provided, one will try to be found.
"""
if pipfile_path is None:
pipfile_path = Pipfile.find()
return Pipfile.load(filename=pipfile_path)
+2376
View File
File diff suppressed because it is too large Load Diff
+449
View File
@@ -0,0 +1,449 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Common objects shared by __init__.py and _ps*.py modules."""
from __future__ import division
import contextlib
import errno
import functools
import os
import socket
import stat
import sys
import warnings
from collections import namedtuple
from socket import AF_INET
from socket import SOCK_DGRAM
from socket import SOCK_STREAM
try:
from socket import AF_INET6
except ImportError:
AF_INET6 = None
try:
from socket import AF_UNIX
except ImportError:
AF_UNIX = None
if sys.version_info >= (3, 4):
import enum
else:
enum = None
__all__ = [
# OS constants
'FREEBSD', 'BSD', 'LINUX', 'NETBSD', 'OPENBSD', 'OSX', 'POSIX', 'SUNOS',
'WINDOWS',
# connection constants
'CONN_CLOSE', 'CONN_CLOSE_WAIT', 'CONN_CLOSING', 'CONN_ESTABLISHED',
'CONN_FIN_WAIT1', 'CONN_FIN_WAIT2', 'CONN_LAST_ACK', 'CONN_LISTEN',
'CONN_NONE', 'CONN_SYN_RECV', 'CONN_SYN_SENT', 'CONN_TIME_WAIT',
# net constants
'NIC_DUPLEX_FULL', 'NIC_DUPLEX_HALF', 'NIC_DUPLEX_UNKNOWN',
# process status constants
'STATUS_DEAD', 'STATUS_DISK_SLEEP', 'STATUS_IDLE', 'STATUS_LOCKED',
'STATUS_RUNNING', 'STATUS_SLEEPING', 'STATUS_STOPPED', 'STATUS_SUSPENDED',
'STATUS_TRACING_STOP', 'STATUS_WAITING', 'STATUS_WAKE_KILL',
'STATUS_WAKING', 'STATUS_ZOMBIE',
# named tuples
'pconn', 'pcputimes', 'pctxsw', 'pgids', 'pio', 'pionice', 'popenfile',
'pthread', 'puids', 'sconn', 'scpustats', 'sdiskio', 'sdiskpart',
'sdiskusage', 'snetio', 'snic', 'snicstats', 'sswap', 'suser',
# utility functions
'conn_tmap', 'deprecated_method', 'isfile_strict', 'memoize',
'parse_environ_block', 'path_exists_strict', 'usage_percent',
'supports_ipv6', 'sockfam_to_enum', 'socktype_to_enum',
]
# ===================================================================
# --- OS constants
# ===================================================================
POSIX = os.name == "posix"
WINDOWS = os.name == "nt"
LINUX = sys.platform.startswith("linux")
OSX = sys.platform.startswith("darwin")
FREEBSD = sys.platform.startswith("freebsd")
OPENBSD = sys.platform.startswith("openbsd")
NETBSD = sys.platform.startswith("netbsd")
BSD = FREEBSD or OPENBSD or NETBSD
SUNOS = sys.platform.startswith("sunos") or sys.platform.startswith("solaris")
# ===================================================================
# --- API constants
# ===================================================================
# Process.status()
STATUS_RUNNING = "running"
STATUS_SLEEPING = "sleeping"
STATUS_DISK_SLEEP = "disk-sleep"
STATUS_STOPPED = "stopped"
STATUS_TRACING_STOP = "tracing-stop"
STATUS_ZOMBIE = "zombie"
STATUS_DEAD = "dead"
STATUS_WAKE_KILL = "wake-kill"
STATUS_WAKING = "waking"
STATUS_IDLE = "idle" # FreeBSD, OSX
STATUS_LOCKED = "locked" # FreeBSD
STATUS_WAITING = "waiting" # FreeBSD
STATUS_SUSPENDED = "suspended" # NetBSD
# Process.connections() and psutil.net_connections()
CONN_ESTABLISHED = "ESTABLISHED"
CONN_SYN_SENT = "SYN_SENT"
CONN_SYN_RECV = "SYN_RECV"
CONN_FIN_WAIT1 = "FIN_WAIT1"
CONN_FIN_WAIT2 = "FIN_WAIT2"
CONN_TIME_WAIT = "TIME_WAIT"
CONN_CLOSE = "CLOSE"
CONN_CLOSE_WAIT = "CLOSE_WAIT"
CONN_LAST_ACK = "LAST_ACK"
CONN_LISTEN = "LISTEN"
CONN_CLOSING = "CLOSING"
CONN_NONE = "NONE"
# net_if_stats()
if enum is None:
NIC_DUPLEX_FULL = 2
NIC_DUPLEX_HALF = 1
NIC_DUPLEX_UNKNOWN = 0
else:
class NicDuplex(enum.IntEnum):
NIC_DUPLEX_FULL = 2
NIC_DUPLEX_HALF = 1
NIC_DUPLEX_UNKNOWN = 0
globals().update(NicDuplex.__members__)
# sensors_battery()
if enum is None:
POWER_TIME_UNKNOWN = -1
POWER_TIME_UNLIMITED = -2
else:
class BatteryTime(enum.IntEnum):
POWER_TIME_UNKNOWN = -1
POWER_TIME_UNLIMITED = -2
globals().update(BatteryTime.__members__)
# ===================================================================
# --- namedtuples
# ===================================================================
# --- for system functions
# psutil.swap_memory()
sswap = namedtuple('sswap', ['total', 'used', 'free', 'percent', 'sin',
'sout'])
# psutil.disk_usage()
sdiskusage = namedtuple('sdiskusage', ['total', 'used', 'free', 'percent'])
# psutil.disk_io_counters()
sdiskio = namedtuple('sdiskio', ['read_count', 'write_count',
'read_bytes', 'write_bytes',
'read_time', 'write_time'])
# psutil.disk_partitions()
sdiskpart = namedtuple('sdiskpart', ['device', 'mountpoint', 'fstype', 'opts'])
# psutil.net_io_counters()
snetio = namedtuple('snetio', ['bytes_sent', 'bytes_recv',
'packets_sent', 'packets_recv',
'errin', 'errout',
'dropin', 'dropout'])
# psutil.users()
suser = namedtuple('suser', ['name', 'terminal', 'host', 'started'])
# psutil.net_connections()
sconn = namedtuple('sconn', ['fd', 'family', 'type', 'laddr', 'raddr',
'status', 'pid'])
# psutil.net_if_addrs()
snic = namedtuple('snic', ['family', 'address', 'netmask', 'broadcast', 'ptp'])
# psutil.net_if_stats()
snicstats = namedtuple('snicstats', ['isup', 'duplex', 'speed', 'mtu'])
# psutil.cpu_stats()
scpustats = namedtuple(
'scpustats', ['ctx_switches', 'interrupts', 'soft_interrupts', 'syscalls'])
# psutil.cpu_freq()
scpufreq = namedtuple('scpufreq', ['current', 'min', 'max'])
# psutil.sensors_temperatures()
shwtemp = namedtuple(
'shwtemp', ['label', 'current', 'high', 'critical'])
# psutil.sensors_battery()
sbattery = namedtuple('sbattery', ['percent', 'secsleft', 'power_plugged'])
# psutil.sensors_battery()
sfan = namedtuple('sfan', ['label', 'current'])
# --- for Process methods
# psutil.Process.cpu_times()
pcputimes = namedtuple('pcputimes',
['user', 'system', 'children_user', 'children_system'])
# psutil.Process.open_files()
popenfile = namedtuple('popenfile', ['path', 'fd'])
# psutil.Process.threads()
pthread = namedtuple('pthread', ['id', 'user_time', 'system_time'])
# psutil.Process.uids()
puids = namedtuple('puids', ['real', 'effective', 'saved'])
# psutil.Process.gids()
pgids = namedtuple('pgids', ['real', 'effective', 'saved'])
# psutil.Process.io_counters()
pio = namedtuple('pio', ['read_count', 'write_count',
'read_bytes', 'write_bytes'])
# psutil.Process.ionice()
pionice = namedtuple('pionice', ['ioclass', 'value'])
# psutil.Process.ctx_switches()
pctxsw = namedtuple('pctxsw', ['voluntary', 'involuntary'])
# psutil.Process.connections()
pconn = namedtuple('pconn', ['fd', 'family', 'type', 'laddr', 'raddr',
'status'])
# ===================================================================
# --- Process.connections() 'kind' parameter mapping
# ===================================================================
conn_tmap = {
"all": ([AF_INET, AF_INET6, AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]),
"tcp": ([AF_INET, AF_INET6], [SOCK_STREAM]),
"tcp4": ([AF_INET], [SOCK_STREAM]),
"udp": ([AF_INET, AF_INET6], [SOCK_DGRAM]),
"udp4": ([AF_INET], [SOCK_DGRAM]),
"inet": ([AF_INET, AF_INET6], [SOCK_STREAM, SOCK_DGRAM]),
"inet4": ([AF_INET], [SOCK_STREAM, SOCK_DGRAM]),
"inet6": ([AF_INET6], [SOCK_STREAM, SOCK_DGRAM]),
}
if AF_INET6 is not None:
conn_tmap.update({
"tcp6": ([AF_INET6], [SOCK_STREAM]),
"udp6": ([AF_INET6], [SOCK_DGRAM]),
})
if AF_UNIX is not None:
conn_tmap.update({
"unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]),
})
del AF_INET, AF_INET6, AF_UNIX, SOCK_STREAM, SOCK_DGRAM
# ===================================================================
# --- utils
# ===================================================================
def usage_percent(used, total, _round=None):
"""Calculate percentage usage of 'used' against 'total'."""
try:
ret = (used / total) * 100
except ZeroDivisionError:
ret = 0.0 if isinstance(used, float) or isinstance(total, float) else 0
if _round is not None:
return round(ret, _round)
else:
return ret
def memoize(fun):
"""A simple memoize decorator for functions supporting (hashable)
positional arguments.
It also provides a cache_clear() function for clearing the cache:
>>> @memoize
... def foo()
... return 1
...
>>> foo()
1
>>> foo.cache_clear()
>>>
"""
@functools.wraps(fun)
def wrapper(*args, **kwargs):
key = (args, frozenset(sorted(kwargs.items())))
try:
return cache[key]
except KeyError:
ret = cache[key] = fun(*args, **kwargs)
return ret
def cache_clear():
"""Clear cache."""
cache.clear()
cache = {}
wrapper.cache_clear = cache_clear
return wrapper
def memoize_when_activated(fun):
"""A memoize decorator which is disabled by default. It can be
activated and deactivated on request.
For efficiency reasons it can be used only against class methods
accepting no arguments.
>>> class Foo:
... @memoize
... def foo()
... print(1)
...
>>> f = Foo()
>>> # deactivated (default)
>>> foo()
1
>>> foo()
1
>>>
>>> # activated
>>> foo.cache_activate()
>>> foo()
1
>>> foo()
>>> foo()
>>>
"""
@functools.wraps(fun)
def wrapper(self):
if not wrapper.cache_activated:
return fun(self)
else:
try:
ret = cache[fun]
except KeyError:
ret = cache[fun] = fun(self)
return ret
def cache_activate():
"""Activate cache."""
wrapper.cache_activated = True
def cache_deactivate():
"""Deactivate and clear cache."""
wrapper.cache_activated = False
cache.clear()
cache = {}
wrapper.cache_activated = False
wrapper.cache_activate = cache_activate
wrapper.cache_deactivate = cache_deactivate
return wrapper
def isfile_strict(path):
"""Same as os.path.isfile() but does not swallow EACCES / EPERM
exceptions, see:
http://mail.python.org/pipermail/python-dev/2012-June/120787.html
"""
try:
st = os.stat(path)
except OSError as err:
if err.errno in (errno.EPERM, errno.EACCES):
raise
return False
else:
return stat.S_ISREG(st.st_mode)
def path_exists_strict(path):
"""Same as os.path.exists() but does not swallow EACCES / EPERM
exceptions, see:
http://mail.python.org/pipermail/python-dev/2012-June/120787.html
"""
try:
os.stat(path)
except OSError as err:
if err.errno in (errno.EPERM, errno.EACCES):
raise
return False
else:
return True
def supports_ipv6():
"""Return True if IPv6 is supported on this platform."""
if not socket.has_ipv6 or not hasattr(socket, "AF_INET6"):
return False
try:
sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
with contextlib.closing(sock):
sock.bind(("::1", 0))
return True
except socket.error:
return False
def parse_environ_block(data):
"""Parse a C environ block of environment variables into a dictionary."""
# The block is usually raw data from the target process. It might contain
# trailing garbage and lines that do not look like assignments.
ret = {}
pos = 0
# localize global variable to speed up access.
WINDOWS_ = WINDOWS
while True:
next_pos = data.find("\0", pos)
# nul byte at the beginning or double nul byte means finish
if next_pos <= pos:
break
# there might not be an equals sign
equal_pos = data.find("=", pos, next_pos)
if equal_pos > pos:
key = data[pos:equal_pos]
value = data[equal_pos + 1:next_pos]
# Windows expects environment variables to be uppercase only
if WINDOWS_:
key = key.upper()
ret[key] = value
pos = next_pos + 1
return ret
def sockfam_to_enum(num):
"""Convert a numeric socket family value to an IntEnum member.
If it's not a known member, return the numeric value itself.
"""
if enum is None:
return num
else: # pragma: no cover
try:
return socket.AddressFamily(num)
except (ValueError, AttributeError):
return num
def socktype_to_enum(num):
"""Convert a numeric socket type value to an IntEnum member.
If it's not a known member, return the numeric value itself.
"""
if enum is None:
return num
else: # pragma: no cover
try:
return socket.AddressType(num)
except (ValueError, AttributeError):
return num
def deprecated_method(replacement):
"""A decorator which can be used to mark a method as deprecated
'replcement' is the method name which will be called instead.
"""
def outer(fun):
msg = "%s() is deprecated; use %s() instead" % (
fun.__name__, replacement)
if fun.__doc__ is None:
fun.__doc__ = msg
@functools.wraps(fun)
def inner(self, *args, **kwargs):
warnings.warn(msg, category=DeprecationWarning, stacklevel=2)
return getattr(self, replacement)(*args, **kwargs)
return inner
return outer
+249
View File
@@ -0,0 +1,249 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Module which provides compatibility with older Python versions."""
import collections
import functools
import os
import sys
__all__ = ["PY3", "long", "xrange", "unicode", "basestring", "u", "b",
"callable", "lru_cache", "which"]
PY3 = sys.version_info[0] == 3
if PY3:
long = int
xrange = range
unicode = str
basestring = str
def u(s):
return s
def b(s):
return s.encode("latin-1")
else:
long = long
xrange = xrange
unicode = unicode
basestring = basestring
def u(s):
return unicode(s, "unicode_escape")
def b(s):
return s
# removed in 3.0, reintroduced in 3.2
try:
callable = callable
except NameError:
def callable(obj):
return any("__call__" in klass.__dict__ for klass in type(obj).__mro__)
# --- stdlib additions
# py 3.2 functools.lru_cache
# Taken from: http://code.activestate.com/recipes/578078
# Credit: Raymond Hettinger
try:
from functools import lru_cache
except ImportError:
try:
from threading import RLock
except ImportError:
from dummy_threading import RLock
_CacheInfo = collections.namedtuple(
"CacheInfo", ["hits", "misses", "maxsize", "currsize"])
class _HashedSeq(list):
__slots__ = 'hashvalue'
def __init__(self, tup, hash=hash):
self[:] = tup
self.hashvalue = hash(tup)
def __hash__(self):
return self.hashvalue
def _make_key(args, kwds, typed,
kwd_mark=(object(), ),
fasttypes=set((int, str, frozenset, type(None))),
sorted=sorted, tuple=tuple, type=type, len=len):
key = args
if kwds:
sorted_items = sorted(kwds.items())
key += kwd_mark
for item in sorted_items:
key += item
if typed:
key += tuple(type(v) for v in args)
if kwds:
key += tuple(type(v) for k, v in sorted_items)
elif len(key) == 1 and type(key[0]) in fasttypes:
return key[0]
return _HashedSeq(key)
def lru_cache(maxsize=100, typed=False):
"""Least-recently-used cache decorator, see:
http://docs.python.org/3/library/functools.html#functools.lru_cache
"""
def decorating_function(user_function):
cache = dict()
stats = [0, 0]
HITS, MISSES = 0, 1
make_key = _make_key
cache_get = cache.get
_len = len
lock = RLock()
root = []
root[:] = [root, root, None, None]
nonlocal_root = [root]
PREV, NEXT, KEY, RESULT = 0, 1, 2, 3
if maxsize == 0:
def wrapper(*args, **kwds):
result = user_function(*args, **kwds)
stats[MISSES] += 1
return result
elif maxsize is None:
def wrapper(*args, **kwds):
key = make_key(args, kwds, typed)
result = cache_get(key, root)
if result is not root:
stats[HITS] += 1
return result
result = user_function(*args, **kwds)
cache[key] = result
stats[MISSES] += 1
return result
else:
def wrapper(*args, **kwds):
if kwds or typed:
key = make_key(args, kwds, typed)
else:
key = args
lock.acquire()
try:
link = cache_get(key)
if link is not None:
root, = nonlocal_root
link_prev, link_next, key, result = link
link_prev[NEXT] = link_next
link_next[PREV] = link_prev
last = root[PREV]
last[NEXT] = root[PREV] = link
link[PREV] = last
link[NEXT] = root
stats[HITS] += 1
return result
finally:
lock.release()
result = user_function(*args, **kwds)
lock.acquire()
try:
root, = nonlocal_root
if key in cache:
pass
elif _len(cache) >= maxsize:
oldroot = root
oldroot[KEY] = key
oldroot[RESULT] = result
root = nonlocal_root[0] = oldroot[NEXT]
oldkey = root[KEY]
root[KEY] = root[RESULT] = None
del cache[oldkey]
cache[key] = oldroot
else:
last = root[PREV]
link = [last, root, key, result]
last[NEXT] = root[PREV] = cache[key] = link
stats[MISSES] += 1
finally:
lock.release()
return result
def cache_info():
"""Report cache statistics"""
lock.acquire()
try:
return _CacheInfo(stats[HITS], stats[MISSES], maxsize,
len(cache))
finally:
lock.release()
def cache_clear():
"""Clear the cache and cache statistics"""
lock.acquire()
try:
cache.clear()
root = nonlocal_root[0]
root[:] = [root, root, None, None]
stats[:] = [0, 0]
finally:
lock.release()
wrapper.__wrapped__ = user_function
wrapper.cache_info = cache_info
wrapper.cache_clear = cache_clear
return functools.update_wrapper(wrapper, user_function)
return decorating_function
# python 3.3
try:
from shutil import which
except ImportError:
def which(cmd, mode=os.F_OK | os.X_OK, path=None):
"""Given a command, mode, and a PATH string, return the path which
conforms to the given mode on the PATH, or None if there is no such
file.
`mode` defaults to os.F_OK | os.X_OK. `path` defaults to the result
of os.environ.get("PATH"), or can be overridden with a custom search
path.
"""
def _access_check(fn, mode):
return (os.path.exists(fn) and os.access(fn, mode) and
not os.path.isdir(fn))
if os.path.dirname(cmd):
if _access_check(cmd, mode):
return cmd
return None
if path is None:
path = os.environ.get("PATH", os.defpath)
if not path:
return None
path = path.split(os.pathsep)
if sys.platform == "win32":
if os.curdir not in path:
path.insert(0, os.curdir)
pathext = os.environ.get("PATHEXT", "").split(os.pathsep)
if any(cmd.lower().endswith(ext.lower()) for ext in pathext):
files = [cmd]
else:
files = [cmd + ext for ext in pathext]
else:
files = [cmd]
seen = set()
for dir in path:
normdir = os.path.normcase(dir)
if normdir not in seen:
seen.add(normdir)
for thefile in files:
name = os.path.join(dir, thefile)
if _access_check(name, mode):
return name
return None
+855
View File
@@ -0,0 +1,855 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""FreeBSD, OpenBSD and NetBSD platforms implementation."""
import contextlib
import errno
import functools
import os
import xml.etree.ElementTree as ET
from collections import namedtuple
from . import _common
from . import _psposix
from . import _psutil_bsd as cext
from . import _psutil_posix as cext_posix
from ._common import conn_tmap
from ._common import FREEBSD
from ._common import memoize
from ._common import memoize_when_activated
from ._common import NETBSD
from ._common import OPENBSD
from ._common import sockfam_to_enum
from ._common import socktype_to_enum
from ._common import usage_percent
from ._compat import which
__extra__all__ = []
# =====================================================================
# --- globals
# =====================================================================
if FREEBSD:
PROC_STATUSES = {
cext.SIDL: _common.STATUS_IDLE,
cext.SRUN: _common.STATUS_RUNNING,
cext.SSLEEP: _common.STATUS_SLEEPING,
cext.SSTOP: _common.STATUS_STOPPED,
cext.SZOMB: _common.STATUS_ZOMBIE,
cext.SWAIT: _common.STATUS_WAITING,
cext.SLOCK: _common.STATUS_LOCKED,
}
elif OPENBSD or NETBSD:
PROC_STATUSES = {
cext.SIDL: _common.STATUS_IDLE,
cext.SSLEEP: _common.STATUS_SLEEPING,
cext.SSTOP: _common.STATUS_STOPPED,
# According to /usr/include/sys/proc.h SZOMB is unused.
# test_zombie_process() shows that SDEAD is the right
# equivalent. Also it appears there's no equivalent of
# psutil.STATUS_DEAD. SDEAD really means STATUS_ZOMBIE.
# cext.SZOMB: _common.STATUS_ZOMBIE,
cext.SDEAD: _common.STATUS_ZOMBIE,
cext.SZOMB: _common.STATUS_ZOMBIE,
# From http://www.eecs.harvard.edu/~margo/cs161/videos/proc.h.txt
# OpenBSD has SRUN and SONPROC: SRUN indicates that a process
# is runnable but *not* yet running, i.e. is on a run queue.
# SONPROC indicates that the process is actually executing on
# a CPU, i.e. it is no longer on a run queue.
# As such we'll map SRUN to STATUS_WAKING and SONPROC to
# STATUS_RUNNING
cext.SRUN: _common.STATUS_WAKING,
cext.SONPROC: _common.STATUS_RUNNING,
}
elif NETBSD:
PROC_STATUSES = {
cext.SIDL: _common.STATUS_IDLE,
cext.SACTIVE: _common.STATUS_RUNNING,
cext.SDYING: _common.STATUS_ZOMBIE,
cext.SSTOP: _common.STATUS_STOPPED,
cext.SZOMB: _common.STATUS_ZOMBIE,
cext.SDEAD: _common.STATUS_DEAD,
cext.SSUSPENDED: _common.STATUS_SUSPENDED, # unique to NetBSD
}
TCP_STATUSES = {
cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED,
cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT,
cext.TCPS_SYN_RECEIVED: _common.CONN_SYN_RECV,
cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1,
cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2,
cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT,
cext.TCPS_CLOSED: _common.CONN_CLOSE,
cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT,
cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK,
cext.TCPS_LISTEN: _common.CONN_LISTEN,
cext.TCPS_CLOSING: _common.CONN_CLOSING,
cext.PSUTIL_CONN_NONE: _common.CONN_NONE,
}
if NETBSD:
PAGESIZE = os.sysconf("SC_PAGESIZE")
else:
PAGESIZE = os.sysconf("SC_PAGE_SIZE")
AF_LINK = cext_posix.AF_LINK
kinfo_proc_map = dict(
ppid=0,
status=1,
real_uid=2,
effective_uid=3,
saved_uid=4,
real_gid=5,
effective_gid=6,
saved_gid=7,
ttynr=8,
create_time=9,
ctx_switches_vol=10,
ctx_switches_unvol=11,
read_io_count=12,
write_io_count=13,
user_time=14,
sys_time=15,
ch_user_time=16,
ch_sys_time=17,
rss=18,
vms=19,
memtext=20,
memdata=21,
memstack=22,
cpunum=23,
name=24,
)
# these get overwritten on "import psutil" from the __init__.py file
NoSuchProcess = None
ZombieProcess = None
AccessDenied = None
TimeoutExpired = None
# =====================================================================
# --- named tuples
# =====================================================================
# psutil.virtual_memory()
svmem = namedtuple(
'svmem', ['total', 'available', 'percent', 'used', 'free',
'active', 'inactive', 'buffers', 'cached', 'shared', 'wired'])
# psutil.cpu_times()
scputimes = namedtuple(
'scputimes', ['user', 'nice', 'system', 'idle', 'irq'])
# psutil.Process.memory_info()
pmem = namedtuple('pmem', ['rss', 'vms', 'text', 'data', 'stack'])
# psutil.Process.memory_full_info()
pfullmem = pmem
# psutil.Process.cpu_times()
pcputimes = namedtuple('pcputimes',
['user', 'system', 'children_user', 'children_system'])
# psutil.Process.memory_maps(grouped=True)
pmmap_grouped = namedtuple(
'pmmap_grouped', 'path rss, private, ref_count, shadow_count')
# psutil.Process.memory_maps(grouped=False)
pmmap_ext = namedtuple(
'pmmap_ext', 'addr, perms path rss, private, ref_count, shadow_count')
# psutil.disk_io_counters()
if FREEBSD:
sdiskio = namedtuple('sdiskio', ['read_count', 'write_count',
'read_bytes', 'write_bytes',
'read_time', 'write_time',
'busy_time'])
else:
sdiskio = namedtuple('sdiskio', ['read_count', 'write_count',
'read_bytes', 'write_bytes'])
# =====================================================================
# --- memory
# =====================================================================
def virtual_memory():
"""System virtual memory as a namedtuple."""
mem = cext.virtual_mem()
total, free, active, inactive, wired, cached, buffers, shared = mem
if NETBSD:
# On NetBSD buffers and shared mem is determined via /proc.
# The C ext set them to 0.
with open('/proc/meminfo', 'rb') as f:
for line in f:
if line.startswith(b'Buffers:'):
buffers = int(line.split()[1]) * 1024
elif line.startswith(b'MemShared:'):
shared = int(line.split()[1]) * 1024
avail = inactive + cached + free
used = active + wired + cached
percent = usage_percent((total - avail), total, _round=1)
return svmem(total, avail, percent, used, free,
active, inactive, buffers, cached, shared, wired)
def swap_memory():
"""System swap memory as (total, used, free, sin, sout) namedtuple."""
total, used, free, sin, sout = cext.swap_mem()
percent = usage_percent(used, total, _round=1)
return _common.sswap(total, used, free, percent, sin, sout)
# =====================================================================
# --- CPU
# =====================================================================
def cpu_times():
"""Return system per-CPU times as a namedtuple"""
user, nice, system, idle, irq = cext.cpu_times()
return scputimes(user, nice, system, idle, irq)
if hasattr(cext, "per_cpu_times"):
def per_cpu_times():
"""Return system CPU times as a namedtuple"""
ret = []
for cpu_t in cext.per_cpu_times():
user, nice, system, idle, irq = cpu_t
item = scputimes(user, nice, system, idle, irq)
ret.append(item)
return ret
else:
# XXX
# Ok, this is very dirty.
# On FreeBSD < 8 we cannot gather per-cpu information, see:
# https://github.com/giampaolo/psutil/issues/226
# If num cpus > 1, on first call we return single cpu times to avoid a
# crash at psutil import time.
# Next calls will fail with NotImplementedError
def per_cpu_times():
"""Return system CPU times as a namedtuple"""
if cpu_count_logical() == 1:
return [cpu_times()]
if per_cpu_times.__called__:
raise NotImplementedError("supported only starting from FreeBSD 8")
per_cpu_times.__called__ = True
return [cpu_times()]
per_cpu_times.__called__ = False
def cpu_count_logical():
"""Return the number of logical CPUs in the system."""
return cext.cpu_count_logical()
if OPENBSD or NETBSD:
def cpu_count_physical():
# OpenBSD and NetBSD do not implement this.
return 1 if cpu_count_logical() == 1 else None
else:
def cpu_count_physical():
"""Return the number of physical CPUs in the system."""
# From the C module we'll get an XML string similar to this:
# http://manpages.ubuntu.com/manpages/precise/man4/smp.4freebsd.html
# We may get None in case "sysctl kern.sched.topology_spec"
# is not supported on this BSD version, in which case we'll mimic
# os.cpu_count() and return None.
ret = None
s = cext.cpu_count_phys()
if s is not None:
# get rid of padding chars appended at the end of the string
index = s.rfind("</groups>")
if index != -1:
s = s[:index + 9]
root = ET.fromstring(s)
try:
ret = len(root.findall('group/children/group/cpu')) or None
finally:
# needed otherwise it will memleak
root.clear()
if not ret:
# If logical CPUs are 1 it's obvious we'll have only 1
# physical CPU.
if cpu_count_logical() == 1:
return 1
return ret
def cpu_stats():
"""Return various CPU stats as a named tuple."""
if FREEBSD:
# Note: the C ext is returning some metrics we are not exposing:
# traps.
ctxsw, intrs, soft_intrs, syscalls, traps = cext.cpu_stats()
elif NETBSD:
# XXX
# Note about intrs: the C extension returns 0. intrs
# can be determined via /proc/stat; it has the same value as
# soft_intrs thought so the kernel is faking it (?).
#
# Note about syscalls: the C extension always sets it to 0 (?).
#
# Note: the C ext is returning some metrics we are not exposing:
# traps, faults and forks.
ctxsw, intrs, soft_intrs, syscalls, traps, faults, forks = \
cext.cpu_stats()
with open('/proc/stat', 'rb') as f:
for line in f:
if line.startswith(b'intr'):
intrs = int(line.split()[1])
elif OPENBSD:
# Note: the C ext is returning some metrics we are not exposing:
# traps, faults and forks.
ctxsw, intrs, soft_intrs, syscalls, traps, faults, forks = \
cext.cpu_stats()
return _common.scpustats(ctxsw, intrs, soft_intrs, syscalls)
# =====================================================================
# --- disks
# =====================================================================
def disk_partitions(all=False):
"""Return mounted disk partitions as a list of namedtuples.
'all' argument is ignored, see:
https://github.com/giampaolo/psutil/issues/906
"""
retlist = []
partitions = cext.disk_partitions()
for partition in partitions:
device, mountpoint, fstype, opts = partition
ntuple = _common.sdiskpart(device, mountpoint, fstype, opts)
retlist.append(ntuple)
return retlist
disk_usage = _psposix.disk_usage
disk_io_counters = cext.disk_io_counters
# =====================================================================
# --- network
# =====================================================================
net_io_counters = cext.net_io_counters
net_if_addrs = cext_posix.net_if_addrs
def net_if_stats():
"""Get NIC stats (isup, duplex, speed, mtu)."""
names = net_io_counters().keys()
ret = {}
for name in names:
mtu = cext_posix.net_if_mtu(name)
isup = cext_posix.net_if_flags(name)
duplex, speed = cext_posix.net_if_duplex_speed(name)
if hasattr(_common, 'NicDuplex'):
duplex = _common.NicDuplex(duplex)
ret[name] = _common.snicstats(isup, duplex, speed, mtu)
return ret
def net_connections(kind):
"""System-wide network connections."""
if OPENBSD:
ret = []
for pid in pids():
try:
cons = Process(pid).connections(kind)
except (NoSuchProcess, ZombieProcess):
continue
else:
for conn in cons:
conn = list(conn)
conn.append(pid)
ret.append(_common.sconn(*conn))
return ret
if kind not in _common.conn_tmap:
raise ValueError("invalid %r kind argument; choose between %s"
% (kind, ', '.join([repr(x) for x in conn_tmap])))
families, types = conn_tmap[kind]
ret = set()
if NETBSD:
rawlist = cext.net_connections(-1)
else:
rawlist = cext.net_connections()
for item in rawlist:
fd, fam, type, laddr, raddr, status, pid = item
# TODO: apply filter at C level
if fam in families and type in types:
try:
status = TCP_STATUSES[status]
except KeyError:
# XXX: Not sure why this happens. I saw this occurring
# with IPv6 sockets opened by 'vim'. Those sockets
# have a very short lifetime so maybe the kernel
# can't initialize their status?
status = TCP_STATUSES[cext.PSUTIL_CONN_NONE]
fam = sockfam_to_enum(fam)
type = socktype_to_enum(type)
nt = _common.sconn(fd, fam, type, laddr, raddr, status, pid)
ret.add(nt)
return list(ret)
# =====================================================================
# --- sensors
# =====================================================================
if FREEBSD:
def sensors_battery():
"""Return battery info."""
percent, minsleft, power_plugged = cext.sensors_battery()
power_plugged = power_plugged == 1
if power_plugged:
secsleft = _common.POWER_TIME_UNLIMITED
elif minsleft == -1:
secsleft = _common.POWER_TIME_UNKNOWN
else:
secsleft = minsleft * 60
return _common.sbattery(percent, secsleft, power_plugged)
# =====================================================================
# --- other system functions
# =====================================================================
def boot_time():
"""The system boot time expressed in seconds since the epoch."""
return cext.boot_time()
def users():
"""Return currently connected users as a list of namedtuples."""
retlist = []
rawlist = cext.users()
for item in rawlist:
user, tty, hostname, tstamp = item
if tty == '~':
continue # reboot or shutdown
nt = _common.suser(user, tty or None, hostname, tstamp)
retlist.append(nt)
return retlist
# =====================================================================
# --- processes
# =====================================================================
@memoize
def _pid_0_exists():
try:
Process(0).name()
except NoSuchProcess:
return False
except AccessDenied:
return True
else:
return True
def pids():
"""Returns a list of PIDs currently running on the system."""
ret = cext.pids()
if OPENBSD and (0 not in ret) and _pid_0_exists():
# On OpenBSD the kernel does not return PID 0 (neither does
# ps) but it's actually querable (Process(0) will succeed).
ret.insert(0, 0)
return ret
if OPENBSD or NETBSD:
def pid_exists(pid):
"""Return True if pid exists."""
exists = _psposix.pid_exists(pid)
if not exists:
# We do this because _psposix.pid_exists() lies in case of
# zombie processes.
return pid in pids()
else:
return True
else:
pid_exists = _psposix.pid_exists
def wrap_exceptions(fun):
"""Decorator which translates bare OSError exceptions into
NoSuchProcess and AccessDenied.
"""
@functools.wraps(fun)
def wrapper(self, *args, **kwargs):
try:
return fun(self, *args, **kwargs)
except OSError as err:
if self.pid == 0:
if 0 in pids():
raise AccessDenied(self.pid, self._name)
else:
raise
if err.errno == errno.ESRCH:
if not pid_exists(self.pid):
raise NoSuchProcess(self.pid, self._name)
else:
raise ZombieProcess(self.pid, self._name, self._ppid)
if err.errno in (errno.EPERM, errno.EACCES):
raise AccessDenied(self.pid, self._name)
raise
return wrapper
@contextlib.contextmanager
def wrap_exceptions_procfs(inst):
"""Same as above, for routines relying on reading /proc fs."""
try:
yield
except EnvironmentError as err:
# ENOENT (no such file or directory) gets raised on open().
# ESRCH (no such process) can get raised on read() if
# process is gone in meantime.
if err.errno in (errno.ENOENT, errno.ESRCH):
if not pid_exists(inst.pid):
raise NoSuchProcess(inst.pid, inst._name)
else:
raise ZombieProcess(inst.pid, inst._name, inst._ppid)
if err.errno in (errno.EPERM, errno.EACCES):
raise AccessDenied(inst.pid, inst._name)
raise
class Process(object):
"""Wrapper class around underlying C implementation."""
__slots__ = ["pid", "_name", "_ppid"]
def __init__(self, pid):
self.pid = pid
self._name = None
self._ppid = None
@memoize_when_activated
def oneshot(self):
"""Retrieves multiple process info in one shot as a raw tuple."""
ret = cext.proc_oneshot_info(self.pid)
assert len(ret) == len(kinfo_proc_map)
return ret
def oneshot_enter(self):
self.oneshot.cache_activate()
def oneshot_exit(self):
self.oneshot.cache_deactivate()
@wrap_exceptions
def name(self):
name = self.oneshot()[kinfo_proc_map['name']]
return name if name is not None else cext.proc_name(self.pid)
@wrap_exceptions
def exe(self):
if FREEBSD:
return cext.proc_exe(self.pid)
elif NETBSD:
if self.pid == 0:
# /proc/0 dir exists but /proc/0/exe doesn't
return ""
with wrap_exceptions_procfs(self):
return os.readlink("/proc/%s/exe" % self.pid)
else:
# OpenBSD: exe cannot be determined; references:
# https://chromium.googlesource.com/chromium/src/base/+/
# master/base_paths_posix.cc
# We try our best guess by using which against the first
# cmdline arg (may return None).
cmdline = self.cmdline()
if cmdline:
return which(cmdline[0])
else:
return ""
@wrap_exceptions
def cmdline(self):
if OPENBSD and self.pid == 0:
return [] # ...else it crashes
elif NETBSD:
# XXX - most of the times the underlying sysctl() call on Net
# and Open BSD returns a truncated string.
# Also /proc/pid/cmdline behaves the same so it looks
# like this is a kernel bug.
try:
return cext.proc_cmdline(self.pid)
except OSError as err:
if err.errno == errno.EINVAL:
if not pid_exists(self.pid):
raise NoSuchProcess(self.pid, self._name)
else:
raise ZombieProcess(self.pid, self._name, self._ppid)
else:
raise
else:
return cext.proc_cmdline(self.pid)
@wrap_exceptions
def terminal(self):
tty_nr = self.oneshot()[kinfo_proc_map['ttynr']]
tmap = _psposix.get_terminal_map()
try:
return tmap[tty_nr]
except KeyError:
return None
@wrap_exceptions
def ppid(self):
self._ppid = self.oneshot()[kinfo_proc_map['ppid']]
return self._ppid
@wrap_exceptions
def uids(self):
rawtuple = self.oneshot()
return _common.puids(
rawtuple[kinfo_proc_map['real_uid']],
rawtuple[kinfo_proc_map['effective_uid']],
rawtuple[kinfo_proc_map['saved_uid']])
@wrap_exceptions
def gids(self):
rawtuple = self.oneshot()
return _common.pgids(
rawtuple[kinfo_proc_map['real_gid']],
rawtuple[kinfo_proc_map['effective_gid']],
rawtuple[kinfo_proc_map['saved_gid']])
@wrap_exceptions
def cpu_times(self):
rawtuple = self.oneshot()
return _common.pcputimes(
rawtuple[kinfo_proc_map['user_time']],
rawtuple[kinfo_proc_map['sys_time']],
rawtuple[kinfo_proc_map['ch_user_time']],
rawtuple[kinfo_proc_map['ch_sys_time']])
if FREEBSD:
@wrap_exceptions
def cpu_num(self):
return self.oneshot()[kinfo_proc_map['cpunum']]
@wrap_exceptions
def memory_info(self):
rawtuple = self.oneshot()
return pmem(
rawtuple[kinfo_proc_map['rss']],
rawtuple[kinfo_proc_map['vms']],
rawtuple[kinfo_proc_map['memtext']],
rawtuple[kinfo_proc_map['memdata']],
rawtuple[kinfo_proc_map['memstack']])
memory_full_info = memory_info
@wrap_exceptions
def create_time(self):
return self.oneshot()[kinfo_proc_map['create_time']]
@wrap_exceptions
def num_threads(self):
if hasattr(cext, "proc_num_threads"):
# FreeBSD
return cext.proc_num_threads(self.pid)
else:
return len(self.threads())
@wrap_exceptions
def num_ctx_switches(self):
rawtuple = self.oneshot()
return _common.pctxsw(
rawtuple[kinfo_proc_map['ctx_switches_vol']],
rawtuple[kinfo_proc_map['ctx_switches_unvol']])
@wrap_exceptions
def threads(self):
# Note: on OpenSBD this (/dev/mem) requires root access.
rawlist = cext.proc_threads(self.pid)
retlist = []
for thread_id, utime, stime in rawlist:
ntuple = _common.pthread(thread_id, utime, stime)
retlist.append(ntuple)
if OPENBSD:
# On OpenBSD the underlying C function does not raise NSP
# in case the process is gone (and the returned list may
# incomplete).
self.name() # raise NSP if the process disappeared on us
return retlist
@wrap_exceptions
def connections(self, kind='inet'):
if kind not in conn_tmap:
raise ValueError("invalid %r kind argument; choose between %s"
% (kind, ', '.join([repr(x) for x in conn_tmap])))
if NETBSD:
families, types = conn_tmap[kind]
ret = set()
rawlist = cext.net_connections(self.pid)
for item in rawlist:
fd, fam, type, laddr, raddr, status, pid = item
assert pid == self.pid
if fam in families and type in types:
try:
status = TCP_STATUSES[status]
except KeyError:
status = TCP_STATUSES[cext.PSUTIL_CONN_NONE]
fam = sockfam_to_enum(fam)
type = socktype_to_enum(type)
nt = _common.pconn(fd, fam, type, laddr, raddr, status)
ret.add(nt)
# On NetBSD the underlying C function does not raise NSP
# in case the process is gone (and the returned list may
# incomplete).
self.name() # raise NSP if the process disappeared on us
return list(ret)
families, types = conn_tmap[kind]
rawlist = cext.proc_connections(self.pid, families, types)
ret = []
for item in rawlist:
fd, fam, type, laddr, raddr, status = item
fam = sockfam_to_enum(fam)
type = socktype_to_enum(type)
status = TCP_STATUSES[status]
nt = _common.pconn(fd, fam, type, laddr, raddr, status)
ret.append(nt)
if OPENBSD:
# On OpenBSD the underlying C function does not raise NSP
# in case the process is gone (and the returned list may
# incomplete).
self.name() # raise NSP if the process disappeared on us
return ret
@wrap_exceptions
def wait(self, timeout=None):
try:
return _psposix.wait_pid(self.pid, timeout)
except _psposix.TimeoutExpired:
raise TimeoutExpired(timeout, self.pid, self._name)
@wrap_exceptions
def nice_get(self):
return cext_posix.getpriority(self.pid)
@wrap_exceptions
def nice_set(self, value):
return cext_posix.setpriority(self.pid, value)
@wrap_exceptions
def status(self):
code = self.oneshot()[kinfo_proc_map['status']]
# XXX is '?' legit? (we're not supposed to return it anyway)
return PROC_STATUSES.get(code, '?')
@wrap_exceptions
def io_counters(self):
rawtuple = self.oneshot()
return _common.pio(
rawtuple[kinfo_proc_map['read_io_count']],
rawtuple[kinfo_proc_map['write_io_count']],
-1,
-1)
@wrap_exceptions
def cwd(self):
"""Return process current working directory."""
# sometimes we get an empty string, in which case we turn
# it into None
if OPENBSD and self.pid == 0:
return None # ...else it would raise EINVAL
elif NETBSD:
with wrap_exceptions_procfs(self):
return os.readlink("/proc/%s/cwd" % self.pid)
elif hasattr(cext, 'proc_open_files'):
# FreeBSD < 8 does not support functions based on
# kinfo_getfile() and kinfo_getvmmap()
return cext.proc_cwd(self.pid) or None
else:
raise NotImplementedError(
"supported only starting from FreeBSD 8" if
FREEBSD else "")
nt_mmap_grouped = namedtuple(
'mmap', 'path rss, private, ref_count, shadow_count')
nt_mmap_ext = namedtuple(
'mmap', 'addr, perms path rss, private, ref_count, shadow_count')
def _not_implemented(self):
raise NotImplementedError
# FreeBSD < 8 does not support functions based on kinfo_getfile()
# and kinfo_getvmmap()
if hasattr(cext, 'proc_open_files'):
@wrap_exceptions
def open_files(self):
"""Return files opened by process as a list of namedtuples."""
rawlist = cext.proc_open_files(self.pid)
return [_common.popenfile(path, fd) for path, fd in rawlist]
else:
open_files = _not_implemented
# FreeBSD < 8 does not support functions based on kinfo_getfile()
# and kinfo_getvmmap()
if hasattr(cext, 'proc_num_fds'):
@wrap_exceptions
def num_fds(self):
"""Return the number of file descriptors opened by this process."""
ret = cext.proc_num_fds(self.pid)
if NETBSD:
# On NetBSD the underlying C function does not raise NSP
# in case the process is gone.
self.name() # raise NSP if the process disappeared on us
return ret
else:
num_fds = _not_implemented
# --- FreeBSD only APIs
if FREEBSD:
@wrap_exceptions
def cpu_affinity_get(self):
return cext.proc_cpu_affinity_get(self.pid)
@wrap_exceptions
def cpu_affinity_set(self, cpus):
# Pre-emptively check if CPUs are valid because the C
# function has a weird behavior in case of invalid CPUs,
# see: https://github.com/giampaolo/psutil/issues/586
allcpus = tuple(range(len(per_cpu_times())))
for cpu in cpus:
if cpu not in allcpus:
raise ValueError("invalid CPU #%i (choose between %s)"
% (cpu, allcpus))
try:
cext.proc_cpu_affinity_set(self.pid, cpus)
except OSError as err:
# 'man cpuset_setaffinity' about EDEADLK:
# <<the call would leave a thread without a valid CPU to run
# on because the set does not overlap with the thread's
# anonymous mask>>
if err.errno in (errno.EINVAL, errno.EDEADLK):
for cpu in cpus:
if cpu not in allcpus:
raise ValueError(
"invalid CPU #%i (choose between %s)" % (
cpu, allcpus))
raise
@wrap_exceptions
def memory_maps(self):
return cext.proc_memory_maps(self.pid)
+1872
View File
File diff suppressed because it is too large Load Diff
+506
View File
@@ -0,0 +1,506 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""OSX platform implementation."""
import errno
import functools
import os
from collections import namedtuple
from . import _common
from . import _psposix
from . import _psutil_osx as cext
from . import _psutil_posix as cext_posix
from ._common import conn_tmap
from ._common import isfile_strict
from ._common import memoize_when_activated
from ._common import parse_environ_block
from ._common import sockfam_to_enum
from ._common import socktype_to_enum
from ._common import usage_percent
__extra__all__ = []
# =====================================================================
# --- globals
# =====================================================================
PAGESIZE = os.sysconf("SC_PAGE_SIZE")
AF_LINK = cext_posix.AF_LINK
# http://students.mimuw.edu.pl/lxr/source/include/net/tcp_states.h
TCP_STATUSES = {
cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED,
cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT,
cext.TCPS_SYN_RECEIVED: _common.CONN_SYN_RECV,
cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1,
cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2,
cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT,
cext.TCPS_CLOSED: _common.CONN_CLOSE,
cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT,
cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK,
cext.TCPS_LISTEN: _common.CONN_LISTEN,
cext.TCPS_CLOSING: _common.CONN_CLOSING,
cext.PSUTIL_CONN_NONE: _common.CONN_NONE,
}
PROC_STATUSES = {
cext.SIDL: _common.STATUS_IDLE,
cext.SRUN: _common.STATUS_RUNNING,
cext.SSLEEP: _common.STATUS_SLEEPING,
cext.SSTOP: _common.STATUS_STOPPED,
cext.SZOMB: _common.STATUS_ZOMBIE,
}
kinfo_proc_map = dict(
ppid=0,
ruid=1,
euid=2,
suid=3,
rgid=4,
egid=5,
sgid=6,
ttynr=7,
ctime=8,
status=9,
name=10,
)
pidtaskinfo_map = dict(
cpuutime=0,
cpustime=1,
rss=2,
vms=3,
pfaults=4,
pageins=5,
numthreads=6,
volctxsw=7,
)
# these get overwritten on "import psutil" from the __init__.py file
NoSuchProcess = None
ZombieProcess = None
AccessDenied = None
TimeoutExpired = None
# =====================================================================
# --- named tuples
# =====================================================================
# psutil.cpu_times()
scputimes = namedtuple('scputimes', ['user', 'nice', 'system', 'idle'])
# psutil.virtual_memory()
svmem = namedtuple(
'svmem', ['total', 'available', 'percent', 'used', 'free',
'active', 'inactive', 'wired'])
# psutil.Process.memory_info()
pmem = namedtuple('pmem', ['rss', 'vms', 'pfaults', 'pageins'])
# psutil.Process.memory_full_info()
pfullmem = namedtuple('pfullmem', pmem._fields + ('uss', ))
# psutil.Process.memory_maps(grouped=True)
pmmap_grouped = namedtuple(
'pmmap_grouped',
'path rss private swapped dirtied ref_count shadow_depth')
# psutil.Process.memory_maps(grouped=False)
pmmap_ext = namedtuple(
'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields))
# =====================================================================
# --- memory
# =====================================================================
def virtual_memory():
"""System virtual memory as a namedtuple."""
total, active, inactive, wired, free = cext.virtual_mem()
avail = inactive + free
used = active + inactive + wired
percent = usage_percent((total - avail), total, _round=1)
return svmem(total, avail, percent, used, free,
active, inactive, wired)
def swap_memory():
"""Swap system memory as a (total, used, free, sin, sout) tuple."""
total, used, free, sin, sout = cext.swap_mem()
percent = usage_percent(used, total, _round=1)
return _common.sswap(total, used, free, percent, sin, sout)
# =====================================================================
# --- CPU
# =====================================================================
def cpu_times():
"""Return system CPU times as a namedtuple."""
user, nice, system, idle = cext.cpu_times()
return scputimes(user, nice, system, idle)
def per_cpu_times():
"""Return system CPU times as a named tuple"""
ret = []
for cpu_t in cext.per_cpu_times():
user, nice, system, idle = cpu_t
item = scputimes(user, nice, system, idle)
ret.append(item)
return ret
def cpu_count_logical():
"""Return the number of logical CPUs in the system."""
return cext.cpu_count_logical()
def cpu_count_physical():
"""Return the number of physical CPUs in the system."""
return cext.cpu_count_phys()
def cpu_stats():
ctx_switches, interrupts, soft_interrupts, syscalls, traps = \
cext.cpu_stats()
return _common.scpustats(
ctx_switches, interrupts, soft_interrupts, syscalls)
def cpu_freq():
"""Return CPU frequency.
On OSX per-cpu frequency is not supported.
Also, the returned frequency never changes, see:
https://arstechnica.com/civis/viewtopic.php?f=19&t=465002
"""
curr, min_, max_ = cext.cpu_freq()
return [_common.scpufreq(curr, min_, max_)]
# =====================================================================
# --- disks
# =====================================================================
disk_usage = _psposix.disk_usage
disk_io_counters = cext.disk_io_counters
def disk_partitions(all=False):
"""Return mounted disk partitions as a list of namedtuples."""
retlist = []
partitions = cext.disk_partitions()
for partition in partitions:
device, mountpoint, fstype, opts = partition
if device == 'none':
device = ''
if not all:
if not os.path.isabs(device) or not os.path.exists(device):
continue
ntuple = _common.sdiskpart(device, mountpoint, fstype, opts)
retlist.append(ntuple)
return retlist
# =====================================================================
# --- network
# =====================================================================
net_io_counters = cext.net_io_counters
net_if_addrs = cext_posix.net_if_addrs
def net_connections(kind='inet'):
"""System-wide network connections."""
# Note: on OSX this will fail with AccessDenied unless
# the process is owned by root.
ret = []
for pid in pids():
try:
cons = Process(pid).connections(kind)
except NoSuchProcess:
continue
else:
if cons:
for c in cons:
c = list(c) + [pid]
ret.append(_common.sconn(*c))
return ret
def net_if_stats():
"""Get NIC stats (isup, duplex, speed, mtu)."""
names = net_io_counters().keys()
ret = {}
for name in names:
mtu = cext_posix.net_if_mtu(name)
isup = cext_posix.net_if_flags(name)
duplex, speed = cext_posix.net_if_duplex_speed(name)
if hasattr(_common, 'NicDuplex'):
duplex = _common.NicDuplex(duplex)
ret[name] = _common.snicstats(isup, duplex, speed, mtu)
return ret
# =====================================================================
# --- other system functions
# =====================================================================
def boot_time():
"""The system boot time expressed in seconds since the epoch."""
return cext.boot_time()
def users():
"""Return currently connected users as a list of namedtuples."""
retlist = []
rawlist = cext.users()
for item in rawlist:
user, tty, hostname, tstamp = item
if tty == '~':
continue # reboot or shutdown
if not tstamp:
continue
nt = _common.suser(user, tty or None, hostname or None, tstamp)
retlist.append(nt)
return retlist
# =====================================================================
# --- processes
# =====================================================================
pids = cext.pids
pid_exists = _psposix.pid_exists
def wrap_exceptions(fun):
"""Decorator which translates bare OSError exceptions into
NoSuchProcess and AccessDenied.
"""
@functools.wraps(fun)
def wrapper(self, *args, **kwargs):
try:
return fun(self, *args, **kwargs)
except OSError as err:
if self.pid == 0:
if 0 in pids():
raise AccessDenied(self.pid, self._name)
else:
raise
if err.errno == errno.ESRCH:
if not pid_exists(self.pid):
raise NoSuchProcess(self.pid, self._name)
else:
raise ZombieProcess(self.pid, self._name, self._ppid)
if err.errno in (errno.EPERM, errno.EACCES):
raise AccessDenied(self.pid, self._name)
raise
return wrapper
class Process(object):
"""Wrapper class around underlying C implementation."""
__slots__ = ["pid", "_name", "_ppid"]
def __init__(self, pid):
self.pid = pid
self._name = None
self._ppid = None
@memoize_when_activated
def _get_kinfo_proc(self):
# Note: should work with all PIDs without permission issues.
ret = cext.proc_kinfo_oneshot(self.pid)
assert len(ret) == len(kinfo_proc_map)
return ret
@memoize_when_activated
def _get_pidtaskinfo(self):
# Note: should work for PIDs owned by user only.
ret = cext.proc_pidtaskinfo_oneshot(self.pid)
assert len(ret) == len(pidtaskinfo_map)
return ret
def oneshot_enter(self):
self._get_kinfo_proc.cache_activate()
self._get_pidtaskinfo.cache_activate()
def oneshot_exit(self):
self._get_kinfo_proc.cache_deactivate()
self._get_pidtaskinfo.cache_deactivate()
@wrap_exceptions
def name(self):
name = self._get_kinfo_proc()[kinfo_proc_map['name']]
return name if name is not None else cext.proc_name(self.pid)
@wrap_exceptions
def exe(self):
return cext.proc_exe(self.pid)
@wrap_exceptions
def cmdline(self):
if not pid_exists(self.pid):
raise NoSuchProcess(self.pid, self._name)
return cext.proc_cmdline(self.pid)
@wrap_exceptions
def environ(self):
if not pid_exists(self.pid):
raise NoSuchProcess(self.pid, self._name)
return parse_environ_block(cext.proc_environ(self.pid))
@wrap_exceptions
def ppid(self):
self._ppid = self._get_kinfo_proc()[kinfo_proc_map['ppid']]
return self._ppid
@wrap_exceptions
def cwd(self):
return cext.proc_cwd(self.pid)
@wrap_exceptions
def uids(self):
rawtuple = self._get_kinfo_proc()
return _common.puids(
rawtuple[kinfo_proc_map['ruid']],
rawtuple[kinfo_proc_map['euid']],
rawtuple[kinfo_proc_map['suid']])
@wrap_exceptions
def gids(self):
rawtuple = self._get_kinfo_proc()
return _common.puids(
rawtuple[kinfo_proc_map['rgid']],
rawtuple[kinfo_proc_map['egid']],
rawtuple[kinfo_proc_map['sgid']])
@wrap_exceptions
def terminal(self):
tty_nr = self._get_kinfo_proc()[kinfo_proc_map['ttynr']]
tmap = _psposix.get_terminal_map()
try:
return tmap[tty_nr]
except KeyError:
return None
@wrap_exceptions
def memory_info(self):
rawtuple = self._get_pidtaskinfo()
return pmem(
rawtuple[pidtaskinfo_map['rss']],
rawtuple[pidtaskinfo_map['vms']],
rawtuple[pidtaskinfo_map['pfaults']],
rawtuple[pidtaskinfo_map['pageins']],
)
@wrap_exceptions
def memory_full_info(self):
basic_mem = self.memory_info()
uss = cext.proc_memory_uss(self.pid)
return pfullmem(*basic_mem + (uss, ))
@wrap_exceptions
def cpu_times(self):
rawtuple = self._get_pidtaskinfo()
return _common.pcputimes(
rawtuple[pidtaskinfo_map['cpuutime']],
rawtuple[pidtaskinfo_map['cpustime']],
# children user / system times are not retrievable (set to 0)
0, 0)
@wrap_exceptions
def create_time(self):
return self._get_kinfo_proc()[kinfo_proc_map['ctime']]
@wrap_exceptions
def num_ctx_switches(self):
# Unvoluntary value seems not to be available;
# getrusage() numbers seems to confirm this theory.
# We set it to 0.
vol = self._get_pidtaskinfo()[pidtaskinfo_map['volctxsw']]
return _common.pctxsw(vol, 0)
@wrap_exceptions
def num_threads(self):
return self._get_pidtaskinfo()[pidtaskinfo_map['numthreads']]
@wrap_exceptions
def open_files(self):
if self.pid == 0:
return []
files = []
rawlist = cext.proc_open_files(self.pid)
for path, fd in rawlist:
if isfile_strict(path):
ntuple = _common.popenfile(path, fd)
files.append(ntuple)
return files
@wrap_exceptions
def connections(self, kind='inet'):
if kind not in conn_tmap:
raise ValueError("invalid %r kind argument; choose between %s"
% (kind, ', '.join([repr(x) for x in conn_tmap])))
families, types = conn_tmap[kind]
rawlist = cext.proc_connections(self.pid, families, types)
ret = []
for item in rawlist:
fd, fam, type, laddr, raddr, status = item
status = TCP_STATUSES[status]
fam = sockfam_to_enum(fam)
type = socktype_to_enum(type)
nt = _common.pconn(fd, fam, type, laddr, raddr, status)
ret.append(nt)
return ret
@wrap_exceptions
def num_fds(self):
if self.pid == 0:
return 0
return cext.proc_num_fds(self.pid)
@wrap_exceptions
def wait(self, timeout=None):
try:
return _psposix.wait_pid(self.pid, timeout)
except _psposix.TimeoutExpired:
raise TimeoutExpired(timeout, self.pid, self._name)
@wrap_exceptions
def nice_get(self):
return cext_posix.getpriority(self.pid)
@wrap_exceptions
def nice_set(self, value):
return cext_posix.setpriority(self.pid, value)
@wrap_exceptions
def status(self):
code = self._get_kinfo_proc()[kinfo_proc_map['status']]
# XXX is '?' legit? (we're not supposed to return it anyway)
return PROC_STATUSES.get(code, '?')
@wrap_exceptions
def threads(self):
rawlist = cext.proc_threads(self.pid)
retlist = []
for thread_id, utime, stime in rawlist:
ntuple = _common.pthread(thread_id, utime, stime)
retlist.append(ntuple)
return retlist
@wrap_exceptions
def memory_maps(self):
return cext.proc_memory_maps(self.pid)
+184
View File
@@ -0,0 +1,184 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Routines common to all posix systems."""
import errno
import glob
import os
import sys
import time
from ._common import memoize
from ._common import sdiskusage
from ._common import usage_percent
from ._compat import PY3
from ._compat import unicode
__all__ = ['TimeoutExpired', 'pid_exists', 'wait_pid', 'disk_usage',
'get_terminal_map']
class TimeoutExpired(Exception):
pass
def pid_exists(pid):
"""Check whether pid exists in the current process table."""
if pid == 0:
# According to "man 2 kill" PID 0 has a special meaning:
# it refers to <<every process in the process group of the
# calling process>> so we don't want to go any further.
# If we get here it means this UNIX platform *does* have
# a process with id 0.
return True
try:
os.kill(pid, 0)
except OSError as err:
if err.errno == errno.ESRCH:
# ESRCH == No such process
return False
elif err.errno == errno.EPERM:
# EPERM clearly means there's a process to deny access to
return True
else:
# According to "man 2 kill" possible error values are
# (EINVAL, EPERM, ESRCH) therefore we should never get
# here. If we do let's be explicit in considering this
# an error.
raise err
else:
return True
def wait_pid(pid, timeout=None):
"""Wait for process with pid 'pid' to terminate and return its
exit status code as an integer.
If pid is not a children of os.getpid() (current process) just
waits until the process disappears and return None.
If pid does not exist at all return None immediately.
Raise TimeoutExpired on timeout expired.
"""
def check_timeout(delay):
if timeout is not None:
if timer() >= stop_at:
raise TimeoutExpired()
time.sleep(delay)
return min(delay * 2, 0.04)
timer = getattr(time, 'monotonic', time.time)
if timeout is not None:
def waitcall():
return os.waitpid(pid, os.WNOHANG)
stop_at = timer() + timeout
else:
def waitcall():
return os.waitpid(pid, 0)
delay = 0.0001
while True:
try:
retpid, status = waitcall()
except OSError as err:
if err.errno == errno.EINTR:
delay = check_timeout(delay)
continue
elif err.errno == errno.ECHILD:
# This has two meanings:
# - pid is not a child of os.getpid() in which case
# we keep polling until it's gone
# - pid never existed in the first place
# In both cases we'll eventually return None as we
# can't determine its exit status code.
while True:
if pid_exists(pid):
delay = check_timeout(delay)
else:
return
else:
raise
else:
if retpid == 0:
# WNOHANG was used, pid is still running
delay = check_timeout(delay)
continue
# process exited due to a signal; return the integer of
# that signal
if os.WIFSIGNALED(status):
return -os.WTERMSIG(status)
# process exited using exit(2) system call; return the
# integer exit(2) system call has been called with
elif os.WIFEXITED(status):
return os.WEXITSTATUS(status)
else:
# should never happen
raise ValueError("unknown process exit status %r" % status)
def disk_usage(path):
"""Return disk usage associated with path.
Note: UNIX usually reserves 5% disk space which is not accessible
by user. In this function "total" and "used" values reflect the
total and used disk space whereas "free" and "percent" represent
the "free" and "used percent" user disk space.
"""
try:
st = os.statvfs(path)
except UnicodeEncodeError:
if not PY3 and isinstance(path, unicode):
# this is a bug with os.statvfs() and unicode on
# Python 2, see:
# - https://github.com/giampaolo/psutil/issues/416
# - http://bugs.python.org/issue18695
try:
path = path.encode(sys.getfilesystemencoding())
except UnicodeEncodeError:
pass
st = os.statvfs(path)
else:
raise
# Total space which is only available to root (unless changed
# at system level).
total = (st.f_blocks * st.f_frsize)
# Remaining free space usable by root.
avail_to_root = (st.f_bfree * st.f_frsize)
# Remaining free space usable by user.
avail_to_user = (st.f_bavail * st.f_frsize)
# Total space being used in general.
used = (total - avail_to_root)
# Total space which is available to user (same as 'total' but
# for the user).
total_user = used + avail_to_user
# User usage percent compared to the total amount of space
# the user can use. This number would be higher if compared
# to root's because the user has less space (usually -5%).
usage_percent_user = usage_percent(used, total_user, _round=1)
# NB: the percentage is -5% than what shown by df due to
# reserved blocks that we are currently not considering:
# https://github.com/giampaolo/psutil/issues/829#issuecomment-223750462
return sdiskusage(
total=total, used=used, free=avail_to_user, percent=usage_percent_user)
@memoize
def get_terminal_map():
"""Get a map of device-id -> path as a dict.
Used by Process.terminal()
"""
ret = {}
ls = glob.glob('/dev/tty*') + glob.glob('/dev/pts/*')
for name in ls:
assert name not in ret, name
try:
ret[os.stat(name).st_rdev] = name
except OSError as err:
if err.errno != errno.ENOENT:
raise
return ret
+702
View File
@@ -0,0 +1,702 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Sun OS Solaris platform implementation."""
import errno
import os
import socket
import subprocess
import sys
from collections import namedtuple
from . import _common
from . import _psposix
from . import _psutil_posix as cext_posix
from . import _psutil_sunos as cext
from ._common import isfile_strict
from ._common import memoize_when_activated
from ._common import sockfam_to_enum
from ._common import socktype_to_enum
from ._common import usage_percent
from ._compat import b
from ._compat import PY3
__extra__all__ = ["CONN_IDLE", "CONN_BOUND", "PROCFS_PATH"]
# =====================================================================
# --- globals
# =====================================================================
PAGE_SIZE = os.sysconf('SC_PAGE_SIZE')
AF_LINK = cext_posix.AF_LINK
IS_64_BIT = sys.maxsize > 2**32
CONN_IDLE = "IDLE"
CONN_BOUND = "BOUND"
PROC_STATUSES = {
cext.SSLEEP: _common.STATUS_SLEEPING,
cext.SRUN: _common.STATUS_RUNNING,
cext.SZOMB: _common.STATUS_ZOMBIE,
cext.SSTOP: _common.STATUS_STOPPED,
cext.SIDL: _common.STATUS_IDLE,
cext.SONPROC: _common.STATUS_RUNNING, # same as run
cext.SWAIT: _common.STATUS_WAITING,
}
TCP_STATUSES = {
cext.TCPS_ESTABLISHED: _common.CONN_ESTABLISHED,
cext.TCPS_SYN_SENT: _common.CONN_SYN_SENT,
cext.TCPS_SYN_RCVD: _common.CONN_SYN_RECV,
cext.TCPS_FIN_WAIT_1: _common.CONN_FIN_WAIT1,
cext.TCPS_FIN_WAIT_2: _common.CONN_FIN_WAIT2,
cext.TCPS_TIME_WAIT: _common.CONN_TIME_WAIT,
cext.TCPS_CLOSED: _common.CONN_CLOSE,
cext.TCPS_CLOSE_WAIT: _common.CONN_CLOSE_WAIT,
cext.TCPS_LAST_ACK: _common.CONN_LAST_ACK,
cext.TCPS_LISTEN: _common.CONN_LISTEN,
cext.TCPS_CLOSING: _common.CONN_CLOSING,
cext.PSUTIL_CONN_NONE: _common.CONN_NONE,
cext.TCPS_IDLE: CONN_IDLE, # sunos specific
cext.TCPS_BOUND: CONN_BOUND, # sunos specific
}
# these get overwritten on "import psutil" from the __init__.py file
NoSuchProcess = None
ZombieProcess = None
AccessDenied = None
TimeoutExpired = None
# =====================================================================
# --- named tuples
# =====================================================================
# psutil.cpu_times()
scputimes = namedtuple('scputimes', ['user', 'system', 'idle', 'iowait'])
# psutil.cpu_times(percpu=True)
pcputimes = namedtuple('pcputimes',
['user', 'system', 'children_user', 'children_system'])
# psutil.virtual_memory()
svmem = namedtuple('svmem', ['total', 'available', 'percent', 'used', 'free'])
# psutil.Process.memory_info()
pmem = namedtuple('pmem', ['rss', 'vms'])
pfullmem = pmem
# psutil.Process.memory_maps(grouped=True)
pmmap_grouped = namedtuple('pmmap_grouped',
['path', 'rss', 'anonymous', 'locked'])
# psutil.Process.memory_maps(grouped=False)
pmmap_ext = namedtuple(
'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields))
# =====================================================================
# --- utils
# =====================================================================
def get_procfs_path():
"""Return updated psutil.PROCFS_PATH constant."""
return sys.modules['psutil'].PROCFS_PATH
# =====================================================================
# --- memory
# =====================================================================
def virtual_memory():
"""Report virtual memory metrics."""
# we could have done this with kstat, but IMHO this is good enough
total = os.sysconf('SC_PHYS_PAGES') * PAGE_SIZE
# note: there's no difference on Solaris
free = avail = os.sysconf('SC_AVPHYS_PAGES') * PAGE_SIZE
used = total - free
percent = usage_percent(used, total, _round=1)
return svmem(total, avail, percent, used, free)
def swap_memory():
"""Report swap memory metrics."""
sin, sout = cext.swap_mem()
# XXX
# we are supposed to get total/free by doing so:
# http://cvs.opensolaris.org/source/xref/onnv/onnv-gate/
# usr/src/cmd/swap/swap.c
# ...nevertheless I can't manage to obtain the same numbers as 'swap'
# cmdline utility, so let's parse its output (sigh!)
p = subprocess.Popen(['/usr/bin/env', 'PATH=/usr/sbin:/sbin:%s' %
os.environ['PATH'], 'swap', '-l'],
stdout=subprocess.PIPE)
stdout, stderr = p.communicate()
if PY3:
stdout = stdout.decode(sys.stdout.encoding)
if p.returncode != 0:
raise RuntimeError("'swap -l' failed (retcode=%s)" % p.returncode)
lines = stdout.strip().split('\n')[1:]
if not lines:
raise RuntimeError('no swap device(s) configured')
total = free = 0
for line in lines:
line = line.split()
t, f = line[-2:]
total += int(int(t) * 512)
free += int(int(f) * 512)
used = total - free
percent = usage_percent(used, total, _round=1)
return _common.sswap(total, used, free, percent,
sin * PAGE_SIZE, sout * PAGE_SIZE)
# =====================================================================
# --- CPU
# =====================================================================
def cpu_times():
"""Return system-wide CPU times as a named tuple"""
ret = cext.per_cpu_times()
return scputimes(*[sum(x) for x in zip(*ret)])
def per_cpu_times():
"""Return system per-CPU times as a list of named tuples"""
ret = cext.per_cpu_times()
return [scputimes(*x) for x in ret]
def cpu_count_logical():
"""Return the number of logical CPUs in the system."""
try:
return os.sysconf("SC_NPROCESSORS_ONLN")
except ValueError:
# mimic os.cpu_count() behavior
return None
def cpu_count_physical():
"""Return the number of physical CPUs in the system."""
return cext.cpu_count_phys()
def cpu_stats():
"""Return various CPU stats as a named tuple."""
ctx_switches, interrupts, syscalls, traps = cext.cpu_stats()
soft_interrupts = 0
return _common.scpustats(ctx_switches, interrupts, soft_interrupts,
syscalls)
# =====================================================================
# --- disks
# =====================================================================
disk_io_counters = cext.disk_io_counters
disk_usage = _psposix.disk_usage
def disk_partitions(all=False):
"""Return system disk partitions."""
# TODO - the filtering logic should be better checked so that
# it tries to reflect 'df' as much as possible
retlist = []
partitions = cext.disk_partitions()
for partition in partitions:
device, mountpoint, fstype, opts = partition
if device == 'none':
device = ''
if not all:
# Differently from, say, Linux, we don't have a list of
# common fs types so the best we can do, AFAIK, is to
# filter by filesystem having a total size > 0.
if not disk_usage(mountpoint).total:
continue
ntuple = _common.sdiskpart(device, mountpoint, fstype, opts)
retlist.append(ntuple)
return retlist
# =====================================================================
# --- network
# =====================================================================
net_io_counters = cext.net_io_counters
net_if_addrs = cext_posix.net_if_addrs
def net_connections(kind, _pid=-1):
"""Return socket connections. If pid == -1 return system-wide
connections (as opposed to connections opened by one process only).
Only INET sockets are returned (UNIX are not).
"""
cmap = _common.conn_tmap.copy()
if _pid == -1:
cmap.pop('unix', 0)
if kind not in cmap:
raise ValueError("invalid %r kind argument; choose between %s"
% (kind, ', '.join([repr(x) for x in cmap])))
families, types = _common.conn_tmap[kind]
rawlist = cext.net_connections(_pid)
ret = set()
for item in rawlist:
fd, fam, type_, laddr, raddr, status, pid = item
if fam not in families:
continue
if type_ not in types:
continue
status = TCP_STATUSES[status]
fam = sockfam_to_enum(fam)
type_ = socktype_to_enum(type_)
if _pid == -1:
nt = _common.sconn(fd, fam, type_, laddr, raddr, status, pid)
else:
nt = _common.pconn(fd, fam, type_, laddr, raddr, status)
ret.add(nt)
return list(ret)
def net_if_stats():
"""Get NIC stats (isup, duplex, speed, mtu)."""
ret = cext.net_if_stats()
for name, items in ret.items():
isup, duplex, speed, mtu = items
if hasattr(_common, 'NicDuplex'):
duplex = _common.NicDuplex(duplex)
ret[name] = _common.snicstats(isup, duplex, speed, mtu)
return ret
# =====================================================================
# --- other system functions
# =====================================================================
def boot_time():
"""The system boot time expressed in seconds since the epoch."""
return cext.boot_time()
def users():
"""Return currently connected users as a list of namedtuples."""
retlist = []
rawlist = cext.users()
localhost = (':0.0', ':0')
for item in rawlist:
user, tty, hostname, tstamp, user_process = item
# note: the underlying C function includes entries about
# system boot, run level and others. We might want
# to use them in the future.
if not user_process:
continue
if hostname in localhost:
hostname = 'localhost'
nt = _common.suser(user, tty, hostname, tstamp)
retlist.append(nt)
return retlist
# =====================================================================
# --- processes
# =====================================================================
def pids():
"""Returns a list of PIDs currently running on the system."""
return [int(x) for x in os.listdir(b(get_procfs_path())) if x.isdigit()]
def pid_exists(pid):
"""Check for the existence of a unix pid."""
return _psposix.pid_exists(pid)
def wrap_exceptions(fun):
"""Call callable into a try/except clause and translate ENOENT,
EACCES and EPERM in NoSuchProcess or AccessDenied exceptions.
"""
def wrapper(self, *args, **kwargs):
try:
return fun(self, *args, **kwargs)
except EnvironmentError as err:
if self.pid == 0:
if 0 in pids():
raise AccessDenied(self.pid, self._name)
else:
raise
# ENOENT (no such file or directory) gets raised on open().
# ESRCH (no such process) can get raised on read() if
# process is gone in meantime.
if err.errno in (errno.ENOENT, errno.ESRCH):
if not pid_exists(self.pid):
raise NoSuchProcess(self.pid, self._name)
else:
raise ZombieProcess(self.pid, self._name, self._ppid)
if err.errno in (errno.EPERM, errno.EACCES):
raise AccessDenied(self.pid, self._name)
raise
return wrapper
class Process(object):
"""Wrapper class around underlying C implementation."""
__slots__ = ["pid", "_name", "_ppid", "_procfs_path"]
def __init__(self, pid):
self.pid = pid
self._name = None
self._ppid = None
self._procfs_path = get_procfs_path()
def oneshot_enter(self):
self._proc_name_and_args.cache_activate()
self._proc_basic_info.cache_activate()
self._proc_cred.cache_activate()
def oneshot_exit(self):
self._proc_name_and_args.cache_deactivate()
self._proc_basic_info.cache_deactivate()
self._proc_cred.cache_deactivate()
@memoize_when_activated
def _proc_name_and_args(self):
return cext.proc_name_and_args(self.pid, self._procfs_path)
@memoize_when_activated
def _proc_basic_info(self):
return cext.proc_basic_info(self.pid, self._procfs_path)
@memoize_when_activated
def _proc_cred(self):
return cext.proc_cred(self.pid, self._procfs_path)
@wrap_exceptions
def name(self):
# note: max len == 15
return self._proc_name_and_args()[0]
@wrap_exceptions
def exe(self):
try:
return os.readlink(
"%s/%s/path/a.out" % (self._procfs_path, self.pid))
except OSError:
pass # continue and guess the exe name from the cmdline
# Will be guessed later from cmdline but we want to explicitly
# invoke cmdline here in order to get an AccessDenied
# exception if the user has not enough privileges.
self.cmdline()
return ""
@wrap_exceptions
def cmdline(self):
return self._proc_name_and_args()[1].split(' ')
@wrap_exceptions
def create_time(self):
return self._proc_basic_info()[3]
@wrap_exceptions
def num_threads(self):
return self._proc_basic_info()[5]
@wrap_exceptions
def nice_get(self):
# For some reason getpriority(3) return ESRCH (no such process)
# for certain low-pid processes, no matter what (even as root).
# The process actually exists though, as it has a name,
# creation time, etc.
# The best thing we can do here appears to be raising AD.
# Note: tested on Solaris 11; on Open Solaris 5 everything is
# fine.
try:
return cext_posix.getpriority(self.pid)
except EnvironmentError as err:
# 48 is 'operation not supported' but errno does not expose
# it. It occurs for low system pids.
if err.errno in (errno.ENOENT, errno.ESRCH, 48):
if pid_exists(self.pid):
raise AccessDenied(self.pid, self._name)
raise
@wrap_exceptions
def nice_set(self, value):
if self.pid in (2, 3):
# Special case PIDs: internally setpriority(3) return ESRCH
# (no such process), no matter what.
# The process actually exists though, as it has a name,
# creation time, etc.
raise AccessDenied(self.pid, self._name)
return cext_posix.setpriority(self.pid, value)
@wrap_exceptions
def ppid(self):
self._ppid = self._proc_basic_info()[0]
return self._ppid
@wrap_exceptions
def uids(self):
real, effective, saved, _, _, _ = self._proc_cred()
return _common.puids(real, effective, saved)
@wrap_exceptions
def gids(self):
_, _, _, real, effective, saved = self._proc_cred()
return _common.puids(real, effective, saved)
@wrap_exceptions
def cpu_times(self):
try:
times = cext.proc_cpu_times(self.pid, self._procfs_path)
except OSError as err:
if err.errno == errno.EOVERFLOW and not IS_64_BIT:
# We may get here if we attempt to query a 64bit process
# with a 32bit python.
# Error originates from read() and also tools like "cat"
# fail in the same way (!).
# Since there simply is no way to determine CPU times we
# return 0.0 as a fallback. See:
# https://github.com/giampaolo/psutil/issues/857
times = (0.0, 0.0, 0.0, 0.0)
else:
raise
return _common.pcputimes(*times)
@wrap_exceptions
def cpu_num(self):
return cext.proc_cpu_num(self.pid, self._procfs_path)
@wrap_exceptions
def terminal(self):
procfs_path = self._procfs_path
hit_enoent = False
tty = wrap_exceptions(
self._proc_basic_info()[0])
if tty != cext.PRNODEV:
for x in (0, 1, 2, 255):
try:
return os.readlink(
'%s/%d/path/%d' % (procfs_path, self.pid, x))
except OSError as err:
if err.errno == errno.ENOENT:
hit_enoent = True
continue
raise
if hit_enoent:
# raise NSP if the process disappeared on us
os.stat('%s/%s' % (procfs_path, self.pid))
@wrap_exceptions
def cwd(self):
# /proc/PID/path/cwd may not be resolved by readlink() even if
# it exists (ls shows it). If that's the case and the process
# is still alive return None (we can return None also on BSD).
# Reference: http://goo.gl/55XgO
procfs_path = self._procfs_path
try:
return os.readlink("%s/%s/path/cwd" % (procfs_path, self.pid))
except OSError as err:
if err.errno == errno.ENOENT:
os.stat("%s/%s" % (procfs_path, self.pid)) # raise NSP or AD
return None
raise
@wrap_exceptions
def memory_info(self):
ret = self._proc_basic_info()
rss, vms = ret[1] * 1024, ret[2] * 1024
return pmem(rss, vms)
memory_full_info = memory_info
@wrap_exceptions
def status(self):
code = self._proc_basic_info()[6]
# XXX is '?' legit? (we're not supposed to return it anyway)
return PROC_STATUSES.get(code, '?')
@wrap_exceptions
def threads(self):
procfs_path = self._procfs_path
ret = []
tids = os.listdir('%s/%d/lwp' % (procfs_path, self.pid))
hit_enoent = False
for tid in tids:
tid = int(tid)
try:
utime, stime = cext.query_process_thread(
self.pid, tid, procfs_path)
except EnvironmentError as err:
if err.errno == errno.EOVERFLOW and not IS_64_BIT:
# We may get here if we attempt to query a 64bit process
# with a 32bit python.
# Error originates from read() and also tools like "cat"
# fail in the same way (!).
# Since there simply is no way to determine CPU times we
# return 0.0 as a fallback. See:
# https://github.com/giampaolo/psutil/issues/857
continue
# ENOENT == thread gone in meantime
if err.errno == errno.ENOENT:
hit_enoent = True
continue
raise
else:
nt = _common.pthread(tid, utime, stime)
ret.append(nt)
if hit_enoent:
# raise NSP if the process disappeared on us
os.stat('%s/%s' % (procfs_path, self.pid))
return ret
@wrap_exceptions
def open_files(self):
retlist = []
hit_enoent = False
procfs_path = self._procfs_path
pathdir = '%s/%d/path' % (procfs_path, self.pid)
for fd in os.listdir('%s/%d/fd' % (procfs_path, self.pid)):
path = os.path.join(pathdir, fd)
if os.path.islink(path):
try:
file = os.readlink(path)
except OSError as err:
# ENOENT == file which is gone in the meantime
if err.errno == errno.ENOENT:
hit_enoent = True
continue
raise
else:
if isfile_strict(file):
retlist.append(_common.popenfile(file, int(fd)))
if hit_enoent:
# raise NSP if the process disappeared on us
os.stat('%s/%s' % (procfs_path, self.pid))
return retlist
def _get_unix_sockets(self, pid):
"""Get UNIX sockets used by process by parsing 'pfiles' output."""
# TODO: rewrite this in C (...but the damn netstat source code
# does not include this part! Argh!!)
cmd = "pfiles %s" % pid
p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
stdout, stderr = p.communicate()
if PY3:
stdout, stderr = [x.decode(sys.stdout.encoding)
for x in (stdout, stderr)]
if p.returncode != 0:
if 'permission denied' in stderr.lower():
raise AccessDenied(self.pid, self._name)
if 'no such process' in stderr.lower():
raise NoSuchProcess(self.pid, self._name)
raise RuntimeError("%r command error\n%s" % (cmd, stderr))
lines = stdout.split('\n')[2:]
for i, line in enumerate(lines):
line = line.lstrip()
if line.startswith('sockname: AF_UNIX'):
path = line.split(' ', 2)[2]
type = lines[i - 2].strip()
if type == 'SOCK_STREAM':
type = socket.SOCK_STREAM
elif type == 'SOCK_DGRAM':
type = socket.SOCK_DGRAM
else:
type = -1
yield (-1, socket.AF_UNIX, type, path, "", _common.CONN_NONE)
@wrap_exceptions
def connections(self, kind='inet'):
ret = net_connections(kind, _pid=self.pid)
# The underlying C implementation retrieves all OS connections
# and filters them by PID. At this point we can't tell whether
# an empty list means there were no connections for process or
# process is no longer active so we force NSP in case the PID
# is no longer there.
if not ret:
# will raise NSP if process is gone
os.stat('%s/%s' % (self._procfs_path, self.pid))
# UNIX sockets
if kind in ('all', 'unix'):
ret.extend([_common.pconn(*conn) for conn in
self._get_unix_sockets(self.pid)])
return ret
nt_mmap_grouped = namedtuple('mmap', 'path rss anon locked')
nt_mmap_ext = namedtuple('mmap', 'addr perms path rss anon locked')
@wrap_exceptions
def memory_maps(self):
def toaddr(start, end):
return '%s-%s' % (hex(start)[2:].strip('L'),
hex(end)[2:].strip('L'))
procfs_path = self._procfs_path
retlist = []
try:
rawlist = cext.proc_memory_maps(self.pid, procfs_path)
except OSError as err:
if err.errno == errno.EOVERFLOW and not IS_64_BIT:
# We may get here if we attempt to query a 64bit process
# with a 32bit python.
# Error originates from read() and also tools like "cat"
# fail in the same way (!).
# Since there simply is no way to determine CPU times we
# return 0.0 as a fallback. See:
# https://github.com/giampaolo/psutil/issues/857
return []
else:
raise
hit_enoent = False
for item in rawlist:
addr, addrsize, perm, name, rss, anon, locked = item
addr = toaddr(addr, addrsize)
if not name.startswith('['):
try:
name = os.readlink(
'%s/%s/path/%s' % (procfs_path, self.pid, name))
except OSError as err:
if err.errno == errno.ENOENT:
# sometimes the link may not be resolved by
# readlink() even if it exists (ls shows it).
# If that's the case we just return the
# unresolved link path.
# This seems an incosistency with /proc similar
# to: http://goo.gl/55XgO
name = '%s/%s/path/%s' % (procfs_path, self.pid, name)
hit_enoent = True
else:
raise
retlist.append((addr, perm, name, rss, anon, locked))
if hit_enoent:
# raise NSP if the process disappeared on us
os.stat('%s/%s' % (procfs_path, self.pid))
return retlist
@wrap_exceptions
def num_fds(self):
return len(os.listdir("%s/%s/fd" % (self._procfs_path, self.pid)))
@wrap_exceptions
def num_ctx_switches(self):
return _common.pctxsw(
*cext.proc_num_ctx_switches(self.pid, self._procfs_path))
@wrap_exceptions
def wait(self, timeout=None):
try:
return _psposix.wait_pid(self.pid, timeout)
except _psposix.TimeoutExpired:
raise TimeoutExpired(timeout, self.pid, self._name)
+1047
View File
File diff suppressed because it is too large Load Diff
+126
View File
@@ -0,0 +1,126 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Routines common to all platforms.
*/
#ifdef PSUTIL_POSIX
#include <sys/types.h>
#include <signal.h>
#endif
#include <Python.h>
/*
* Set OSError(errno=ESRCH, strerror="No such process") Python exception.
*/
PyObject *
NoSuchProcess(void) {
PyObject *exc;
char *msg = strerror(ESRCH);
exc = PyObject_CallFunction(PyExc_OSError, "(is)", ESRCH, msg);
PyErr_SetObject(PyExc_OSError, exc);
Py_XDECREF(exc);
return NULL;
}
/*
* Set OSError(errno=EACCES, strerror="Permission denied") Python exception.
*/
PyObject *
AccessDenied(void) {
PyObject *exc;
char *msg = strerror(EACCES);
exc = PyObject_CallFunction(PyExc_OSError, "(is)", EACCES, msg);
PyErr_SetObject(PyExc_OSError, exc);
Py_XDECREF(exc);
return NULL;
}
#ifdef PSUTIL_POSIX
/*
* Check if PID exists. Return values:
* 1: exists
* 0: does not exist
* -1: error (Python exception is set)
*/
int
psutil_pid_exists(long pid) {
int ret;
// No negative PID exists, plus -1 is an alias for sending signal
// too all processes except system ones. Not what we want.
if (pid < 0)
return 0;
// As per "man 2 kill" PID 0 is an alias for sending the signal to
// every process in the process group of the calling process.
// Not what we want. Some platforms have PID 0, some do not.
// We decide that at runtime.
if (pid == 0) {
#if defined(PSUTIL_LINUX) || defined(PSUTIL_FREEBSD)
return 0;
#else
return 1;
#endif
}
#if defined(PSUTIL_OSX)
ret = kill((pid_t)pid , 0);
#else
ret = kill(pid , 0);
#endif
if (ret == 0)
return 1;
else {
if (errno == ESRCH) {
// ESRCH == No such process
return 0;
}
else if (errno == EPERM) {
// EPERM clearly indicates there's a process to deny
// access to.
return 1;
}
else {
// According to "man 2 kill" possible error values are
// (EINVAL, EPERM, ESRCH) therefore we should never get
// here. If we do let's be explicit in considering this
// an error.
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
}
}
/*
* Utility used for those syscalls which do not return a meaningful
* error that we can translate into an exception which makes sense.
* As such, we'll have to guess.
* On UNIX, if errno is set, we return that one (OSError).
* Else, if PID does not exist we assume the syscall failed because
* of that so we raise NoSuchProcess.
* If none of this is true we giveup and raise RuntimeError(msg).
* This will always set a Python exception and return NULL.
*/
int
psutil_raise_for_pid(long pid, char *msg) {
// Set exception to AccessDenied if pid exists else NoSuchProcess.
if (errno != 0) {
PyErr_SetFromErrno(PyExc_OSError);
return 0;
}
if (psutil_pid_exists(pid) == 0)
NoSuchProcess();
else
PyErr_SetString(PyExc_RuntimeError, msg);
return 0;
}
#endif
+15
View File
@@ -0,0 +1,15 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
PyObject* AccessDenied(void);
PyObject* NoSuchProcess(void);
#ifdef PSUTIL_POSIX
int psutil_pid_exists(long pid);
void psutil_raise_for_pid(long pid, char *msg);
#endif
+688
View File
@@ -0,0 +1,688 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Linux-specific functions.
*/
#ifndef _GNU_SOURCE
#define _GNU_SOURCE 1
#endif
#include <Python.h>
#include <errno.h>
#include <stdlib.h>
#include <mntent.h>
#include <features.h>
#include <utmp.h>
#include <sched.h>
#include <linux/version.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <linux/sockios.h>
#include <linux/if.h>
// see: https://github.com/giampaolo/psutil/issues/659
#ifdef PSUTIL_ETHTOOL_MISSING_TYPES
#include <linux/types.h>
typedef __u64 u64;
typedef __u32 u32;
typedef __u16 u16;
typedef __u8 u8;
#endif
/* Avoid redefinition of struct sysinfo with musl libc */
#define _LINUX_SYSINFO_H
#include <linux/ethtool.h>
/* The minimum number of CPUs allocated in a cpu_set_t */
static const int NCPUS_START = sizeof(unsigned long) * CHAR_BIT;
// Linux >= 2.6.13
#define PSUTIL_HAVE_IOPRIO defined(__NR_ioprio_get) && defined(__NR_ioprio_set)
// Linux >= 2.6.36 (supposedly) and glibc >= 13
#define PSUTIL_HAVE_PRLIMIT \
(LINUX_VERSION_CODE >= KERNEL_VERSION(2, 6, 36)) && \
(__GLIBC__ >= 2 && __GLIBC_MINOR__ >= 13) && \
defined(__NR_prlimit64)
#if PSUTIL_HAVE_PRLIMIT
#define _FILE_OFFSET_BITS 64
#include <time.h>
#include <sys/resource.h>
#endif
// May happen on old RedHat versions, see:
// https://github.com/giampaolo/psutil/issues/607
#ifndef DUPLEX_UNKNOWN
#define DUPLEX_UNKNOWN 0xff
#endif
#if PSUTIL_HAVE_IOPRIO
enum {
IOPRIO_WHO_PROCESS = 1,
};
static inline int
ioprio_get(int which, int who) {
return syscall(__NR_ioprio_get, which, who);
}
static inline int
ioprio_set(int which, int who, int ioprio) {
return syscall(__NR_ioprio_set, which, who, ioprio);
}
#define IOPRIO_CLASS_SHIFT 13
#define IOPRIO_PRIO_MASK ((1UL << IOPRIO_CLASS_SHIFT) - 1)
#define IOPRIO_PRIO_CLASS(mask) ((mask) >> IOPRIO_CLASS_SHIFT)
#define IOPRIO_PRIO_DATA(mask) ((mask) & IOPRIO_PRIO_MASK)
#define IOPRIO_PRIO_VALUE(class, data) (((class) << IOPRIO_CLASS_SHIFT) | data)
/*
* Return a (ioclass, iodata) Python tuple representing process I/O priority.
*/
static PyObject *
psutil_proc_ioprio_get(PyObject *self, PyObject *args) {
long pid;
int ioprio, ioclass, iodata;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
ioprio = ioprio_get(IOPRIO_WHO_PROCESS, pid);
if (ioprio == -1)
return PyErr_SetFromErrno(PyExc_OSError);
ioclass = IOPRIO_PRIO_CLASS(ioprio);
iodata = IOPRIO_PRIO_DATA(ioprio);
return Py_BuildValue("ii", ioclass, iodata);
}
/*
* A wrapper around ioprio_set(); sets process I/O priority.
* ioclass can be either IOPRIO_CLASS_RT, IOPRIO_CLASS_BE, IOPRIO_CLASS_IDLE
* or 0. iodata goes from 0 to 7 depending on ioclass specified.
*/
static PyObject *
psutil_proc_ioprio_set(PyObject *self, PyObject *args) {
long pid;
int ioprio, ioclass, iodata;
int retval;
if (! PyArg_ParseTuple(args, "lii", &pid, &ioclass, &iodata))
return NULL;
ioprio = IOPRIO_PRIO_VALUE(ioclass, iodata);
retval = ioprio_set(IOPRIO_WHO_PROCESS, pid, ioprio);
if (retval == -1)
return PyErr_SetFromErrno(PyExc_OSError);
Py_RETURN_NONE;
}
#endif
#if PSUTIL_HAVE_PRLIMIT
/*
* A wrapper around prlimit(2); sets process resource limits.
* This can be used for both get and set, in which case extra
* 'soft' and 'hard' args must be provided.
*/
static PyObject *
psutil_linux_prlimit(PyObject *self, PyObject *args) {
long pid;
int ret, resource;
struct rlimit old, new;
struct rlimit *newp = NULL;
PyObject *py_soft = NULL;
PyObject *py_hard = NULL;
if (! PyArg_ParseTuple(args, "li|OO", &pid, &resource, &py_soft, &py_hard))
return NULL;
// get
if (py_soft == NULL && py_hard == NULL) {
ret = prlimit(pid, resource, NULL, &old);
if (ret == -1)
return PyErr_SetFromErrno(PyExc_OSError);
#if defined(PSUTIL_HAVE_LONG_LONG)
if (sizeof(old.rlim_cur) > sizeof(long)) {
return Py_BuildValue("LL",
(PY_LONG_LONG)old.rlim_cur,
(PY_LONG_LONG)old.rlim_max);
}
#endif
return Py_BuildValue("ll", (long)old.rlim_cur, (long)old.rlim_max);
}
// set
else {
#if defined(PSUTIL_HAVE_LARGEFILE_SUPPORT)
new.rlim_cur = PyLong_AsLongLong(py_soft);
if (new.rlim_cur == (rlim_t) - 1 && PyErr_Occurred())
return NULL;
new.rlim_max = PyLong_AsLongLong(py_hard);
if (new.rlim_max == (rlim_t) - 1 && PyErr_Occurred())
return NULL;
#else
new.rlim_cur = PyLong_AsLong(py_soft);
if (new.rlim_cur == (rlim_t) - 1 && PyErr_Occurred())
return NULL;
new.rlim_max = PyLong_AsLong(py_hard);
if (new.rlim_max == (rlim_t) - 1 && PyErr_Occurred())
return NULL;
#endif
newp = &new;
ret = prlimit(pid, resource, newp, &old);
if (ret == -1)
return PyErr_SetFromErrno(PyExc_OSError);
Py_RETURN_NONE;
}
}
#endif
/*
* Return disk mounted partitions as a list of tuples including device,
* mount point and filesystem type
*/
static PyObject *
psutil_disk_partitions(PyObject *self, PyObject *args) {
FILE *file = NULL;
struct mntent *entry;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
if (py_retlist == NULL)
return NULL;
// MOUNTED constant comes from mntent.h and it's == '/etc/mtab'
Py_BEGIN_ALLOW_THREADS
file = setmntent(MOUNTED, "r");
Py_END_ALLOW_THREADS
if ((file == 0) || (file == NULL)) {
PyErr_SetFromErrnoWithFilename(PyExc_OSError, MOUNTED);
goto error;
}
while ((entry = getmntent(file))) {
if (entry == NULL) {
PyErr_Format(PyExc_RuntimeError, "getmntent() syscall failed");
goto error;
}
py_tuple = Py_BuildValue("(ssss)",
entry->mnt_fsname, // device
entry->mnt_dir, // mount point
entry->mnt_type, // fs type
entry->mnt_opts); // options
if (! py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
endmntent(file);
return py_retlist;
error:
if (file != NULL)
endmntent(file);
Py_XDECREF(py_tuple);
Py_DECREF(py_retlist);
return NULL;
}
/*
* A wrapper around sysinfo(), return system memory usage statistics.
*/
static PyObject *
psutil_linux_sysinfo(PyObject *self, PyObject *args) {
struct sysinfo info;
if (sysinfo(&info) != 0)
return PyErr_SetFromErrno(PyExc_OSError);
// note: boot time might also be determined from here
return Py_BuildValue(
"(kkkkkkI)",
info.totalram, // total
info.freeram, // free
info.bufferram, // buffer
info.sharedram, // shared
info.totalswap, // swap tot
info.freeswap, // swap free
info.mem_unit // multiplier
);
}
/*
* Return process CPU affinity as a Python list
* The dual implementation exists because of:
* https://github.com/giampaolo/psutil/issues/536
*/
#ifdef CPU_ALLOC
static PyObject *
psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args) {
int cpu, ncpus, count, cpucount_s;
long pid;
size_t setsize;
cpu_set_t *mask = NULL;
PyObject *py_list = NULL;
if (!PyArg_ParseTuple(args, "l", &pid))
return NULL;
ncpus = NCPUS_START;
while (1) {
setsize = CPU_ALLOC_SIZE(ncpus);
mask = CPU_ALLOC(ncpus);
if (mask == NULL)
return PyErr_NoMemory();
if (sched_getaffinity(pid, setsize, mask) == 0)
break;
CPU_FREE(mask);
if (errno != EINVAL)
return PyErr_SetFromErrno(PyExc_OSError);
if (ncpus > INT_MAX / 2) {
PyErr_SetString(PyExc_OverflowError, "could not allocate "
"a large enough CPU set");
return NULL;
}
ncpus = ncpus * 2;
}
py_list = PyList_New(0);
if (py_list == NULL)
goto error;
cpucount_s = CPU_COUNT_S(setsize, mask);
for (cpu = 0, count = cpucount_s; count; cpu++) {
if (CPU_ISSET_S(cpu, setsize, mask)) {
#if PY_MAJOR_VERSION >= 3
PyObject *cpu_num = PyLong_FromLong(cpu);
#else
PyObject *cpu_num = PyInt_FromLong(cpu);
#endif
if (cpu_num == NULL)
goto error;
if (PyList_Append(py_list, cpu_num)) {
Py_DECREF(cpu_num);
goto error;
}
Py_DECREF(cpu_num);
--count;
}
}
CPU_FREE(mask);
return py_list;
error:
if (mask)
CPU_FREE(mask);
Py_XDECREF(py_list);
return NULL;
}
#else
/*
* Alternative implementation in case CPU_ALLOC is not defined.
*/
static PyObject *
psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args) {
cpu_set_t cpuset;
unsigned int len = sizeof(cpu_set_t);
long pid;
int i;
PyObject* py_retlist = NULL;
PyObject *py_cpu_num = NULL;
if (!PyArg_ParseTuple(args, "l", &pid))
return NULL;
CPU_ZERO(&cpuset);
if (sched_getaffinity(pid, len, &cpuset) < 0)
return PyErr_SetFromErrno(PyExc_OSError);
py_retlist = PyList_New(0);
if (py_retlist == NULL)
goto error;
for (i = 0; i < CPU_SETSIZE; ++i) {
if (CPU_ISSET(i, &cpuset)) {
py_cpu_num = Py_BuildValue("i", i);
if (py_cpu_num == NULL)
goto error;
if (PyList_Append(py_retlist, py_cpu_num))
goto error;
Py_DECREF(py_cpu_num);
}
}
return py_retlist;
error:
Py_XDECREF(py_cpu_num);
Py_XDECREF(py_retlist);
return NULL;
}
#endif
/*
* Set process CPU affinity; expects a bitmask
*/
static PyObject *
psutil_proc_cpu_affinity_set(PyObject *self, PyObject *args) {
cpu_set_t cpu_set;
size_t len;
long pid;
int i, seq_len;
PyObject *py_cpu_set;
PyObject *py_cpu_seq = NULL;
if (!PyArg_ParseTuple(args, "lO", &pid, &py_cpu_set))
return NULL;
if (!PySequence_Check(py_cpu_set)) {
PyErr_Format(PyExc_TypeError, "sequence argument expected, got %s",
Py_TYPE(py_cpu_set)->tp_name);
goto error;
}
py_cpu_seq = PySequence_Fast(py_cpu_set, "expected a sequence or integer");
if (!py_cpu_seq)
goto error;
seq_len = PySequence_Fast_GET_SIZE(py_cpu_seq);
CPU_ZERO(&cpu_set);
for (i = 0; i < seq_len; i++) {
PyObject *item = PySequence_Fast_GET_ITEM(py_cpu_seq, i);
#if PY_MAJOR_VERSION >= 3
long value = PyLong_AsLong(item);
#else
long value = PyInt_AsLong(item);
#endif
if ((value == -1) || PyErr_Occurred()) {
if (!PyErr_Occurred())
PyErr_SetString(PyExc_ValueError, "invalid CPU value");
goto error;
}
CPU_SET(value, &cpu_set);
}
len = sizeof(cpu_set);
if (sched_setaffinity(pid, len, &cpu_set)) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
Py_DECREF(py_cpu_seq);
Py_RETURN_NONE;
error:
if (py_cpu_seq != NULL)
Py_DECREF(py_cpu_seq);
return NULL;
}
/*
* Return currently connected users as a list of tuples.
*/
static PyObject *
psutil_users(PyObject *self, PyObject *args) {
struct utmp *ut;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
PyObject *py_user_proc = NULL;
if (py_retlist == NULL)
return NULL;
setutent();
while (NULL != (ut = getutent())) {
py_tuple = NULL;
py_user_proc = NULL;
if (ut->ut_type == USER_PROCESS)
py_user_proc = Py_True;
else
py_user_proc = Py_False;
py_tuple = Py_BuildValue(
"(sssfO)",
ut->ut_user, // username
ut->ut_line, // tty
ut->ut_host, // hostname
(float)ut->ut_tv.tv_sec, // tstamp
py_user_proc // (bool) user process
);
if (! py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
endutent();
return py_retlist;
error:
Py_XDECREF(py_tuple);
Py_XDECREF(py_user_proc);
Py_DECREF(py_retlist);
endutent();
return NULL;
}
/*
* Return stats about a particular network
* interface. References:
* https://github.com/dpaleino/wicd/blob/master/wicd/backends/be-ioctl.py
* http://www.i-scream.org/libstatgrab/
*/
static PyObject*
psutil_net_if_duplex_speed(PyObject* self, PyObject* args) {
char *nic_name;
int sock = 0;
int ret;
int duplex;
int speed;
struct ifreq ifr;
struct ethtool_cmd ethcmd;
PyObject *py_retlist = NULL;
if (! PyArg_ParseTuple(args, "s", &nic_name))
return NULL;
sock = socket(AF_INET, SOCK_DGRAM, 0);
if (sock == -1)
goto error;
strncpy(ifr.ifr_name, nic_name, sizeof(ifr.ifr_name));
// duplex and speed
memset(&ethcmd, 0, sizeof ethcmd);
ethcmd.cmd = ETHTOOL_GSET;
ifr.ifr_data = (void *)&ethcmd;
ret = ioctl(sock, SIOCETHTOOL, &ifr);
if (ret != -1) {
duplex = ethcmd.duplex;
speed = ethcmd.speed;
}
else {
if ((errno == EOPNOTSUPP) || (errno == EINVAL)) {
// EOPNOTSUPP may occur in case of wi-fi cards.
// For EINVAL see:
// https://github.com/giampaolo/psutil/issues/797
// #issuecomment-202999532
duplex = DUPLEX_UNKNOWN;
speed = 0;
}
else {
goto error;
}
}
close(sock);
py_retlist = Py_BuildValue("[ii]", duplex, speed);
if (!py_retlist)
goto error;
return py_retlist;
error:
if (sock != -1)
close(sock);
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
/*
* Define the psutil C module methods and initialize the module.
*/
static PyMethodDef
PsutilMethods[] = {
// --- per-process functions
#if PSUTIL_HAVE_IOPRIO
{"proc_ioprio_get", psutil_proc_ioprio_get, METH_VARARGS,
"Get process I/O priority"},
{"proc_ioprio_set", psutil_proc_ioprio_set, METH_VARARGS,
"Set process I/O priority"},
#endif
{"proc_cpu_affinity_get", psutil_proc_cpu_affinity_get, METH_VARARGS,
"Return process CPU affinity as a Python long (the bitmask)."},
{"proc_cpu_affinity_set", psutil_proc_cpu_affinity_set, METH_VARARGS,
"Set process CPU affinity; expects a bitmask."},
// --- system related functions
{"disk_partitions", psutil_disk_partitions, METH_VARARGS,
"Return disk mounted partitions as a list of tuples including "
"device, mount point and filesystem type"},
{"users", psutil_users, METH_VARARGS,
"Return currently connected users as a list of tuples"},
{"net_if_duplex_speed", psutil_net_if_duplex_speed, METH_VARARGS,
"Return duplex and speed info about a NIC"},
// --- linux specific
{"linux_sysinfo", psutil_linux_sysinfo, METH_VARARGS,
"A wrapper around sysinfo(), return system memory usage statistics"},
#if PSUTIL_HAVE_PRLIMIT
{"linux_prlimit", psutil_linux_prlimit, METH_VARARGS,
"Get or set process resource limits."},
#endif
{NULL, NULL, 0, NULL}
};
struct module_state {
PyObject *error;
};
#if PY_MAJOR_VERSION >= 3
#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))
#else
#define GETSTATE(m) (&_state)
#endif
#if PY_MAJOR_VERSION >= 3
static int
psutil_linux_traverse(PyObject *m, visitproc visit, void *arg) {
Py_VISIT(GETSTATE(m)->error);
return 0;
}
static int
psutil_linux_clear(PyObject *m) {
Py_CLEAR(GETSTATE(m)->error);
return 0;
}
static struct PyModuleDef
moduledef = {
PyModuleDef_HEAD_INIT,
"psutil_linux",
NULL,
sizeof(struct module_state),
PsutilMethods,
NULL,
psutil_linux_traverse,
psutil_linux_clear,
NULL
};
#define INITERROR return NULL
PyMODINIT_FUNC PyInit__psutil_linux(void)
#else
#define INITERROR return
void init_psutil_linux(void)
#endif
{
PyObject *v;
#if PY_MAJOR_VERSION >= 3
PyObject *module = PyModule_Create(&moduledef);
#else
PyObject *module = Py_InitModule("_psutil_linux", PsutilMethods);
#endif
PyModule_AddIntConstant(module, "version", PSUTIL_VERSION);
#if PSUTIL_HAVE_PRLIMIT
PyModule_AddIntConstant(module, "RLIMIT_AS", RLIMIT_AS);
PyModule_AddIntConstant(module, "RLIMIT_CORE", RLIMIT_CORE);
PyModule_AddIntConstant(module, "RLIMIT_CPU", RLIMIT_CPU);
PyModule_AddIntConstant(module, "RLIMIT_DATA", RLIMIT_DATA);
PyModule_AddIntConstant(module, "RLIMIT_FSIZE", RLIMIT_FSIZE);
PyModule_AddIntConstant(module, "RLIMIT_LOCKS", RLIMIT_LOCKS);
PyModule_AddIntConstant(module, "RLIMIT_MEMLOCK", RLIMIT_MEMLOCK);
PyModule_AddIntConstant(module, "RLIMIT_NOFILE", RLIMIT_NOFILE);
PyModule_AddIntConstant(module, "RLIMIT_NPROC", RLIMIT_NPROC);
PyModule_AddIntConstant(module, "RLIMIT_RSS", RLIMIT_RSS);
PyModule_AddIntConstant(module, "RLIMIT_STACK", RLIMIT_STACK);
#if defined(HAVE_LONG_LONG)
if (sizeof(RLIM_INFINITY) > sizeof(long)) {
v = PyLong_FromLongLong((PY_LONG_LONG) RLIM_INFINITY);
} else
#endif
{
v = PyLong_FromLong((long) RLIM_INFINITY);
}
if (v) {
PyModule_AddObject(module, "RLIM_INFINITY", v);
}
#ifdef RLIMIT_MSGQUEUE
PyModule_AddIntConstant(module, "RLIMIT_MSGQUEUE", RLIMIT_MSGQUEUE);
#endif
#ifdef RLIMIT_NICE
PyModule_AddIntConstant(module, "RLIMIT_NICE", RLIMIT_NICE);
#endif
#ifdef RLIMIT_RTPRIO
PyModule_AddIntConstant(module, "RLIMIT_RTPRIO", RLIMIT_RTPRIO);
#endif
#ifdef RLIMIT_RTTIME
PyModule_AddIntConstant(module, "RLIMIT_RTTIME", RLIMIT_RTTIME);
#endif
#ifdef RLIMIT_SIGPENDING
PyModule_AddIntConstant(module, "RLIMIT_SIGPENDING", RLIMIT_SIGPENDING);
#endif
#endif
PyModule_AddIntConstant(module, "DUPLEX_HALF", DUPLEX_HALF);
PyModule_AddIntConstant(module, "DUPLEX_FULL", DUPLEX_FULL);
PyModule_AddIntConstant(module, "DUPLEX_UNKNOWN", DUPLEX_UNKNOWN);
if (module == NULL)
INITERROR;
#if PY_MAJOR_VERSION >= 3
return module;
#endif
}
+1910
View File
File diff suppressed because it is too large Load Diff
+619
View File
@@ -0,0 +1,619 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Functions specific to all POSIX compliant platforms.
*/
#include <Python.h>
#include <errno.h>
#include <stdlib.h>
#include <sys/resource.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <net/if.h>
#ifdef PSUTIL_SUNOS10
#include "arch/solaris/v10/ifaddrs.h"
#else
#include <ifaddrs.h>
#endif
#if defined(PSUTIL_LINUX)
#include <netdb.h>
#include <linux/if_packet.h>
#elif defined(PSUTIL_BSD) || defined(PSUTIL_OSX)
#include <netdb.h>
#include <netinet/in.h>
#include <net/if_dl.h>
#include <sys/sockio.h>
#include <net/if_media.h>
#include <net/if.h>
#elif defined(PSUTIL_SUNOS)
#include <netdb.h>
#include <sys/sockio.h>
#endif
/*
* Given a PID return process priority as a Python integer.
*/
static PyObject *
psutil_posix_getpriority(PyObject *self, PyObject *args) {
long pid;
int priority;
errno = 0;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
#ifdef PSUTIL_OSX
priority = getpriority(PRIO_PROCESS, (id_t)pid);
#else
priority = getpriority(PRIO_PROCESS, pid);
#endif
if (errno != 0)
return PyErr_SetFromErrno(PyExc_OSError);
return Py_BuildValue("i", priority);
}
/*
* Given a PID and a value change process priority.
*/
static PyObject *
psutil_posix_setpriority(PyObject *self, PyObject *args) {
long pid;
int priority;
int retval;
if (! PyArg_ParseTuple(args, "li", &pid, &priority))
return NULL;
#ifdef PSUTIL_OSX
retval = setpriority(PRIO_PROCESS, (id_t)pid, priority);
#else
retval = setpriority(PRIO_PROCESS, pid, priority);
#endif
if (retval == -1)
return PyErr_SetFromErrno(PyExc_OSError);
Py_RETURN_NONE;
}
/*
* Translate a sockaddr struct into a Python string.
* Return None if address family is not AF_INET* or AF_PACKET.
*/
static PyObject *
psutil_convert_ipaddr(struct sockaddr *addr, int family) {
char buf[NI_MAXHOST];
int err;
int addrlen;
size_t n;
size_t len;
const char *data;
char *ptr;
if (addr == NULL) {
Py_INCREF(Py_None);
return Py_None;
}
else if (family == AF_INET || family == AF_INET6) {
if (family == AF_INET)
addrlen = sizeof(struct sockaddr_in);
else
addrlen = sizeof(struct sockaddr_in6);
err = getnameinfo(addr, addrlen, buf, sizeof(buf), NULL, 0,
NI_NUMERICHOST);
if (err != 0) {
// XXX we get here on FreeBSD when processing 'lo' / AF_INET6
// broadcast. Not sure what to do other than returning None.
// ifconfig does not show anything BTW.
//PyErr_Format(PyExc_RuntimeError, gai_strerror(err));
//return NULL;
Py_INCREF(Py_None);
return Py_None;
}
else {
return Py_BuildValue("s", buf);
}
}
#ifdef PSUTIL_LINUX
else if (family == AF_PACKET) {
struct sockaddr_ll *lladdr = (struct sockaddr_ll *)addr;
len = lladdr->sll_halen;
data = (const char *)lladdr->sll_addr;
}
#elif defined(PSUTIL_BSD) || defined(PSUTIL_OSX)
else if (addr->sa_family == AF_LINK) {
// Note: prior to Python 3.4 socket module does not expose
// AF_LINK so we'll do.
struct sockaddr_dl *dladdr = (struct sockaddr_dl *)addr;
len = dladdr->sdl_alen;
data = LLADDR(dladdr);
}
#endif
else {
// unknown family
Py_INCREF(Py_None);
return Py_None;
}
// AF_PACKET or AF_LINK
if (len > 0) {
ptr = buf;
for (n = 0; n < len; ++n) {
sprintf(ptr, "%02x:", data[n] & 0xff);
ptr += 3;
}
*--ptr = '\0';
return Py_BuildValue("s", buf);
}
else {
Py_INCREF(Py_None);
return Py_None;
}
}
/*
* Return NICs information a-la ifconfig as a list of tuples.
* TODO: on Solaris we won't get any MAC address.
*/
static PyObject*
psutil_net_if_addrs(PyObject* self, PyObject* args) {
struct ifaddrs *ifaddr, *ifa;
int family;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
PyObject *py_address = NULL;
PyObject *py_netmask = NULL;
PyObject *py_broadcast = NULL;
PyObject *py_ptp = NULL;
if (py_retlist == NULL)
return NULL;
if (getifaddrs(&ifaddr) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
for (ifa = ifaddr; ifa != NULL; ifa = ifa->ifa_next) {
if (!ifa->ifa_addr)
continue;
family = ifa->ifa_addr->sa_family;
py_address = psutil_convert_ipaddr(ifa->ifa_addr, family);
// If the primary address can't be determined just skip it.
// I've never seen this happen on Linux but I did on FreeBSD.
if (py_address == Py_None)
continue;
if (py_address == NULL)
goto error;
py_netmask = psutil_convert_ipaddr(ifa->ifa_netmask, family);
if (py_netmask == NULL)
goto error;
if (ifa->ifa_flags & IFF_BROADCAST) {
py_broadcast = psutil_convert_ipaddr(ifa->ifa_broadaddr, family);
Py_INCREF(Py_None);
py_ptp = Py_None;
}
else if (ifa->ifa_flags & IFF_POINTOPOINT) {
py_ptp = psutil_convert_ipaddr(ifa->ifa_dstaddr, family);
Py_INCREF(Py_None);
py_broadcast = Py_None;
}
else {
Py_INCREF(Py_None);
Py_INCREF(Py_None);
py_broadcast = Py_None;
py_ptp = Py_None;
}
if ((py_broadcast == NULL) || (py_ptp == NULL))
goto error;
py_tuple = Py_BuildValue(
"(siOOOO)",
ifa->ifa_name,
family,
py_address,
py_netmask,
py_broadcast,
py_ptp
);
if (! py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
Py_DECREF(py_address);
Py_DECREF(py_netmask);
Py_DECREF(py_broadcast);
Py_DECREF(py_ptp);
}
freeifaddrs(ifaddr);
return py_retlist;
error:
if (ifaddr != NULL)
freeifaddrs(ifaddr);
Py_DECREF(py_retlist);
Py_XDECREF(py_tuple);
Py_XDECREF(py_address);
Py_XDECREF(py_netmask);
Py_XDECREF(py_broadcast);
Py_XDECREF(py_ptp);
return NULL;
}
/*
* Return NIC MTU. References:
* http://www.i-scream.org/libstatgrab/
*/
static PyObject *
psutil_net_if_mtu(PyObject *self, PyObject *args) {
char *nic_name;
int sock = 0;
int ret;
#ifdef PSUTIL_SUNOS10
struct lifreq lifr;
#else
struct ifreq ifr;
#endif
if (! PyArg_ParseTuple(args, "s", &nic_name))
return NULL;
sock = socket(AF_INET, SOCK_DGRAM, 0);
if (sock == -1)
goto error;
#ifdef PSUTIL_SUNOS10
strncpy(lifr.lifr_name, nic_name, sizeof(lifr.lifr_name));
ret = ioctl(sock, SIOCGIFMTU, &lifr);
#else
strncpy(ifr.ifr_name, nic_name, sizeof(ifr.ifr_name));
ret = ioctl(sock, SIOCGIFMTU, &ifr);
#endif
if (ret == -1)
goto error;
close(sock);
#ifdef PSUTIL_SUNOS10
return Py_BuildValue("i", lifr.lifr_mtu);
#else
return Py_BuildValue("i", ifr.ifr_mtu);
#endif
error:
if (sock != 0)
close(sock);
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
/*
* Inspect NIC flags, returns a bool indicating whether the NIC is
* running. References:
* http://www.i-scream.org/libstatgrab/
*/
static PyObject *
psutil_net_if_flags(PyObject *self, PyObject *args) {
char *nic_name;
int sock = 0;
int ret;
struct ifreq ifr;
if (! PyArg_ParseTuple(args, "s", &nic_name))
return NULL;
sock = socket(AF_INET, SOCK_DGRAM, 0);
if (sock == -1)
goto error;
strncpy(ifr.ifr_name, nic_name, sizeof(ifr.ifr_name));
ret = ioctl(sock, SIOCGIFFLAGS, &ifr);
if (ret == -1)
goto error;
close(sock);
if ((ifr.ifr_flags & IFF_UP) != 0)
return Py_BuildValue("O", Py_True);
else
return Py_BuildValue("O", Py_False);
error:
if (sock != 0)
close(sock);
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
/*
* net_if_stats() OSX/BSD implementation.
*/
#if defined(PSUTIL_BSD) || defined(PSUTIL_OSX)
int psutil_get_nic_speed(int ifm_active) {
// Determine NIC speed. Taken from:
// http://www.i-scream.org/libstatgrab/
// Assuming only ETHER devices
switch(IFM_TYPE(ifm_active)) {
case IFM_ETHER:
switch(IFM_SUBTYPE(ifm_active)) {
#if defined(IFM_HPNA_1) && ((!defined(IFM_10G_LR)) \
|| (IFM_10G_LR != IFM_HPNA_1))
// HomePNA 1.0 (1Mb/s)
case(IFM_HPNA_1):
return 1;
#endif
// 10 Mbit
case(IFM_10_T): // 10BaseT - RJ45
case(IFM_10_2): // 10Base2 - Thinnet
case(IFM_10_5): // 10Base5 - AUI
case(IFM_10_STP): // 10BaseT over shielded TP
case(IFM_10_FL): // 10baseFL - Fiber
return 10;
// 100 Mbit
case(IFM_100_TX): // 100BaseTX - RJ45
case(IFM_100_FX): // 100BaseFX - Fiber
case(IFM_100_T4): // 100BaseT4 - 4 pair cat 3
case(IFM_100_VG): // 100VG-AnyLAN
case(IFM_100_T2): // 100BaseT2
return 100;
// 1000 Mbit
case(IFM_1000_SX): // 1000BaseSX - multi-mode fiber
case(IFM_1000_LX): // 1000baseLX - single-mode fiber
case(IFM_1000_CX): // 1000baseCX - 150ohm STP
#if defined(IFM_1000_TX) && !defined(PSUTIL_OPENBSD)
// FreeBSD 4 and others (but NOT OpenBSD) -> #define IFM_1000_T in net/if_media.h
case(IFM_1000_TX):
#endif
#ifdef IFM_1000_FX
case(IFM_1000_FX):
#endif
#ifdef IFM_1000_T
case(IFM_1000_T):
#endif
return 1000;
#if defined(IFM_10G_SR) || defined(IFM_10G_LR) || defined(IFM_10G_CX4) \
|| defined(IFM_10G_T)
#ifdef IFM_10G_SR
case(IFM_10G_SR):
#endif
#ifdef IFM_10G_LR
case(IFM_10G_LR):
#endif
#ifdef IFM_10G_CX4
case(IFM_10G_CX4):
#endif
#ifdef IFM_10G_TWINAX
case(IFM_10G_TWINAX):
#endif
#ifdef IFM_10G_TWINAX_LONG
case(IFM_10G_TWINAX_LONG):
#endif
#ifdef IFM_10G_T
case(IFM_10G_T):
#endif
return 10000;
#endif
#if defined(IFM_2500_SX)
#ifdef IFM_2500_SX
case(IFM_2500_SX):
#endif
return 2500;
#endif // any 2.5GBit stuff...
// We don't know what it is
default:
return 0;
}
break;
#ifdef IFM_TOKEN
case IFM_TOKEN:
switch(IFM_SUBTYPE(ifm_active)) {
case IFM_TOK_STP4: // Shielded twisted pair 4m - DB9
case IFM_TOK_UTP4: // Unshielded twisted pair 4m - RJ45
return 4;
case IFM_TOK_STP16: // Shielded twisted pair 16m - DB9
case IFM_TOK_UTP16: // Unshielded twisted pair 16m - RJ45
return 16;
#if defined(IFM_TOK_STP100) || defined(IFM_TOK_UTP100)
#ifdef IFM_TOK_STP100
case IFM_TOK_STP100: // Shielded twisted pair 100m - DB9
#endif
#ifdef IFM_TOK_UTP100
case IFM_TOK_UTP100: // Unshielded twisted pair 100m - RJ45
#endif
return 100;
#endif
// We don't know what it is
default:
return 0;
}
break;
#endif
#ifdef IFM_FDDI
case IFM_FDDI:
switch(IFM_SUBTYPE(ifm_active)) {
// We don't know what it is
default:
return 0;
}
break;
#endif
case IFM_IEEE80211:
switch(IFM_SUBTYPE(ifm_active)) {
case IFM_IEEE80211_FH1: // Frequency Hopping 1Mbps
case IFM_IEEE80211_DS1: // Direct Sequence 1Mbps
return 1;
case IFM_IEEE80211_FH2: // Frequency Hopping 2Mbps
case IFM_IEEE80211_DS2: // Direct Sequence 2Mbps
return 2;
case IFM_IEEE80211_DS5: // Direct Sequence 5Mbps
return 5;
case IFM_IEEE80211_DS11: // Direct Sequence 11Mbps
return 11;
case IFM_IEEE80211_DS22: // Direct Sequence 22Mbps
return 22;
// We don't know what it is
default:
return 0;
}
break;
default:
return 0;
}
}
/*
* Return stats about a particular network interface.
* References:
* http://www.i-scream.org/libstatgrab/
*/
static PyObject *
psutil_net_if_duplex_speed(PyObject *self, PyObject *args) {
char *nic_name;
int sock = 0;
int ret;
int duplex;
int speed;
struct ifreq ifr;
struct ifmediareq ifmed;
if (! PyArg_ParseTuple(args, "s", &nic_name))
return NULL;
sock = socket(AF_INET, SOCK_DGRAM, 0);
if (sock == -1)
goto error;
strncpy(ifr.ifr_name, nic_name, sizeof(ifr.ifr_name));
// speed / duplex
memset(&ifmed, 0, sizeof(struct ifmediareq));
strlcpy(ifmed.ifm_name, nic_name, sizeof(ifmed.ifm_name));
ret = ioctl(sock, SIOCGIFMEDIA, (caddr_t)&ifmed);
if (ret == -1) {
speed = 0;
duplex = 0;
}
else {
speed = psutil_get_nic_speed(ifmed.ifm_active);
if ((ifmed.ifm_active | IFM_FDX) == ifmed.ifm_active)
duplex = 2;
else if ((ifmed.ifm_active | IFM_HDX) == ifmed.ifm_active)
duplex = 1;
else
duplex = 0;
}
close(sock);
return Py_BuildValue("[ii]", duplex, speed);
error:
if (sock != 0)
close(sock);
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
#endif // net_if_stats() OSX/BSD implementation
/*
* define the psutil C module methods and initialize the module.
*/
static PyMethodDef
PsutilMethods[] = {
{"getpriority", psutil_posix_getpriority, METH_VARARGS,
"Return process priority"},
{"setpriority", psutil_posix_setpriority, METH_VARARGS,
"Set process priority"},
{"net_if_addrs", psutil_net_if_addrs, METH_VARARGS,
"Retrieve NICs information"},
{"net_if_mtu", psutil_net_if_mtu, METH_VARARGS,
"Retrieve NIC MTU"},
{"net_if_flags", psutil_net_if_flags, METH_VARARGS,
"Retrieve NIC flags"},
#if defined(PSUTIL_BSD) || defined(PSUTIL_OSX)
{"net_if_duplex_speed", psutil_net_if_duplex_speed, METH_VARARGS,
"Return NIC stats."},
#endif
{NULL, NULL, 0, NULL}
};
struct module_state {
PyObject *error;
};
#if PY_MAJOR_VERSION >= 3
#define GETSTATE(m) ((struct module_state*)PyModule_GetState(m))
#else
#define GETSTATE(m) (&_state)
#endif
#if PY_MAJOR_VERSION >= 3
static int
psutil_posix_traverse(PyObject *m, visitproc visit, void *arg) {
Py_VISIT(GETSTATE(m)->error);
return 0;
}
static int
psutil_posix_clear(PyObject *m) {
Py_CLEAR(GETSTATE(m)->error);
return 0;
}
static struct PyModuleDef moduledef = {
PyModuleDef_HEAD_INIT,
"psutil_posix",
NULL,
sizeof(struct module_state),
PsutilMethods,
NULL,
psutil_posix_traverse,
psutil_posix_clear,
NULL
};
#define INITERROR return NULL
PyMODINIT_FUNC PyInit__psutil_posix(void)
#else
#define INITERROR return
void init_psutil_posix(void)
#endif
{
#if PY_MAJOR_VERSION >= 3
PyObject *module = PyModule_Create(&moduledef);
#else
PyObject *module = Py_InitModule("_psutil_posix", PsutilMethods);
#endif
#if defined(PSUTIL_BSD) || defined(PSUTIL_OSX) || defined(PSUTIL_SUNOS)
PyModule_AddIntConstant(module, "AF_LINK", AF_LINK);
#endif
if (module == NULL)
INITERROR;
#if PY_MAJOR_VERSION >= 3
return module;
#endif
}
+5
View File
@@ -0,0 +1,5 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+965
View File
@@ -0,0 +1,965 @@
# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Windows platform implementation."""
import contextlib
import errno
import functools
import os
import sys
from collections import namedtuple
from . import _common
try:
from . import _psutil_windows as cext
except ImportError as err:
if str(err).lower().startswith("dll load failed") and \
sys.getwindowsversion()[0] < 6:
# We may get here if:
# 1) we are on an old Windows version
# 2) psutil was installed via pip + wheel
# See: https://github.com/giampaolo/psutil/issues/811
# It must be noted that psutil can still (kind of) work
# on outdated systems if compiled / installed from sources,
# but if we get here it means this this was a wheel (or exe).
msg = "this Windows version is too old (< Windows Vista); "
msg += "psutil 3.4.2 is the latest version which supports Windows "
msg += "2000, XP and 2003 server"
raise RuntimeError(msg)
else:
raise
from ._common import conn_tmap
from ._common import isfile_strict
from ._common import parse_environ_block
from ._common import sockfam_to_enum
from ._common import socktype_to_enum
from ._common import usage_percent
from ._common import memoize_when_activated
from ._compat import long
from ._compat import lru_cache
from ._compat import PY3
from ._compat import xrange
from ._psutil_windows import ABOVE_NORMAL_PRIORITY_CLASS
from ._psutil_windows import BELOW_NORMAL_PRIORITY_CLASS
from ._psutil_windows import HIGH_PRIORITY_CLASS
from ._psutil_windows import IDLE_PRIORITY_CLASS
from ._psutil_windows import NORMAL_PRIORITY_CLASS
from ._psutil_windows import REALTIME_PRIORITY_CLASS
if sys.version_info >= (3, 4):
import enum
else:
enum = None
# process priority constants, import from __init__.py:
# http://msdn.microsoft.com/en-us/library/ms686219(v=vs.85).aspx
__extra__all__ = [
"win_service_iter", "win_service_get",
"ABOVE_NORMAL_PRIORITY_CLASS", "BELOW_NORMAL_PRIORITY_CLASS",
"HIGH_PRIORITY_CLASS", "IDLE_PRIORITY_CLASS",
"NORMAL_PRIORITY_CLASS", "REALTIME_PRIORITY_CLASS",
"CONN_DELETE_TCB",
"AF_LINK",
]
# =====================================================================
# --- globals
# =====================================================================
CONN_DELETE_TCB = "DELETE_TCB"
WAIT_TIMEOUT = 0x00000102 # 258 in decimal
ACCESS_DENIED_SET = frozenset([errno.EPERM, errno.EACCES,
cext.ERROR_ACCESS_DENIED])
if enum is None:
AF_LINK = -1
else:
AddressFamily = enum.IntEnum('AddressFamily', {'AF_LINK': -1})
AF_LINK = AddressFamily.AF_LINK
TCP_STATUSES = {
cext.MIB_TCP_STATE_ESTAB: _common.CONN_ESTABLISHED,
cext.MIB_TCP_STATE_SYN_SENT: _common.CONN_SYN_SENT,
cext.MIB_TCP_STATE_SYN_RCVD: _common.CONN_SYN_RECV,
cext.MIB_TCP_STATE_FIN_WAIT1: _common.CONN_FIN_WAIT1,
cext.MIB_TCP_STATE_FIN_WAIT2: _common.CONN_FIN_WAIT2,
cext.MIB_TCP_STATE_TIME_WAIT: _common.CONN_TIME_WAIT,
cext.MIB_TCP_STATE_CLOSED: _common.CONN_CLOSE,
cext.MIB_TCP_STATE_CLOSE_WAIT: _common.CONN_CLOSE_WAIT,
cext.MIB_TCP_STATE_LAST_ACK: _common.CONN_LAST_ACK,
cext.MIB_TCP_STATE_LISTEN: _common.CONN_LISTEN,
cext.MIB_TCP_STATE_CLOSING: _common.CONN_CLOSING,
cext.MIB_TCP_STATE_DELETE_TCB: CONN_DELETE_TCB,
cext.PSUTIL_CONN_NONE: _common.CONN_NONE,
}
if enum is not None:
class Priority(enum.IntEnum):
ABOVE_NORMAL_PRIORITY_CLASS = ABOVE_NORMAL_PRIORITY_CLASS
BELOW_NORMAL_PRIORITY_CLASS = BELOW_NORMAL_PRIORITY_CLASS
HIGH_PRIORITY_CLASS = HIGH_PRIORITY_CLASS
IDLE_PRIORITY_CLASS = IDLE_PRIORITY_CLASS
NORMAL_PRIORITY_CLASS = NORMAL_PRIORITY_CLASS
REALTIME_PRIORITY_CLASS = REALTIME_PRIORITY_CLASS
globals().update(Priority.__members__)
pinfo_map = dict(
num_handles=0,
ctx_switches=1,
user_time=2,
kernel_time=3,
create_time=4,
num_threads=5,
io_rcount=6,
io_wcount=7,
io_rbytes=8,
io_wbytes=9,
io_count_others=10,
io_bytes_others=11,
num_page_faults=12,
peak_wset=13,
wset=14,
peak_paged_pool=15,
paged_pool=16,
peak_non_paged_pool=17,
non_paged_pool=18,
pagefile=19,
peak_pagefile=20,
mem_private=21,
)
# these get overwritten on "import psutil" from the __init__.py file
NoSuchProcess = None
AccessDenied = None
TimeoutExpired = None
# =====================================================================
# --- named tuples
# =====================================================================
# psutil.cpu_times()
scputimes = namedtuple('scputimes',
['user', 'system', 'idle', 'interrupt', 'dpc'])
# psutil.virtual_memory()
svmem = namedtuple('svmem', ['total', 'available', 'percent', 'used', 'free'])
# psutil.Process.memory_info()
pmem = namedtuple(
'pmem', ['rss', 'vms',
'num_page_faults', 'peak_wset', 'wset', 'peak_paged_pool',
'paged_pool', 'peak_nonpaged_pool', 'nonpaged_pool',
'pagefile', 'peak_pagefile', 'private'])
# psutil.Process.memory_full_info()
pfullmem = namedtuple('pfullmem', pmem._fields + ('uss', ))
# psutil.Process.memory_maps(grouped=True)
pmmap_grouped = namedtuple('pmmap_grouped', ['path', 'rss'])
# psutil.Process.memory_maps(grouped=False)
pmmap_ext = namedtuple(
'pmmap_ext', 'addr perms ' + ' '.join(pmmap_grouped._fields))
# psutil.Process.io_counters()
pio = namedtuple('pio', ['read_count', 'write_count',
'read_bytes', 'write_bytes',
'other_count', 'other_bytes'])
# =====================================================================
# --- utils
# =====================================================================
@lru_cache(maxsize=512)
def convert_dos_path(s):
"""Convert paths using native DOS format like:
"\Device\HarddiskVolume1\Windows\systemew\file.txt"
into:
"C:\Windows\systemew\file.txt"
"""
if PY3 and not isinstance(s, str):
s = s.decode('utf8')
rawdrive = '\\'.join(s.split('\\')[:3])
driveletter = cext.win32_QueryDosDevice(rawdrive)
return os.path.join(driveletter, s[len(rawdrive):])
def py2_strencode(s, encoding=sys.getfilesystemencoding()):
"""Encode a string in the given encoding. Falls back on returning
the string as is if it can't be encoded.
"""
if PY3 or isinstance(s, str):
return s
else:
try:
return s.encode(encoding)
except UnicodeEncodeError:
# Filesystem codec failed, return the plain unicode
# string (this should never happen).
return s
# =====================================================================
# --- memory
# =====================================================================
def virtual_memory():
"""System virtual memory as a namedtuple."""
mem = cext.virtual_mem()
totphys, availphys, totpagef, availpagef, totvirt, freevirt = mem
#
total = totphys
avail = availphys
free = availphys
used = total - avail
percent = usage_percent((total - avail), total, _round=1)
return svmem(total, avail, percent, used, free)
def swap_memory():
"""Swap system memory as a (total, used, free, sin, sout) tuple."""
mem = cext.virtual_mem()
total = mem[2]
free = mem[3]
used = total - free
percent = usage_percent(used, total, _round=1)
return _common.sswap(total, used, free, percent, 0, 0)
# =====================================================================
# --- disk
# =====================================================================
disk_io_counters = cext.disk_io_counters
def disk_usage(path):
"""Return disk usage associated with path."""
try:
total, free = cext.disk_usage(path)
except WindowsError:
if not os.path.exists(path):
msg = "No such file or directory: '%s'" % path
raise OSError(errno.ENOENT, msg)
raise
used = total - free
percent = usage_percent(used, total, _round=1)
return _common.sdiskusage(total, used, free, percent)
def disk_partitions(all):
"""Return disk partitions."""
rawlist = cext.disk_partitions(all)
return [_common.sdiskpart(*x) for x in rawlist]
# =====================================================================
# --- CPU
# =====================================================================
def cpu_times():
"""Return system CPU times as a named tuple."""
user, system, idle = cext.cpu_times()
# Internally, GetSystemTimes() is used, and it doesn't return
# interrupt and dpc times. cext.per_cpu_times() does, so we
# rely on it to get those only.
percpu_summed = scputimes(*[sum(n) for n in zip(*cext.per_cpu_times())])
return scputimes(user, system, idle,
percpu_summed.interrupt, percpu_summed.dpc)
def per_cpu_times():
"""Return system per-CPU times as a list of named tuples."""
ret = []
for user, system, idle, interrupt, dpc in cext.per_cpu_times():
item = scputimes(user, system, idle, interrupt, dpc)
ret.append(item)
return ret
def cpu_count_logical():
"""Return the number of logical CPUs in the system."""
return cext.cpu_count_logical()
def cpu_count_physical():
"""Return the number of physical CPUs in the system."""
return cext.cpu_count_phys()
def cpu_stats():
"""Return CPU statistics."""
ctx_switches, interrupts, dpcs, syscalls = cext.cpu_stats()
soft_interrupts = 0
return _common.scpustats(ctx_switches, interrupts, soft_interrupts,
syscalls)
def cpu_freq():
"""Return CPU frequency.
On Windows per-cpu frequency is not supported.
"""
curr, max_ = cext.cpu_freq()
min_ = 0.0
return [_common.scpufreq(float(curr), min_, float(max_))]
# =====================================================================
# --- network
# =====================================================================
def net_connections(kind, _pid=-1):
"""Return socket connections. If pid == -1 return system-wide
connections (as opposed to connections opened by one process only).
"""
if kind not in conn_tmap:
raise ValueError("invalid %r kind argument; choose between %s"
% (kind, ', '.join([repr(x) for x in conn_tmap])))
families, types = conn_tmap[kind]
rawlist = cext.net_connections(_pid, families, types)
ret = set()
for item in rawlist:
fd, fam, type, laddr, raddr, status, pid = item
status = TCP_STATUSES[status]
fam = sockfam_to_enum(fam)
type = socktype_to_enum(type)
if _pid == -1:
nt = _common.sconn(fd, fam, type, laddr, raddr, status, pid)
else:
nt = _common.pconn(fd, fam, type, laddr, raddr, status)
ret.add(nt)
return list(ret)
def net_if_stats():
"""Get NIC stats (isup, duplex, speed, mtu)."""
ret = cext.net_if_stats()
for name, items in ret.items():
name = py2_strencode(name)
isup, duplex, speed, mtu = items
if hasattr(_common, 'NicDuplex'):
duplex = _common.NicDuplex(duplex)
ret[name] = _common.snicstats(isup, duplex, speed, mtu)
return ret
def net_io_counters():
"""Return network I/O statistics for every network interface
installed on the system as a dict of raw tuples.
"""
ret = cext.net_io_counters()
return dict([(py2_strencode(k), v) for k, v in ret.items()])
def net_if_addrs():
"""Return the addresses associated to each NIC."""
ret = []
for items in cext.net_if_addrs():
items = list(items)
items[0] = py2_strencode(items[0])
ret.append(items)
return ret
# =====================================================================
# --- sensors
# =====================================================================
def sensors_battery():
"""Return battery information."""
# For constants meaning see:
# https://msdn.microsoft.com/en-us/library/windows/desktop/
# aa373232(v=vs.85).aspx
acline_status, flags, percent, secsleft = cext.sensors_battery()
power_plugged = acline_status == 1
no_battery = bool(flags & 128)
charging = bool(flags & 8)
if no_battery:
return None
if power_plugged or charging:
secsleft = _common.POWER_TIME_UNLIMITED
elif secsleft == -1:
secsleft = _common.POWER_TIME_UNKNOWN
return _common.sbattery(percent, secsleft, power_plugged)
# =====================================================================
# --- other system functions
# =====================================================================
def boot_time():
"""The system boot time expressed in seconds since the epoch."""
return cext.boot_time()
def users():
"""Return currently connected users as a list of namedtuples."""
retlist = []
rawlist = cext.users()
for item in rawlist:
user, hostname, tstamp = item
user = py2_strencode(user)
nt = _common.suser(user, None, hostname, tstamp)
retlist.append(nt)
return retlist
# =====================================================================
# --- Windows services
# =====================================================================
def win_service_iter():
"""Return a list of WindowsService instances."""
for name, display_name in cext.winservice_enumerate():
yield WindowsService(name, display_name)
def win_service_get(name):
"""Open a Windows service and return it as a WindowsService instance."""
service = WindowsService(name, None)
service._display_name = service._query_config()['display_name']
return service
class WindowsService(object):
"""Represents an installed Windows service."""
def __init__(self, name, display_name):
self._name = name
self._display_name = display_name
def __str__(self):
details = "(name=%r, display_name=%r)" % (
self._name, self._display_name)
return "%s%s" % (self.__class__.__name__, details)
def __repr__(self):
return "<%s at %s>" % (self.__str__(), id(self))
def __eq__(self, other):
# Test for equality with another WindosService object based
# on name.
if not isinstance(other, WindowsService):
return NotImplemented
return self._name == other._name
def __ne__(self, other):
return not self == other
def _query_config(self):
with self._wrap_exceptions():
display_name, binpath, username, start_type = \
cext.winservice_query_config(self._name)
# XXX - update _self.display_name?
return dict(
display_name=display_name,
binpath=binpath,
username=username,
start_type=start_type)
def _query_status(self):
with self._wrap_exceptions():
status, pid = cext.winservice_query_status(self._name)
if pid == 0:
pid = None
return dict(status=status, pid=pid)
@contextlib.contextmanager
def _wrap_exceptions(self):
"""Ctx manager which translates bare OSError and WindowsError
exceptions into NoSuchProcess and AccessDenied.
"""
try:
yield
except WindowsError as err:
NO_SUCH_SERVICE_SET = (cext.ERROR_INVALID_NAME,
cext.ERROR_SERVICE_DOES_NOT_EXIST)
if err.errno in ACCESS_DENIED_SET:
raise AccessDenied(
pid=None, name=self._name,
msg="service %r is not querable (not enough privileges)" %
self._name)
elif err.errno in NO_SUCH_SERVICE_SET or \
err.winerror in NO_SUCH_SERVICE_SET:
raise NoSuchProcess(
pid=None, name=self._name,
msg="service %r does not exist)" % self._name)
else:
raise
# config query
def name(self):
"""The service name. This string is how a service is referenced
and can be passed to win_service_get() to get a new
WindowsService instance.
"""
return self._name
def display_name(self):
"""The service display name. The value is cached when this class
is instantiated.
"""
return self._display_name
def binpath(self):
"""The fully qualified path to the service binary/exe file as
a string, including command line arguments.
"""
return self._query_config()['binpath']
def username(self):
"""The name of the user that owns this service."""
return self._query_config()['username']
def start_type(self):
"""A string which can either be "automatic", "manual" or
"disabled".
"""
return self._query_config()['start_type']
# status query
def pid(self):
"""The process PID, if any, else None. This can be passed
to Process class to control the service's process.
"""
return self._query_status()['pid']
def status(self):
"""Service status as a string."""
return self._query_status()['status']
def description(self):
"""Service long description."""
return cext.winservice_query_descr(self.name())
# utils
def as_dict(self):
"""Utility method retrieving all the information above as a
dictionary.
"""
d = self._query_config()
d.update(self._query_status())
d['name'] = self.name()
d['display_name'] = self.display_name()
d['description'] = self.description()
return d
# actions
# XXX: the necessary C bindings for start() and stop() are
# implemented but for now I prefer not to expose them.
# I may change my mind in the future. Reasons:
# - they require Administrator privileges
# - can't implement a timeout for stop() (unless by using a thread,
# which sucks)
# - would require adding ServiceAlreadyStarted and
# ServiceAlreadyStopped exceptions, adding two new APIs.
# - we might also want to have modify(), which would basically mean
# rewriting win32serviceutil.ChangeServiceConfig, which involves a
# lot of stuff (and API constants which would pollute the API), see:
# http://pyxr.sourceforge.net/PyXR/c/python24/lib/site-packages/
# win32/lib/win32serviceutil.py.html#0175
# - psutil is typically about "read only" monitoring stuff;
# win_service_* APIs should only be used to retrieve a service and
# check whether it's running
# def start(self, timeout=None):
# with self._wrap_exceptions():
# cext.winservice_start(self.name())
# if timeout:
# giveup_at = time.time() + timeout
# while True:
# if self.status() == "running":
# return
# else:
# if time.time() > giveup_at:
# raise TimeoutExpired(timeout)
# else:
# time.sleep(.1)
# def stop(self):
# # Note: timeout is not implemented because it's just not
# # possible, see:
# # http://stackoverflow.com/questions/11973228/
# with self._wrap_exceptions():
# return cext.winservice_stop(self.name())
# =====================================================================
# --- processes
# =====================================================================
pids = cext.pids
pid_exists = cext.pid_exists
ppid_map = cext.ppid_map # used internally by Process.children()
def wrap_exceptions(fun):
"""Decorator which translates bare OSError and WindowsError
exceptions into NoSuchProcess and AccessDenied.
"""
@functools.wraps(fun)
def wrapper(self, *args, **kwargs):
try:
return fun(self, *args, **kwargs)
except OSError as err:
if err.errno in ACCESS_DENIED_SET:
raise AccessDenied(self.pid, self._name)
if err.errno == errno.ESRCH:
raise NoSuchProcess(self.pid, self._name)
raise
return wrapper
class Process(object):
"""Wrapper class around underlying C implementation."""
__slots__ = ["pid", "_name", "_ppid"]
def __init__(self, pid):
self.pid = pid
self._name = None
self._ppid = None
# --- oneshot() stuff
def oneshot_enter(self):
self.oneshot_info.cache_activate()
def oneshot_exit(self):
self.oneshot_info.cache_deactivate()
@memoize_when_activated
def oneshot_info(self):
"""Return multiple information about this process as a
raw tuple.
"""
ret = cext.proc_info(self.pid)
assert len(ret) == len(pinfo_map)
return ret
@wrap_exceptions
def name(self):
"""Return process name, which on Windows is always the final
part of the executable.
"""
# This is how PIDs 0 and 4 are always represented in taskmgr
# and process-hacker.
if self.pid == 0:
return "System Idle Process"
elif self.pid == 4:
return "System"
else:
try:
# Note: this will fail with AD for most PIDs owned
# by another user but it's faster.
return py2_strencode(os.path.basename(self.exe()))
except AccessDenied:
return py2_strencode(cext.proc_name(self.pid))
@wrap_exceptions
def exe(self):
# Note: os.path.exists(path) may return False even if the file
# is there, see:
# http://stackoverflow.com/questions/3112546/os-path-exists-lies
# see https://github.com/giampaolo/psutil/issues/414
# see https://github.com/giampaolo/psutil/issues/528
if self.pid in (0, 4):
raise AccessDenied(self.pid, self._name)
return py2_strencode(convert_dos_path(cext.proc_exe(self.pid)))
@wrap_exceptions
def cmdline(self):
ret = cext.proc_cmdline(self.pid)
if PY3:
return ret
else:
return [py2_strencode(s) for s in ret]
@wrap_exceptions
def environ(self):
return parse_environ_block(cext.proc_environ(self.pid))
def ppid(self):
try:
return ppid_map()[self.pid]
except KeyError:
raise NoSuchProcess(self.pid, self._name)
def _get_raw_meminfo(self):
try:
return cext.proc_memory_info(self.pid)
except OSError as err:
if err.errno in ACCESS_DENIED_SET:
# TODO: the C ext can probably be refactored in order
# to get this from cext.proc_info()
info = self.oneshot_info()
return (
info[pinfo_map['num_page_faults']],
info[pinfo_map['peak_wset']],
info[pinfo_map['wset']],
info[pinfo_map['peak_paged_pool']],
info[pinfo_map['paged_pool']],
info[pinfo_map['peak_non_paged_pool']],
info[pinfo_map['non_paged_pool']],
info[pinfo_map['pagefile']],
info[pinfo_map['peak_pagefile']],
info[pinfo_map['mem_private']],
)
raise
@wrap_exceptions
def memory_info(self):
# on Windows RSS == WorkingSetSize and VSM == PagefileUsage.
# Underlying C function returns fields of PROCESS_MEMORY_COUNTERS
# struct.
t = self._get_raw_meminfo()
rss = t[2] # wset
vms = t[7] # pagefile
return pmem(*(rss, vms, ) + t)
@wrap_exceptions
def memory_full_info(self):
basic_mem = self.memory_info()
uss = cext.proc_memory_uss(self.pid)
return pfullmem(*basic_mem + (uss, ))
def memory_maps(self):
try:
raw = cext.proc_memory_maps(self.pid)
except OSError as err:
# XXX - can't use wrap_exceptions decorator as we're
# returning a generator; probably needs refactoring.
if err.errno in ACCESS_DENIED_SET:
raise AccessDenied(self.pid, self._name)
if err.errno == errno.ESRCH:
raise NoSuchProcess(self.pid, self._name)
raise
else:
for addr, perm, path, rss in raw:
path = convert_dos_path(path)
addr = hex(addr)
yield (addr, perm, path, rss)
@wrap_exceptions
def kill(self):
return cext.proc_kill(self.pid)
@wrap_exceptions
def send_signal(self, sig):
os.kill(self.pid, sig)
@wrap_exceptions
def wait(self, timeout=None):
if timeout is None:
cext_timeout = cext.INFINITE
else:
# WaitForSingleObject() expects time in milliseconds
cext_timeout = int(timeout * 1000)
ret = cext.proc_wait(self.pid, cext_timeout)
if ret == WAIT_TIMEOUT:
raise TimeoutExpired(timeout, self.pid, self._name)
return ret
@wrap_exceptions
def username(self):
if self.pid in (0, 4):
return 'NT AUTHORITY\\SYSTEM'
return cext.proc_username(self.pid)
@wrap_exceptions
def create_time(self):
# special case for kernel process PIDs; return system boot time
if self.pid in (0, 4):
return boot_time()
try:
return cext.proc_create_time(self.pid)
except OSError as err:
if err.errno in ACCESS_DENIED_SET:
return self.oneshot_info()[pinfo_map['create_time']]
raise
@wrap_exceptions
def num_threads(self):
return self.oneshot_info()[pinfo_map['num_threads']]
@wrap_exceptions
def threads(self):
rawlist = cext.proc_threads(self.pid)
retlist = []
for thread_id, utime, stime in rawlist:
ntuple = _common.pthread(thread_id, utime, stime)
retlist.append(ntuple)
return retlist
@wrap_exceptions
def cpu_times(self):
try:
user, system = cext.proc_cpu_times(self.pid)
except OSError as err:
if err.errno in ACCESS_DENIED_SET:
info = self.oneshot_info()
user = info[pinfo_map['user_time']]
system = info[pinfo_map['kernel_time']]
else:
raise
# Children user/system times are not retrievable (set to 0).
return _common.pcputimes(user, system, 0, 0)
@wrap_exceptions
def suspend(self):
return cext.proc_suspend(self.pid)
@wrap_exceptions
def resume(self):
return cext.proc_resume(self.pid)
@wrap_exceptions
def cwd(self):
if self.pid in (0, 4):
raise AccessDenied(self.pid, self._name)
# return a normalized pathname since the native C function appends
# "\\" at the and of the path
path = cext.proc_cwd(self.pid)
return py2_strencode(os.path.normpath(path))
@wrap_exceptions
def open_files(self):
if self.pid in (0, 4):
return []
ret = set()
# Filenames come in in native format like:
# "\Device\HarddiskVolume1\Windows\systemew\file.txt"
# Convert the first part in the corresponding drive letter
# (e.g. "C:\") by using Windows's QueryDosDevice()
raw_file_names = cext.proc_open_files(self.pid)
for _file in raw_file_names:
_file = convert_dos_path(_file)
if isfile_strict(_file):
if not PY3:
_file = py2_strencode(_file)
ntuple = _common.popenfile(_file, -1)
ret.add(ntuple)
return list(ret)
@wrap_exceptions
def connections(self, kind='inet'):
return net_connections(kind, _pid=self.pid)
@wrap_exceptions
def nice_get(self):
value = cext.proc_priority_get(self.pid)
if enum is not None:
value = Priority(value)
return value
@wrap_exceptions
def nice_set(self, value):
return cext.proc_priority_set(self.pid, value)
# available on Windows >= Vista
if hasattr(cext, "proc_io_priority_get"):
@wrap_exceptions
def ionice_get(self):
return cext.proc_io_priority_get(self.pid)
@wrap_exceptions
def ionice_set(self, value, _):
if _:
raise TypeError("set_proc_ionice() on Windows takes only "
"1 argument (2 given)")
if value not in (2, 1, 0):
raise ValueError("value must be 2 (normal), 1 (low) or 0 "
"(very low); got %r" % value)
return cext.proc_io_priority_set(self.pid, value)
@wrap_exceptions
def io_counters(self):
try:
ret = cext.proc_io_counters(self.pid)
except OSError as err:
if err.errno in ACCESS_DENIED_SET:
info = self.oneshot_info()
ret = (
info[pinfo_map['io_rcount']],
info[pinfo_map['io_wcount']],
info[pinfo_map['io_rbytes']],
info[pinfo_map['io_wbytes']],
info[pinfo_map['io_count_others']],
info[pinfo_map['io_bytes_others']],
)
else:
raise
return pio(*ret)
@wrap_exceptions
def status(self):
suspended = cext.proc_is_suspended(self.pid)
if suspended:
return _common.STATUS_STOPPED
else:
return _common.STATUS_RUNNING
@wrap_exceptions
def cpu_affinity_get(self):
def from_bitmask(x):
return [i for i in xrange(64) if (1 << i) & x]
bitmask = cext.proc_cpu_affinity_get(self.pid)
return from_bitmask(bitmask)
@wrap_exceptions
def cpu_affinity_set(self, value):
def to_bitmask(l):
if not l:
raise ValueError("invalid argument %r" % l)
out = 0
for b in l:
out |= 2 ** b
return out
# SetProcessAffinityMask() states that ERROR_INVALID_PARAMETER
# is returned for an invalid CPU but this seems not to be true,
# therefore we check CPUs validy beforehand.
allcpus = list(range(len(per_cpu_times())))
for cpu in value:
if cpu not in allcpus:
if not isinstance(cpu, (int, long)):
raise TypeError(
"invalid CPU %r; an integer is required" % cpu)
else:
raise ValueError("invalid CPU %r" % cpu)
bitmask = to_bitmask(value)
cext.proc_cpu_affinity_set(self.pid, bitmask)
@wrap_exceptions
def num_handles(self):
try:
return cext.proc_num_handles(self.pid)
except OSError as err:
if err.errno in ACCESS_DENIED_SET:
return self.oneshot_info()[pinfo_map['num_handles']]
raise
@wrap_exceptions
def num_ctx_switches(self):
ctx_switches = self.oneshot_info()[pinfo_map['ctx_switches']]
# only voluntary ctx switches are supported
return _common.pctxsw(ctx_switches, 0)
File diff suppressed because it is too large Load Diff
+32
View File
@@ -0,0 +1,32 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
typedef struct kinfo_proc kinfo_proc;
int psutil_get_proc_list(struct kinfo_proc **procList, size_t *procCount);
int psutil_kinfo_proc(const pid_t pid, struct kinfo_proc *proc);
//
PyObject* psutil_cpu_count_phys(PyObject* self, PyObject* args);
PyObject* psutil_disk_io_counters(PyObject* self, PyObject* args);
PyObject* psutil_get_cmdline(long pid);
PyObject* psutil_per_cpu_times(PyObject* self, PyObject* args);
PyObject* psutil_proc_cpu_affinity_get(PyObject* self, PyObject* args);
PyObject* psutil_proc_cpu_affinity_set(PyObject* self, PyObject* args);
PyObject* psutil_proc_cwd(PyObject* self, PyObject* args);
PyObject* psutil_proc_exe(PyObject* self, PyObject* args);
PyObject* psutil_proc_memory_maps(PyObject* self, PyObject* args);
PyObject* psutil_proc_num_fds(PyObject* self, PyObject* args);
PyObject* psutil_proc_num_threads(PyObject* self, PyObject* args);
PyObject* psutil_proc_threads(PyObject* self, PyObject* args);
PyObject* psutil_swap_mem(PyObject* self, PyObject* args);
PyObject* psutil_virtual_mem(PyObject* self, PyObject* args);
PyObject* psutil_cpu_stats(PyObject* self, PyObject* args);
#if defined(PSUTIL_FREEBSD)
PyObject* psutil_sensors_battery(PyObject* self, PyObject* args);
#endif
+631
View File
@@ -0,0 +1,631 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
#include <sys/user.h>
#include <sys/file.h>
#include <sys/socketvar.h> // for struct xsocket
#include <sys/un.h>
#include <sys/unpcb.h>
#include <sys/types.h>
#include <sys/sysctl.h>
#include <netinet/in.h> // for xinpcb struct
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <netinet/in_pcb.h>
#include <netinet/tcp.h>
#include <netinet/tcp_timer.h>
#include <netinet/tcp_var.h> // for struct xtcpcb
#include <netinet/tcp_fsm.h> // for TCP connection states
#include <arpa/inet.h> // for inet_ntop()
#include <net/if_media.h>
#include <libutil.h>
#include "../../_psutil_common.h"
#define HASHSIZE 1009
// a signaler for connections without an actual status
static int PSUTIL_CONN_NONE = 128;
static struct xfile *psutil_xfiles;
static int psutil_nxfiles;
// The tcplist fetching and walking is borrowed from netstat/inet.c.
static char *
psutil_fetch_tcplist(void) {
char *buf;
size_t len;
for (;;) {
if (sysctlbyname("net.inet.tcp.pcblist", NULL, &len, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
buf = malloc(len);
if (buf == NULL) {
PyErr_NoMemory();
return NULL;
}
if (sysctlbyname("net.inet.tcp.pcblist", buf, &len, NULL, 0) < 0) {
free(buf);
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
return buf;
}
}
static int
psutil_sockaddr_port(int family, struct sockaddr_storage *ss) {
struct sockaddr_in6 *sin6;
struct sockaddr_in *sin;
if (family == AF_INET) {
sin = (struct sockaddr_in *)ss;
return (sin->sin_port);
}
else {
sin6 = (struct sockaddr_in6 *)ss;
return (sin6->sin6_port);
}
}
static void *
psutil_sockaddr_addr(int family, struct sockaddr_storage *ss) {
struct sockaddr_in6 *sin6;
struct sockaddr_in *sin;
if (family == AF_INET) {
sin = (struct sockaddr_in *)ss;
return (&sin->sin_addr);
}
else {
sin6 = (struct sockaddr_in6 *)ss;
return (&sin6->sin6_addr);
}
}
static socklen_t
psutil_sockaddr_addrlen(int family) {
if (family == AF_INET)
return (sizeof(struct in_addr));
else
return (sizeof(struct in6_addr));
}
static int
psutil_sockaddr_matches(int family, int port, void *pcb_addr,
struct sockaddr_storage *ss) {
if (psutil_sockaddr_port(family, ss) != port)
return (0);
return (memcmp(psutil_sockaddr_addr(family, ss), pcb_addr,
psutil_sockaddr_addrlen(family)) == 0);
}
static struct tcpcb *
psutil_search_tcplist(char *buf, struct kinfo_file *kif) {
struct tcpcb *tp;
struct inpcb *inp;
struct xinpgen *xig, *oxig;
struct xsocket *so;
oxig = xig = (struct xinpgen *)buf;
for (xig = (struct xinpgen *)((char *)xig + xig->xig_len);
xig->xig_len > sizeof(struct xinpgen);
xig = (struct xinpgen *)((char *)xig + xig->xig_len)) {
tp = &((struct xtcpcb *)xig)->xt_tp;
inp = &((struct xtcpcb *)xig)->xt_inp;
so = &((struct xtcpcb *)xig)->xt_socket;
if (so->so_type != kif->kf_sock_type ||
so->xso_family != kif->kf_sock_domain ||
so->xso_protocol != kif->kf_sock_protocol)
continue;
if (kif->kf_sock_domain == AF_INET) {
if (!psutil_sockaddr_matches(
AF_INET, inp->inp_lport, &inp->inp_laddr,
&kif->kf_sa_local))
continue;
if (!psutil_sockaddr_matches(
AF_INET, inp->inp_fport, &inp->inp_faddr,
&kif->kf_sa_peer))
continue;
} else {
if (!psutil_sockaddr_matches(
AF_INET6, inp->inp_lport, &inp->in6p_laddr,
&kif->kf_sa_local))
continue;
if (!psutil_sockaddr_matches(
AF_INET6, inp->inp_fport, &inp->in6p_faddr,
&kif->kf_sa_peer))
continue;
}
return (tp);
}
return NULL;
}
int
psutil_populate_xfiles() {
size_t len;
if ((psutil_xfiles = malloc(len = sizeof *psutil_xfiles)) == NULL) {
PyErr_NoMemory();
return 0;
}
while (sysctlbyname("kern.file", psutil_xfiles, &len, 0, 0) == -1) {
if (errno != ENOMEM) {
PyErr_SetFromErrno(0);
return 0;
}
len *= 2;
if ((psutil_xfiles = realloc(psutil_xfiles, len)) == NULL) {
PyErr_NoMemory();
return 0;
}
}
if (len > 0 && psutil_xfiles->xf_size != sizeof *psutil_xfiles) {
PyErr_Format(PyExc_RuntimeError, "struct xfile size mismatch");
return 0;
}
psutil_nxfiles = len / sizeof *psutil_xfiles;
return 1;
}
int
psutil_get_pid_from_sock(int sock_hash) {
struct xfile *xf;
int hash, n;
for (xf = psutil_xfiles, n = 0; n < psutil_nxfiles; ++n, ++xf) {
if (xf->xf_data == NULL)
continue;
hash = (int)((uintptr_t)xf->xf_data % HASHSIZE);
if (sock_hash == hash)
return xf->xf_pid;
}
return -1;
}
// Reference:
// https://gitorious.org/freebsd/freebsd/source/
// f1d6f4778d2044502209708bc167c05f9aa48615:usr.bin/sockstat/sockstat.c
int psutil_gather_inet(int proto, PyObject *py_retlist) {
struct xinpgen *xig, *exig;
struct xinpcb *xip;
struct xtcpcb *xtp;
struct inpcb *inp;
struct xsocket *so;
const char *varname = NULL;
size_t len, bufsize;
void *buf;
int hash;
int retry;
int type;
PyObject *py_tuple = NULL;
PyObject *py_laddr = NULL;
PyObject *py_raddr = NULL;
switch (proto) {
case IPPROTO_TCP:
varname = "net.inet.tcp.pcblist";
type = SOCK_STREAM;
break;
case IPPROTO_UDP:
varname = "net.inet.udp.pcblist";
type = SOCK_DGRAM;
break;
}
buf = NULL;
bufsize = 8192;
retry = 5;
do {
for (;;) {
buf = realloc(buf, bufsize);
if (buf == NULL)
continue; // XXX
len = bufsize;
if (sysctlbyname(varname, buf, &len, NULL, 0) == 0)
break;
if (errno != ENOMEM) {
PyErr_SetFromErrno(0);
goto error;
}
bufsize *= 2;
}
xig = (struct xinpgen *)buf;
exig = (struct xinpgen *)(void *)((char *)buf + len - sizeof *exig);
if (xig->xig_len != sizeof *xig || exig->xig_len != sizeof *exig) {
PyErr_Format(PyExc_RuntimeError, "struct xinpgen size mismatch");
goto error;
}
} while (xig->xig_gen != exig->xig_gen && retry--);
for (;;) {
int lport, rport, pid, status, family;
xig = (struct xinpgen *)(void *)((char *)xig + xig->xig_len);
if (xig >= exig)
break;
switch (proto) {
case IPPROTO_TCP:
xtp = (struct xtcpcb *)xig;
if (xtp->xt_len != sizeof *xtp) {
PyErr_Format(PyExc_RuntimeError,
"struct xtcpcb size mismatch");
goto error;
}
inp = &xtp->xt_inp;
so = &xtp->xt_socket;
status = xtp->xt_tp.t_state;
break;
case IPPROTO_UDP:
xip = (struct xinpcb *)xig;
if (xip->xi_len != sizeof *xip) {
PyErr_Format(PyExc_RuntimeError,
"struct xinpcb size mismatch");
goto error;
}
inp = &xip->xi_inp;
so = &xip->xi_socket;
status = PSUTIL_CONN_NONE;
break;
default:
PyErr_Format(PyExc_RuntimeError, "invalid proto");
goto error;
}
char lip[200], rip[200];
hash = (int)((uintptr_t)so->xso_so % HASHSIZE);
pid = psutil_get_pid_from_sock(hash);
if (pid < 0)
continue;
lport = ntohs(inp->inp_lport);
rport = ntohs(inp->inp_fport);
if (inp->inp_vflag & INP_IPV4) {
family = AF_INET;
inet_ntop(AF_INET, &inp->inp_laddr.s_addr, lip, sizeof(lip));
inet_ntop(AF_INET, &inp->inp_faddr.s_addr, rip, sizeof(rip));
}
else if (inp->inp_vflag & INP_IPV6) {
family = AF_INET6;
inet_ntop(AF_INET6, &inp->in6p_laddr.s6_addr, lip, sizeof(lip));
inet_ntop(AF_INET6, &inp->in6p_faddr.s6_addr, rip, sizeof(rip));
}
// construct python tuple/list
py_laddr = Py_BuildValue("(si)", lip, lport);
if (!py_laddr)
goto error;
if (rport != 0)
py_raddr = Py_BuildValue("(si)", rip, rport);
else
py_raddr = Py_BuildValue("()");
if (!py_raddr)
goto error;
py_tuple = Py_BuildValue("(iiiNNii)", -1, family, type, py_laddr,
py_raddr, status, pid);
if (!py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
free(buf);
return 1;
error:
Py_XDECREF(py_tuple);
Py_XDECREF(py_laddr);
Py_XDECREF(py_raddr);
free(buf);
return 0;
}
int psutil_gather_unix(int proto, PyObject *py_retlist) {
struct xunpgen *xug, *exug;
struct xunpcb *xup;
const char *varname = NULL;
const char *protoname = NULL;
size_t len;
size_t bufsize;
void *buf;
int hash;
int retry;
int pid;
struct sockaddr_un *sun;
char path[PATH_MAX];
PyObject *py_tuple = NULL;
PyObject *py_laddr = NULL;
PyObject *py_raddr = NULL;
switch (proto) {
case SOCK_STREAM:
varname = "net.local.stream.pcblist";
protoname = "stream";
break;
case SOCK_DGRAM:
varname = "net.local.dgram.pcblist";
protoname = "dgram";
break;
}
buf = NULL;
bufsize = 8192;
retry = 5;
do {
for (;;) {
buf = realloc(buf, bufsize);
if (buf == NULL) {
PyErr_NoMemory();
goto error;
}
len = bufsize;
if (sysctlbyname(varname, buf, &len, NULL, 0) == 0)
break;
if (errno != ENOMEM) {
PyErr_SetFromErrno(0);
goto error;
}
bufsize *= 2;
}
xug = (struct xunpgen *)buf;
exug = (struct xunpgen *)(void *)
((char *)buf + len - sizeof *exug);
if (xug->xug_len != sizeof *xug || exug->xug_len != sizeof *exug) {
PyErr_Format(PyExc_RuntimeError, "struct xinpgen size mismatch");
goto error;
}
} while (xug->xug_gen != exug->xug_gen && retry--);
for (;;) {
xug = (struct xunpgen *)(void *)((char *)xug + xug->xug_len);
if (xug >= exug)
break;
xup = (struct xunpcb *)xug;
if (xup->xu_len != sizeof *xup)
goto error;
hash = (int)((uintptr_t) xup->xu_socket.xso_so % HASHSIZE);
pid = psutil_get_pid_from_sock(hash);
if (pid < 0)
continue;
sun = (struct sockaddr_un *)&xup->xu_addr;
snprintf(path, sizeof(path), "%.*s",
(int)(sun->sun_len - (sizeof(*sun) - sizeof(sun->sun_path))),
sun->sun_path);
py_tuple = Py_BuildValue("(iiisOii)", -1, AF_UNIX, proto, path,
Py_None, PSUTIL_CONN_NONE, pid);
if (!py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
Py_INCREF(Py_None);
}
free(buf);
return 1;
error:
Py_XDECREF(py_tuple);
Py_XDECREF(py_laddr);
Py_XDECREF(py_raddr);
free(buf);
return 0;
}
PyObject*
psutil_net_connections(PyObject* self, PyObject* args) {
// Return system-wide open connections.
PyObject *py_retlist = PyList_New(0);
if (py_retlist == NULL)
return NULL;
if (psutil_populate_xfiles() != 1)
goto error;
if (psutil_gather_inet(IPPROTO_TCP, py_retlist) == 0)
goto error;
if (psutil_gather_inet(IPPROTO_UDP, py_retlist) == 0)
goto error;
if (psutil_gather_unix(SOCK_STREAM, py_retlist) == 0)
goto error;
if (psutil_gather_unix(SOCK_DGRAM, py_retlist) == 0)
goto error;
free(psutil_xfiles);
return py_retlist;
error:
Py_DECREF(py_retlist);
free(psutil_xfiles);
return NULL;
}
PyObject *
psutil_proc_connections(PyObject *self, PyObject *args) {
// Return connections opened by process.
long pid;
int i, cnt;
struct kinfo_file *freep = NULL;
struct kinfo_file *kif;
char *tcplist = NULL;
struct tcpcb *tcp;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
PyObject *py_laddr = NULL;
PyObject *py_raddr = NULL;
PyObject *py_af_filter = NULL;
PyObject *py_type_filter = NULL;
PyObject *py_family = NULL;
PyObject *py_type = NULL;
if (py_retlist == NULL)
return NULL;
if (! PyArg_ParseTuple(args, "lOO", &pid, &py_af_filter, &py_type_filter))
goto error;
if (!PySequence_Check(py_af_filter) || !PySequence_Check(py_type_filter)) {
PyErr_SetString(PyExc_TypeError, "arg 2 or 3 is not a sequence");
goto error;
}
errno = 0;
freep = kinfo_getfile(pid, &cnt);
if (freep == NULL) {
psutil_raise_for_pid(pid, "kinfo_getfile() failed");
goto error;
}
tcplist = psutil_fetch_tcplist();
if (tcplist == NULL) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
for (i = 0; i < cnt; i++) {
int lport, rport, state;
char lip[200], rip[200];
char path[PATH_MAX];
int inseq;
py_tuple = NULL;
py_laddr = NULL;
py_raddr = NULL;
kif = &freep[i];
if (kif->kf_type == KF_TYPE_SOCKET) {
// apply filters
py_family = PyLong_FromLong((long)kif->kf_sock_domain);
inseq = PySequence_Contains(py_af_filter, py_family);
Py_DECREF(py_family);
if (inseq == 0)
continue;
py_type = PyLong_FromLong((long)kif->kf_sock_type);
inseq = PySequence_Contains(py_type_filter, py_type);
Py_DECREF(py_type);
if (inseq == 0)
continue;
// IPv4 / IPv6 socket
if ((kif->kf_sock_domain == AF_INET) ||
(kif->kf_sock_domain == AF_INET6)) {
// fill status
state = PSUTIL_CONN_NONE;
if (kif->kf_sock_type == SOCK_STREAM) {
tcp = psutil_search_tcplist(tcplist, kif);
if (tcp != NULL)
state = (int)tcp->t_state;
}
// build addr and port
inet_ntop(
kif->kf_sock_domain,
psutil_sockaddr_addr(kif->kf_sock_domain,
&kif->kf_sa_local),
lip,
sizeof(lip));
inet_ntop(
kif->kf_sock_domain,
psutil_sockaddr_addr(kif->kf_sock_domain,
&kif->kf_sa_peer),
rip,
sizeof(rip));
lport = htons(psutil_sockaddr_port(kif->kf_sock_domain,
&kif->kf_sa_local));
rport = htons(psutil_sockaddr_port(kif->kf_sock_domain,
&kif->kf_sa_peer));
// construct python tuple/list
py_laddr = Py_BuildValue("(si)", lip, lport);
if (!py_laddr)
goto error;
if (rport != 0)
py_raddr = Py_BuildValue("(si)", rip, rport);
else
py_raddr = Py_BuildValue("()");
if (!py_raddr)
goto error;
py_tuple = Py_BuildValue(
"(iiiNNi)",
kif->kf_fd,
kif->kf_sock_domain,
kif->kf_sock_type,
py_laddr,
py_raddr,
state
);
if (!py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
// UNIX socket
else if (kif->kf_sock_domain == AF_UNIX) {
struct sockaddr_un *sun;
sun = (struct sockaddr_un *)&kif->kf_sa_local;
snprintf(
path, sizeof(path), "%.*s",
(int)(sun->sun_len - (sizeof(*sun) - sizeof(sun->sun_path))),
sun->sun_path);
py_tuple = Py_BuildValue(
"(iiisOi)",
kif->kf_fd,
kif->kf_sock_domain,
kif->kf_sock_type,
path,
Py_None,
PSUTIL_CONN_NONE
);
if (!py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
Py_INCREF(Py_None);
}
}
}
free(freep);
free(tcplist);
return py_retlist;
error:
Py_XDECREF(py_tuple);
Py_XDECREF(py_laddr);
Py_XDECREF(py_raddr);
Py_DECREF(py_retlist);
if (freep != NULL)
free(freep);
if (tcplist != NULL)
free(tcplist);
return NULL;
}
+11
View File
@@ -0,0 +1,11 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
PyObject* psutil_proc_connections(PyObject* self, PyObject* args);
PyObject* psutil_net_connections(PyObject* self, PyObject* args);
+664
View File
@@ -0,0 +1,664 @@
/*
* Copyright (c) 2009, Giampaolo Rodola', Landry Breuil.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Platform-specific module methods for NetBSD.
*/
#if defined(PSUTIL_NETBSD)
#define _KMEMUSER
#endif
#include <Python.h>
#include <assert.h>
#include <err.h>
#include <errno.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/param.h>
#include <sys/sysctl.h>
#include <sys/user.h>
#include <sys/proc.h>
#include <sys/swap.h> // for swap_mem
#include <signal.h>
#include <kvm.h>
// connection stuff
#include <netdb.h> // for NI_MAXHOST
#include <sys/socket.h>
#include <sys/sched.h> // for CPUSTATES & CP_*
#define _KERNEL // for DTYPE_*
#include <sys/file.h>
#undef _KERNEL
#include <sys/disk.h> // struct diskstats
#include <netinet/in.h>
#include <arpa/inet.h>
#include "netbsd_socks.h"
#include "netbsd.h"
#include "../../_psutil_common.h"
#define PSUTIL_KPT2DOUBLE(t) (t ## _sec + t ## _usec / 1000000.0)
#define PSUTIL_TV2DOUBLE(t) ((t).tv_sec + (t).tv_usec / 1000000.0)
// ============================================================================
// Utility functions
// ============================================================================
int
psutil_kinfo_proc(pid_t pid, kinfo_proc *proc) {
// Fills a kinfo_proc struct based on process pid.
int ret;
int mib[6];
size_t size = sizeof(kinfo_proc);
mib[0] = CTL_KERN;
mib[1] = KERN_PROC2;
mib[2] = KERN_PROC_PID;
mib[3] = pid;
mib[4] = size;
mib[5] = 1;
ret = sysctl((int*)mib, 6, proc, &size, NULL, 0);
if (ret == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
// sysctl stores 0 in the size if we can't find the process information.
if (size == 0) {
NoSuchProcess();
return -1;
}
return 0;
}
struct kinfo_file *
kinfo_getfile(pid_t pid, int* cnt) {
// Mimic's FreeBSD kinfo_file call, taking a pid and a ptr to an
// int as arg and returns an array with cnt struct kinfo_file.
int mib[6];
size_t len;
struct kinfo_file* kf;
mib[0] = CTL_KERN;
mib[1] = KERN_FILE2;
mib[2] = KERN_FILE_BYPID;
mib[3] = (int) pid;
mib[4] = sizeof(struct kinfo_file);
mib[5] = 0;
// get the size of what would be returned
if (sysctl(mib, 6, NULL, &len, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
if ((kf = malloc(len)) == NULL) {
PyErr_NoMemory();
return NULL;
}
mib[5] = (int)(len / sizeof(struct kinfo_file));
if (sysctl(mib, 6, kf, &len, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
*cnt = (int)(len / sizeof(struct kinfo_file));
return kf;
}
// XXX: This is no longer used as per
// https://github.com/giampaolo/psutil/pull/557#issuecomment-171912820
// Current implementation uses /proc instead.
// Left here just in case.
PyObject *
psutil_proc_exe(PyObject *self, PyObject *args) {
#if __NetBSD_Version__ >= 799000000
pid_t pid;
char pathname[MAXPATHLEN];
int error;
int mib[4];
int ret;
size_t size;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
if (pid == 0) {
// else returns ENOENT
return Py_BuildValue("s", "");
}
mib[0] = CTL_KERN;
mib[1] = KERN_PROC_ARGS;
mib[2] = pid;
mib[3] = KERN_PROC_PATHNAME;
size = sizeof(pathname);
error = sysctl(mib, 4, NULL, &size, NULL, 0);
if (error == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
error = sysctl(mib, 4, pathname, &size, NULL, 0);
if (error == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
if (size == 0 || strlen(pathname) == 0) {
ret = psutil_pid_exists(pid);
if (ret == -1)
return NULL;
else if (ret == 0)
return NoSuchProcess();
else
strcpy(pathname, "");
}
#if PY_MAJOR_VERSION >= 3
return PyUnicode_DecodeFSDefault(pathname);
#else
return Py_BuildValue("s", pathname);
#endif
#else
return Py_BuildValue("s", "");
#endif
}
PyObject *
psutil_proc_num_threads(PyObject *self, PyObject *args) {
// Return number of threads used by process as a Python integer.
long pid;
kinfo_proc kp;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
if (psutil_kinfo_proc(pid, &kp) == -1)
return NULL;
return Py_BuildValue("l", (long)kp.p_nlwps);
}
PyObject *
psutil_proc_threads(PyObject *self, PyObject *args) {
pid_t pid;
int mib[5];
int i, nlwps;
ssize_t st;
size_t size;
struct kinfo_lwp *kl = NULL;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
if (py_retlist == NULL)
return NULL;
if (! PyArg_ParseTuple(args, "l", &pid))
goto error;
mib[0] = CTL_KERN;
mib[1] = KERN_LWP;
mib[2] = pid;
mib[3] = sizeof(struct kinfo_lwp);
mib[4] = 0;
st = sysctl(mib, 5, NULL, &size, NULL, 0);
if (st == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
if (size == 0) {
NoSuchProcess();
goto error;
}
mib[4] = size / sizeof(size_t);
kl = malloc(size);
if (kl == NULL) {
PyErr_NoMemory();
goto error;
}
st = sysctl(mib, 5, kl, &size, NULL, 0);
if (st == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
if (size == 0) {
NoSuchProcess();
goto error;
}
nlwps = (int)(size / sizeof(struct kinfo_lwp));
for (i = 0; i < nlwps; i++) {
py_tuple = Py_BuildValue("idd",
(&kl[i])->l_lid,
PSUTIL_KPT2DOUBLE((&kl[i])->l_rtime),
PSUTIL_KPT2DOUBLE((&kl[i])->l_rtime));
if (py_tuple == NULL)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
free(kl);
return py_retlist;
error:
Py_XDECREF(py_tuple);
Py_DECREF(py_retlist);
if (kl != NULL)
free(kl);
return NULL;
}
// ============================================================================
// APIS
// ============================================================================
int
psutil_get_proc_list(kinfo_proc **procList, size_t *procCount) {
// Returns a list of all BSD processes on the system. This routine
// allocates the list and puts it in *procList and a count of the
// number of entries in *procCount. You are responsible for freeing
// this list (use "free" from System framework).
// On success, the function returns 0.
// On error, the function returns a BSD errno value.
kinfo_proc *result;
int done;
static const int name[] = { CTL_KERN, KERN_PROC, KERN_PROC, 0 };
// Declaring name as const requires us to cast it when passing it to
// sysctl because the prototype doesn't include the const modifier.
size_t length;
char errbuf[_POSIX2_LINE_MAX];
kinfo_proc *x;
int cnt;
kvm_t *kd;
assert( procList != NULL);
assert(*procList == NULL);
assert(procCount != NULL);
kd = kvm_openfiles(NULL, NULL, NULL, KVM_NO_FILES, errbuf);
if (kd == NULL) {
PyErr_Format(
PyExc_RuntimeError, "kvm_openfiles() syscall failed: %s", errbuf);
return errno;
}
result = kvm_getproc2(kd, KERN_PROC_ALL, 0, sizeof(kinfo_proc), &cnt);
if (result == NULL) {
PyErr_Format(PyExc_RuntimeError, "kvm_getproc2() syscall failed");
kvm_close(kd);
return errno;
}
*procCount = (size_t)cnt;
size_t mlen = cnt * sizeof(kinfo_proc);
if ((*procList = malloc(mlen)) == NULL) {
PyErr_NoMemory();
kvm_close(kd);
return errno;
}
memcpy(*procList, result, mlen);
assert(*procList != NULL);
kvm_close(kd);
return 0;
}
char *
psutil_get_cmd_args(pid_t pid, size_t *argsize) {
int mib[4];
ssize_t st;
size_t argmax;
size_t size;
char *procargs = NULL;
mib[0] = CTL_KERN;
mib[1] = KERN_ARGMAX;
size = sizeof(argmax);
st = sysctl(mib, 2, &argmax, &size, NULL, 0);
if (st == -1) {
warn("failed to get kern.argmax");
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
procargs = (char *)malloc(argmax);
if (procargs == NULL) {
PyErr_NoMemory();
return NULL;
}
mib[0] = CTL_KERN;
mib[1] = KERN_PROC_ARGS;
mib[2] = pid;
mib[3] = KERN_PROC_ARGV;
st = sysctl(mib, 4, procargs, &argmax, NULL, 0);
if (st == -1) {
warn("failed to get kern.procargs");
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
*argsize = argmax;
return procargs;
}
// Return the command line as a python list object.
// XXX - most of the times sysctl() returns a truncated string.
// Also /proc/pid/cmdline behaves the same so it looks like this
// is a kernel bug.
PyObject *
psutil_get_cmdline(pid_t pid) {
char *argstr = NULL;
int pos = 0;
size_t argsize = 0;
PyObject *py_arg = NULL;
PyObject *py_retlist = PyList_New(0);
if (py_retlist == NULL)
return NULL;
if (pid == 0)
return py_retlist;
argstr = psutil_get_cmd_args(pid, &argsize);
if (argstr == NULL)
goto error;
// args are returned as a flattened string with \0 separators between
// arguments add each string to the list then step forward to the next
// separator
if (argsize > 0) {
while (pos < argsize) {
#if PY_MAJOR_VERSION >= 3
py_arg = PyUnicode_DecodeFSDefault(&argstr[pos]);
#else
py_arg = Py_BuildValue("s", &argstr[pos]);
#endif
if (!py_arg)
goto error;
if (PyList_Append(py_retlist, py_arg))
goto error;
Py_DECREF(py_arg);
pos = pos + strlen(&argstr[pos]) + 1;
}
}
free(argstr);
return py_retlist;
error:
Py_XDECREF(py_arg);
Py_DECREF(py_retlist);
if (argstr != NULL)
free(argstr);
return NULL;
}
/*
* Virtual memory stats, taken from:
* https://github.com/satterly/zabbix-stats/blob/master/src/libs/zbxsysinfo/
* netbsd/memory.c
*/
PyObject *
psutil_virtual_mem(PyObject *self, PyObject *args) {
size_t size;
struct uvmexp_sysctl uv;
int mib[] = {CTL_VM, VM_UVMEXP2};
long pagesize = getpagesize();
size = sizeof(uv);
if (sysctl(mib, 2, &uv, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
return Py_BuildValue("KKKKKKKK",
(unsigned long long) uv.npages << uv.pageshift, // total
(unsigned long long) uv.free << uv.pageshift, // free
(unsigned long long) uv.active << uv.pageshift, // active
(unsigned long long) uv.inactive << uv.pageshift, // inactive
(unsigned long long) uv.wired << uv.pageshift, // wired
(unsigned long long) uv.filepages + uv.execpages * pagesize, // cached
// These are determined from /proc/meminfo in Python.
(unsigned long long) 0, // buffers
(unsigned long long) 0 // shared
);
}
PyObject *
psutil_swap_mem(PyObject *self, PyObject *args) {
uint64_t swap_total, swap_free;
struct swapent *swdev;
int nswap, i;
nswap = swapctl(SWAP_NSWAP, 0, 0);
if (nswap == 0) {
// This means there's no swap partition.
return Py_BuildValue("(iiiii)", 0, 0, 0, 0, 0);
}
swdev = calloc(nswap, sizeof(*swdev));
if (swdev == NULL) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
if (swapctl(SWAP_STATS, swdev, nswap) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
// Total things up.
swap_total = swap_free = 0;
for (i = 0; i < nswap; i++) {
if (swdev[i].se_flags & SWF_ENABLE) {
swap_total += swdev[i].se_nblks * DEV_BSIZE;
swap_free += (swdev[i].se_nblks - swdev[i].se_inuse) * DEV_BSIZE;
}
}
free(swdev);
// Get swap in/out
unsigned int total;
size_t size = sizeof(total);
struct uvmexp_sysctl uv;
int mib[] = {CTL_VM, VM_UVMEXP2};
long pagesize = getpagesize();
size = sizeof(uv);
if (sysctl(mib, 2, &uv, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
return Py_BuildValue("(LLLll)",
swap_total,
(swap_total - swap_free),
swap_free,
(long) uv.pgswapin * pagesize, // swap in
(long) uv.pgswapout * pagesize); // swap out
error:
free(swdev);
}
PyObject *
psutil_proc_num_fds(PyObject *self, PyObject *args) {
long pid;
int cnt;
struct kinfo_file *freep;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
errno = 0;
freep = kinfo_getfile(pid, &cnt);
if (freep == NULL) {
psutil_raise_for_pid(pid, "kinfo_getfile() failed");
return NULL;
}
free(freep);
return Py_BuildValue("i", cnt);
}
PyObject *
psutil_per_cpu_times(PyObject *self, PyObject *args) {
// XXX: why static?
static int maxcpus;
int mib[3];
int ncpu;
size_t len;
size_t size;
int i;
PyObject *py_cputime = NULL;
PyObject *py_retlist = PyList_New(0);
if (py_retlist == NULL)
return NULL;
// retrieve the number of cpus
mib[0] = CTL_HW;
mib[1] = HW_NCPU;
len = sizeof(ncpu);
if (sysctl(mib, 2, &ncpu, &len, NULL, 0) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
uint64_t cpu_time[CPUSTATES];
for (i = 0; i < ncpu; i++) {
// per-cpu info
mib[0] = CTL_KERN;
mib[1] = KERN_CP_TIME;
mib[2] = i;
size = sizeof(cpu_time);
if (sysctl(mib, 3, &cpu_time, &size, NULL, 0) == -1) {
warn("failed to get kern.cptime2");
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
py_cputime = Py_BuildValue(
"(ddddd)",
(double)cpu_time[CP_USER] / CLOCKS_PER_SEC,
(double)cpu_time[CP_NICE] / CLOCKS_PER_SEC,
(double)cpu_time[CP_SYS] / CLOCKS_PER_SEC,
(double)cpu_time[CP_IDLE] / CLOCKS_PER_SEC,
(double)cpu_time[CP_INTR] / CLOCKS_PER_SEC);
if (!py_cputime)
goto error;
if (PyList_Append(py_retlist, py_cputime))
goto error;
Py_DECREF(py_cputime);
}
return py_retlist;
error:
Py_XDECREF(py_cputime);
Py_DECREF(py_retlist);
return NULL;
}
PyObject *
psutil_disk_io_counters(PyObject *self, PyObject *args) {
int i, dk_ndrive, mib[3];
size_t len;
struct io_sysctl *stats;
PyObject *py_disk_info = NULL;
PyObject *py_retdict = PyDict_New();
if (py_retdict == NULL)
return NULL;
mib[0] = CTL_HW;
mib[1] = HW_IOSTATS;
mib[2] = sizeof(struct io_sysctl);
len = 0;
if (sysctl(mib, 3, NULL, &len, NULL, 0) < 0) {
warn("can't get HW_IOSTATS");
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
dk_ndrive = (int)(len / sizeof(struct io_sysctl));
stats = malloc(len);
if (stats == NULL) {
PyErr_NoMemory();
goto error;
}
if (sysctl(mib, 3, stats, &len, NULL, 0) < 0 ) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
for (i = 0; i < dk_ndrive; i++) {
py_disk_info = Py_BuildValue(
"(KKKK)",
stats[i].rxfer,
stats[i].wxfer,
stats[i].rbytes,
stats[i].wbytes
);
if (!py_disk_info)
goto error;
if (PyDict_SetItemString(py_retdict, stats[i].name, py_disk_info))
goto error;
Py_DECREF(py_disk_info);
}
free(stats);
return py_retdict;
error:
Py_XDECREF(py_disk_info);
Py_DECREF(py_retdict);
if (stats != NULL)
free(stats);
return NULL;
}
PyObject *
psutil_cpu_stats(PyObject *self, PyObject *args) {
size_t size;
struct uvmexp_sysctl uv;
int uvmexp_mib[] = {CTL_VM, VM_UVMEXP2};
size = sizeof(uv);
if (sysctl(uvmexp_mib, 2, &uv, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
return Py_BuildValue(
"IIIIIII",
uv.swtch, // ctx switches
uv.intrs, // interrupts - XXX always 0, will be determined via /proc
uv.softs, // soft interrupts
uv.syscalls, // syscalls - XXX always 0
uv.traps, // traps
uv.faults, // faults
uv.forks // forks
);
}
+28
View File
@@ -0,0 +1,28 @@
/*
* Copyright (c) 2009, Giampaolo Rodola', Landry Breuil.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
typedef struct kinfo_proc2 kinfo_proc;
int psutil_kinfo_proc(pid_t pid, kinfo_proc *proc);
struct kinfo_file * kinfo_getfile(pid_t pid, int* cnt);
int psutil_get_proc_list(kinfo_proc **procList, size_t *procCount);
char *psutil_get_cmd_args(pid_t pid, size_t *argsize);
//
PyObject *psutil_get_cmdline(pid_t pid);
PyObject *psutil_proc_threads(PyObject *self, PyObject *args);
PyObject *psutil_virtual_mem(PyObject *self, PyObject *args);
PyObject *psutil_swap_mem(PyObject *self, PyObject *args);
PyObject *psutil_proc_num_fds(PyObject *self, PyObject *args);
PyObject *psutil_proc_connections(PyObject *self, PyObject *args);
PyObject *psutil_per_cpu_times(PyObject *self, PyObject *args);
PyObject* psutil_disk_io_counters(PyObject* self, PyObject* args);
PyObject* psutil_proc_exe(PyObject* self, PyObject* args);
PyObject* psutil_proc_num_threads(PyObject* self, PyObject* args);
PyObject* psutil_cpu_stats(PyObject* self, PyObject* args);
+443
View File
@@ -0,0 +1,443 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'.
* Copyright (c) 2015, Ryo ONODERA.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/sysctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <string.h>
#include <sys/cdefs.h>
#include <arpa/inet.h>
#include <sys/queue.h>
#include <sys/un.h>
#include <sys/file.h>
// a signaler for connections without an actual status
int PSUTIL_CONN_NONE = 128;
// address family filter
enum af_filter {
INET,
INET4,
INET6,
TCP,
TCP4,
TCP6,
UDP,
UDP4,
UDP6,
UNIX,
ALL,
};
// kinfo_file results
struct kif {
SLIST_ENTRY(kif) kifs;
struct kinfo_file *kif;
};
// kinfo_file results list
SLIST_HEAD(kifhead, kif) kihead = SLIST_HEAD_INITIALIZER(kihead);
// kinfo_pcb results
struct kpcb {
SLIST_ENTRY(kpcb) kpcbs;
struct kinfo_pcb *kpcb;
};
// kinfo_pcb results list
SLIST_HEAD(kpcbhead, kpcb) kpcbhead = SLIST_HEAD_INITIALIZER(kpcbhead);
static void psutil_kiflist_init(void);
static void psutil_kiflist_clear(void);
static void psutil_kpcblist_init(void);
static void psutil_kpcblist_clear(void);
static int psutil_get_files(void);
static int psutil_get_sockets(const char *name);
static int psutil_get_info(int aff);
// Initialize kinfo_file results list.
static void
psutil_kiflist_init(void) {
SLIST_INIT(&kihead);
return;
}
// Clear kinfo_file results list.
static void
psutil_kiflist_clear(void) {
while (!SLIST_EMPTY(&kihead)) {
SLIST_REMOVE_HEAD(&kihead, kifs);
}
return;
}
// Initialize kinof_pcb result list.
static void
psutil_kpcblist_init(void) {
SLIST_INIT(&kpcbhead);
return;
}
// Clear kinof_pcb result list.
static void
psutil_kpcblist_clear(void) {
while (!SLIST_EMPTY(&kpcbhead)) {
SLIST_REMOVE_HEAD(&kpcbhead, kpcbs);
}
return;
}
// Get all open files including socket.
static int
psutil_get_files(void) {
size_t len;
int mib[6];
char *buf;
off_t offset;
int j;
mib[0] = CTL_KERN;
mib[1] = KERN_FILE2;
mib[2] = KERN_FILE_BYFILE;
mib[3] = 0;
mib[4] = sizeof(struct kinfo_file);
mib[5] = 0;
if (sysctl(mib, 6, NULL, &len, NULL, 0) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
offset = len % sizeof(off_t);
mib[5] = len / sizeof(struct kinfo_file);
if ((buf = malloc(len + offset)) == NULL) {
PyErr_NoMemory();
return -1;
}
if (sysctl(mib, 6, buf + offset, &len, NULL, 0) == -1) {
free(buf);
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
len /= sizeof(struct kinfo_file);
struct kinfo_file *ki = (struct kinfo_file *)(buf + offset);
for (j = 0; j < len; j++) {
struct kif *kif = malloc(sizeof(struct kif));
kif->kif = &ki[j];
SLIST_INSERT_HEAD(&kihead, kif, kifs);
}
/*
// debug
struct kif *k;
SLIST_FOREACH(k, &kihead, kifs) {
printf("%d\n", k->kif->ki_pid);
}
*/
return 0;
}
// Get open sockets.
static int
psutil_get_sockets(const char *name) {
size_t namelen;
int mib[8];
int ret, j;
struct kinfo_pcb *pcb;
size_t len;
memset(mib, 0, sizeof(mib));
if (sysctlnametomib(name, mib, &namelen) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
if (sysctl(mib, __arraycount(mib), NULL, &len, NULL, 0) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
if ((pcb = malloc(len)) == NULL) {
PyErr_NoMemory();
return -1;
}
memset(pcb, 0, len);
mib[6] = sizeof(*pcb);
mib[7] = len / sizeof(*pcb);
if (sysctl(mib, __arraycount(mib), pcb, &len, NULL, 0) == -1) {
free(pcb);
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
len /= sizeof(struct kinfo_pcb);
struct kinfo_pcb *kp = (struct kinfo_pcb *)pcb;
for (j = 0; j < len; j++) {
struct kpcb *kpcb = malloc(sizeof(struct kpcb));
kpcb->kpcb = &kp[j];
SLIST_INSERT_HEAD(&kpcbhead, kpcb, kpcbs);
}
/*
// debug
struct kif *k;
struct kpcb *k;
SLIST_FOREACH(k, &kpcbhead, kpcbs) {
printf("ki_type: %d\n", k->kpcb->ki_type);
printf("ki_family: %d\n", k->kpcb->ki_family);
}
*/
return 0;
}
// Collect open file and connections.
static int
psutil_get_info(int aff) {
switch (aff) {
case INET:
if (psutil_get_sockets("net.inet.tcp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet.udp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.tcp6.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.udp6.pcblist") != 0)
return -1;
break;
case INET4:
if (psutil_get_sockets("net.inet.tcp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet.udp.pcblist") != 0)
return -1;
break;
case INET6:
if (psutil_get_sockets("net.inet6.tcp6.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.udp6.pcblist") != 0)
return -1;
break;
case TCP:
if (psutil_get_sockets("net.inet.tcp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.tcp6.pcblist") != 0)
return -1;
break;
case TCP4:
if (psutil_get_sockets("net.inet.tcp.pcblist") != 0)
return -1;
break;
case TCP6:
if (psutil_get_sockets("net.inet6.tcp6.pcblist") != 0)
return -1;
break;
case UDP:
if (psutil_get_sockets("net.inet.udp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.udp6.pcblist") != 0)
return -1;
break;
case UDP4:
if (psutil_get_sockets("net.inet.udp.pcblist") != 0)
return -1;
break;
case UDP6:
if (psutil_get_sockets("net.inet6.udp6.pcblist") != 0)
return -1;
break;
case UNIX:
if (psutil_get_sockets("net.local.stream.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.local.seqpacket.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.local.dgram.pcblist") != 0)
return -1;
break;
case ALL:
if (psutil_get_sockets("net.inet.tcp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet.udp.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.tcp6.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.inet6.udp6.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.local.stream.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.local.seqpacket.pcblist") != 0)
return -1;
if (psutil_get_sockets("net.local.dgram.pcblist") != 0)
return -1;
break;
}
return 0;
}
/*
* Return system-wide connections (unless a pid != -1 is passed).
*/
PyObject *
psutil_net_connections(PyObject *self, PyObject *args) {
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
PyObject *py_laddr = NULL;
PyObject *py_raddr = NULL;
pid_t pid;
if (py_retlist == NULL)
return NULL;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
psutil_kiflist_init();
psutil_kpcblist_init();
if (psutil_get_files() != 0)
goto error;
if (psutil_get_info(ALL) != 0)
goto error;
struct kif *k;
SLIST_FOREACH(k, &kihead, kifs) {
struct kpcb *kp;
if ((pid != -1) && (k->kif->ki_pid != pid))
continue;
SLIST_FOREACH(kp, &kpcbhead, kpcbs) {
if (k->kif->ki_fdata != kp->kpcb->ki_sockaddr)
continue;
char laddr[PATH_MAX];
char raddr[PATH_MAX];
int32_t lport;
int32_t rport;
int32_t status;
// IPv4 or IPv6
if ((kp->kpcb->ki_family == AF_INET) ||
(kp->kpcb->ki_family == AF_INET6)) {
if (kp->kpcb->ki_family == AF_INET) {
// IPv4
struct sockaddr_in *sin_src =
(struct sockaddr_in *)&kp->kpcb->ki_src;
struct sockaddr_in *sin_dst =
(struct sockaddr_in *)&kp->kpcb->ki_dst;
// source addr and port
inet_ntop(AF_INET, &sin_src->sin_addr, laddr,
sizeof(laddr));
lport = ntohs(sin_src->sin_port);
// remote addr and port
inet_ntop(AF_INET, &sin_dst->sin_addr, raddr,
sizeof(raddr));
rport = ntohs(sin_dst->sin_port);
}
else {
// IPv6
struct sockaddr_in6 *sin6_src =
(struct sockaddr_in6 *)&kp->kpcb->ki_src;
struct sockaddr_in6 *sin6_dst =
(struct sockaddr_in6 *)&kp->kpcb->ki_dst;
// local addr and port
inet_ntop(AF_INET6, &sin6_src->sin6_addr, laddr,
sizeof(laddr));
lport = ntohs(sin6_src->sin6_port);
// remote addr and port
inet_ntop(AF_INET6, &sin6_dst->sin6_addr, raddr,
sizeof(raddr));
rport = ntohs(sin6_dst->sin6_port);
}
// status
if (kp->kpcb->ki_type == SOCK_STREAM)
status = kp->kpcb->ki_tstate;
else
status = PSUTIL_CONN_NONE;
// build addr tuple
py_laddr = Py_BuildValue("(si)", laddr, lport);
if (! py_laddr)
goto error;
if (rport != 0)
py_raddr = Py_BuildValue("(si)", raddr, rport);
else
py_raddr = Py_BuildValue("()");
if (! py_raddr)
goto error;
}
else if (kp->kpcb->ki_family == AF_UNIX) {
// UNIX sockets
struct sockaddr_un *sun_src =
(struct sockaddr_un *)&kp->kpcb->ki_src;
struct sockaddr_un *sun_dst =
(struct sockaddr_un *)&kp->kpcb->ki_dst;
strcpy(laddr, sun_src->sun_path);
strcpy(raddr, sun_dst->sun_path);
status = PSUTIL_CONN_NONE;
// TODO: handle unicode
py_laddr = Py_BuildValue("s", laddr);
if (! py_laddr)
goto error;
// TODO: handle unicode
py_raddr = Py_BuildValue("s", raddr);
if (! py_raddr)
goto error;
}
// append tuple to list
py_tuple = Py_BuildValue(
"(iiiNNii)",
k->kif->ki_fd,
kp->kpcb->ki_family,
kp->kpcb->ki_type,
py_laddr,
py_raddr,
status,
k->kif->ki_pid);
if (! py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
}
psutil_kiflist_clear();
psutil_kpcblist_clear();
return py_retlist;
error:
Py_XDECREF(py_tuple);
Py_XDECREF(py_laddr);
Py_XDECREF(py_raddr);
return 0;
}
+10
View File
@@ -0,0 +1,10 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'.
* Copyright (c) 2015, Ryo ONODERA.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
PyObject *psutil_proc_connections(PyObject *, PyObject *);
PyObject *psutil_net_connections(PyObject *, PyObject *);
+788
View File
@@ -0,0 +1,788 @@
/*
* Copyright (c) 2009, Giampaolo Rodola', Landry Breuil.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Platform-specific module methods for OpenBSD.
*/
#include <Python.h>
#include <assert.h>
#include <errno.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/param.h>
#include <sys/sysctl.h>
#include <sys/user.h>
#include <sys/proc.h>
#include <sys/mount.h> // for VFS_*
#include <sys/swap.h> // for swap_mem
#include <sys/vmmeter.h> // for vmtotal struct
#include <signal.h>
#include <kvm.h>
// connection stuff
#include <netdb.h> // for NI_MAXHOST
#include <sys/socket.h>
#include <sys/sched.h> // for CPUSTATES & CP_*
#define _KERNEL // for DTYPE_*
#include <sys/file.h>
#undef _KERNEL
#include <sys/disk.h> // struct diskstats
#include <arpa/inet.h> // for inet_ntoa()
#include <err.h> // for warn() & err()
#include "../../_psutil_common.h"
#define PSUTIL_KPT2DOUBLE(t) (t ## _sec + t ## _usec / 1000000.0)
// #define PSUTIL_TV2DOUBLE(t) ((t).tv_sec + (t).tv_usec / 1000000.0)
// a signaler for connections without an actual status
int PSUTIL_CONN_NONE = 128;
// ============================================================================
// Utility functions
// ============================================================================
int
psutil_kinfo_proc(pid_t pid, struct kinfo_proc *proc) {
// Fills a kinfo_proc struct based on process pid.
int ret;
int mib[6];
size_t size = sizeof(struct kinfo_proc);
mib[0] = CTL_KERN;
mib[1] = KERN_PROC;
mib[2] = KERN_PROC_PID;
mib[3] = pid;
mib[4] = size;
mib[5] = 1;
ret = sysctl((int*)mib, 6, proc, &size, NULL, 0);
if (ret == -1) {
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
// sysctl stores 0 in the size if we can't find the process information.
if (size == 0) {
NoSuchProcess();
return -1;
}
return 0;
}
struct kinfo_file *
kinfo_getfile(long pid, int* cnt) {
// Mimic's FreeBSD kinfo_file call, taking a pid and a ptr to an
// int as arg and returns an array with cnt struct kinfo_file.
int mib[6];
size_t len;
struct kinfo_file* kf;
mib[0] = CTL_KERN;
mib[1] = KERN_FILE;
mib[2] = KERN_FILE_BYPID;
mib[3] = (int) pid;
mib[4] = sizeof(struct kinfo_file);
mib[5] = 0;
/* get the size of what would be returned */
if (sysctl(mib, 6, NULL, &len, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
if ((kf = malloc(len)) == NULL) {
PyErr_NoMemory();
return NULL;
}
mib[5] = (int)(len / sizeof(struct kinfo_file));
if (sysctl(mib, 6, kf, &len, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
*cnt = (int)(len / sizeof(struct kinfo_file));
return kf;
}
// ============================================================================
// APIS
// ============================================================================
int
psutil_get_proc_list(struct kinfo_proc **procList, size_t *procCount) {
// Returns a list of all BSD processes on the system. This routine
// allocates the list and puts it in *procList and a count of the
// number of entries in *procCount. You are responsible for freeing
// this list (use "free" from System framework).
// On success, the function returns 0.
// On error, the function returns a BSD errno value.
struct kinfo_proc *result;
// Declaring name as const requires us to cast it when passing it to
// sysctl because the prototype doesn't include the const modifier.
char errbuf[_POSIX2_LINE_MAX];
int cnt;
kvm_t *kd;
assert(procList != NULL);
assert(*procList == NULL);
assert(procCount != NULL);
kd = kvm_openfiles(NULL, NULL, NULL, KVM_NO_FILES, errbuf);
if (kd == NULL) {
return errno;
}
result = kvm_getprocs(kd, KERN_PROC_ALL, 0, sizeof(struct kinfo_proc), &cnt);
if (result == NULL) {
kvm_close(kd);
err(1, NULL);
return errno;
}
*procCount = (size_t)cnt;
size_t mlen = cnt * sizeof(struct kinfo_proc);
if ((*procList = malloc(mlen)) == NULL) {
kvm_close(kd);
err(1, NULL);
return errno;
}
memcpy(*procList, result, mlen);
assert(*procList != NULL);
kvm_close(kd);
return 0;
}
char **
_psutil_get_argv(long pid) {
static char **argv;
int argv_mib[] = {CTL_KERN, KERN_PROC_ARGS, pid, KERN_PROC_ARGV};
size_t argv_size = 128;
/* Loop and reallocate until we have enough space to fit argv. */
for (;; argv_size *= 2) {
if ((argv = realloc(argv, argv_size)) == NULL)
err(1, NULL);
if (sysctl(argv_mib, 4, argv, &argv_size, NULL, 0) == 0)
return argv;
if (errno == ESRCH) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
if (errno != ENOMEM)
err(1, NULL);
}
}
// returns the command line as a python list object
PyObject *
psutil_get_cmdline(long pid) {
static char **argv;
char **p;
PyObject *py_arg = NULL;
PyObject *py_retlist = Py_BuildValue("[]");
if (!py_retlist)
return NULL;
if (pid < 0)
return py_retlist;
if ((argv = _psutil_get_argv(pid)) == NULL)
goto error;
for (p = argv; *p != NULL; p++) {
#if PY_MAJOR_VERSION >= 3
py_arg = PyUnicode_DecodeFSDefault(*p);
#else
py_arg = Py_BuildValue("s", *p);
#endif
if (!py_arg)
goto error;
if (PyList_Append(py_retlist, py_arg))
goto error;
Py_DECREF(py_arg);
}
return py_retlist;
error:
Py_XDECREF(py_arg);
Py_DECREF(py_retlist);
return NULL;
}
PyObject *
psutil_proc_threads(PyObject *self, PyObject *args) {
// OpenBSD reference:
// https://github.com/janmojzis/pstree/blob/master/proc_kvm.c
// Note: this requires root access, else it will fail trying
// to access /dev/kmem.
long pid;
kvm_t *kd = NULL;
int nentries, i;
char errbuf[4096];
struct kinfo_proc *kp;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
if (py_retlist == NULL)
return NULL;
if (! PyArg_ParseTuple(args, "l", &pid))
goto error;
kd = kvm_openfiles(0, 0, 0, O_RDONLY, errbuf);
if (! kd) {
if (strstr(errbuf, "Permission denied") != NULL)
AccessDenied();
else
PyErr_Format(PyExc_RuntimeError, "kvm_openfiles() syscall failed");
goto error;
}
kp = kvm_getprocs(
kd, KERN_PROC_PID | KERN_PROC_SHOW_THREADS | KERN_PROC_KTHREAD, pid,
sizeof(*kp), &nentries);
if (! kp) {
if (strstr(errbuf, "Permission denied") != NULL)
AccessDenied();
else
PyErr_Format(PyExc_RuntimeError, "kvm_getprocs() syscall failed");
goto error;
}
for (i = 0; i < nentries; i++) {
if (kp[i].p_tid < 0)
continue;
if (kp[i].p_pid == pid) {
py_tuple = Py_BuildValue(
"Idd",
kp[i].p_tid,
PSUTIL_KPT2DOUBLE(kp[i].p_uutime),
PSUTIL_KPT2DOUBLE(kp[i].p_ustime));
if (py_tuple == NULL)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
}
kvm_close(kd);
return py_retlist;
error:
Py_XDECREF(py_tuple);
Py_DECREF(py_retlist);
if (kd != NULL)
kvm_close(kd);
return NULL;
}
PyObject *
psutil_virtual_mem(PyObject *self, PyObject *args) {
int64_t total_physmem;
int uvmexp_mib[] = {CTL_VM, VM_UVMEXP};
int bcstats_mib[] = {CTL_VFS, VFS_GENERIC, VFS_BCACHESTAT};
int physmem_mib[] = {CTL_HW, HW_PHYSMEM64};
int vmmeter_mib[] = {CTL_VM, VM_METER};
size_t size;
struct uvmexp uvmexp;
struct bcachestats bcstats;
struct vmtotal vmdata;
long pagesize = getpagesize();
size = sizeof(total_physmem);
if (sysctl(physmem_mib, 2, &total_physmem, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
size = sizeof(uvmexp);
if (sysctl(uvmexp_mib, 2, &uvmexp, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
size = sizeof(bcstats);
if (sysctl(bcstats_mib, 3, &bcstats, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
size = sizeof(vmdata);
if (sysctl(vmmeter_mib, 2, &vmdata, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
return Py_BuildValue("KKKKKKKK",
// Note: many programs calculate total memory as
// "uvmexp.npages * pagesize" but this is incorrect and does not
// match "sysctl | grep hw.physmem".
(unsigned long long) total_physmem,
(unsigned long long) uvmexp.free * pagesize,
(unsigned long long) uvmexp.active * pagesize,
(unsigned long long) uvmexp.inactive * pagesize,
(unsigned long long) uvmexp.wired * pagesize,
// this is how "top" determines it
(unsigned long long) bcstats.numbufpages * pagesize, // cached
(unsigned long long) 0, // buffers
(unsigned long long) vmdata.t_vmshr + vmdata.t_rmshr // shared
);
}
PyObject *
psutil_swap_mem(PyObject *self, PyObject *args) {
uint64_t swap_total, swap_free;
struct swapent *swdev;
int nswap, i;
if ((nswap = swapctl(SWAP_NSWAP, 0, 0)) == 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
if ((swdev = calloc(nswap, sizeof(*swdev))) == NULL) {
PyErr_NoMemory();
return NULL;
}
if (swapctl(SWAP_STATS, swdev, nswap) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
// Total things up.
swap_total = swap_free = 0;
for (i = 0; i < nswap; i++) {
if (swdev[i].se_flags & SWF_ENABLE) {
swap_free += (swdev[i].se_nblks - swdev[i].se_inuse);
swap_total += swdev[i].se_nblks;
}
}
free(swdev);
return Py_BuildValue("(LLLII)",
swap_total * DEV_BSIZE,
(swap_total - swap_free) * DEV_BSIZE,
swap_free * DEV_BSIZE,
// swap in / swap out is not supported as the
// swapent struct does not provide any info
// about it.
0, 0);
error:
free(swdev);
return NULL;
}
PyObject *
psutil_proc_num_fds(PyObject *self, PyObject *args) {
long pid;
int cnt;
struct kinfo_file *freep;
struct kinfo_proc kipp;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
if (psutil_kinfo_proc(pid, &kipp) == -1)
return NULL;
errno = 0;
freep = kinfo_getfile(pid, &cnt);
if (freep == NULL) {
psutil_raise_for_pid(pid, "kinfo_getfile() failed");
return NULL;
}
free(freep);
return Py_BuildValue("i", cnt);
}
PyObject *
psutil_proc_cwd(PyObject *self, PyObject *args) {
// Reference:
// http://anoncvs.spacehopper.org/openbsd-src/tree/bin/ps/print.c#n179
long pid;
struct kinfo_proc kp;
char path[MAXPATHLEN];
size_t pathlen = sizeof path;
if (! PyArg_ParseTuple(args, "l", &pid))
return NULL;
if (psutil_kinfo_proc(pid, &kp) == -1)
return NULL;
int name[] = { CTL_KERN, KERN_PROC_CWD, pid };
if (sysctl(name, 3, path, &pathlen, NULL, 0) != 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
#if PY_MAJOR_VERSION >= 3
return PyUnicode_DecodeFSDefault(path);
#else
return Py_BuildValue("s", path);
#endif
}
// see sys/kern/kern_sysctl.c lines 1100 and
// usr.bin/fstat/fstat.c print_inet_details()
static char *
psutil_convert_ipv4(int family, uint32_t addr[4]) {
struct in_addr a;
memcpy(&a, addr, sizeof(a));
return inet_ntoa(a);
}
static char *
psutil_inet6_addrstr(struct in6_addr *p)
{
struct sockaddr_in6 sin6;
static char hbuf[NI_MAXHOST];
const int niflags = NI_NUMERICHOST;
memset(&sin6, 0, sizeof(sin6));
sin6.sin6_family = AF_INET6;
sin6.sin6_len = sizeof(struct sockaddr_in6);
sin6.sin6_addr = *p;
if (IN6_IS_ADDR_LINKLOCAL(p) &&
*(u_int16_t *)&sin6.sin6_addr.s6_addr[2] != 0) {
sin6.sin6_scope_id =
ntohs(*(u_int16_t *)&sin6.sin6_addr.s6_addr[2]);
sin6.sin6_addr.s6_addr[2] = sin6.sin6_addr.s6_addr[3] = 0;
}
if (getnameinfo((struct sockaddr *)&sin6, sin6.sin6_len,
hbuf, sizeof(hbuf), NULL, 0, niflags))
return "invalid";
return hbuf;
}
PyObject *
psutil_proc_connections(PyObject *self, PyObject *args) {
long pid;
int i, cnt;
struct kinfo_file *freep = NULL;
struct kinfo_file *kif;
char *tcplist = NULL;
PyObject *py_retlist = PyList_New(0);
PyObject *py_tuple = NULL;
PyObject *py_laddr = NULL;
PyObject *py_raddr = NULL;
PyObject *py_af_filter = NULL;
PyObject *py_type_filter = NULL;
PyObject *py_family = NULL;
PyObject *_type = NULL;
if (py_retlist == NULL)
return NULL;
if (! PyArg_ParseTuple(args, "lOO", &pid, &py_af_filter, &py_type_filter))
goto error;
if (!PySequence_Check(py_af_filter) || !PySequence_Check(py_type_filter)) {
PyErr_SetString(PyExc_TypeError, "arg 2 or 3 is not a sequence");
goto error;
}
errno = 0;
freep = kinfo_getfile(pid, &cnt);
if (freep == NULL) {
psutil_raise_for_pid(pid, "kinfo_getfile() failed");
goto error;
}
for (i = 0; i < cnt; i++) {
int state;
int lport;
int rport;
char addrbuf[NI_MAXHOST + 2];
int inseq;
struct in6_addr laddr6;
py_tuple = NULL;
py_laddr = NULL;
py_raddr = NULL;
kif = &freep[i];
if (kif->f_type == DTYPE_SOCKET) {
// apply filters
py_family = PyLong_FromLong((long)kif->so_family);
inseq = PySequence_Contains(py_af_filter, py_family);
Py_DECREF(py_family);
if (inseq == 0)
continue;
_type = PyLong_FromLong((long)kif->so_type);
inseq = PySequence_Contains(py_type_filter, _type);
Py_DECREF(_type);
if (inseq == 0)
continue;
// IPv4 / IPv6 socket
if ((kif->so_family == AF_INET) || (kif->so_family == AF_INET6)) {
// fill status
if (kif->so_type == SOCK_STREAM)
state = kif->t_state;
else
state = PSUTIL_CONN_NONE;
// ports
lport = ntohs(kif->inp_lport);
rport = ntohs(kif->inp_fport);
// local address, IPv4
if (kif->so_family == AF_INET) {
py_laddr = Py_BuildValue(
"(si)",
psutil_convert_ipv4(kif->so_family, kif->inp_laddru),
lport);
if (!py_laddr)
goto error;
}
else {
// local address, IPv6
memcpy(&laddr6, kif->inp_laddru, sizeof(laddr6));
snprintf(addrbuf, sizeof(addrbuf), "%s",
psutil_inet6_addrstr(&laddr6));
py_laddr = Py_BuildValue("(si)", addrbuf, lport);
if (!py_laddr)
goto error;
}
if (rport != 0) {
// remote address, IPv4
if (kif->so_family == AF_INET) {
py_raddr = Py_BuildValue(
"(si)",
psutil_convert_ipv4(
kif->so_family, kif->inp_faddru),
rport);
}
else {
// remote address, IPv6
memcpy(&laddr6, kif->inp_faddru, sizeof(laddr6));
snprintf(addrbuf, sizeof(addrbuf), "%s",
psutil_inet6_addrstr(&laddr6));
py_raddr = Py_BuildValue("(si)", addrbuf, rport);
if (!py_raddr)
goto error;
}
}
else {
py_raddr = Py_BuildValue("()");
}
if (!py_raddr)
goto error;
py_tuple = Py_BuildValue(
"(iiiNNi)",
kif->fd_fd,
kif->so_family,
kif->so_type,
py_laddr,
py_raddr,
state);
if (!py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
}
// UNIX socket
else if (kif->so_family == AF_UNIX) {
py_tuple = Py_BuildValue(
"(iiisOi)",
kif->fd_fd,
kif->so_family,
kif->so_type,
kif->unp_path,
Py_None,
PSUTIL_CONN_NONE);
if (!py_tuple)
goto error;
if (PyList_Append(py_retlist, py_tuple))
goto error;
Py_DECREF(py_tuple);
Py_INCREF(Py_None);
}
}
}
free(freep);
free(tcplist);
return py_retlist;
error:
Py_XDECREF(py_tuple);
Py_XDECREF(py_laddr);
Py_XDECREF(py_raddr);
Py_DECREF(py_retlist);
if (freep != NULL)
free(freep);
if (tcplist != NULL)
free(tcplist);
return NULL;
}
PyObject *
psutil_per_cpu_times(PyObject *self, PyObject *args) {
int mib[3];
int ncpu;
size_t len;
size_t size;
int i;
PyObject *py_retlist = PyList_New(0);
PyObject *py_cputime = NULL;
if (py_retlist == NULL)
return NULL;
// retrieve the number of cpus
mib[0] = CTL_HW;
mib[1] = HW_NCPU;
len = sizeof(ncpu);
if (sysctl(mib, 2, &ncpu, &len, NULL, 0) == -1) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
uint64_t cpu_time[CPUSTATES];
for (i = 0; i < ncpu; i++) {
// per-cpu info
mib[0] = CTL_KERN;
mib[1] = KERN_CPTIME2;
mib[2] = i;
size = sizeof(cpu_time);
if (sysctl(mib, 3, &cpu_time, &size, NULL, 0) == -1) {
warn("failed to get kern.cptime2");
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
py_cputime = Py_BuildValue(
"(ddddd)",
(double)cpu_time[CP_USER] / CLOCKS_PER_SEC,
(double)cpu_time[CP_NICE] / CLOCKS_PER_SEC,
(double)cpu_time[CP_SYS] / CLOCKS_PER_SEC,
(double)cpu_time[CP_IDLE] / CLOCKS_PER_SEC,
(double)cpu_time[CP_INTR] / CLOCKS_PER_SEC);
if (!py_cputime)
goto error;
if (PyList_Append(py_retlist, py_cputime))
goto error;
Py_DECREF(py_cputime);
}
return py_retlist;
error:
Py_XDECREF(py_cputime);
Py_DECREF(py_retlist);
return NULL;
}
PyObject *
psutil_disk_io_counters(PyObject *self, PyObject *args) {
int i, dk_ndrive, mib[3];
size_t len;
struct diskstats *stats;
PyObject *py_retdict = PyDict_New();
PyObject *py_disk_info = NULL;
if (py_retdict == NULL)
return NULL;
mib[0] = CTL_HW;
mib[1] = HW_DISKSTATS;
len = 0;
if (sysctl(mib, 2, NULL, &len, NULL, 0) < 0) {
warn("can't get hw.diskstats size");
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
dk_ndrive = (int)(len / sizeof(struct diskstats));
stats = malloc(len);
if (stats == NULL) {
warn("can't malloc");
PyErr_NoMemory();
goto error;
}
if (sysctl(mib, 2, stats, &len, NULL, 0) < 0 ) {
warn("could not read hw.diskstats");
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
for (i = 0; i < dk_ndrive; i++) {
py_disk_info = Py_BuildValue(
"(KKKK)",
stats[i].ds_rxfer, // num reads
stats[i].ds_wxfer, // num writes
stats[i].ds_rbytes, // read bytes
stats[i].ds_wbytes // write bytes
);
if (!py_disk_info)
goto error;
if (PyDict_SetItemString(py_retdict, stats[i].ds_name, py_disk_info))
goto error;
Py_DECREF(py_disk_info);
}
free(stats);
return py_retdict;
error:
Py_XDECREF(py_disk_info);
Py_DECREF(py_retdict);
if (stats != NULL)
free(stats);
return NULL;
}
PyObject *
psutil_cpu_stats(PyObject *self, PyObject *args) {
size_t size;
struct uvmexp uv;
int uvmexp_mib[] = {CTL_VM, VM_UVMEXP};
size = sizeof(uv);
if (sysctl(uvmexp_mib, 2, &uv, &size, NULL, 0) < 0) {
PyErr_SetFromErrno(PyExc_OSError);
return NULL;
}
return Py_BuildValue(
"IIIIIII",
uv.swtch, // ctx switches
uv.intrs, // interrupts - XXX always 0, will be determined via /proc
uv.softs, // soft interrupts
uv.syscalls, // syscalls - XXX always 0
uv.traps, // traps
uv.faults, // faults
uv.forks // forks
);
}
+27
View File
@@ -0,0 +1,27 @@
/*
* Copyright (c) 2009, Giampaolo Rodola', Landry Breuil.
* All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
typedef struct kinfo_proc kinfo_proc;
int psutil_kinfo_proc(pid_t pid, struct kinfo_proc *proc);
struct kinfo_file * kinfo_getfile(long pid, int* cnt);
int psutil_get_proc_list(struct kinfo_proc **procList, size_t *procCount);
char **_psutil_get_argv(long pid);
PyObject * psutil_get_cmdline(long pid);
//
PyObject *psutil_proc_threads(PyObject *self, PyObject *args);
PyObject *psutil_virtual_mem(PyObject *self, PyObject *args);
PyObject *psutil_swap_mem(PyObject *self, PyObject *args);
PyObject *psutil_proc_num_fds(PyObject *self, PyObject *args);
PyObject *psutil_proc_cwd(PyObject *self, PyObject *args);
PyObject *psutil_proc_connections(PyObject *self, PyObject *args);
PyObject *psutil_per_cpu_times(PyObject *self, PyObject *args);
PyObject* psutil_disk_io_counters(PyObject* self, PyObject* args);
PyObject* psutil_cpu_stats(PyObject* self, PyObject* args);
+372
View File
@@ -0,0 +1,372 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Helper functions related to fetching process information.
* Used by _psutil_osx module methods.
*/
#include <Python.h>
#include <assert.h>
#include <errno.h>
#include <limits.h> // for INT_MAX
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <sys/sysctl.h>
#include <libproc.h>
#include "process_info.h"
#include "../../_psutil_common.h"
/*
* Returns a list of all BSD processes on the system. This routine
* allocates the list and puts it in *procList and a count of the
* number of entries in *procCount. You are responsible for freeing
* this list (use "free" from System framework).
* On success, the function returns 0.
* On error, the function returns a BSD errno value.
*/
int
psutil_get_proc_list(kinfo_proc **procList, size_t *procCount) {
int mib3[3] = { CTL_KERN, KERN_PROC, KERN_PROC_ALL };
size_t size, size2;
void *ptr;
int err;
int lim = 8; // some limit
assert( procList != NULL);
assert(*procList == NULL);
assert(procCount != NULL);
*procCount = 0;
/*
* We start by calling sysctl with ptr == NULL and size == 0.
* That will succeed, and set size to the appropriate length.
* We then allocate a buffer of at least that size and call
* sysctl with that buffer. If that succeeds, we're done.
* If that call fails with ENOMEM, we throw the buffer away
* and try again.
* Note that the loop calls sysctl with NULL again. This is
* is necessary because the ENOMEM failure case sets size to
* the amount of data returned, not the amount of data that
* could have been returned.
*/
while (lim-- > 0) {
size = 0;
if (sysctl((int *)mib3, 3, NULL, &size, NULL, 0) == -1)
return errno;
size2 = size + (size >> 3); // add some
if (size2 > size) {
ptr = malloc(size2);
if (ptr == NULL)
ptr = malloc(size);
else
size = size2;
}
else {
ptr = malloc(size);
}
if (ptr == NULL)
return ENOMEM;
if (sysctl((int *)mib3, 3, ptr, &size, NULL, 0) == -1) {
err = errno;
free(ptr);
if (err != ENOMEM)
return err;
}
else {
*procList = (kinfo_proc *)ptr;
*procCount = size / sizeof(kinfo_proc);
return 0;
}
}
return ENOMEM;
}
// Read the maximum argument size for processes
int
psutil_get_argmax() {
int argmax;
int mib[] = { CTL_KERN, KERN_ARGMAX };
size_t size = sizeof(argmax);
if (sysctl(mib, 2, &argmax, &size, NULL, 0) == 0)
return argmax;
return 0;
}
// return process args as a python list
PyObject *
psutil_get_cmdline(long pid) {
int mib[3];
int nargs;
size_t len;
char *procargs = NULL;
char *arg_ptr;
char *arg_end;
char *curr_arg;
size_t argmax;
PyObject *py_arg = NULL;
PyObject *py_retlist = NULL;
// special case for PID 0 (kernel_task) where cmdline cannot be fetched
if (pid == 0)
return Py_BuildValue("[]");
// read argmax and allocate memory for argument space.
argmax = psutil_get_argmax();
if (! argmax) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
procargs = (char *)malloc(argmax);
if (NULL == procargs) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
// read argument space
mib[0] = CTL_KERN;
mib[1] = KERN_PROCARGS2;
mib[2] = (pid_t)pid;
if (sysctl(mib, 3, procargs, &argmax, NULL, 0) < 0) {
if (EINVAL == errno) {
// EINVAL == access denied OR nonexistent PID
if (psutil_pid_exists(pid))
AccessDenied();
else
NoSuchProcess();
}
goto error;
}
arg_end = &procargs[argmax];
// copy the number of arguments to nargs
memcpy(&nargs, procargs, sizeof(nargs));
arg_ptr = procargs + sizeof(nargs);
len = strlen(arg_ptr);
arg_ptr += len + 1;
if (arg_ptr == arg_end) {
free(procargs);
return Py_BuildValue("[]");
}
// skip ahead to the first argument
for (; arg_ptr < arg_end; arg_ptr++) {
if (*arg_ptr != '\0')
break;
}
// iterate through arguments
curr_arg = arg_ptr;
py_retlist = Py_BuildValue("[]");
if (!py_retlist)
goto error;
while (arg_ptr < arg_end && nargs > 0) {
if (*arg_ptr++ == '\0') {
#if PY_MAJOR_VERSION >= 3
py_arg = PyUnicode_DecodeFSDefault(curr_arg);
#else
py_arg = Py_BuildValue("s", curr_arg);
#endif
if (!py_arg)
goto error;
if (PyList_Append(py_retlist, py_arg))
goto error;
Py_DECREF(py_arg);
// iterate to next arg and decrement # of args
curr_arg = arg_ptr;
nargs--;
}
}
free(procargs);
return py_retlist;
error:
Py_XDECREF(py_arg);
Py_XDECREF(py_retlist);
if (procargs != NULL)
free(procargs);
return NULL;
}
// return process environment as a python string
PyObject *
psutil_get_environ(long pid) {
int mib[3];
int nargs;
char *procargs = NULL;
char *procenv = NULL;
char *arg_ptr;
char *arg_end;
char *env_start;
size_t argmax;
PyObject *py_ret = NULL;
// special case for PID 0 (kernel_task) where cmdline cannot be fetched
if (pid == 0)
goto empty;
// read argmax and allocate memory for argument space.
argmax = psutil_get_argmax();
if (! argmax) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
procargs = (char *)malloc(argmax);
if (NULL == procargs) {
PyErr_SetFromErrno(PyExc_OSError);
goto error;
}
// read argument space
mib[0] = CTL_KERN;
mib[1] = KERN_PROCARGS2;
mib[2] = (pid_t)pid;
if (sysctl(mib, 3, procargs, &argmax, NULL, 0) < 0) {
if (EINVAL == errno) {
// EINVAL == access denied OR nonexistent PID
if (psutil_pid_exists(pid))
AccessDenied();
else
NoSuchProcess();
}
goto error;
}
arg_end = &procargs[argmax];
// copy the number of arguments to nargs
memcpy(&nargs, procargs, sizeof(nargs));
// skip executable path
arg_ptr = procargs + sizeof(nargs);
arg_ptr = memchr(arg_ptr, '\0', arg_end - arg_ptr);
if (arg_ptr == NULL || arg_ptr == arg_end)
goto empty;
// skip ahead to the first argument
for (; arg_ptr < arg_end; arg_ptr++) {
if (*arg_ptr != '\0')
break;
}
// iterate through arguments
while (arg_ptr < arg_end && nargs > 0) {
if (*arg_ptr++ == '\0')
nargs--;
}
// build an environment variable block
env_start = arg_ptr;
procenv = calloc(1, arg_end - arg_ptr);
if (procenv == NULL) {
PyErr_NoMemory();
goto error;
}
while (*arg_ptr != '\0' && arg_ptr < arg_end) {
char *s = memchr(arg_ptr + 1, '\0', arg_end - arg_ptr);
if (s == NULL)
break;
memcpy(procenv + (arg_ptr - env_start), arg_ptr, s - arg_ptr);
arg_ptr = s + 1;
}
#if PY_MAJOR_VERSION >= 3
py_ret = PyUnicode_DecodeFSDefaultAndSize(
procenv, arg_ptr - env_start + 1);
#else
py_ret = PyString_FromStringAndSize(procenv, arg_ptr - env_start + 1);
#endif
if (!py_ret) {
// XXX: don't want to free() this as per:
// https://github.com/giampaolo/psutil/issues/926
// It sucks but not sure what else to do.
procargs = NULL;
goto error;
}
free(procargs);
free(procenv);
return py_ret;
empty:
if (procargs != NULL)
free(procargs);
return Py_BuildValue("s", "");
error:
Py_XDECREF(py_ret);
if (procargs != NULL)
free(procargs);
if (procenv != NULL)
free(procargs);
return NULL;
}
int
psutil_get_kinfo_proc(long pid, struct kinfo_proc *kp) {
int mib[4];
size_t len;
mib[0] = CTL_KERN;
mib[1] = KERN_PROC;
mib[2] = KERN_PROC_PID;
mib[3] = (pid_t)pid;
// fetch the info with sysctl()
len = sizeof(struct kinfo_proc);
// now read the data from sysctl
if (sysctl(mib, 4, kp, &len, NULL, 0) == -1) {
// raise an exception and throw errno as the error
PyErr_SetFromErrno(PyExc_OSError);
return -1;
}
// sysctl succeeds but len is zero, happens when process has gone away
if (len == 0) {
NoSuchProcess();
return -1;
}
return 0;
}
/*
* A wrapper around proc_pidinfo().
* Returns 0 on failure (and Python exception gets already set).
*/
int
psutil_proc_pidinfo(long pid, int flavor, uint64_t arg, void *pti, int size) {
errno = 0;
int ret = proc_pidinfo((int)pid, flavor, arg, pti, size);
if ((ret <= 0) || ((unsigned long)ret < sizeof(pti))) {
psutil_raise_for_pid(pid, "proc_pidinfo() syscall failed");
return 0;
}
return ret;
}
+17
View File
@@ -0,0 +1,17 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#include <Python.h>
typedef struct kinfo_proc kinfo_proc;
int psutil_get_argmax(void);
int psutil_get_kinfo_proc(long pid, struct kinfo_proc *kp);
int psutil_get_proc_list(kinfo_proc **procList, size_t *procCount);
int psutil_proc_pidinfo(
long pid, int flavor, uint64_t arg, void *pti, int size);
PyObject* psutil_get_cmdline(long pid);
PyObject* psutil_get_environ(long pid);
+124
View File
@@ -0,0 +1,124 @@
/* Refrences:
* https://lists.samba.org/archive/samba-technical/2009-February/063079.html
* http://stackoverflow.com/questions/4139405/#4139811
* https://code.google.com/p/openpgm/source/browse/trunk/openpgm/pgm/getifaddrs.c
*/
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <net/if.h>
#include <netinet/in.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/sockio.h>
#include "ifaddrs.h"
#define MAX(x,y) ((x)>(y)?(x):(y))
#define SIZE(p) MAX((p).ss_len,sizeof(p))
static struct sockaddr *
sa_dup (struct sockaddr *sa1)
{
struct sockaddr *sa2;
size_t sz = sizeof(sa1);
sa2 = (struct sockaddr *) calloc(1,sz);
memcpy(sa2,sa1,sz);
return(sa2);
}
void freeifaddrs (struct ifaddrs *ifp)
{
if (NULL == ifp) return;
free(ifp->ifa_name);
free(ifp->ifa_addr);
free(ifp->ifa_netmask);
free(ifp->ifa_dstaddr);
freeifaddrs(ifp->ifa_next);
free(ifp);
}
int getifaddrs (struct ifaddrs **ifap)
{
int sd = -1;
char *ccp, *ecp;
struct lifconf ifc;
struct lifreq *ifr;
struct lifnum lifn;
struct ifaddrs *cifa = NULL; /* current */
struct ifaddrs *pifa = NULL; /* previous */
const size_t IFREQSZ = sizeof(struct lifreq);
sd = socket(AF_INET, SOCK_STREAM, 0);
if (sd < 0)
goto error;
ifc.lifc_buf = NULL;
*ifap = NULL;
/* find how much memory to allocate for the SIOCGLIFCONF call */
lifn.lifn_family = AF_UNSPEC;
lifn.lifn_flags = 0;
if (ioctl(sd, SIOCGLIFNUM, &lifn) < 0)
goto error;
/* Sun and Apple code likes to pad the interface count here in case interfaces
* are coming up between calls */
lifn.lifn_count += 4;
ifc.lifc_family = AF_UNSPEC;
ifc.lifc_len = lifn.lifn_count * sizeof(struct lifreq);
ifc.lifc_buf = calloc(1, ifc.lifc_len);
if (ioctl(sd, SIOCGLIFCONF, &ifc) < 0)
goto error;
ccp = (char *)ifc.lifc_req;
ecp = ccp + ifc.lifc_len;
while (ccp < ecp) {
ifr = (struct lifreq *) ccp;
cifa = (struct ifaddrs *) calloc(1, sizeof(struct ifaddrs));
cifa->ifa_next = NULL;
cifa->ifa_name = strdup(ifr->lifr_name);
if (pifa == NULL) *ifap = cifa; /* first one */
else pifa->ifa_next = cifa;
if (ioctl(sd, SIOCGLIFADDR, ifr, IFREQSZ) < 0)
goto error;
cifa->ifa_addr = sa_dup((struct sockaddr*)&ifr->lifr_addr);
if (ioctl(sd, SIOCGLIFNETMASK, ifr, IFREQSZ) < 0)
goto error;
cifa->ifa_netmask = sa_dup((struct sockaddr*)&ifr->lifr_addr);
cifa->ifa_flags = 0;
cifa->ifa_dstaddr = NULL;
if (0 == ioctl(sd, SIOCGLIFFLAGS, ifr)) /* optional */
cifa->ifa_flags = ifr->lifr_flags;
if (ioctl(sd, SIOCGLIFDSTADDR, ifr, IFREQSZ) < 0) {
if (0 == ioctl(sd, SIOCGLIFBRDADDR, ifr, IFREQSZ))
cifa->ifa_dstaddr = sa_dup((struct sockaddr*)&ifr->lifr_addr);
}
else cifa->ifa_dstaddr = sa_dup((struct sockaddr*)&ifr->lifr_addr);
pifa = cifa;
ccp += IFREQSZ;
}
free(ifc.lifc_buf);
close(sd);
return 0;
error:
if (ifc.lifc_buf != NULL)
free(ifc.lifc_buf);
if (sd != -1)
close(sd);
return (-1);
}
+26
View File
@@ -0,0 +1,26 @@
/* Reference: https://lists.samba.org/archive/samba-technical/2009-February/063079.html */
#ifndef __IFADDRS_H___
#define __IFADDRS_H___
#include <sys/socket.h>
#include <net/if.h>
#undef ifa_dstaddr
#undef ifa_broadaddr
#define ifa_broadaddr ifa_dstaddr
struct ifaddrs {
struct ifaddrs *ifa_next;
char *ifa_name;
unsigned int ifa_flags;
struct sockaddr *ifa_addr;
struct sockaddr *ifa_netmask;
struct sockaddr *ifa_dstaddr;
};
extern int getifaddrs(struct ifaddrs **);
extern void freeifaddrs(struct ifaddrs *);
#endif
+41
View File
@@ -0,0 +1,41 @@
// mingw headers are missing this
typedef enum _LOGICAL_PROCESSOR_RELATIONSHIP {
RelationProcessorCore,
RelationNumaNode,
RelationCache,
RelationProcessorPackage,
RelationGroup,
RelationAll=0xffff
} LOGICAL_PROCESSOR_RELATIONSHIP;
typedef enum _PROCESSOR_CACHE_TYPE {
CacheUnified,CacheInstruction,CacheData,CacheTrace
} PROCESSOR_CACHE_TYPE;
typedef struct _CACHE_DESCRIPTOR {
BYTE Level;
BYTE Associativity;
WORD LineSize;
DWORD Size;
PROCESSOR_CACHE_TYPE Type;
} CACHE_DESCRIPTOR,*PCACHE_DESCRIPTOR;
typedef struct _SYSTEM_LOGICAL_PROCESSOR_INFORMATION {
ULONG_PTR ProcessorMask;
LOGICAL_PROCESSOR_RELATIONSHIP Relationship;
union {
struct {
BYTE Flags;
} ProcessorCore;
struct {
DWORD NodeNumber;
} NumaNode;
CACHE_DESCRIPTOR Cache;
ULONGLONG Reserved[2];
};
} SYSTEM_LOGICAL_PROCESSOR_INFORMATION,*PSYSTEM_LOGICAL_PROCESSOR_INFORMATION;
WINBASEAPI WINBOOL WINAPI
GetLogicalProcessorInformation(PSYSTEM_LOGICAL_PROCESSOR_INFORMATION Buffer,
PDWORD ReturnedLength);
+38
View File
@@ -0,0 +1,38 @@
#include "inet_ntop.h"
// From: https://memset.wordpress.com/2010/10/09/inet_ntop-for-win32/
PCSTR
WSAAPI
inet_ntop(__in INT Family,
__in PVOID pAddr,
__out_ecount(StringBufSize) PSTR pStringBuf,
__in size_t StringBufSize) {
DWORD dwAddressLength = 0;
struct sockaddr_storage srcaddr;
struct sockaddr_in *srcaddr4 = (struct sockaddr_in*) &srcaddr;
struct sockaddr_in6 *srcaddr6 = (struct sockaddr_in6*) &srcaddr;
memset(&srcaddr, 0, sizeof(struct sockaddr_storage));
srcaddr.ss_family = Family;
if (Family == AF_INET)
{
dwAddressLength = sizeof(struct sockaddr_in);
memcpy(&(srcaddr4->sin_addr), pAddr, sizeof(struct in_addr));
} else if (Family == AF_INET6)
{
dwAddressLength = sizeof(struct sockaddr_in6);
memcpy(&(srcaddr6->sin6_addr), pAddr, sizeof(struct in6_addr));
} else {
return NULL;
}
if (WSAAddressToString((LPSOCKADDR) &srcaddr,
dwAddressLength,
0,
pStringBuf,
(LPDWORD) &StringBufSize) != 0) {
return NULL;
}
return pStringBuf;
}
+10
View File
@@ -0,0 +1,10 @@
#include <ws2tcpip.h>
PCSTR
WSAAPI
inet_ntop(
__in INT Family,
__in PVOID pAddr,
__out_ecount(StringBufSize) PSTR pStringBuf,
__in size_t StringBufSize
);
+341
View File
@@ -0,0 +1,341 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#if !defined(__NTEXTAPI_H__)
#define __NTEXTAPI_H__
#include <winternl.h>
typedef struct {
LARGE_INTEGER IdleTime;
LARGE_INTEGER KernelTime;
LARGE_INTEGER UserTime;
LARGE_INTEGER DpcTime;
LARGE_INTEGER InterruptTime;
ULONG InterruptCount;
} _SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION;
typedef struct {
LARGE_INTEGER IdleProcessTime;
LARGE_INTEGER IoReadTransferCount;
LARGE_INTEGER IoWriteTransferCount;
LARGE_INTEGER IoOtherTransferCount;
ULONG IoReadOperationCount;
ULONG IoWriteOperationCount;
ULONG IoOtherOperationCount;
ULONG AvailablePages;
ULONG CommittedPages;
ULONG CommitLimit;
ULONG PeakCommitment;
ULONG PageFaultCount;
ULONG CopyOnWriteCount;
ULONG TransitionCount;
ULONG CacheTransitionCount;
ULONG DemandZeroCount;
ULONG PageReadCount;
ULONG PageReadIoCount;
ULONG CacheReadCount;
ULONG CacheIoCount;
ULONG DirtyPagesWriteCount;
ULONG DirtyWriteIoCount;
ULONG MappedPagesWriteCount;
ULONG MappedWriteIoCount;
ULONG PagedPoolPages;
ULONG NonPagedPoolPages;
ULONG PagedPoolAllocs;
ULONG PagedPoolFrees;
ULONG NonPagedPoolAllocs;
ULONG NonPagedPoolFrees;
ULONG FreeSystemPtes;
ULONG ResidentSystemCodePage;
ULONG TotalSystemDriverPages;
ULONG TotalSystemCodePages;
ULONG NonPagedPoolLookasideHits;
ULONG PagedPoolLookasideHits;
ULONG AvailablePagedPoolPages;
ULONG ResidentSystemCachePage;
ULONG ResidentPagedPoolPage;
ULONG ResidentSystemDriverPage;
ULONG CcFastReadNoWait;
ULONG CcFastReadWait;
ULONG CcFastReadResourceMiss;
ULONG CcFastReadNotPossible;
ULONG CcFastMdlReadNoWait;
ULONG CcFastMdlReadWait;
ULONG CcFastMdlReadResourceMiss;
ULONG CcFastMdlReadNotPossible;
ULONG CcMapDataNoWait;
ULONG CcMapDataWait;
ULONG CcMapDataNoWaitMiss;
ULONG CcMapDataWaitMiss;
ULONG CcPinMappedDataCount;
ULONG CcPinReadNoWait;
ULONG CcPinReadWait;
ULONG CcPinReadNoWaitMiss;
ULONG CcPinReadWaitMiss;
ULONG CcCopyReadNoWait;
ULONG CcCopyReadWait;
ULONG CcCopyReadNoWaitMiss;
ULONG CcCopyReadWaitMiss;
ULONG CcMdlReadNoWait;
ULONG CcMdlReadWait;
ULONG CcMdlReadNoWaitMiss;
ULONG CcMdlReadWaitMiss;
ULONG CcReadAheadIos;
ULONG CcLazyWriteIos;
ULONG CcLazyWritePages;
ULONG CcDataFlushes;
ULONG CcDataPages;
ULONG ContextSwitches;
ULONG FirstLevelTbFills;
ULONG SecondLevelTbFills;
ULONG SystemCalls;
} _SYSTEM_PERFORMANCE_INFORMATION;
typedef struct {
ULONG ContextSwitches;
ULONG DpcCount;
ULONG DpcRate;
ULONG TimeIncrement;
ULONG DpcBypassCount;
ULONG ApcBypassCount;
} _SYSTEM_INTERRUPT_INFORMATION;
typedef enum _KTHREAD_STATE {
Initialized,
Ready,
Running,
Standby,
Terminated,
Waiting,
Transition,
DeferredReady,
GateWait,
MaximumThreadState
} KTHREAD_STATE, *PKTHREAD_STATE;
typedef enum _KWAIT_REASON {
Executive = 0,
FreePage = 1,
PageIn = 2,
PoolAllocation = 3,
DelayExecution = 4,
Suspended = 5,
UserRequest = 6,
WrExecutive = 7,
WrFreePage = 8,
WrPageIn = 9,
WrPoolAllocation = 10,
WrDelayExecution = 11,
WrSuspended = 12,
WrUserRequest = 13,
WrEventPair = 14,
WrQueue = 15,
WrLpcReceive = 16,
WrLpcReply = 17,
WrVirtualMemory = 18,
WrPageOut = 19,
WrRendezvous = 20,
Spare2 = 21,
Spare3 = 22,
Spare4 = 23,
Spare5 = 24,
WrCalloutStack = 25,
WrKernel = 26,
WrResource = 27,
WrPushLock = 28,
WrMutex = 29,
WrQuantumEnd = 30,
WrDispatchInt = 31,
WrPreempted = 32,
WrYieldExecution = 33,
WrFastMutex = 34,
WrGuardedMutex = 35,
WrRundown = 36,
MaximumWaitReason = 37
} KWAIT_REASON, *PKWAIT_REASON;
typedef struct _CLIENT_ID {
HANDLE UniqueProcess;
HANDLE UniqueThread;
} CLIENT_ID, *PCLIENT_ID;
typedef struct _SYSTEM_THREAD_INFORMATION {
LARGE_INTEGER KernelTime;
LARGE_INTEGER UserTime;
LARGE_INTEGER CreateTime;
ULONG WaitTime;
PVOID StartAddress;
CLIENT_ID ClientId;
LONG Priority;
LONG BasePriority;
ULONG ContextSwitches;
ULONG ThreadState;
KWAIT_REASON WaitReason;
} SYSTEM_THREAD_INFORMATION, *PSYSTEM_THREAD_INFORMATION;
typedef struct _TEB *PTEB;
// private
typedef struct _SYSTEM_EXTENDED_THREAD_INFORMATION {
SYSTEM_THREAD_INFORMATION ThreadInfo;
PVOID StackBase;
PVOID StackLimit;
PVOID Win32StartAddress;
PTEB TebBase;
ULONG_PTR Reserved2;
ULONG_PTR Reserved3;
ULONG_PTR Reserved4;
} SYSTEM_EXTENDED_THREAD_INFORMATION, *PSYSTEM_EXTENDED_THREAD_INFORMATION;
typedef struct _SYSTEM_PROCESS_INFORMATION2 {
ULONG NextEntryOffset;
ULONG NumberOfThreads;
LARGE_INTEGER SpareLi1;
LARGE_INTEGER SpareLi2;
LARGE_INTEGER SpareLi3;
LARGE_INTEGER CreateTime;
LARGE_INTEGER UserTime;
LARGE_INTEGER KernelTime;
UNICODE_STRING ImageName;
LONG BasePriority;
HANDLE UniqueProcessId;
HANDLE InheritedFromUniqueProcessId;
ULONG HandleCount;
ULONG SessionId;
ULONG_PTR PageDirectoryBase;
SIZE_T PeakVirtualSize;
SIZE_T VirtualSize;
DWORD PageFaultCount;
SIZE_T PeakWorkingSetSize;
SIZE_T WorkingSetSize;
SIZE_T QuotaPeakPagedPoolUsage;
SIZE_T QuotaPagedPoolUsage;
SIZE_T QuotaPeakNonPagedPoolUsage;
SIZE_T QuotaNonPagedPoolUsage;
SIZE_T PagefileUsage;
SIZE_T PeakPagefileUsage;
SIZE_T PrivatePageCount;
LARGE_INTEGER ReadOperationCount;
LARGE_INTEGER WriteOperationCount;
LARGE_INTEGER OtherOperationCount;
LARGE_INTEGER ReadTransferCount;
LARGE_INTEGER WriteTransferCount;
LARGE_INTEGER OtherTransferCount;
SYSTEM_THREAD_INFORMATION Threads[1];
} SYSTEM_PROCESS_INFORMATION2, *PSYSTEM_PROCESS_INFORMATION2;
#define SYSTEM_PROCESS_INFORMATION SYSTEM_PROCESS_INFORMATION2
#define PSYSTEM_PROCESS_INFORMATION PSYSTEM_PROCESS_INFORMATION2
// ================================================
// psutil.users() support
// ================================================
typedef struct _WINSTATION_INFO {
BYTE Reserved1[72];
ULONG SessionId;
BYTE Reserved2[4];
FILETIME ConnectTime;
FILETIME DisconnectTime;
FILETIME LastInputTime;
FILETIME LoginTime;
BYTE Reserved3[1096];
FILETIME CurrentTime;
} WINSTATION_INFO, *PWINSTATION_INFO;
typedef BOOLEAN (WINAPI * PWINSTATIONQUERYINFORMATIONW)
(HANDLE,ULONG,WINSTATIONINFOCLASS,PVOID,ULONG,PULONG);
/*
* NtQueryInformationProcess code taken from
* http://wj32.wordpress.com/2009/01/24/howto-get-the-command-line-of-processes/
* typedefs needed to compile against ntdll functions not exposted in the API
*/
typedef LONG NTSTATUS;
typedef NTSTATUS (NTAPI *_NtQueryInformationProcess)(
HANDLE ProcessHandle,
DWORD ProcessInformationClass,
PVOID ProcessInformation,
DWORD ProcessInformationLength,
PDWORD ReturnLength
);
typedef NTSTATUS (NTAPI *_NtSetInformationProcess)(
HANDLE ProcessHandle,
DWORD ProcessInformationClass,
PVOID ProcessInformation,
DWORD ProcessInformationLength
);
typedef enum _PROCESSINFOCLASS2 {
_ProcessBasicInformation,
ProcessQuotaLimits,
ProcessIoCounters,
ProcessVmCounters,
ProcessTimes,
ProcessBasePriority,
ProcessRaisePriority,
_ProcessDebugPort,
ProcessExceptionPort,
ProcessAccessToken,
ProcessLdtInformation,
ProcessLdtSize,
ProcessDefaultHardErrorMode,
ProcessIoPortHandlers,
ProcessPooledUsageAndLimits,
ProcessWorkingSetWatch,
ProcessUserModeIOPL,
ProcessEnableAlignmentFaultFixup,
ProcessPriorityClass,
ProcessWx86Information,
ProcessHandleCount,
ProcessAffinityMask,
ProcessPriorityBoost,
ProcessDeviceMap,
ProcessSessionInformation,
ProcessForegroundInformation,
_ProcessWow64Information,
/* added after XP+ */
_ProcessImageFileName,
ProcessLUIDDeviceMapsEnabled,
_ProcessBreakOnTermination,
ProcessDebugObjectHandle,
ProcessDebugFlags,
ProcessHandleTracing,
ProcessIoPriority,
ProcessExecuteFlags,
ProcessResourceManagement,
ProcessCookie,
ProcessImageInformation,
MaxProcessInfoClass
} PROCESSINFOCLASS2;
#define PROCESSINFOCLASS PROCESSINFOCLASS2
#define ProcessBasicInformation _ProcessBasicInformation
#define ProcessWow64Information _ProcessWow64Information
#define ProcessDebugPort _ProcessDebugPort
#define ProcessImageFileName _ProcessImageFileName
#define ProcessBreakOnTermination _ProcessBreakOnTermination
#endif // __NTEXTAPI_H__
+522
View File
@@ -0,0 +1,522 @@
/*
* Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
*/
#include "process_handles.h"
static _NtQuerySystemInformation __NtQuerySystemInformation = NULL;
static _NtQueryObject __NtQueryObject = NULL;
CRITICAL_SECTION g_cs;
BOOL g_initialized = FALSE;
NTSTATUS g_status;
HANDLE g_hFile = NULL;
HANDLE g_hEvtStart = NULL;
HANDLE g_hEvtFinish = NULL;
HANDLE g_hThread = NULL;
PUNICODE_STRING g_pNameBuffer = NULL;
ULONG g_dwSize = 0;
ULONG g_dwLength = 0;
PVOID
GetLibraryProcAddress(PSTR LibraryName, PSTR ProcName) {
return GetProcAddress(GetModuleHandleA(LibraryName), ProcName);
}
PyObject *
psutil_get_open_files(long dwPid, HANDLE hProcess) {
OSVERSIONINFO osvi;
ZeroMemory(&osvi, sizeof(OSVERSIONINFO));
osvi.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);
GetVersionEx(&osvi);
// Threaded version only works for Vista+
if (osvi.dwMajorVersion >= 6)
return psutil_get_open_files_ntqueryobject(dwPid, hProcess);
else
return psutil_get_open_files_getmappedfilename(dwPid, hProcess);
}
VOID
psutil_get_open_files_init(BOOL threaded) {
if (g_initialized == TRUE)
return;
// Resolve the Windows API calls
__NtQuerySystemInformation =
GetLibraryProcAddress("ntdll.dll", "NtQuerySystemInformation");
__NtQueryObject = GetLibraryProcAddress("ntdll.dll", "NtQueryObject");
// Create events for signalling work between threads
if (threaded == TRUE) {
g_hEvtStart = CreateEvent(NULL, FALSE, FALSE, NULL);
g_hEvtFinish = CreateEvent(NULL, FALSE, FALSE, NULL);
InitializeCriticalSection(&g_cs);
}
g_initialized = TRUE;
}
PyObject *
psutil_get_open_files_ntqueryobject(long dwPid, HANDLE hProcess) {
NTSTATUS status;
PSYSTEM_HANDLE_INFORMATION_EX pHandleInfo = NULL;
DWORD dwInfoSize = 0x10000;
DWORD dwRet = 0;
PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX hHandle = NULL;
DWORD i = 0;
BOOLEAN error = FALSE;
DWORD dwWait = 0;
PyObject* py_retlist = NULL;
PyObject* py_path = NULL;
if (g_initialized == FALSE)
psutil_get_open_files_init(TRUE);
// Due to the use of global variables, ensure only 1 call
// to psutil_get_open_files() is running
EnterCriticalSection(&g_cs);
if (__NtQuerySystemInformation == NULL ||
__NtQueryObject == NULL ||
g_hEvtStart == NULL ||
g_hEvtFinish == NULL)
{
PyErr_SetFromWindowsErr(0);
error = TRUE;
goto cleanup;
}
// Py_BuildValue raises an exception if NULL is returned
py_retlist = PyList_New(0);
if (py_retlist == NULL) {
error = TRUE;
goto cleanup;
}
do {
if (pHandleInfo != NULL) {
HeapFree(GetProcessHeap(), 0, pHandleInfo);
pHandleInfo = NULL;
}
// NtQuerySystemInformation won't give us the correct buffer size,
// so we guess by doubling the buffer size.
dwInfoSize *= 2;
pHandleInfo = HeapAlloc(GetProcessHeap(),
HEAP_ZERO_MEMORY,
dwInfoSize);
if (pHandleInfo == NULL) {
PyErr_NoMemory();
error = TRUE;
goto cleanup;
}
} while ((status = __NtQuerySystemInformation(
SystemExtendedHandleInformation,
pHandleInfo,
dwInfoSize,
&dwRet)) == STATUS_INFO_LENGTH_MISMATCH);
// NtQuerySystemInformation stopped giving us STATUS_INFO_LENGTH_MISMATCH
if (!NT_SUCCESS(status)) {
PyErr_SetFromWindowsErr(HRESULT_FROM_NT(status));
error = TRUE;
goto cleanup;
}
for (i = 0; i < pHandleInfo->NumberOfHandles; i++) {
hHandle = &pHandleInfo->Handles[i];
// Check if this hHandle belongs to the PID the user specified.
if (hHandle->UniqueProcessId != (HANDLE)dwPid)
goto loop_cleanup;
if (!DuplicateHandle(hProcess,
hHandle->HandleValue,
GetCurrentProcess(),
&g_hFile,
0,
TRUE,
DUPLICATE_SAME_ACCESS))
{
/*
printf("[%d] DuplicateHandle (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
goto loop_cleanup;
}
// Guess buffer size is MAX_PATH + 1
g_dwLength = (MAX_PATH+1) * sizeof(WCHAR);
do {
// Release any previously allocated buffer
if (g_pNameBuffer != NULL) {
HeapFree(GetProcessHeap(), 0, g_pNameBuffer);
g_pNameBuffer = NULL;
g_dwSize = 0;
}
// NtQueryObject puts the required buffer size in g_dwLength
// WinXP edge case puts g_dwLength == 0, just skip this handle
if (g_dwLength == 0)
goto loop_cleanup;
g_dwSize = g_dwLength;
if (g_dwSize > 0) {
g_pNameBuffer = HeapAlloc(GetProcessHeap(),
HEAP_ZERO_MEMORY,
g_dwSize);
if (g_pNameBuffer == NULL)
goto loop_cleanup;
}
dwWait = psutil_NtQueryObject();
// If the call does not return, skip this handle
if (dwWait != WAIT_OBJECT_0)
goto loop_cleanup;
} while (g_status == STATUS_INFO_LENGTH_MISMATCH);
// NtQueryObject stopped returning STATUS_INFO_LENGTH_MISMATCH
if (!NT_SUCCESS(g_status))
goto loop_cleanup;
// Convert to PyUnicode and append it to the return list
if (g_pNameBuffer->Length > 0) {
/*
printf("[%d] Filename (%#x) %#d bytes: %S\n",
dwPid,
hHandle->HandleValue,
g_pNameBuffer->Length,
g_pNameBuffer->Buffer);
*/
py_path = PyUnicode_FromWideChar(g_pNameBuffer->Buffer,
g_pNameBuffer->Length/2);
if (py_path == NULL) {
/*
printf("[%d] PyUnicode_FromWideChar (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
error = TRUE;
goto loop_cleanup;
}
if (PyList_Append(py_retlist, py_path)) {
/*
printf("[%d] PyList_Append (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
error = TRUE;
goto loop_cleanup;
}
}
loop_cleanup:
Py_XDECREF(py_path);
py_path = NULL;
if (g_pNameBuffer != NULL)
HeapFree(GetProcessHeap(), 0, g_pNameBuffer);
g_pNameBuffer = NULL;
g_dwSize = 0;
g_dwLength = 0;
if (g_hFile != NULL)
CloseHandle(g_hFile);
g_hFile = NULL;
}
cleanup:
if (g_pNameBuffer != NULL)
HeapFree(GetProcessHeap(), 0, g_pNameBuffer);
g_pNameBuffer = NULL;
g_dwSize = 0;
g_dwLength = 0;
if (g_hFile != NULL)
CloseHandle(g_hFile);
g_hFile = NULL;
if (pHandleInfo != NULL)
HeapFree(GetProcessHeap(), 0, pHandleInfo);
pHandleInfo = NULL;
if (error) {
Py_XDECREF(py_retlist);
py_retlist = NULL;
}
LeaveCriticalSection(&g_cs);
return py_retlist;
}
DWORD
psutil_NtQueryObject() {
DWORD dwWait = 0;
if (g_hThread == NULL)
g_hThread = CreateThread(
NULL,
0,
psutil_NtQueryObjectThread,
NULL,
0,
NULL);
if (g_hThread == NULL)
return GetLastError();
// Signal the worker thread to start
SetEvent(g_hEvtStart);
// Wait for the worker thread to finish
dwWait = WaitForSingleObject(g_hEvtFinish, NTQO_TIMEOUT);
// If the thread hangs, kill it and cleanup
if (dwWait == WAIT_TIMEOUT) {
SuspendThread(g_hThread);
TerminateThread(g_hThread, 1);
WaitForSingleObject(g_hThread, INFINITE);
CloseHandle(g_hThread);
g_hThread = NULL;
}
return dwWait;
}
DWORD WINAPI
psutil_NtQueryObjectThread(LPVOID lpvParam) {
// Loop infinitely waiting for work
while (TRUE) {
WaitForSingleObject(g_hEvtStart, INFINITE);
g_status = __NtQueryObject(g_hFile,
ObjectNameInformation,
g_pNameBuffer,
g_dwSize,
&g_dwLength);
SetEvent(g_hEvtFinish);
}
}
PyObject *
psutil_get_open_files_getmappedfilename(long dwPid, HANDLE hProcess) {
NTSTATUS status;
PSYSTEM_HANDLE_INFORMATION_EX pHandleInfo = NULL;
DWORD dwInfoSize = 0x10000;
DWORD dwRet = 0;
PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX hHandle = NULL;
HANDLE hFile = NULL;
HANDLE hMap = NULL;
DWORD i = 0;
BOOLEAN error = FALSE;
PyObject* py_retlist = NULL;
PyObject* py_path = NULL;
ULONG dwSize = 0;
LPVOID pMem = NULL;
TCHAR pszFilename[MAX_PATH+1];
if (g_initialized == FALSE)
psutil_get_open_files_init(FALSE);
if (__NtQuerySystemInformation == NULL || __NtQueryObject == NULL) {
PyErr_SetFromWindowsErr(0);
error = TRUE;
goto cleanup;
}
// Py_BuildValue raises an exception if NULL is returned
py_retlist = PyList_New(0);
if (py_retlist == NULL) {
error = TRUE;
goto cleanup;
}
do {
if (pHandleInfo != NULL) {
HeapFree(GetProcessHeap(), 0, pHandleInfo);
pHandleInfo = NULL;
}
// NtQuerySystemInformation won't give us the correct buffer size,
// so we guess by doubling the buffer size.
dwInfoSize *= 2;
pHandleInfo = HeapAlloc(GetProcessHeap(),
HEAP_ZERO_MEMORY,
dwInfoSize);
if (pHandleInfo == NULL) {
PyErr_NoMemory();
error = TRUE;
goto cleanup;
}
} while ((status = __NtQuerySystemInformation(
SystemExtendedHandleInformation,
pHandleInfo,
dwInfoSize,
&dwRet)) == STATUS_INFO_LENGTH_MISMATCH);
// NtQuerySystemInformation stopped giving us STATUS_INFO_LENGTH_MISMATCH
if (!NT_SUCCESS(status)) {
PyErr_SetFromWindowsErr(HRESULT_FROM_NT(status));
error = TRUE;
goto cleanup;
}
for (i = 0; i < pHandleInfo->NumberOfHandles; i++) {
hHandle = &pHandleInfo->Handles[i];
// Check if this hHandle belongs to the PID the user specified.
if (hHandle->UniqueProcessId != (HANDLE)dwPid)
goto loop_cleanup;
if (!DuplicateHandle(hProcess,
hHandle->HandleValue,
GetCurrentProcess(),
&hFile,
0,
TRUE,
DUPLICATE_SAME_ACCESS))
{
/*
printf("[%d] DuplicateHandle (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
goto loop_cleanup;
}
hMap = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
if (hMap == NULL) {
/*
printf("[%d] CreateFileMapping (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
goto loop_cleanup;
}
pMem = MapViewOfFile(hMap, FILE_MAP_READ, 0, 0, 1);
if (pMem == NULL) {
/*
printf("[%d] MapViewOfFile (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
goto loop_cleanup;
}
dwSize = GetMappedFileName(
GetCurrentProcess(), pMem, pszFilename, MAX_PATH);
if (dwSize == 0) {
/*
printf("[%d] GetMappedFileName (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
goto loop_cleanup;
}
pszFilename[dwSize] = '\0';
/*
printf("[%d] Filename (%#x) %#d bytes: %S\n",
dwPid,
hHandle->HandleValue,
dwSize,
pszFilename);
*/
py_path = PyUnicode_FromWideChar(pszFilename, dwSize);
if (py_path == NULL) {
/*
printf("[%d] PyUnicode_FromStringAndSize (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
error = TRUE;
goto loop_cleanup;
}
if (PyList_Append(py_retlist, py_path)) {
/*
printf("[%d] PyList_Append (%#x): %#x \n",
dwPid,
hHandle->HandleValue,
GetLastError());
*/
error = TRUE;
goto loop_cleanup;
}
loop_cleanup:
Py_XDECREF(py_path);
py_path = NULL;
if (pMem != NULL)
UnmapViewOfFile(pMem);
pMem = NULL;
if (hMap != NULL)
CloseHandle(hMap);
hMap = NULL;
if (hFile != NULL)
CloseHandle(hFile);
hFile = NULL;
dwSize = 0;
}
cleanup:
if (pMem != NULL)
UnmapViewOfFile(pMem);
pMem = NULL;
if (hMap != NULL)
CloseHandle(hMap);
hMap = NULL;
if (hFile != NULL)
CloseHandle(hFile);
hFile = NULL;
if (pHandleInfo != NULL)
HeapFree(GetProcessHeap(), 0, pHandleInfo);
pHandleInfo = NULL;
if (error) {
Py_XDECREF(py_retlist);
py_retlist = NULL;
}
return py_retlist;
}
+111
View File
@@ -0,0 +1,111 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#ifndef __PROCESS_HANDLES_H__
#define __PROCESS_HANDLES_H__
#ifndef UNICODE
#define UNICODE
#endif
#include <Python.h>
#include <stdio.h>
#include <windows.h>
#include <strsafe.h>
#include <winternl.h>
#include <psapi.h>
#ifndef NT_SUCCESS
#define NT_SUCCESS(x) ((x) >= 0)
#endif
#define STATUS_INFO_LENGTH_MISMATCH 0xc0000004
#define ObjectBasicInformation 0
#define ObjectNameInformation 1
#define ObjectTypeInformation 2
#define HANDLE_TYPE_FILE 28
#define NTQO_TIMEOUT 100
typedef NTSTATUS (NTAPI *_NtQuerySystemInformation)(
ULONG SystemInformationClass,
PVOID SystemInformation,
ULONG SystemInformationLength,
PULONG ReturnLength
);
typedef NTSTATUS (NTAPI *_NtQueryObject)(
HANDLE ObjectHandle,
ULONG ObjectInformationClass,
PVOID ObjectInformation,
ULONG ObjectInformationLength,
PULONG ReturnLength
);
// Undocumented FILE_INFORMATION_CLASS: FileNameInformation
static const SYSTEM_INFORMATION_CLASS SystemExtendedHandleInformation = (SYSTEM_INFORMATION_CLASS)64;
typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX {
PVOID Object;
HANDLE UniqueProcessId;
HANDLE HandleValue;
ULONG GrantedAccess;
USHORT CreatorBackTraceIndex;
USHORT ObjectTypeIndex;
ULONG HandleAttributes;
ULONG Reserved;
} SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX, *PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX;
typedef struct _SYSTEM_HANDLE_INFORMATION_EX {
ULONG_PTR NumberOfHandles;
ULONG_PTR Reserved;
SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX Handles[1];
} SYSTEM_HANDLE_INFORMATION_EX, *PSYSTEM_HANDLE_INFORMATION_EX;
typedef enum _POOL_TYPE {
NonPagedPool,
PagedPool,
NonPagedPoolMustSucceed,
DontUseThisType,
NonPagedPoolCacheAligned,
PagedPoolCacheAligned,
NonPagedPoolCacheAlignedMustS
} POOL_TYPE, *PPOOL_TYPE;
typedef struct _OBJECT_TYPE_INFORMATION {
UNICODE_STRING Name;
ULONG TotalNumberOfObjects;
ULONG TotalNumberOfHandles;
ULONG TotalPagedPoolUsage;
ULONG TotalNonPagedPoolUsage;
ULONG TotalNamePoolUsage;
ULONG TotalHandleTableUsage;
ULONG HighWaterNumberOfObjects;
ULONG HighWaterNumberOfHandles;
ULONG HighWaterPagedPoolUsage;
ULONG HighWaterNonPagedPoolUsage;
ULONG HighWaterNamePoolUsage;
ULONG HighWaterHandleTableUsage;
ULONG InvalidAttributes;
GENERIC_MAPPING GenericMapping;
ULONG ValidAccess;
BOOLEAN SecurityRequired;
BOOLEAN MaintainHandleCount;
USHORT MaintainTypeList;
POOL_TYPE PoolType;
ULONG PagedPoolUsage;
ULONG NonPagedPoolUsage;
} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;
PVOID GetLibraryProcAddress(PSTR LibraryName, PSTR ProcName);
VOID psutil_get_open_files_init(BOOL threaded);
PyObject* psutil_get_open_files(long pid, HANDLE processHandle);
PyObject* psutil_get_open_files_ntqueryobject(long dwPid, HANDLE hProcess);
PyObject* psutil_get_open_files_getmappedfilename(long dwPid, HANDLE hProcess);
DWORD psutil_NtQueryObject(void);
DWORD WINAPI psutil_NtQueryObjectThread(LPVOID lpvParam);
#endif // __PROCESS_HANDLES_H__
+873
View File
@@ -0,0 +1,873 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Helper functions related to fetching process information. Used by
* _psutil_windows module methods.
*/
#include <Python.h>
#include <windows.h>
#include <Psapi.h>
#include <tlhelp32.h>
#include "security.h"
#include "process_info.h"
#include "ntextapi.h"
#include "../../_psutil_common.h"
/*
* A wrapper around OpenProcess setting NSP exception if process
* no longer exists.
* "pid" is the process pid, "dwDesiredAccess" is the first argument
* exptected by OpenProcess.
* Return a process handle or NULL.
*/
HANDLE
psutil_handle_from_pid_waccess(DWORD pid, DWORD dwDesiredAccess) {
HANDLE hProcess;
DWORD processExitCode = 0;
if (pid == 0) {
// otherwise we'd get NoSuchProcess
return AccessDenied();
}
hProcess = OpenProcess(dwDesiredAccess, FALSE, pid);
if (hProcess == NULL) {
if (GetLastError() == ERROR_INVALID_PARAMETER)
NoSuchProcess();
else
PyErr_SetFromWindowsErr(0);
return NULL;
}
// make sure the process is running
GetExitCodeProcess(hProcess, &processExitCode);
if (processExitCode == 0) {
NoSuchProcess();
CloseHandle(hProcess);
return NULL;
}
return hProcess;
}
/*
* Same as psutil_handle_from_pid_waccess but implicitly uses
* PROCESS_QUERY_INFORMATION | PROCESS_VM_READ as dwDesiredAccess
* parameter for OpenProcess.
*/
HANDLE
psutil_handle_from_pid(DWORD pid) {
DWORD dwDesiredAccess = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
return psutil_handle_from_pid_waccess(pid, dwDesiredAccess);
}
/*
* Given a Python int referencing a process handle close the process handle.
*/
PyObject *
psutil_win32_CloseHandle(PyObject *self, PyObject *args) {
unsigned long handle;
if (! PyArg_ParseTuple(args, "k", &handle))
return NULL;
// TODO: may want to check return value;
CloseHandle((HANDLE)handle);
Py_RETURN_NONE;
}
DWORD *
psutil_get_pids(DWORD *numberOfReturnedPIDs) {
// Win32 SDK says the only way to know if our process array
// wasn't large enough is to check the returned size and make
// sure that it doesn't match the size of the array.
// If it does we allocate a larger array and try again
// Stores the actual array
DWORD *procArray = NULL;
DWORD procArrayByteSz;
int procArraySz = 0;
// Stores the byte size of the returned array from enumprocesses
DWORD enumReturnSz = 0;
do {
procArraySz += 1024;
free(procArray);
procArrayByteSz = procArraySz * sizeof(DWORD);
procArray = malloc(procArrayByteSz);
if (procArray == NULL) {
PyErr_NoMemory();
return NULL;
}
if (! EnumProcesses(procArray, procArrayByteSz, &enumReturnSz)) {
free(procArray);
PyErr_SetFromWindowsErr(0);
return NULL;
}
} while (enumReturnSz == procArraySz * sizeof(DWORD));
// The number of elements is the returned size / size of each element
*numberOfReturnedPIDs = enumReturnSz / sizeof(DWORD);
return procArray;
}
int
psutil_pid_is_running(DWORD pid) {
HANDLE hProcess;
DWORD exitCode;
// Special case for PID 0 System Idle Process
if (pid == 0)
return 1;
if (pid < 0)
return 0;
hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ,
FALSE, pid);
if (NULL == hProcess) {
// invalid parameter is no such process
if (GetLastError() == ERROR_INVALID_PARAMETER) {
CloseHandle(hProcess);
return 0;
}
// access denied obviously means there's a process to deny access to...
if (GetLastError() == ERROR_ACCESS_DENIED) {
CloseHandle(hProcess);
return 1;
}
CloseHandle(hProcess);
PyErr_SetFromWindowsErr(0);
return -1;
}
if (GetExitCodeProcess(hProcess, &exitCode)) {
CloseHandle(hProcess);
return (exitCode == STILL_ACTIVE);
}
// access denied means there's a process there so we'll assume
// it's running
if (GetLastError() == ERROR_ACCESS_DENIED) {
CloseHandle(hProcess);
return 1;
}
PyErr_SetFromWindowsErr(0);
CloseHandle(hProcess);
return -1;
}
int
psutil_pid_in_proclist(DWORD pid) {
DWORD *proclist = NULL;
DWORD numberOfReturnedPIDs;
DWORD i;
proclist = psutil_get_pids(&numberOfReturnedPIDs);
if (proclist == NULL)
return -1;
for (i = 0; i < numberOfReturnedPIDs; i++) {
if (pid == proclist[i]) {
free(proclist);
return 1;
}
}
free(proclist);
return 0;
}
// Check exit code from a process handle. Return FALSE on an error also
// XXX - not used anymore
int
handlep_is_running(HANDLE hProcess) {
DWORD dwCode;
if (NULL == hProcess)
return 0;
if (GetExitCodeProcess(hProcess, &dwCode)) {
if (dwCode == STILL_ACTIVE)
return 1;
}
return 0;
}
// Helper structures to access the memory correctly. Some of these might also
// be defined in the winternl.h header file but unfortunately not in a usable
// way.
// see http://msdn2.microsoft.com/en-us/library/aa489609.aspx
#ifndef NT_SUCCESS
#define NT_SUCCESS(Status) ((NTSTATUS)(Status) >= 0)
#endif
// http://msdn.microsoft.com/en-us/library/aa813741(VS.85).aspx
typedef struct {
BYTE Reserved1[16];
PVOID Reserved2[5];
UNICODE_STRING CurrentDirectoryPath;
PVOID CurrentDirectoryHandle;
UNICODE_STRING DllPath;
UNICODE_STRING ImagePathName;
UNICODE_STRING CommandLine;
LPCWSTR env;
} RTL_USER_PROCESS_PARAMETERS_, *PRTL_USER_PROCESS_PARAMETERS_;
// https://msdn.microsoft.com/en-us/library/aa813706(v=vs.85).aspx
#ifdef _WIN64
typedef struct {
BYTE Reserved1[2];
BYTE BeingDebugged;
BYTE Reserved2[21];
PVOID LoaderData;
PRTL_USER_PROCESS_PARAMETERS_ ProcessParameters;
/* More fields ... */
} PEB_;
#else
typedef struct {
BYTE Reserved1[2];
BYTE BeingDebugged;
BYTE Reserved2[1];
PVOID Reserved3[2];
PVOID Ldr;
PRTL_USER_PROCESS_PARAMETERS_ ProcessParameters;
/* More fields ... */
} PEB_;
#endif
#ifdef _WIN64
/* When we are a 64 bit process accessing a 32 bit (WoW64) process we need to
use the 32 bit structure layout. */
typedef struct {
USHORT Length;
USHORT MaxLength;
DWORD Buffer;
} UNICODE_STRING32;
typedef struct {
BYTE Reserved1[16];
DWORD Reserved2[5];
UNICODE_STRING32 CurrentDirectoryPath;
DWORD CurrentDirectoryHandle;
UNICODE_STRING32 DllPath;
UNICODE_STRING32 ImagePathName;
UNICODE_STRING32 CommandLine;
DWORD env;
} RTL_USER_PROCESS_PARAMETERS32;
typedef struct {
BYTE Reserved1[2];
BYTE BeingDebugged;
BYTE Reserved2[1];
DWORD Reserved3[2];
DWORD Ldr;
DWORD ProcessParameters;
/* More fields ... */
} PEB32;
#else
/* When we are a 32 bit (WoW64) process accessing a 64 bit process we need to
use the 64 bit structure layout and a special function to read its memory.
*/
typedef NTSTATUS (NTAPI *_NtWow64ReadVirtualMemory64)(
IN HANDLE ProcessHandle,
IN PVOID64 BaseAddress,
OUT PVOID Buffer,
IN ULONG64 Size,
OUT PULONG64 NumberOfBytesRead);
typedef enum {
MemoryInformationBasic
} MEMORY_INFORMATION_CLASS;
typedef NTSTATUS (NTAPI *_NtWow64QueryVirtualMemory64)(
IN HANDLE ProcessHandle,
IN PVOID64 BaseAddress,
IN MEMORY_INFORMATION_CLASS MemoryInformationClass,
OUT PMEMORY_BASIC_INFORMATION64 MemoryInformation,
IN ULONG64 Size,
OUT PULONG64 ReturnLength OPTIONAL);
typedef struct {
PVOID Reserved1[2];
PVOID64 PebBaseAddress;
PVOID Reserved2[4];
PVOID UniqueProcessId[2];
PVOID Reserved3[2];
} PROCESS_BASIC_INFORMATION64;
typedef struct {
USHORT Length;
USHORT MaxLength;
PVOID64 Buffer;
} UNICODE_STRING64;
typedef struct {
BYTE Reserved1[16];
PVOID64 Reserved2[5];
UNICODE_STRING64 CurrentDirectoryPath;
PVOID64 CurrentDirectoryHandle;
UNICODE_STRING64 DllPath;
UNICODE_STRING64 ImagePathName;
UNICODE_STRING64 CommandLine;
PVOID64 env;
} RTL_USER_PROCESS_PARAMETERS64;
typedef struct {
BYTE Reserved1[2];
BYTE BeingDebugged;
BYTE Reserved2[21];
PVOID64 LoaderData;
PVOID64 ProcessParameters;
/* More fields ... */
} PEB64;
#endif
/* Given a pointer into a process's memory, figure out how much data can be
* read from it. */
static int psutil_get_process_region_size(HANDLE hProcess,
LPCVOID src,
SIZE_T *psize) {
MEMORY_BASIC_INFORMATION info;
if (!VirtualQueryEx(hProcess, src, &info, sizeof(info))) {
PyErr_SetFromWindowsErr(0);
return -1;
}
*psize = info.RegionSize - ((char*)src - (char*)info.BaseAddress);
return 0;
}
#ifndef _WIN64
/* Given a pointer into a process's memory, figure out how much data can be
* read from it. */
static int psutil_get_process_region_size64(HANDLE hProcess,
const PVOID64 src64,
PULONG64 psize) {
static _NtWow64QueryVirtualMemory64 NtWow64QueryVirtualMemory64 = NULL;
MEMORY_BASIC_INFORMATION64 info64;
if (NtWow64QueryVirtualMemory64 == NULL) {
NtWow64QueryVirtualMemory64 =
(_NtWow64QueryVirtualMemory64)GetProcAddress(
GetModuleHandleA("ntdll.dll"),
"NtWow64QueryVirtualMemory64");
if (NtWow64QueryVirtualMemory64 == NULL) {
PyErr_SetString(PyExc_NotImplementedError,
"NtWow64QueryVirtualMemory64 missing");
return -1;
}
}
if (!NT_SUCCESS(NtWow64QueryVirtualMemory64(
hProcess,
src64,
0,
&info64,
sizeof(info64),
NULL))) {
PyErr_SetFromWindowsErr(0);
return -1;
}
*psize = info64.RegionSize - ((char*)src64 - (char*)info64.BaseAddress);
return 0;
}
#endif
enum psutil_process_data_kind {
KIND_CMDLINE,
KIND_CWD,
KIND_ENVIRON,
};
/* Get data from the process with the given pid. The data is returned in the
pdata output member as a nul terminated string which must be freed on
success.
On success 0 is returned. On error the output parameter is not touched, -1
is returned, and an appropriate Python exception is set. */
static int psutil_get_process_data(long pid,
enum psutil_process_data_kind kind,
WCHAR **pdata,
SIZE_T *psize) {
/* This function is quite complex because there are several cases to be
considered:
Two cases are really simple: we (i.e. the python interpreter) and the
target process are both 32 bit or both 64 bit. In that case the memory
layout of the structures matches up and all is well.
When we are 64 bit and the target process is 32 bit we need to use
custom 32 bit versions of the structures.
When we are 32 bit and the target process is 64 bit we need to use
custom 64 bit version of the structures. Also we need to use separate
Wow64 functions to get the information.
A few helper structs are defined above so that the compiler can handle
calculating the correct offsets.
Additional help also came from the following sources:
https://github.com/kohsuke/winp and
http://wj32.org/wp/2009/01/24/howto-get-the-command-line-of-processes/
http://stackoverflow.com/a/14012919
http://www.drdobbs.com/embracing-64-bit-windows/184401966
*/
static _NtQueryInformationProcess NtQueryInformationProcess = NULL;
#ifndef _WIN64
static _NtQueryInformationProcess NtWow64QueryInformationProcess64 = NULL;
static _NtWow64ReadVirtualMemory64 NtWow64ReadVirtualMemory64 = NULL;
#endif
HANDLE hProcess = NULL;
LPCVOID src;
SIZE_T size;
WCHAR *buffer = NULL;
#ifdef _WIN64
LPVOID ppeb32 = NULL;
#else
PVOID64 src64;
BOOL weAreWow64;
BOOL theyAreWow64;
#endif
hProcess = psutil_handle_from_pid(pid);
if (hProcess == NULL)
return -1;
if (NtQueryInformationProcess == NULL) {
NtQueryInformationProcess = (_NtQueryInformationProcess)GetProcAddress(
GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess");
}
#ifdef _WIN64
/* 64 bit case. Check if the target is a 32 bit process running in WoW64
* mode. */
if (!NT_SUCCESS(NtQueryInformationProcess(hProcess,
ProcessWow64Information,
&ppeb32,
sizeof(LPVOID),
NULL))) {
PyErr_SetFromWindowsErr(0);
goto error;
}
if (ppeb32 != NULL) {
/* We are 64 bit. Target process is 32 bit running in WoW64 mode. */
PEB32 peb32;
RTL_USER_PROCESS_PARAMETERS32 procParameters32;
// read PEB
if (!ReadProcessMemory(hProcess, ppeb32, &peb32, sizeof(peb32), NULL)) {
PyErr_SetFromWindowsErr(0);
goto error;
}
// read process parameters
if (!ReadProcessMemory(hProcess,
UlongToPtr(peb32.ProcessParameters),
&procParameters32,
sizeof(procParameters32),
NULL)) {
PyErr_SetFromWindowsErr(0);
goto error;
}
switch (kind) {
case KIND_CMDLINE:
src = UlongToPtr(procParameters32.CommandLine.Buffer),
size = procParameters32.CommandLine.Length;
break;
case KIND_CWD:
src = UlongToPtr(procParameters32.CurrentDirectoryPath.Buffer);
size = procParameters32.CurrentDirectoryPath.Length;
break;
case KIND_ENVIRON:
src = UlongToPtr(procParameters32.env);
break;
}
} else
#else
/* 32 bit case. Check if the target is also 32 bit. */
if (!IsWow64Process(GetCurrentProcess(), &weAreWow64) ||
!IsWow64Process(hProcess, &theyAreWow64)) {
PyErr_SetFromWindowsErr(0);
goto error;
}
if (weAreWow64 && !theyAreWow64) {
/* We are 32 bit running in WoW64 mode. Target process is 64 bit. */
PROCESS_BASIC_INFORMATION64 pbi64;
PEB64 peb64;
RTL_USER_PROCESS_PARAMETERS64 procParameters64;
if (NtWow64QueryInformationProcess64 == NULL) {
NtWow64QueryInformationProcess64 =
(_NtQueryInformationProcess)GetProcAddress(
GetModuleHandleA("ntdll.dll"),
"NtWow64QueryInformationProcess64");
if (NtWow64QueryInformationProcess64 == NULL) {
PyErr_SetString(PyExc_NotImplementedError,
"NtWow64QueryInformationProcess64 missing");
goto error;
}
}
if (!NT_SUCCESS(NtWow64QueryInformationProcess64(
hProcess,
ProcessBasicInformation,
&pbi64,
sizeof(pbi64),
NULL))) {
PyErr_SetFromWindowsErr(0);
goto error;
}
// read peb
if (NtWow64ReadVirtualMemory64 == NULL) {
NtWow64ReadVirtualMemory64 =
(_NtWow64ReadVirtualMemory64)GetProcAddress(
GetModuleHandleA("ntdll.dll"),
"NtWow64ReadVirtualMemory64");
if (NtWow64ReadVirtualMemory64 == NULL) {
PyErr_SetString(PyExc_NotImplementedError,
"NtWow64ReadVirtualMemory64 missing");
goto error;
}
}
if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(hProcess,
pbi64.PebBaseAddress,
&peb64,
sizeof(peb64),
NULL))) {
PyErr_SetFromWindowsErr(0);
goto error;
}
// read process parameters
if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(hProcess,
peb64.ProcessParameters,
&procParameters64,
sizeof(procParameters64),
NULL))) {
PyErr_SetFromWindowsErr(0);
goto error;
}
switch (kind) {
case KIND_CMDLINE:
src64 = procParameters64.CommandLine.Buffer;
size = procParameters64.CommandLine.Length;
break;
case KIND_CWD:
src64 = procParameters64.CurrentDirectoryPath.Buffer,
size = procParameters64.CurrentDirectoryPath.Length;
break;
case KIND_ENVIRON:
src64 = procParameters64.env;
break;
}
} else
#endif
/* Target process is of the same bitness as us. */
{
PROCESS_BASIC_INFORMATION pbi;
PEB_ peb;
RTL_USER_PROCESS_PARAMETERS_ procParameters;
if (!NT_SUCCESS(NtQueryInformationProcess(hProcess,
ProcessBasicInformation,
&pbi,
sizeof(pbi),
NULL))) {
PyErr_SetFromWindowsErr(0);
goto error;
}
// read peb
if (!ReadProcessMemory(hProcess,
pbi.PebBaseAddress,
&peb,
sizeof(peb),
NULL)) {
PyErr_SetFromWindowsErr(0);
goto error;
}
// read process parameters
if (!ReadProcessMemory(hProcess,
peb.ProcessParameters,
&procParameters,
sizeof(procParameters),
NULL)) {
PyErr_SetFromWindowsErr(0);
goto error;
}
switch (kind) {
case KIND_CMDLINE:
src = procParameters.CommandLine.Buffer;
size = procParameters.CommandLine.Length;
break;
case KIND_CWD:
src = procParameters.CurrentDirectoryPath.Buffer;
size = procParameters.CurrentDirectoryPath.Length;
break;
case KIND_ENVIRON:
src = procParameters.env;
break;
}
}
if (kind == KIND_ENVIRON) {
#ifndef _WIN64
if (weAreWow64 && !theyAreWow64) {
ULONG64 size64;
if (psutil_get_process_region_size64(hProcess, src64, &size64) != 0)
goto error;
size = (SIZE_T)size64;
}
else
#endif
if (psutil_get_process_region_size(hProcess, src, &size) != 0)
goto error;
}
buffer = calloc(size + 2, 1);
if (buffer == NULL) {
PyErr_NoMemory();
goto error;
}
#ifndef _WIN64
if (weAreWow64 && !theyAreWow64) {
if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(hProcess,
src64,
buffer,
size,
NULL))) {
PyErr_SetFromWindowsErr(0);
goto error;
}
} else
#endif
if (!ReadProcessMemory(hProcess, src, buffer, size, NULL)) {
PyErr_SetFromWindowsErr(0);
goto error;
}
CloseHandle(hProcess);
*pdata = buffer;
*psize = size;
return 0;
error:
if (hProcess != NULL)
CloseHandle(hProcess);
if (buffer != NULL)
free(buffer);
return -1;
}
/*
* returns a Python list representing the arguments for the process
* with given pid or NULL on error.
*/
PyObject *
psutil_get_cmdline(long pid) {
PyObject *ret = NULL;
WCHAR *data = NULL;
SIZE_T size;
PyObject *py_retlist = NULL;
PyObject *py_unicode = NULL;
LPWSTR *szArglist = NULL;
int nArgs, i;
if (psutil_get_process_data(pid, KIND_CMDLINE, &data, &size) != 0)
goto out;
// attempt to parse the command line using Win32 API
szArglist = CommandLineToArgvW(data, &nArgs);
if (szArglist == NULL) {
PyErr_SetFromWindowsErr(0);
goto out;
}
// arglist parsed as array of UNICODE_STRING, so convert each to
// Python string object and add to arg list
py_retlist = PyList_New(nArgs);
if (py_retlist == NULL)
goto out;
for (i = 0; i < nArgs; i++) {
py_unicode = PyUnicode_FromWideChar(szArglist[i],
wcslen(szArglist[i]));
if (py_unicode == NULL)
goto out;
PyList_SET_ITEM(py_retlist, i, py_unicode);
py_unicode = NULL;
}
ret = py_retlist;
py_retlist = NULL;
out:
LocalFree(szArglist);
free(data);
Py_XDECREF(py_unicode);
Py_XDECREF(py_retlist);
return ret;
}
PyObject *psutil_get_cwd(long pid) {
PyObject *ret = NULL;
WCHAR *data = NULL;
SIZE_T size;
if (psutil_get_process_data(pid, KIND_CWD, &data, &size) != 0)
goto out;
// convert wchar array to a Python unicode string
ret = PyUnicode_FromWideChar(data, wcslen(data));
out:
free(data);
return ret;
}
/*
* returns a Python string containing the environment variable data for the
* process with given pid or NULL on error.
*/
PyObject *psutil_get_environ(long pid) {
PyObject *ret = NULL;
WCHAR *data = NULL;
SIZE_T size;
if (psutil_get_process_data(pid, KIND_ENVIRON, &data, &size) != 0)
goto out;
// convert wchar array to a Python unicode string
ret = PyUnicode_FromWideChar(data, size / 2);
out:
free(data);
return ret;
}
#define PH_FIRST_PROCESS(Processes) ((PSYSTEM_PROCESS_INFORMATION)(Processes))
#define PH_NEXT_PROCESS(Process) ( \
((PSYSTEM_PROCESS_INFORMATION)(Process))->NextEntryOffset ? \
(PSYSTEM_PROCESS_INFORMATION)((PCHAR)(Process) + \
((PSYSTEM_PROCESS_INFORMATION)(Process))->NextEntryOffset) : \
NULL)
const int STATUS_INFO_LENGTH_MISMATCH = 0xC0000004;
const int STATUS_BUFFER_TOO_SMALL = 0xC0000023L;
/*
* Given a process PID and a PSYSTEM_PROCESS_INFORMATION structure
* fills the structure with various process information by using
* NtQuerySystemInformation.
* We use this as a fallback when faster functions fail with access
* denied. This is slower because it iterates over all processes.
* On success return 1, else 0 with Python exception already set.
*/
int
psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess,
PVOID *retBuffer) {
static ULONG initialBufferSize = 0x4000;
NTSTATUS status;
PVOID buffer;
ULONG bufferSize;
PSYSTEM_PROCESS_INFORMATION process;
// get NtQuerySystemInformation
typedef DWORD (_stdcall * NTQSI_PROC) (int, PVOID, ULONG, PULONG);
NTQSI_PROC NtQuerySystemInformation;
HINSTANCE hNtDll;
hNtDll = LoadLibrary(TEXT("ntdll.dll"));
NtQuerySystemInformation = (NTQSI_PROC)GetProcAddress(
hNtDll, "NtQuerySystemInformation");
bufferSize = initialBufferSize;
buffer = malloc(bufferSize);
if (buffer == NULL) {
PyErr_NoMemory();
goto error;
}
while (TRUE) {
status = NtQuerySystemInformation(SystemProcessInformation, buffer,
bufferSize, &bufferSize);
if (status == STATUS_BUFFER_TOO_SMALL ||
status == STATUS_INFO_LENGTH_MISMATCH)
{
free(buffer);
buffer = malloc(bufferSize);
if (buffer == NULL) {
PyErr_NoMemory();
goto error;
}
}
else {
break;
}
}
if (status != 0) {
PyErr_Format(
PyExc_RuntimeError, "NtQuerySystemInformation() syscall failed");
goto error;
}
if (bufferSize <= 0x20000)
initialBufferSize = bufferSize;
process = PH_FIRST_PROCESS(buffer);
do {
if (process->UniqueProcessId == (HANDLE)pid) {
*retProcess = process;
*retBuffer = buffer;
return 1;
}
} while ( (process = PH_NEXT_PROCESS(process)) );
NoSuchProcess();
goto error;
error:
FreeLibrary(hNtDll);
if (buffer != NULL)
free(buffer);
return 0;
}
+34
View File
@@ -0,0 +1,34 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*/
#if !defined(__PROCESS_INFO_H)
#define __PROCESS_INFO_H
#include <Python.h>
#include <windows.h>
#include "security.h"
#include "ntextapi.h"
#define HANDLE_TO_PYNUM(handle) PyLong_FromUnsignedLong((unsigned long) handle)
#define PYNUM_TO_HANDLE(obj) ((HANDLE)PyLong_AsUnsignedLong(obj))
DWORD* psutil_get_pids(DWORD *numberOfReturnedPIDs);
HANDLE psutil_handle_from_pid(DWORD pid);
HANDLE psutil_handle_from_pid_waccess(DWORD pid, DWORD dwDesiredAccess);
PyObject* psutil_win32_OpenProcess(PyObject *self, PyObject *args);
PyObject* psutil_win32_CloseHandle(PyObject *self, PyObject *args);
int psutil_handlep_is_running(HANDLE hProcess);
int psutil_pid_in_proclist(DWORD pid);
int psutil_pid_is_running(DWORD pid);
PyObject* psutil_get_cmdline(long pid);
PyObject* psutil_get_cwd(long pid);
PyObject* psutil_get_environ(long pid);
int psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess,
PVOID *retBuffer);
#endif
+225
View File
@@ -0,0 +1,225 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Security related functions for Windows platform (Set privileges such as
* SeDebug), as well as security helper functions.
*/
#include <windows.h>
#include <Python.h>
/*
* Convert a process handle to a process token handle.
*/
HANDLE
psutil_token_from_handle(HANDLE hProcess) {
HANDLE hToken = NULL;
if (! OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
return PyErr_SetFromWindowsErr(0);
return hToken;
}
/*
* http://www.ddj.com/windows/184405986
*
* There's a way to determine whether we're running under the Local System
* account. However (you guessed it), we have to call more Win32 functions to
* determine this. Backing up through the code listing, we need to make another
* call to GetTokenInformation, but instead of passing through the TOKEN_USER
* constant, we pass through the TOKEN_PRIVILEGES constant. This value returns
* an array of privileges that the account has in the environment. Iterating
* through the array, we call the function LookupPrivilegeName looking for the
* string SeTcbPrivilege. If the function returns this string, then this
* account has Local System privileges
*/
int
psutil_has_system_privilege(HANDLE hProcess) {
DWORD i;
DWORD dwSize = 0;
DWORD dwRetval = 0;
TCHAR privName[256];
DWORD dwNameSize = 256;
// PTOKEN_PRIVILEGES tp = NULL;
BYTE *pBuffer = NULL;
TOKEN_PRIVILEGES *tp = NULL;
HANDLE hToken = psutil_token_from_handle(hProcess);
if (NULL == hToken)
return -1;
// call GetTokenInformation first to get the buffer size
if (! GetTokenInformation(hToken, TokenPrivileges, NULL, 0, &dwSize)) {
dwRetval = GetLastError();
// if it failed for a reason other than the buffer, bail out
if (dwRetval != ERROR_INSUFFICIENT_BUFFER ) {
PyErr_SetFromWindowsErr(dwRetval);
return 0;
}
}
// allocate buffer and call GetTokenInformation again
// tp = (PTOKEN_PRIVILEGES) GlobalAlloc(GPTR, dwSize);
pBuffer = (BYTE *) malloc(dwSize);
if (pBuffer == NULL) {
PyErr_NoMemory();
return -1;
}
if (! GetTokenInformation(hToken, TokenPrivileges, pBuffer,
dwSize, &dwSize))
{
PyErr_SetFromWindowsErr(0);
free(pBuffer);
return -1;
}
// convert the BYTE buffer to a TOKEN_PRIVILEGES struct pointer
tp = (TOKEN_PRIVILEGES *)pBuffer;
// check all the privileges looking for SeTcbPrivilege
for (i = 0; i < tp->PrivilegeCount; i++) {
// reset the buffer contents and the buffer size
strcpy(privName, "");
dwNameSize = sizeof(privName) / sizeof(TCHAR);
if (! LookupPrivilegeName(NULL,
&tp->Privileges[i].Luid,
(LPTSTR)privName,
&dwNameSize))
{
PyErr_SetFromWindowsErr(0);
free(pBuffer);
return -1;
}
// if we find the SeTcbPrivilege then it's a LocalSystem process
if (! lstrcmpi(privName, TEXT("SeTcbPrivilege"))) {
free(pBuffer);
return 1;
}
}
free(pBuffer);
return 0;
}
BOOL
psutil_set_privilege(HANDLE hToken, LPCTSTR Privilege, BOOL bEnablePrivilege) {
TOKEN_PRIVILEGES tp;
LUID luid;
TOKEN_PRIVILEGES tpPrevious;
DWORD cbPrevious = sizeof(TOKEN_PRIVILEGES);
if (!LookupPrivilegeValue( NULL, Privilege, &luid )) return FALSE;
// first pass. get current privilege setting
tp.PrivilegeCount = 1;
tp.Privileges[0].Luid = luid;
tp.Privileges[0].Attributes = 0;
AdjustTokenPrivileges(
hToken,
FALSE,
&tp,
sizeof(TOKEN_PRIVILEGES),
&tpPrevious,
&cbPrevious
);
if (GetLastError() != ERROR_SUCCESS) return FALSE;
// second pass. set privilege based on previous setting
tpPrevious.PrivilegeCount = 1;
tpPrevious.Privileges[0].Luid = luid;
if (bEnablePrivilege)
tpPrevious.Privileges[0].Attributes |= (SE_PRIVILEGE_ENABLED);
else
tpPrevious.Privileges[0].Attributes ^=
(SE_PRIVILEGE_ENABLED & tpPrevious.Privileges[0].Attributes);
AdjustTokenPrivileges(
hToken,
FALSE,
&tpPrevious,
cbPrevious,
NULL,
NULL
);
if (GetLastError() != ERROR_SUCCESS) return FALSE;
return TRUE;
}
int
psutil_set_se_debug() {
HANDLE hToken;
if (! OpenThreadToken(GetCurrentThread(),
TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
FALSE,
&hToken)
) {
if (GetLastError() == ERROR_NO_TOKEN) {
if (!ImpersonateSelf(SecurityImpersonation)) {
CloseHandle(hToken);
return 0;
}
if (!OpenThreadToken(GetCurrentThread(),
TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
FALSE,
&hToken)
) {
RevertToSelf();
CloseHandle(hToken);
return 0;
}
}
}
// enable SeDebugPrivilege (open any process)
if (! psutil_set_privilege(hToken, SE_DEBUG_NAME, TRUE)) {
RevertToSelf();
CloseHandle(hToken);
return 0;
}
RevertToSelf();
CloseHandle(hToken);
return 1;
}
int
psutil_unset_se_debug() {
HANDLE hToken;
if (! OpenThreadToken(GetCurrentThread(),
TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
FALSE,
&hToken)
) {
if (GetLastError() == ERROR_NO_TOKEN) {
if (! ImpersonateSelf(SecurityImpersonation))
return 0;
if (!OpenThreadToken(GetCurrentThread(),
TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
FALSE,
&hToken))
{
return 0;
}
}
}
// now disable SeDebug
if (! psutil_set_privilege(hToken, SE_DEBUG_NAME, FALSE))
return 0;
CloseHandle(hToken);
return 1;
}
+17
View File
@@ -0,0 +1,17 @@
/*
* Copyright (c) 2009, Jay Loden, Giampaolo Rodola'. All rights reserved.
* Use of this source code is governed by a BSD-style license that can be
* found in the LICENSE file.
*
* Security related functions for Windows platform (Set privileges such as
* SeDebug), as well as security helper functions.
*/
#include <windows.h>
BOOL psutil_set_privilege(HANDLE hToken, LPCTSTR Privilege, BOOL bEnablePrivilege);
HANDLE psutil_token_from_handle(HANDLE hProcess);
int psutil_has_system_privilege(HANDLE hProcess);
int psutil_set_se_debug();
int psutil_unset_se_debug();

Some files were not shown because too many files have changed in this diff Show More