end-to-end: split cli into multiple commands (#1042)

This allows a more step-wise approach, which is useful for e.g. generate
builds in one session, and generating harnesses in other sessions.

Signed-off-by: David Korczynski <[email protected]>

Author: DavidKorczynski

experimental/end_to_end/cli.py

@@ -102,13 +102,14 @@ def _run_introspector_collection(runner_script, project, wd, semaphore):
   semaphore.release()
 
 
-def extract_introspector_reports_for_benchmarks(projects_to_run, workdir, args):
+def extract_introspector_reports_for_benchmarks(projects_to_run, workdir,
+                                                parallel_build_jobs):
   """Runs introspector through each report to collect program analysis data."""
   oss_fuzz_dir = os.path.join(workdir, 'oss-fuzz')
   runner_script = os.path.join(workdir, 'fuzz-introspector',
                                'oss_fuzz_integration', 'runner.py')
 
-  semaphore = threading.Semaphore(args.build_jobs)
+  semaphore = threading.Semaphore(parallel_build_jobs)
   jobs = []
 
   for project in projects_to_run:
@@ -392,25 +393,43 @@ def _create_data_dir(workdir):
   return dst_dir
 
 
-def run_harness_generation(out_gen, workdir, args):
-  """Runs harness generation based on the projects in `out_gen`"""
-
+def prepare_fuzz_introspector_db(out_gen, workdir, parallel_introspector_jobs):
+  # Run introspector collection on the generated projects
   projects_to_run = copy_generated_projects_to_harness_gen(out_gen, workdir)
-  extract_introspector_reports_for_benchmarks(projects_to_run, workdir, args)
+  extract_introspector_reports_for_benchmarks(projects_to_run, workdir,
+                                              parallel_introspector_jobs)
+
+  # Create a fuzz introspector database based on the projects in
+  # the working directory's OSS-Fuzz.
   shutdown_fi_webapp()
   create_fi_db(workdir)
-  if args.until_fi_db:
-    logger.info('Fuzz Introspector webapp created. Exiting.')
-    sys.exit(0)
+
+
+def run_harness_generation(out_gen, workdir, args):
+  """Runs harness generation based on the projects in `out_gen`"""
+
+  # Read the json file from FI to get all current projects.
+  fi_project_json = os.path.join(workdir, 'fuzz-introspector', 'tools',
+                                 'web-fuzzing-introspection', 'app', 'static',
+                                 'assets', 'db', 'all-project-current.json')
+  if not os.path.isfile(fi_project_json):
+    logger.info('Did not find FI DB file.')
+    set()
+
+  projects_to_run = []
+  with open(fi_project_json, 'r') as f:
+    json_content = json.load(f)
+  for elem in json_content:
+    projects_to_run.append(elem['project_name'])
+
+  # Launch the fuzz introspector webapp so it's ready for OFG core
   shutdown_fi_webapp()
   launch_fi_webapp(workdir)
   wait_until_fi_webapp_is_launched()
   dst_data_dir = _create_data_dir(workdir)
   logger.info('Wrote data directory for OFG experiments in %s', dst_data_dir)
-  if args.all_but_ofg_core:
-    # do a prompt exist
-    sys.exit(0)
 
+  # Run OFG core using local OSS-Fuzz and local Fuzz Introspector.
   run_ofg_generation(projects_to_run, workdir, args)
 
   create_merged_oss_fuzz_projects(out_gen)
@@ -442,20 +461,9 @@ def _get_next_data_dst_dir():
   return _get_next_folder_in_idx('data-dir')
 
 
-def run_analysis(args):
-  """Generates builds and harnesses for repositories in input."""
-  workdir = setup_workdirs(args.workdir)
-
-  abs_workdir = os.path.abspath(workdir)
-  if not args.out:
-    out_folder = get_next_out_folder()
-  else:
-    out_folder = args.out
-
-  #if os.path.isdir('results'):
-  #  shutil.rmtree('results')
-
-  oss_fuzz_dir = os.path.join(abs_workdir, 'oss-fuzz-1')
+def _run_build_generation(workdir, out_folder, args):
+  """ Build script generation. """
+  oss_fuzz_dir = os.path.join(workdir, 'oss-fuzz-1')
   target_repositories = runner.extract_target_repositories(args.input)
   if args.agent:
     # Prepare arguments used deeper in OFG core.
@@ -472,14 +480,72 @@ def run_analysis(args):
                          parallel_jobs=args.build_jobs,
                          max_timeout=args.build_timeout)
 
-  # Exit if only builds are required.
-  if args.build_only:
-    logger.info('Finished analysis')
-    logger.info('Results in %s', out_folder)
-    return
 
+def run_fuzz_introspector_db_creation(args):
+  """Entrypoint for fuzz introspector database creation."""
+  args.workdir = os.path.abspath(args.workdir)
+  prepare_fuzz_introspector_db(args.generated_builds, args.workdir,
+                               args.parallel_build_jobs)
+
+
+def run_build_generation(args):
+  """Generates builds and harnesses for repositories in input."""
+
+  # Prepare working directory.
+  workdir = setup_workdirs(args.workdir)
+
+  abs_workdir = os.path.abspath(workdir)
+  if not args.out:
+    out_folder = get_next_out_folder()
+  else:
+    out_folder = args.out
+
+  _run_build_generation(abs_workdir, out_folder, args)
+
+
+def run_cmd_harness_generation(args):
+  """Entrypoint for command for harness generation."""
+
+  # Prepare working directory.
+  abs_workdir = os.path.abspath(args.workdir)
+
+  out_folder = args.out
+
+  # Run harness generation.
   projects_run = run_harness_generation(out_folder, abs_workdir, args)
 
+  # Log results.
+  logger.info('Finished analysis')
+  logger.info('Results in %s', out_folder)
+  logger.info('Projects generated (%d): ', len(projects_run))
+  for project in projects_run:
+    logger.info('- %s', project)
+
+  if os.path.isdir('results'):
+    shutil.copytree('results', os.path.join(out_folder, 'harness-results'))
+
+
+def run_full(args):
+  """Generates builds and harnesses for repositories in input."""
+
+  # Prepare working directory.
+  workdir = setup_workdirs(args.workdir)
+
+  abs_workdir = os.path.abspath(workdir)
+  if not args.out:
+    out_folder = get_next_out_folder()
+  else:
+    out_folder = args.out
+
+  _run_build_generation(abs_workdir, out_folder, args)
+
+  # Prepare fuzz introspector database.
+  prepare_fuzz_introspector_db(out_folder, abs_workdir, args.build_jobs)
+
+  # Run harness generation.
+  projects_run = run_harness_generation(out_folder, abs_workdir, args)
+
+  # Log results.
   logger.info('Finished analysis')
   logger.info('Results in %s', out_folder)
   logger.info('Projects generated (%d): ', len(projects_run))
@@ -493,56 +559,162 @@ def run_analysis(args):
 def parse_commandline():
   """Parse the commandline."""
   parser = argparse.ArgumentParser()
-  parser.add_argument('--input', '-i', help='Input to analyze')
-  parser.add_argument('--out',
-                      '-o',
-                      help='Directory to store output.',
-                      default='oss-fuzz-generated')
-  parser.add_argument('--silent',
-                      '-s',
-                      help='Disable logging in subprocess.',
-                      action='store_true')
-  parser.add_argument('--model',
-                      '-m',
-                      help=('Models available: '
-                            f'{", ".join(models.LLM.all_llm_names())}.'),
-                      type=str)
-  parser.add_argument('--agent',
-                      '-a',
-                      help='Enable agent workflow',
-                      action='store_true')
-  parser.add_argument('--hg-agent',
-                      '-ha',
-                      help='Enable agent harness generation',
-                      action='store_true')
-  parser.add_argument('-gm',
-                      '--generate-benchmarks-max',
-                      help='Max targets to generate per benchmark heuristic.',
-                      type=int,
-                      default=5)
-  parser.add_argument('-mr',
-                      '--max-round',
-                      type=int,
-                      default=5,
-                      help='Max trial round for agents.')
-  parser.add_argument('--build-only',
-                      help='Only generated builds',
-                      action='store_true')
-  parser.add_argument('--build-jobs',
-                      help='Parallel build-generator jobs to run.',
-                      default=2,
-                      type=int)
-  parser.add_argument('--all-but-ofg-core', action='store_true')
-  parser.add_argument(
+  subparsers = parser.add_subparsers(dest='command')
+
+  # Run build generation.
+  run_build_gen = subparsers.add_parser(
+      'generate-builds',
+      help='Generate OSS-Fuzz projects with build scripts but empty fuzzers.')
+  run_build_gen.add_argument('--input', '-i', help='Input to analyze')
+  run_build_gen.add_argument('--out',
+                             '-o',
+                             help='Directory to store output.',
+                             default='oss-fuzz-generated')
+  run_build_gen.add_argument('--silent',
+                             '-s',
+                             help='Disable logging in subprocess.',
+                             action='store_true')
+  run_build_gen.add_argument('--model',
+                             '-m',
+                             help=('Models available: '
+                                   f'{", ".join(models.LLM.all_llm_names())}.'),
+                             type=str)
+  run_build_gen.add_argument('--agent',
+                             '-a',
+                             help='Enable agent workflow',
+                             action='store_true')
+  run_build_gen.add_argument(
+      '-gm',
+      '--generate-benchmarks-max',
+      help='Max targets to generate per benchmark heuristic.',
+      type=int,
+      default=5)
+  run_build_gen.add_argument('-mr',
+                             '--max-round',
+                             type=int,
+                             default=5,
+                             help='Max trial round for agents.')
+  run_build_gen.add_argument('--build-jobs',
+                             help='Parallel build-generator jobs to run.',
+                             default=2,
+                             type=int)
+  run_build_gen.add_argument(
       '--build-timeout',
       help='Timeout for build generation per project, in seconds.',
       default=0,
       type=int)
-  parser.add_argument(
-      '--until-fi-db',
-      help='Run until Fuzz Introspector DB creation and then exit.',
+  run_build_gen.add_argument('-w', '--workdir', help='Work directory to use')
+
+  # Generate fuzz introspector database.
+  run_generate_fi_db_parser = subparsers.add_parser(
+      'generate-fuzz-introspector-database',
+      help='Generates a fuzz introspector database from auto build projects.')
+
+  run_generate_fi_db_parser.add_argument('--generated-builds', required=True)
+  run_generate_fi_db_parser.add_argument('--workdir', required=True)
+  run_generate_fi_db_parser.add_argument('--parallel-build-jobs',
+                                         type=int,
+                                         default=5)
+  run_generate_fi_db_parser.add_argument('--silent',
+                                         '-s',
+                                         help='Disable logging in subprocess.',
+                                         action='store_true')
+
+  # Run harness generation
+  run_harness_generation_parser = subparsers.add_parser(
+      'generate-harnesses',
+      help="Harness generation of OSS-Fuzz projects.",
+  )
+
+  run_harness_generation_parser.add_argument('--out',
+                                             '-o',
+                                             help='Directory to store output.',
+                                             default='oss-fuzz-generated')
+  run_harness_generation_parser.add_argument(
+      '--silent',
+      '-s',
+      help='Disable logging in subprocess.',
+      action='store_true')
+  run_harness_generation_parser.add_argument(
+      '--model',
+      '-m',
+      help=('Models available: '
+            f'{", ".join(models.LLM.all_llm_names())}.'),
+      type=str)
+  run_harness_generation_parser.add_argument('--agent',
+                                             '-a',
+                                             help='Enable agent workflow',
+                                             action='store_true')
+  run_harness_generation_parser.add_argument(
+      '--hg-agent',
+      '-ha',
+      help='Enable agent harness generation',
       action='store_true')
-  parser.add_argument('-w', '--workdir', help='Work directory to use')
+  run_harness_generation_parser.add_argument(
+      '-gm',
+      '--generate-benchmarks-max',
+      help='Max targets to generate per benchmark heuristic.',
+      type=int,
+      default=5)
+  run_harness_generation_parser.add_argument('-mr',
+                                             '--max-round',
+                                             type=int,
+                                             default=5,
+                                             help='Max trial round for agents.')
+  run_harness_generation_parser.add_argument('-w',
+                                             '--workdir',
+                                             help='Work directory to use')
+
+  # Run a full end to end generation.
+  run_full_parser = subparsers.add_parser(
+      'generate-full',
+      help="End to end generation of OSS-Fuzz projects.",
+  )
+  run_full_parser.add_argument('--input', '-i', help='Input to analyze')
+  run_full_parser.add_argument('--out',
+                               '-o',
+                               help='Directory to store output.',
+                               default='oss-fuzz-generated')
+  run_full_parser.add_argument('--silent',
+                               '-s',
+                               help='Disable logging in subprocess.',
+                               action='store_true')
+  run_full_parser.add_argument(
+      '--model',
+      '-m',
+      help=('Models available: '
+            f'{", ".join(models.LLM.all_llm_names())}.'),
+      type=str)
+  run_full_parser.add_argument('--agent',
+                               '-a',
+                               help='Enable agent workflow',
+                               action='store_true')
+  run_full_parser.add_argument('--hg-agent',
+                               '-ha',
+                               help='Enable agent harness generation',
+                               action='store_true')
+  run_full_parser.add_argument(
+      '-gm',
+      '--generate-benchmarks-max',
+      help='Max targets to generate per benchmark heuristic.',
+      type=int,
+      default=5)
+  run_full_parser.add_argument('-mr',
+                               '--max-round',
+                               type=int,
+                               default=5,
+                               help='Max trial round for agents.')
+  run_full_parser.add_argument('--build-jobs',
+                               help='Parallel build-generator jobs to run.',
+                               default=2,
+                               type=int)
+  run_full_parser.add_argument(
+      '--build-timeout',
+      help='Timeout for build generation per project, in seconds.',
+      default=0,
+      type=int)
+  run_full_parser.add_argument('-w', '--workdir', help='Work directory to use')
+
   return parser.parse_args()
 
 
@@ -551,7 +723,15 @@ def main():
   args = parse_commandline()
   setup_logging()
   silent_global = args.silent
-  run_analysis(args)
+
+  if args.command == 'generate-full':
+    run_full(args)
+  if args.command == 'generate-fuzz-introspector-database':
+    run_fuzz_introspector_db_creation(args)
+  if args.command == 'generate-builds':
+    run_build_generation(args)
+  if args.command == 'generate-harnesses':
+    run_cmd_harness_generation(args)
 
 
 if __name__ == '__main__':