Fix and improve AFL linux example

examples/fuzzing/linux_x8664/fuzz_x8664_linux.py

@@ -17,35 +17,44 @@ Steps:
     $ rm -fr afl_outputs/default/
 """
 
-# No more need for importing unicornafl, try afl.ql_afl_fuzz instead!
-
 import os
 import sys
 
-from typing import Optional
+from typing import Sequence
 
-sys.path.append("../../..")
+QLHOME = os.path.realpath(r'../../..')
+
+sys.path.append(QLHOME)
 from qiling import Qiling
 from qiling.const import QL_VERBOSE
 from qiling.extensions import pipe
 from qiling.extensions import afl
 
-def main(input_file: str):
-    ql = Qiling(["./x8664_fuzz"], "../../rootfs/x8664_linux",
-        verbose=QL_VERBOSE.OFF, # keep qiling logging off
-        console=False)          # thwart program output
+
+def main(argv: Sequence[str], rootfs: str, infilename: str):
+    # initialize a qiling instance.
+    # note we keep verbosity off and thwart the program's output to gain some speed-up
+    ql = Qiling(argv, rootfs, verbose=QL_VERBOSE.OFF, console=False)
+
+    # get the image base address
+    img = ql.loader.get_image_by_name('x8664_fuzz')
+    assert img is not None
+
+    # fuzzing scope: the main function
+    main_begins = img.base + 0x1275
+    main_ends = img.base + 0x1293
 
     # redirect stdin to our mock to feed it with incoming fuzzed keystrokes
     ql.os.stdin = pipe.SimpleInStream(sys.stdin.fileno())
 
-    def place_input_callback(ql: Qiling, input: bytes, persistent_round: int) -> Optional[bool]:
+    def place_input_callback(ql: Qiling, feed: bytes, round: int) -> bool:
         """Feed generated stimuli to the fuzzed target.
 
         This method is called with every fuzzing iteration.
         """
 
         # feed fuzzed input to our mock stdin
-        ql.os.stdin.write(input)
+        ql.os.stdin.write(feed)
 
         # signal afl to proceed with this input
         return True
@@ -54,23 +63,24 @@ def main(input_file: str):
         """Have Unicorn fork and start instrumentation.
         """
 
-        afl.ql_afl_fuzz(ql, input_file=input_file, place_input_callback=place_input_callback, exits=[ql.os.exit_point])
+        afl.ql_afl_fuzz(ql, infilename, place_input_callback, [main_ends])
 
-    # get image base address
-    ba = ql.loader.images[0].base
+    # set afl instrumentation [re]starting point
+    ql.hook_address(start_afl, main_begins)
+
+    def __crash(ql: Qiling) -> None:
+        os.abort()
 
     # make the process crash whenever __stack_chk_fail@plt is about to be called.
     # this way afl will count stack protection violations as crashes
-    ql.hook_address(callback=lambda x: os.abort(), address=ba + 0x126e)
-
-    # set afl instrumentation [re]starting point. we set it to 'main'
-    ql.hook_address(callback=start_afl, address=ba + 0x1275)
+    ql.hook_address(__crash, img.base + 0x126e)
 
     # okay, ready to roll
     ql.run()
 
 if __name__ == "__main__":
-    if len(sys.argv) == 1:
-        raise ValueError("No input file provided.")
-
-    main(sys.argv[1])
+    main(
+        rf'{QLHOME}/examples/fuzzing/linux_x8664/x8664_fuzz'.split(),
+        rf'{QLHOME}/examples/rootfs/x8664_linux',
+        rf'{QLHOME}/examples/fuzzing/linux_x8664/afl_inputs/a'
+    )