Merge pull request #121 from cbisnett/add_with_statement_support_to_b…

…locks

Add with statement support to blocks

boofuzz/__init__.py

@@ -126,13 +126,12 @@ def s_switch(name):
 # ## BLOCK MANAGEMENT
 
 
-def s_block_start(name, group=None, encoder=None, dep=None, dep_value=None, dep_values=(), dep_compare="=="):
-    # TODO: Either convert this to a with() statement, or add a new one that is compatible.
+def s_block(name, group=None, encoder=None, dep=None, dep_value=None, dep_values=(), dep_compare="=="):
     """
-    Open a new block under the current request. This routine always returns True so you can make your fuzzer pretty
-    with indenting::
+    Open a new block under the current request. The returned instance supports the "with" interface so it will
+    be automatically closed for you::
 
-        if s_block_start("header"):
+        with s_block("header"):
             s_static("\\x00\\x01")
             if s_block_start("body"):
                 ...
@@ -152,11 +151,45 @@ def s_block_start(name, group=None, encoder=None, dep=None, dep_value=None, dep_
     @type  dep_compare: str
     @param dep_compare: (Optional, def="==") Comparison method to use on dependency (==, !=, >, >=, <, <=)
     """
+    class ScopedBlock(Block):
+        def __init__(self, block):
+            self.block = block
+
+        def __enter__(self):
+            """
+            Setup before entering the "with" statement body
+            """
+            return self.block
+
+        def __exit__(self, type, value, traceback):
+            """
+            Cleanup after executing the "with" statement body
+            """
+            # Automagically close the block when exiting the "with" statement
+            s_block_end()
 
-    block = Block(name, blocks.CURRENT, group, encoder, dep, dep_value, dep_values, dep_compare)
+    block = s_block_start(name, group, encoder, dep, dep_value, dep_values, dep_compare)
+
+    return ScopedBlock(block)
+
+def s_block_start(name, *args, **kwargs):
+    """
+    Open a new block under the current request. This routine always returns an instance so you can make your fuzzer pretty
+    with indenting::
+
+        if s_block_start("header"):
+            s_static("\\x00\\x01")
+            if s_block_start("body"):
+                ...
+        s_block_close()
+
+    @note Prefer using s_block to this function directly
+    @see s_block
+    """
+    block = Block(name, blocks.CURRENT, *args, **kwargs)
     blocks.CURRENT.push(block)
 
-    return True
+    return block
 
 
 # noinspection PyUnusedLocal

unit_tests/blocks.py

@@ -6,6 +6,7 @@ def run():
     dependencies()
     repeaters()
     return_current_mutant()
+    with_statements()
 
     # clear out the requests.
     blocks.REQUESTS = {}
@@ -176,3 +177,14 @@ def return_current_mutant():
         req1.mutate()
     assert (req1.mutant.name == "uhntiss")
     req1.reset()
+
+def with_statements():
+    s_initialize("WITH TEST")
+
+    with s_block("BLOCK1") as b:
+        assert (b.name == "BLOCK1")
+        s_static("test")
+
+    req = s_get("WITH TEST")
+    assert (req.num_mutations() == 0)
+    assert (req.render() == "test")