pwntools/pwnlib/args.py

"""
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 :mod:`pwnlib.args.args`, which is exposed as the global variable
:data:`args`.  Arguments which ``pwntools`` reserves internally are not exposed
this way.  These reserved arguments (such as ``DEBUG`` and ``NOASLR``) instead
adjust :data:`pwnlib.context.context`, so their effect is read from ``context``
(e.g. ``context.log_level``, ``context.aslr``) rather than from ``args`` --
``args.NOASLR`` is always the empty string.

.. 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')

Arguments can also be accessed directly with the dot operator, e.g.:

.. code-block:: python

    if args.REMOTE:
        ...

Any undefined arguments evaluate to an empty string, ``''``.

The full list of supported "magic arguments" and their effects are listed
below.

"""
import collections
import logging
import os
import string
import sys

from pwnlib import term
from pwnlib.context import context

class PwnlibArgs(collections.defaultdict):
    def __getattr__(self, attr):
        if attr.startswith('_'):
            raise AttributeError(attr)
        return self[attr]

args = PwnlibArgs(str)
term_mode  = True
env_prefix = 'PWNLIB_'
free_form  = True

# Check to see if we were invoked as one of the 'pwn xxx' scripts.
# If so, we don't want to remove e.g. "SYS_" from the end of the command
# line, as this breaks things like constgrep.
import pwnlib.commandline
basename = os.path.basename(sys.argv[0])

if basename == 'pwn' or basename in pwnlib.commandline.__all__:
    free_form = False


def isident(s):
    """
    Helper function to check whether a string is a valid identifier,
    as passed in on the command-line.
    """
    first = string.ascii_uppercase + '_'
    body = string.digits + first
    if not s:
        return False
    if s[0] not in first:
        return False
    if not all(c in body for c in s[1:]):
        return False
    return True

def asbool(s):
    """
    Convert a string to its boolean value
    """
    if   s.lower() == 'true':
        return True
    elif s.lower() == 'false':
        return False
    elif s.isdigit():
        return bool(int(s))
    else:
        raise ValueError('must be integer or boolean: %r' % s)

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 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 SILENT(x):
    """Sets the logging verbosity to ``error`` which silences most
    output."""
    LOG_LEVEL('error')

def DEBUG(x):
    """Sets the logging verbosity to ``debug`` which displays much
    more information, including logging each byte sent by tubes."""
    LOG_LEVEL('debug')

def NOTERM(v):
    """Disables pretty terminal settings and animations."""
    if asbool(v):
        global term_mode
        term_mode = False

def TIMEOUT(v):
    """Sets a timeout for tube operations (in seconds) via
    ``context.timeout``, e.g. ``TIMEOUT=30``"""
    context.defaults['timeout'] = int(v)

def RANDOMIZE(v):
    """Enables randomization of various pieces via ``context.randomize``"""
    context.defaults['randomize'] = asbool(v)

def NOASLR(v):
    """Disables ASLR via ``context.aslr``"""
    context.defaults['aslr'] = not asbool(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

def LOCAL_LIBCDB(v):
    """Sets path to local libc-database via ``context.local_libcdb``, e.g. 
    ``LOCAL_LIBCDB='/path/to/libc-databse'``"""
    context.local_libcdb = v

hooks = {
    'LOG_LEVEL': LOG_LEVEL,
    'LOG_FILE': LOG_FILE,
    'DEBUG': DEBUG,
    'NOTERM': NOTERM,
    'SILENT': SILENT,
    'RANDOMIZE': RANDOMIZE,
    'TIMEOUT': TIMEOUT,
    'NOASLR': NOASLR,
    'NOPTRACE': NOPTRACE,
    'STDERR': STDERR,
    'LOCAL_LIBCDB': LOCAL_LIBCDB,
}

def initialize():
    global args, term_mode

    # Hack for readthedocs.org
    if 'READTHEDOCS' in os.environ:
        os.environ['PWNLIB_NOTERM'] = '1'

    for k, v in os.environ.items():
        if not k.startswith(env_prefix):
            continue
        k = k[len(env_prefix):]

        if k in hooks:
            hooks[k](v)
        elif isident(k):
            args[k] = v

    argv = sys.argv[:]
    for arg in sys.argv[:]:
        orig  = arg
        value = 'True'

        if '=' in arg:
            arg, value = arg.split('=', 1)

        if arg in hooks:
            sys.argv.remove(orig)
            hooks[arg](value)

        elif free_form and isident(arg):
            sys.argv.remove(orig)
            args[arg] = value

    if term_mode:
        term.init()