Fix args passing for LifterBF & optimize tests

.gitignore

@@ -95,3 +95,5 @@ ENV/
 .ropeproject
 
 # End of https://www.gitignore.io/api/python
+
+.vscode/

angr_platforms/bf/__init__.py

@@ -1 +1,5 @@
-from . import arch_bf, load_bf, lift_bf, engine_bf, simos_bf
+from .arch_bf import ArchBF
+from .load_bf import BF
+from .lift_bf import LifterBF
+from .engine_bf import UberEngineWithBF
+from .simos_bf import SimBF, SimBFSyscall

angr_platforms/bf/lift_bf.py

@@ -1,16 +1,7 @@
-import archinfo
-import pyvex
+import bitstring
 from pyvex.lifting.util import *
 from pyvex.lifting import register
-from .arch_bf import ArchBF
-import bitstring
-import sys
-import os
-import claripy
-from angr import SimValueError
-import logging
 
-log = logging.getLogger("LifterBF")
 
 # This is actually a BrainFuck lifter for pyVEX.  I'm not joking.
 # Created by edg on 1/14/2017
@@ -246,19 +237,3 @@ class LifterBF(GymratLifter):
 
 # Tell PyVEX that this lifter exists.
 register(LifterBF, 'BF')
-
-if __name__ == '__main__':
-    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-    import logging
-    logging.getLogger('pyvex').setLevel(logging.DEBUG)
-    logging.basicConfig()
-
-    test1 = b'<>+-[].,'
-    test2 = b'<>+-[].,'
-    lifter = LifterBF(arch=archinfo.arch_from_id('bf'), addr=0)
-    lifter.lift(data=test1)
-    lifter.irsb.pp()
-
-    lifter = LifterBF(arch=ArchBF(), addr=0)
-    lifter.lift(data=test2)
-    lifter.irsb.pp()

tests/test_engine_bf.py

@@ -1,42 +1,47 @@
 import os
-import logging
-import pyvex
+import unittest
 import angr
+import pyvex
+
+from angr_platforms.bf import UberEngineWithBF
 
-from angr_platforms.bf.engine_bf import UberEngineWithBF
 
-def test_hello():
-    lifters = pyvex.lifting.lifters['BF']
-    pyvex.lifting.lifters['BF'] = []
-    try:
-        hellobf = str(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../test_programs/bf/hello.bf'))
-        p = angr.Project(hellobf, engine=UberEngineWithBF)
-        entry = p.factory.entry_state()
-        smgr = p.factory.simulation_manager(entry)
-        smgr.explore()
-        assert smgr.deadended[0].posix.dumps(1) == b'Hello World!\n'
-    finally:
-        pyvex.lifting.lifters['BF'] = lifters
+class TestBFEngine(unittest.TestCase):
+    # pylint:disable=missing-class-docstring
+    def test_hello(self):
+        lifters = pyvex.lifting.lifters['BF']
+        pyvex.lifting.lifters['BF'] = []
+        try:
+            hellobf = str(os.path.join(os.path.dirname(os.path.realpath(__file__)),'../test_programs/bf/hello.bf'))
+            p = angr.Project(hellobf, engine=UberEngineWithBF)
+            entry = p.factory.entry_state()
+            smgr = p.factory.simulation_manager(entry)
+            smgr.explore()
+            assert smgr.deadended[0].posix.dumps(1) == b'Hello World!\n'
+        finally:
+            # It's designed to only have "finally" block, no "except" blocks.
+            # We want to make sure lifters['BF'] is restored after the test,
+            # so that other code won't complain about it, while still being
+            # able to detect any test failure in the "try" block
+            pyvex.lifting.lifters['BF'] = lifters
 
-def test_1bytecrackme_good():
-    lifters = pyvex.lifting.lifters['BF']
-    pyvex.lifting.lifters['BF'] = []
-    try:
-        crackme = str(os.path.join(os.path.dirname(os.path.realpath(__file__)),
-            '../test_programs/bf/1bytecrackme-good.bf'))
-        bad_states = lambda state: b"-" in state.posix.dumps(1)
-        p = angr.Project(crackme, engine=UberEngineWithBF)
-        p.arch.vex_arch = None  # force test with engine
-        entry = p.factory.entry_state(remove_options={angr.options.LAZY_SOLVES})
-        smgr = p.factory.simulation_manager(entry)
-        smgr.run(until=lambda lsmgr: len(lsmgr.active) == 0)
-        smgr.stash(from_stash="deadended", to_stash="bad", filter_func=bad_states)
-        assert b"\n" == smgr.deadended[0].posix.dumps(0)
-    finally:
-        pyvex.lifting.lifters['BF'] = lifters
+    def test_1bytecrackme_good(self):
+        lifters = pyvex.lifting.lifters['BF']
+        pyvex.lifting.lifters['BF'] = []
+        try:
+            crackme = str(os.path.join(os.path.dirname(os.path.realpath(__file__)),
+                '../test_programs/bf/1bytecrackme-good.bf'))
+            p = angr.Project(crackme, engine=UberEngineWithBF)
+            p.arch.vex_arch = None  # force test with engine
+            entry = p.factory.entry_state(remove_options={angr.options.LAZY_SOLVES})
+            smgr = p.factory.simulation_manager(entry)
+            smgr.run(until=lambda lsmgr: len(lsmgr.active) == 0)
+            smgr.stash(from_stash="deadended", to_stash="bad", filter_func=lambda state: b"-" in state.posix.dumps(1))
+            assert b"\n" == smgr.deadended[0].posix.dumps(0)
+        finally:
+            # "finally" block only, no "except" blocks. See `test_hello()`
+            pyvex.lifting.lifters['BF'] = lifters
 
 
 if __name__ == '__main__':
-    logging.basicConfig(level=logging.INFO)
-    test_hello()
-    test_1bytecrackme_good()
+    unittest.main()

tests/test_lifter_bf.py

@@ -1,29 +1,47 @@
-import logging
 import os
+import unittest
 import angr
+import archinfo
+
+from angr_platforms.bf import ArchBF, LifterBF
+
+
+class TestBFLifter(unittest.TestCase):
+    # pylint:disable=missing-class-docstring
+    def test_lifter_bf(self):
+        # import logging
+        # logging.getLogger('pyvex').setLevel(logging.DEBUG)
+        # logging.basicConfig()
+        test1 = b'<>+-[].,'
+        test2 = b'<>+-[].,'
+        lifter = LifterBF(archinfo.arch_from_id('bf'), 0)
+        lifter.lift(data=test1)
+        lifter.irsb.pp()
+
+        lifter = LifterBF(ArchBF(), 0)
+        lifter.lift(data=test2)
+        lifter.irsb.pp()
+
+    def test_hello(self):
+        hellobf = str(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../test_programs/bf/hello.bf'))
+        p = angr.Project(hellobf)
+        entry = p.factory.entry_state()
+        smgr = p.factory.simulation_manager(entry)
+        smgr.explore()
+        assert smgr.deadended[0].posix.dumps(1) == b'Hello World!\n'
+
+    def test_1bytecrackme_good(self):
+        crackme = str(
+            os.path.join(os.path.dirname(os.path.realpath(__file__)), '../test_programs/bf/1bytecrackme-good.bf')
+        )
+        bad_states = lambda state: b"-" in state.posix.dumps(1)
+        p = angr.Project(crackme)
+        entry = p.factory.entry_state(remove_options={angr.options.LAZY_SOLVES})
+        smgr = p.factory.simulation_manager(entry)
+        smgr.run(until=lambda lsmgr: len(lsmgr.active) == 0)
+        smgr.stash(from_stash="deadended", to_stash="bad", filter_func=bad_states)
+        assert b"\n" == smgr.deadended[0].posix.dumps(0)
 
-import angr_platforms.bf
-
-def test_hello():
-    hellobf = str(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../test_programs/bf/hello.bf'))
-    p = angr.Project(hellobf)
-    entry = p.factory.entry_state()
-    smgr = p.factory.simulation_manager(entry)
-    smgr.explore()
-    assert smgr.deadended[0].posix.dumps(1) == b'Hello World!\n'
-
-def test_1bytecrackme_good():
-    crackme = str(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../test_programs/bf/1bytecrackme-good.bf'))
-    bad_states = lambda state: b"-" in state.posix.dumps(1)
-    p = angr.Project(crackme)
-    entry = p.factory.entry_state(remove_options={angr.options.LAZY_SOLVES})
-    smgr = p.factory.simulation_manager(entry)
-    smgr.run(until=lambda lsmgr: len(lsmgr.active) == 0)
-    smgr.stash(from_stash="deadended", to_stash="bad", filter_func=bad_states)
-    assert b"\n" == smgr.deadended[0].posix.dumps(0)
 
 if __name__ == '__main__':
-    logging.getLogger('pyvex').setLevel(logging.DEBUG)
-    logging.basicConfig()
-    test_hello()
-    test_1bytecrackme_good()
+    unittest.main()