Merge pull request #7 from quarkslab/broker-replay-mechanism

Broker replay mechanism

bin/pastis-benchmark

@@ -158,20 +158,41 @@ def showmap(bins: str):
 @click.option("--allow-remote", type=bool, is_flag=True, default=False, help="Enable remote connection")
 @click.option('--probe', type=str, help="Probe to load as a python module (should contain a ProbeInterface)", multiple=True)
 @click.option('--skip-cpufreq', is_flag=True, type=bool, default=False, help="Skip CPU frequency scaling check")
+@click.option('--mem-threshold', type=int, default=85, help="RAM consumption limit", show_default=True)
 @click.option('--start-quorum', type=int, default=0, help="Number of client connection to receive before triggering startup", show_default=True)
-@click.option('--proxy', type=str, default="", help="Run the broker as a proxy to another broker")
+@click.option('--filter-inputs', type=bool, is_flag=True, default=False, help="Filter inputs that do not generate coverage", show_default=True)
+@click.option('--stream', type=bool, is_flag=True, default=False, help="Stream input and coverage info in the given file", show_default=True)
+@click.option('--replay-threads', type=int, default=4, help="number of threads to use for input replay", show_default=True)
+@click.option('--proxy', type=str, default="", help="Run the broker as a proxy to another broker: pymodule@ip:port")
+@click.argument('pargs', nargs=-1)
 def run(workspace: str, bins: str, seeds: str, mode: str, injloc: str, aflpp: bool, hfuzz: bool, triton: bool,
         debug: bool, timeout: Optional[int], port: int, hfuzz_path: str, hfuzz_threads: int, spawn: bool,
-        allow_remote: bool, probe: Tuple[str], skip_cpufreq: bool, start_quorum: int, proxy: str):
+        allow_remote: bool, probe: Tuple[str], skip_cpufreq: bool,  mem_threshold: int, start_quorum: int,  proxy: str,
+        filter_inputs: bool, stream: bool, replay_threads: int, pargs: Tuple[str]):
 
     configure_logging(logging.DEBUG if debug else logging.INFO, "%(asctime)s %(name)s [%(levelname)s] %(message)s")
 
-    broker = PastisBroker(workspace, bins, BrokingMode[mode], CheckMode.CHECK_ALL, SeedInjectLoc[injloc], None, [], start_quorum=start_quorum)
+    broker = PastisBroker(workspace,
+                          bins,
+                          BrokingMode[mode],
+                          CheckMode.CHECK_ALL,
+                          SeedInjectLoc[injloc],
+                          None,
+                          list(pargs),
+                          mem_threshold,
+                          start_quorum,
+                          filter_inputs,
+                          stream,
+                          replay_threads)
 
     if proxy:  # proxy format should be:  IP:port@py_module
-        url, py_module = proxy.split("@")
-        proxy_ip, proxy_port = url.split(":")
-        broker.set_proxy(proxy_ip, int(proxy_port), py_module)
+        try:
+            py_module, url = proxy.split("@")
+            proxy_ip, proxy_port = url.split(":")
+            broker.set_proxy(proxy_ip, int(proxy_port), py_module)
+        except ValueError:
+            logging.error(f"Cannot parse proxy: {proxy}, format need to be: py_module@ip:port")
+            return
 
     # Add all given seeds as initial seed
     if seeds is None:

bin/pastis-broker

@@ -54,10 +54,14 @@ def iterate_file(file):
 @click.option('-p', '--port', type=int, default=5555, help="Port to bind to", multiple=False, show_default=True)
 @click.option('--mem-threshold', type=int, default=85, help="RAM consumption limit", show_default=True)
 @click.option('--start-quorum', type=int, default=0, help="Number of client connection to receive before triggering startup", show_default=True)
+@click.option('--filter-inputs', type=bool, is_flag=True, default=False, help="Filter inputs that do not generate coverage", show_default=True)
+@click.option('--stream', type=bool, is_flag=True, default=False, help="Stream input and coverage info in the given file", show_default=True)
+@click.option('--replay-threads', type=int, default=4, help="number of threads to use for input replay", show_default=True)
 @click.argument('pargvs', nargs=-1)
 def main(workspace: str, sast_report: Optional[str], bins: str, mode: str, chkmode: str, injloc: str, engine: Tuple[str],
          tt_config: Optional[str], hf_config: Optional[str], seed: Tuple[str], timeout: Optional[int],
-         port: Optional[int], pargvs: Tuple[str], mem_threshold: int, start_quorum: int):
+         port: Optional[int], pargvs: Tuple[str], mem_threshold: int, start_quorum: int, filter_inputs: bool,
+         stream: bool, replay_threads: int):
     global broker
     # Instanciate the broker
 
@@ -66,7 +70,18 @@ def main(workspace: str, sast_report: Optional[str], bins: str, mode: str, chkmo
         logging.error(f"Check mode {chkmode.name} requires a SAST report (use -r) to provide it")
         sys.exit(1)
 
-    broker = PastisBroker(workspace, bins, BrokingMode[mode], chkmode, SeedInjectLoc[injloc], sast_report, list(pargvs), mem_threshold, start_quorum)
+    broker = PastisBroker(workspace,
+                          bins,
+                          BrokingMode[mode],
+                          chkmode,
+                          SeedInjectLoc[injloc],
+                          sast_report,
+                          list(pargvs),
+                          mem_threshold,
+                          start_quorum,
+                          filter_inputs,
+                          stream,
+                          replay_threads)
 
     # Preload all Fuzzing engine if needed
     for eng in engine:

engines/pastis-aflpp/pastisaflpp/aflpp.py

@@ -46,9 +46,14 @@ def find_alfpp_binary(root_dir: Union[Path, str]) -> Optional[Path]:
             aflpp_path = os.environ.get(AFLPPProcess.AFLPP_ENV_VAR)
             return Path(aflpp_path) / 'afl-fuzz' if aflpp_path else shutil.which(AFLPPProcess.BINARY)
 
-    def start(self, target: str, target_arguments: str, workspace: Workspace, exmode: ExecMode, fuzzmode: FuzzMode, stdin: bool, engine_args: str, cmplog: Optional[str] = None, dictionary: Optional[str] = None):
+    def start(self, target: str, target_arguments: list[str], workspace: Workspace, exmode: ExecMode, fuzzmode: FuzzMode, stdin: bool, engine_args: str, cmplog: Optional[str] = None, dictionary: Optional[str] = None):
+        # Check that we have '@@' if input provided via argv
+        if not stdin:
+            if "@@" not in target_arguments:
+                logging.error(f"seed provided via ARGV but can't find '@@' on program argv")
+                return
         # Build target command line.
-        target_cmdline = f"{target} {target_arguments}"
+        target_cmdline = f"{target} {' '.join(target_arguments)}"
 
         # Build fuzzer arguments.
         # NOTE: Assuming the target receives inputs from stdin.
@@ -78,10 +83,6 @@ def start(self, target: str, target_arguments: str, workspace: Workspace, exmode
         # Build fuzzer command line.
         aflpp_cmdline = f'{self.__path} {aflpp_arguments} -- {target_cmdline}'
 
-        # NOTE: Assuming fixed location for the input file.
-        if not stdin:
-            aflpp_cmdline += " @@"
-
         logging.info(f"Run AFL++: {aflpp_cmdline}")
         logging.debug(f"\tWorkspace: {workspace.root_dir}")
 

engines/pastis-aflpp/pastisaflpp/driver.py

@@ -78,7 +78,7 @@ def start(self, package: BinaryPackage, argv: List[str], exmode: ExecMode, fuzzm
 
         logging.info(f"Start process (injectloc: {seed_inj.name})")
         self.aflpp.start(str(package.executable_path.absolute()),
-                         " ".join(argv),
+                         argv,
                          self.workspace,
                          exmode,
                          fuzzmode,

engines/pastis-honggfuzz/pastishf/driver.py

@@ -78,15 +78,15 @@ def start(self, package: BinaryPackage, argv: List[str], exmode: ExecMode, fuzzm
         self.workspace.start()  # Start looking at directories
 
         logging.info("Start process")
-        if not self.honggfuzz.start(self.__package.executable_path.absolute(),
-                             " ".join(argv),
+        if not self.honggfuzz.start(str(self.__package.executable_path.absolute()),
+                             argv,
                              self.workspace,
                              exmode,
                              fuzzmode,
                              seed_inj == SeedInjectLoc.STDIN,
                              engine_args,
                              str(package.dictionary.absolute()) if package.dictionary else None):
-            self._agent.send_log(LogLevel.ERROR, "Cannot start target, HFQBDIPRELAOD not found")
+            self._agent.send_log(LogLevel.ERROR, "Cannot start target")
         self._started = True
 
         # Start the replay worker (note that the queue might already have started to be filled by agent thread)

engines/pastis-honggfuzz/pastishf/honggfuzz.py

@@ -44,10 +44,19 @@ def __init__(self, path: str = None):
 
         self.__process = None
 
-    def start(self, target: str, target_arguments: str, workspace: Workspace, exmode: ExecMode, fuzzmode: FuzzMode,
+    def start(self, target: str, target_arguments: list[str], workspace: Workspace, exmode: ExecMode, fuzzmode: FuzzMode,
               stdin: bool, engine_args: str, dictionary: Optional[str] = None) -> bool:
+        if not stdin:
+            if "@@" in target_arguments:  # Change '@@' for ___FILE___
+                idx = target_arguments.index("@@")
+                target_arguments[idx] = "___FILE___"
+            else:
+                if "___FILE___" not in target_arguments:
+                    logging.error(f"seed provided via ARGV but can't find '@@'/___FILE___ on program argv")
+                    return False
+
         # Build target command line.
-        target_cmdline = f"{target} {target_arguments}"
+        target_cmdline = f"{target} {' '.join(target_arguments)}"
 
         HFQBDI_LIB_PATH = os.getenv('HFQBDI_LIB_PATH')
 
@@ -78,10 +87,6 @@ def start(self, target: str, target_arguments: str, workspace: Workspace, exmode
         # Build fuzzer command line.
         hfuzz_cmdline = f'{self.__path} {hfuzz_arguments} -- {target_cmdline}'
 
-        # NOTE: Assuming fixed location for the input file.
-        if not stdin:
-            hfuzz_cmdline += " ___FILE___"
-
         logging.info(f"Run Honggfuzz with: {hfuzz_cmdline}")
         logging.debug(f"\tWorkspace: {workspace.root_dir}")
 

engines/pastis-triton/bin/pastis-triton

@@ -58,18 +58,21 @@ def cli():
 @cli.command()
 @click.option('-h', '--host', type=str, default='localhost', help='Host to connect to')
 @click.option('-p', '--port', type=int, default=5555, help='Port to connect to')
+@click.option('--debug', type=bool,  is_flag=True, show_default=True, default=False, help='Enable debug logs')
 @click.option('--probe', type=str, help="Probe to load as a python module (should contain a ProbeInterface)", multiple=True)
-def online(host: str, port: int, probe: Tuple[str]):
+def online(host: str, port: int, debug: bool, probe: Tuple[str]):
     """
     This is the online mode of the pastis-triton exploration. With this mode,
     the client (pastis-triton) will try to connect to the broker. Then, the broker
     will send us the binary to explore, the configuration and initiale seeds.
 
     :param host: The remote host to connect
     :param port: The remote host's port to connect
+    :param debug: Configure debugging logs
+    :param probe: Probes to enable (Python modules imported with importlib)
     """
 
-    configure_logs(logging.INFO)
+    configure_logs(logging.DEBUG if debug else logging.INFO)
 
     # Create the network agent and connect to the broker
     agent = ClientAgent()

engines/pastis-triton/pastisdse/pastisdse.py

@@ -314,10 +314,10 @@ def start_received(self, fname: str, binary: bytes, engine: FuzzingEngineInfo, e
             self.config = Config.from_json(engine_args)
         else:
             self.config = Config()  # Empty configuration
-
-        # Override config argv if provided
-        if argv:
-            self.config.program_argv = argv
+            # Use argv ONLY if no configuration provided
+            self.config.program_argv = [f"./{fname}"]
+            if argv:
+                self.config.program_argv.extend(argv) # Preprend the binary to argv
 
         """
         Actions taken depending on seed format & co:
@@ -563,7 +563,7 @@ def _process_seed_received(self, typ: SeedType, seed: Seed):
                     self.dse.add_input_seed(seed)
                 else:
                     # Check whether the seed improves the current coverage.
-                    if self.dse.coverage.improves_coverage(coverage):
+                    if self.dse.coverage.improve_coverage(coverage):
                         logging.info(f"seed added {seed.hash} [{typ.name}] (coverage merged)")
                         self.seeds_merged += 1
                         self.dse.coverage.merge(coverage)

libpastis/proto/message.proto

@@ -67,7 +67,7 @@ message StartMsg {
     string coverage_mode   = 8;       // coverage strategy for Triton (ignored by Triton)
     SeedInjectLoc seed_location  = 9;
     string engine_args           = 10; // Serialized JSON of engine parameters
-    repeated string program_argv = 11;
+    repeated string program_argv = 11;  // Arguments (without program name (argv[0]))
 }
 
 message StopMsg {