run_all_tests.sh and harness.py

[zachriggle]

run_all_tests.sh should be replaced with a simple scandown for unit tests. For example, pytest or nosetests scans for all files which start with test, and executes all functions whose name start with test.

harness.py can then be renamed to test.py and replaced with a simple wrapper script that does something like what's shown below.

Example test.py

#!/usr/bin/env python2
from pwn import *

def run_exploit(**kwargs):
    # set up the flag and target file
    write('flag', randoms(20, string.ascii_letters))
    saveflag = tempfile.NamedTemporaryFile()

    # Set up arguments
    global args
    args['SAVEFLAG'] = saveflag.name
    args['FLAG'] = 'flag'
    args.update(**kwargs)

    exploit = __import__('exploit', level=0)
    del sys.modules['exploit']
    del exploit

    # verify
    assert read(saveflag.name) == read('flag')

def test_local():
    'Run the exploit locally'
    run_exploit()

def test_remote():
    l = listen(0)
    l.spawn_process('pwnme')
    run_exploit(REMOTE='localhost', PORT=l.lport)

if __name__ == '__main__':
    test_local()
    test_remote()

py.test and nosetest example

Given this input script

pip install pytest nose
cat > test_foo.py <<EOF
from pwn import *

def test_normal_success():
    print "Lol"

def test_normal_error():
    print "Shucks!"
    raise Exception()

def test_success():
    log.info("Hurray!")

def test_failure():
    log.error("Oh no!")
EOF
PWNLIB_NOTERM=1 py.test
PWNLIB_NOTERM=1 nosetests

py.test output

=============================================================================== test session starts ===============================================================================
platform linux2 -- Python 2.7.8 -- py-1.4.26 -- pytest-2.6.4
collected 4 items

test_foo.py .F.F

==================================================================================== FAILURES =====================================================================================
________________________________________________________________________________ test_normal_error ________________________________________________________________________________

    def test_normal_error():
        print "Shucks!"
>       raise Exception()
E       Exception

test_foo.py:8: Exception
------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------
Shucks!
__________________________________________________________________________________ test_failure ___________________________________________________________________________________

    def test_failure():
>       log.error("Oh no!")

test_foo.py:14:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <pwnlib.log.Logger object at 0x7fc1eea42810>, m = 'Oh no!', a = (), kw = {'extra': {'pwnlib_stop': False, 'pwnlib_symbol': 'ERROR'}}

    def error(self, m, *a, **kw):
        """error(message)

            Logs an error message, and raises an ``Exception``.
            """
        self.__log(logging.ERROR, m, a, kw, text.on_red('ERROR'))
>       raise Exception(m)
E       Exception: Oh no!

/home/riggle/pwntools/pwnlib/log.py:134: Exception
------------------------------------------------------------------------------ Captured stdout call -------------------------------------------------------------------------------
[ERROR] Oh no!
======================================================================= 2 failed, 2 passed in 0.23 seconds ========================================================================

nosetests output

.E.E
======================================================================
ERROR: test_foo.test_normal_error
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/riggle/.pyenv/versions/2.7.8/lib/python2.7/site-packages/nose/case.py", line 197, in runTest
    self.test(*self.arg)
  File "/media/SSD1T/riggle/yyy/test_foo.py", line 8, in test_normal_error
    raise Exception()
Exception:
-------------------- >> begin captured stdout << ---------------------
Shucks!

--------------------- >> end captured stdout << ----------------------

======================================================================
ERROR: test_foo.test_failure
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/riggle/.pyenv/versions/2.7.8/lib/python2.7/site-packages/nose/case.py", line 197, in runTest
    self.test(*self.arg)
  File "/media/SSD1T/riggle/yyy/test_foo.py", line 14, in test_failure
    log.error("Oh no!")
  File "/home/riggle/pwntools/pwnlib/log.py", line 134, in error
    raise Exception(m)
Exception: Oh no!
-------------------- >> begin captured stdout << ---------------------
[ERROR] Oh no!

--------------------- >> end captured stdout << ----------------------
-------------------- >> begin captured logging << --------------------
pwnlib.exploit: ERROR: Oh no!
--------------------- >> end captured logging << ---------------------

----------------------------------------------------------------------
Ran 4 tests in 0.126s

FAILED (errors=2)