build.py

import argparse
import subprocess
from pathlib import Path
import yaml

import config
import fuzzers

def build_base():
    print("Building base image")

    subprocess.check_output(
        ['docker', 'build', '-t', 'fs_base', '-f', 'Dockerfile', '.'],
        cwd='./docker/base'
    )

def build_benchmark(benchmark):
    print(f"Building benchmark: {benchmark}")

    subprocess.check_output(
        ['docker', 'build', '-t', 'fs_' + benchmark, '-f', 'Dockerfile', '.'],
        cwd=f'./benchmarks/{benchmark}'
    )

def build_fuzzer(fuzzer, benchmark):
    print(f"Building fuzzer: {fuzzer} for benchmark: {benchmark}")

    subprocess.check_output(
        [
            'docker', 'build',
            '--build-arg', f'parent_image=fs_{benchmark}',
            '-t', 'fs_' + fuzzer + '_' + benchmark,
            '-f', 'builder.Dockerfile', '.'
        ],
        cwd=f'./fuzzers/{fuzzer}'
    )

def build_fuzzer_builder(fuzzer, benchmark, debug=False):
    print(f"Building fuzzer builder: {fuzzer} for benchmark: {benchmark}")

    cmd = []
    cmd += [
        'docker', 'build',
        '--build-arg', f'parent_image=fs_{fuzzer}_{benchmark}',
        '--build-arg', f'benchmark={benchmark}',
        '--build-arg', f'fuzzer={fuzzer}',
    ]
    if debug:
        cmd += ['--build-arg', 'debug_builder=1']
    cmd += [
        '-t', 'fs_' + fuzzer + '_' + benchmark + '_builder',
        '-f', 'docker/builder/Dockerfile',
        '.'
    ]
    subprocess.check_output(cmd)

def export_fuzzer(fuzzer, benchmark, build_dir: Path):
    if build_dir.exists():
        print('Clearing existing build directory')
        for child in build_dir.iterdir():
            if child.is_dir():
                subprocess.check_output(['rm', '-rf', child])
            else:
                child.unlink()

    build_dir.mkdir(parents=True, exist_ok=True)

    print(f"Exporting fuzzer: {fuzzer} for benchmark: {benchmark}")

    subprocess.check_output(
        [
            'docker', 'build',
            '--build-arg', f'parent_image=fs_{fuzzer}_{benchmark}_builder',
            '-t', 'fs_' + fuzzer + '_' + benchmark + '_export',
            '-f', 'docker/exporter/Dockerfile',
            '--output=type=local,dest=' + str(build_dir),
            '.'
        ]
    )

def main(args):
    cfg = config.Config('config.yaml')

    fuzzer = args.fuzzer
    benchmark = args.benchmark
    print(f"Fuzzer: {fuzzer}")
    print(f"Benchmark: {benchmark}")

    build_base()
    build_benchmark(benchmark)
    build_fuzzer(fuzzer, benchmark)
    build_fuzzer_builder(fuzzer, benchmark, debug=args.debug)

    if args.debug:
        # run the debug container
        subprocess.run(['docker', 'run', '-it', '--rm', 'fs_' + fuzzer + '_' + benchmark + '_builder'])
        return
    
    export_fuzzer(fuzzer, benchmark, build_dir=cfg.build_dir(fuzzer, benchmark))

    # Check for postprocess.py
    if (Path(f'./fuzzers/{fuzzer}/postprocess.py')).exists():
        print('Running postprocess.py')
        __import__('fuzzers', fromlist=[fuzzer])
        postprocess_script = getattr(__import__(f'fuzzers.{fuzzer}', fromlist=['postprocess']), 'postprocess')

        binfo = yaml.safe_load(open(f'benchmarks/{benchmark}/benchmark.yaml'))

        postprocess_script.run(cfg, cfg.build_dir(fuzzer, benchmark), binfo['fuzz_target'])

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--fuzzer', '-f', type=str, required=True, help='Fuzzer name')
    parser.add_argument('--benchmark', '-b', type=str, required=True, help='Benchmark name')
    parser.add_argument('--debug', '-d', action='store_true', help='Debug builder')
    args = parser.parse_args()
    main(args)