Add STDERR argument, and write docs for args module (#760)

Add STDERR argument, and write docs for args module

Fixes #759

Author: zachriggle

docs/source/args.rst

@@ -0,0 +1,9 @@
+.. testsetup:: *
+
+   from pwn import *
+
+:mod:`pwnlib.args` --- Magic Command-Line Arguments
+=====================================================
+
+.. automodule:: pwnlib.args
+   :members:

docs/source/conf.py

@@ -55,6 +55,16 @@
 pwnlib.term.text.when = 'never'
 pwnlib.log.install_default_handler()
 pwnlib.log.rootlogger.setLevel(1)
+
+# Sphinx modifies sys.stdout, and context.log_terminal has
+# a reference to the original instance.  We need to update
+# it for logging to be captured.
+class stdout(object):
+    def __getattr__(self, name):
+        return getattr(sys.stdout, name)
+    def __setattr__(self, name, value):
+        return setattr(sys.stdout, name, value)
+pwnlib.context.ContextType.defaults['log_console'] = stdout()
 '''
 
 autodoc_member_order = 'alphabetical'

docs/source/index.rst

@@ -43,6 +43,7 @@ Each of the ``pwntools`` modules is documented here.
    :glob:
 
    adb
+   args
    asm
    atexception
    atexit

pwnlib/args.py

@@ -1,5 +1,43 @@
 #!/usr/bin/env python2
 """
+Pwntools exposes several magic command-line arguments and environment
+variables when operating in `from pwn import *` mode.
+
+The arguments extracted from the command-line and removed from ``sys.argv``.
+
+Arguments can be set by appending them to the command-line, or setting
+them in the environment prefixed by ``PWNLIB_``.
+
+The easiest example is to enable more verbose debugging.  Just set ``DEBUG``.
+
+.. code-block:: bash
+
+    $ PWNLIB_DEBUG=1 python exploit.py
+    $ python exploit.py DEBUG
+
+These arguments are automatically extracted, regardless of their name, and
+exposed via ``pwnlib.args.args``, which is exposed as the global variable
+``args``.  Arguments which ``pwntools`` reserves internally are not exposed
+this way.
+
+.. code-block:: bash
+
+    $ python -c 'from pwn import *; print args' A=1 B=Hello HOST=1.2.3.4 DEBUG
+    defaultdict(<type 'str'>, {'A': '1', 'HOST': '1.2.3.4', 'B': 'Hello'})
+
+This is very useful for conditional code, for example determining whether to
+run an exploit locally or to connect to a remote server.  Arguments which are
+not specified evaluate to an empty string.
+
+.. code-block:: python
+
+    if args['REMOTE']:
+        io = remote('exploitme.com', 4141)
+    else:
+        io = process('./pwnable')
+
+The full list of supported "magic arguments" and their effects are listed
+below.
 """
 import collections
 import logging
@@ -42,46 +80,67 @@ def asbool(s):
     else:
         raise ValueError('must be integer or boolean: %r' % s)
 
-def set_log_level(x):
+def LOG_LEVEL(x):
+    """Sets the logging verbosity used via ``context.log_level``,
+    e.g. ``LOG_LEVEL=debug``.
+    """
     with context.local(log_level=x):
         context.defaults['log_level']=context.log_level
 
-def set_log_file(x):
+def LOG_FILE(x):
+    """Sets a log file to be used via ``context.log_file``, e.g.
+    ``LOG_FILE=./log.txt``"""
     context.log_file=x
 
-def set_log_level_error(x):
-    set_log_level('error')
+def SILENT(x):
+    """Sets the logging verbosity to ``error`` which silences most
+    output."""
+    LOG_FILE('error')
 
-def set_log_level_debug(x):
-    set_log_level('debug')
+def DEBUG(x):
+    """Sets the logging verbosity to ``debug`` which displays much
+    more information, including logging each byte sent by tubes."""
+    LOG_FILE('debug')
 
-def set_noterm(v):
+def NOTERM(v):
+    """Disables pretty terminal settings and animations."""
     if asbool(v):
         global term_mode
         term_mode = False
 
-def set_timeout(v):
+def TIMEOUT(v):
+    """Sets a timeout for tube operations (in seconds) via
+    ``context.timeout``, e.g. ``TIMEOUT=30``"""
     context.defaults['timeout'] = int(v)
 
-def set_randomize(v):
+def RANDOMIZE(v):
+    """Enables randomization of various pieces via ``context.randomize``"""
     context.defaults['randomize'] = asbool(v)
 
-def set_aslr(v):
+def NOASLR(v):
+    """Disables ASLR via ``context.aslr``"""
     context.defaults['aslr'] = not asbool(v)
 
-def set_noptrace(v):
+def NOPTRACE(v):
+    """Disables facilities which require ``ptrace`` such as ``gdb.attach()``
+    statements, via ``context.noptrace``."""
     context.defaults['noptrace'] = asbool(v)
 
+def STDERR(v):
+    """Sends logging to ``stderr`` by default, instead of ``stdout``"""
+    context.log_console = sys.stderr
+
 hooks = {
-    'LOG_LEVEL': set_log_level,
-    'LOG_FILE': set_log_file,
-    'DEBUG': set_log_level_debug,
-    'NOTERM': set_noterm,
-    'SILENT': set_log_level_error,
-    'RANDOMIZE': set_randomize,
-    'TIMEOUT': set_timeout,
-    'NOASLR': set_aslr,
-    'NOPTRACE': set_noptrace,
+    'LOG_LEVEL': LOG_LEVEL,
+    'LOG_FILE': LOG_FILE,
+    'DEBUG': DEBUG,
+    'NOTERM': NOTERM,
+    'SILENT': SILENT,
+    'RANDOMIZE': RANDOMIZE,
+    'TIMEOUT': TIMEOUT,
+    'NOASLR': NOASLR,
+    'NOPTRACE': NOPTRACE,
+    'STDERR': STDERR,
 }
 
 def initialize():

pwnlib/commandline/common.py

@@ -5,8 +5,6 @@
 import pwnlib
 from pwnlib.context import context
 
-pwnlib.log.console.stream = sys.stderr
-
 choices = map(str, [16,32,64])
 choices += list(context.oses)
 choices += list(context.architectures)

pwnlib/commandline/main.py

@@ -1,3 +1,5 @@
+import sys
+
 from . import asm
 from . import checksec
 from . import common
@@ -15,6 +17,7 @@
 from . import unhex
 from . import update
 from .common import parser
+from ..context import context
 
 commands = {
     'asm': asm.main,
@@ -36,7 +39,8 @@
 
 def main():
     args = parser.parse_args()
-    commands[args.command](args)
+    with context.local(log_console = sys.stderr):
+        commands[args.command](args)
 
 if __name__ == '__main__':
     main()

pwnlib/context/__init__.py

@@ -334,6 +334,7 @@ class ContextType(object):
         'kernel': None,
         'log_level': logging.INFO,
         'log_file': _devnull(),
+        'log_console': sys.stdout,
         'randomize': False,
         'newline': '\n',
         'noptrace': False,
@@ -922,6 +923,24 @@ def log_file(self, value):
         value.flush()
         return value
 
+    @_validator
+    def log_console(self, stream):
+        """
+        Sets the default logging console target.
+
+        Examples:
+
+            >>> context.log_level = 'warn'
+            >>> log.warn("Hello")
+            [!] Hello
+            >>> context.log_console=open('/dev/null', 'w')
+            >>> log.warn("Hello")
+            >>> context.clear()
+        """
+        if isinstance(stream, str):
+            stream = open(stream, 'wt')
+        return stream
+
     @property
     def mask(self):
         return (1 << self.bits) - 1

pwnlib/log.py

@@ -113,7 +113,7 @@
 from .term import text
 
 # list of prefixes to use for the different message types.  note that the `text`
-# module won't add any escape codes if `sys.stderr.isatty()` is `False`
+# module won't add any escape codes if `pwnlib.context.log_console.isatty()` is `False`
 _msgtype_prefixes = {
     'status'       : [text.magenta, 'x'],
     'success'      : [text.bold_green, '+'],
@@ -490,10 +490,14 @@ class Handler(logging.StreamHandler):
     logger will not be emitted but rather an animated progress line will be
     created.
 
-    This handler outputs to ``sys.stderr``.
-
     An instance of this handler is added to the ``'pwnlib'`` logger.
     """
+    @property
+    def stream(self):
+        return context.log_console
+    @stream.setter
+    def stream(self, value):
+        pass
     def emit(self, record):
         """
         Emit a log record or create/update an animated progress logger
@@ -652,7 +656,7 @@ def handle(self, *a, **kw):
 #     logger.addHandler(myCoolPitchingHandler)
 #
 rootlogger = getLogger('pwnlib')
-console   = Handler(sys.stdout)
+console   = Handler()
 formatter = Formatter()
 console.setFormatter(formatter)
 
@@ -663,8 +667,6 @@ def install_default_handler():
     the ``pwnlib`` root logger.  This function is automatically called from when
     importing :mod:`pwn`.
     '''
-    console.stream = sys.stdout
-
     logger         = logging.getLogger('pwnlib')
 
     if console not in logger.handlers: