1#!/usr/bin/env python
2
3# ################################################################
4# Copyright (c) 2016-present, Facebook, Inc.
5# All rights reserved.
6#
7# This source code is licensed under both the BSD-style license (found in the
8# LICENSE file in the root directory of this source tree) and the GPLv2 (found
9# in the COPYING file in the root directory of this source tree).
10# ##########################################################################
11
12import argparse
13import contextlib
14import os
15import re
16import shutil
17import subprocess
18import sys
19import tempfile
20
21
22def abs_join(a, *p):
23    return os.path.abspath(os.path.join(a, *p))
24
25
26# Constants
27FUZZ_DIR = os.path.abspath(os.path.dirname(__file__))
28CORPORA_DIR = abs_join(FUZZ_DIR, 'corpora')
29TARGETS = [
30    'simple_round_trip',
31    'stream_round_trip',
32    'block_round_trip',
33    'simple_decompress',
34    'stream_decompress',
35    'block_decompress',
36]
37ALL_TARGETS = TARGETS + ['all']
38FUZZ_RNG_SEED_SIZE = 4
39
40# Standard environment variables
41CC = os.environ.get('CC', 'cc')
42CXX = os.environ.get('CXX', 'c++')
43CPPFLAGS = os.environ.get('CPPFLAGS', '')
44CFLAGS = os.environ.get('CFLAGS', '-O3')
45CXXFLAGS = os.environ.get('CXXFLAGS', CFLAGS)
46LDFLAGS = os.environ.get('LDFLAGS', '')
47MFLAGS = os.environ.get('MFLAGS', '-j')
48
49# Fuzzing environment variables
50LIB_FUZZING_ENGINE = os.environ.get('LIB_FUZZING_ENGINE', 'libregression.a')
51AFL_FUZZ = os.environ.get('AFL_FUZZ', 'afl-fuzz')
52DECODECORPUS = os.environ.get('DECODECORPUS',
53                              abs_join(FUZZ_DIR, '..', 'decodecorpus'))
54
55# Sanitizer environment variables
56MSAN_EXTRA_CPPFLAGS = os.environ.get('MSAN_EXTRA_CPPFLAGS', '')
57MSAN_EXTRA_CFLAGS = os.environ.get('MSAN_EXTRA_CFLAGS', '')
58MSAN_EXTRA_CXXFLAGS = os.environ.get('MSAN_EXTRA_CXXFLAGS', '')
59MSAN_EXTRA_LDFLAGS = os.environ.get('MSAN_EXTRA_LDFLAGS', '')
60
61
62def create(r):
63    d = os.path.abspath(r)
64    if not os.path.isdir(d):
65        os.mkdir(d)
66    return d
67
68
69def check(r):
70    d = os.path.abspath(r)
71    if not os.path.isdir(d):
72        return None
73    return d
74
75
76@contextlib.contextmanager
77def tmpdir():
78    dirpath = tempfile.mkdtemp()
79    try:
80        yield dirpath
81    finally:
82        shutil.rmtree(dirpath, ignore_errors=True)
83
84
85def parse_targets(in_targets):
86    targets = set()
87    for target in in_targets:
88        if not target:
89            continue
90        if target == 'all':
91            targets = targets.union(TARGETS)
92        elif target in TARGETS:
93            targets.add(target)
94        else:
95            raise RuntimeError('{} is not a valid target'.format(target))
96    return list(targets)
97
98
99def targets_parser(args, description):
100    parser = argparse.ArgumentParser(prog=args.pop(0), description=description)
101    parser.add_argument(
102        'TARGET',
103        nargs='*',
104        type=str,
105        help='Fuzz target(s) to build {{{}}}'.format(', '.join(ALL_TARGETS)))
106    args, extra = parser.parse_known_args(args)
107    args.extra = extra
108
109    args.TARGET = parse_targets(args.TARGET)
110
111    return args
112
113
114def parse_env_flags(args, flags):
115    """
116    Look for flags set by environment variables.
117    """
118    san_flags = ','.join(re.findall('-fsanitize=((?:[a-z]+,?)+)', flags))
119    nosan_flags = ','.join(re.findall('-fno-sanitize=((?:[a-z]+,?)+)', flags))
120
121    def set_sanitizer(sanitizer, default, san, nosan):
122        if sanitizer in san and sanitizer in nosan:
123            raise RuntimeError('-fno-sanitize={s} and -fsanitize={s} passed'.
124                               format(s=sanitizer))
125        if sanitizer in san:
126            return True
127        if sanitizer in nosan:
128            return False
129        return default
130
131    san = set(san_flags.split(','))
132    nosan = set(nosan_flags.split(','))
133
134    args.asan = set_sanitizer('address', args.asan, san, nosan)
135    args.msan = set_sanitizer('memory', args.msan, san, nosan)
136    args.ubsan = set_sanitizer('undefined', args.ubsan, san, nosan)
137
138    args.sanitize = args.asan or args.msan or args.ubsan
139
140    return args
141
142
143def compiler_version(cc, cxx):
144    """
145    Determines the compiler and version.
146    Only works for clang and gcc.
147    """
148    cc_version_bytes = subprocess.check_output([cc, "--version"])
149    cxx_version_bytes = subprocess.check_output([cxx, "--version"])
150    if cc_version_bytes.startswith(b'clang'):
151        assert(cxx_version_bytes.startswith(b'clang'))
152        compiler = 'clang'
153    if cc_version_bytes.startswith(b'gcc'):
154        assert(cxx_version_bytes.startswith(b'g++'))
155        compiler = 'gcc'
156    version_regex = b'([0-9])+\.([0-9])+\.([0-9])+'
157    version_match = re.search(version_regex, cc_version_bytes)
158    version = tuple(int(version_match.group(i)) for i in range(1, 4))
159    return compiler, version
160
161
162def overflow_ubsan_flags(cc, cxx):
163    compiler, version = compiler_version(cc, cxx)
164    if compiler == 'gcc':
165        return ['-fno-sanitize=signed-integer-overflow']
166    if compiler == 'clang' and version >= (5, 0, 0):
167        return ['-fno-sanitize=pointer-overflow']
168    return []
169
170
171def build_parser(args):
172    description = """
173    Cleans the repository and builds a fuzz target (or all).
174    Many flags default to environment variables (default says $X='y').
175    Options that aren't enabling features default to the correct values for
176    zstd.
177    Enable sanitizers with --enable-*san.
178    For regression testing just build.
179    For libFuzzer set LIB_FUZZING_ENGINE and pass --enable-coverage.
180    For AFL set CC and CXX to AFL's compilers and set
181    LIB_FUZZING_ENGINE='libregression.a'.
182    """
183    parser = argparse.ArgumentParser(prog=args.pop(0), description=description)
184    parser.add_argument(
185        '--lib-fuzzing-engine',
186        dest='lib_fuzzing_engine',
187        type=str,
188        default=LIB_FUZZING_ENGINE,
189        help=('The fuzzing engine to use e.g. /path/to/libFuzzer.a '
190              "(default: $LIB_FUZZING_ENGINE='{})".format(LIB_FUZZING_ENGINE)))
191    parser.add_argument(
192        '--enable-coverage',
193        dest='coverage',
194        action='store_true',
195        help='Enable coverage instrumentation (-fsanitize-coverage)')
196    parser.add_argument(
197        '--enable-asan', dest='asan', action='store_true', help='Enable UBSAN')
198    parser.add_argument(
199        '--enable-ubsan',
200        dest='ubsan',
201        action='store_true',
202        help='Enable UBSAN')
203    parser.add_argument(
204        '--enable-ubsan-pointer-overflow',
205        dest='ubsan_pointer_overflow',
206        action='store_true',
207        help='Enable UBSAN pointer overflow check (known failure)')
208    parser.add_argument(
209        '--enable-msan', dest='msan', action='store_true', help='Enable MSAN')
210    parser.add_argument(
211        '--enable-msan-track-origins', dest='msan_track_origins',
212        action='store_true', help='Enable MSAN origin tracking')
213    parser.add_argument(
214        '--msan-extra-cppflags',
215        dest='msan_extra_cppflags',
216        type=str,
217        default=MSAN_EXTRA_CPPFLAGS,
218        help="Extra CPPFLAGS for MSAN (default: $MSAN_EXTRA_CPPFLAGS='{}')".
219        format(MSAN_EXTRA_CPPFLAGS))
220    parser.add_argument(
221        '--msan-extra-cflags',
222        dest='msan_extra_cflags',
223        type=str,
224        default=MSAN_EXTRA_CFLAGS,
225        help="Extra CFLAGS for MSAN (default: $MSAN_EXTRA_CFLAGS='{}')".format(
226            MSAN_EXTRA_CFLAGS))
227    parser.add_argument(
228        '--msan-extra-cxxflags',
229        dest='msan_extra_cxxflags',
230        type=str,
231        default=MSAN_EXTRA_CXXFLAGS,
232        help="Extra CXXFLAGS for MSAN (default: $MSAN_EXTRA_CXXFLAGS='{}')".
233        format(MSAN_EXTRA_CXXFLAGS))
234    parser.add_argument(
235        '--msan-extra-ldflags',
236        dest='msan_extra_ldflags',
237        type=str,
238        default=MSAN_EXTRA_LDFLAGS,
239        help="Extra LDFLAGS for MSAN (default: $MSAN_EXTRA_LDFLAGS='{}')".
240        format(MSAN_EXTRA_LDFLAGS))
241    parser.add_argument(
242        '--enable-sanitize-recover',
243        dest='sanitize_recover',
244        action='store_true',
245        help='Non-fatal sanitizer errors where possible')
246    parser.add_argument(
247        '--debug',
248        dest='debug',
249        type=int,
250        default=1,
251        help='Set ZSTD_DEBUG (default: 1)')
252    parser.add_argument(
253        '--force-memory-access',
254        dest='memory_access',
255        type=int,
256        default=0,
257        help='Set MEM_FORCE_MEMORY_ACCESS (default: 0)')
258    parser.add_argument(
259        '--fuzz-rng-seed-size',
260        dest='fuzz_rng_seed_size',
261        type=int,
262        default=4,
263        help='Set FUZZ_RNG_SEED_SIZE (default: 4)')
264    parser.add_argument(
265        '--disable-fuzzing-mode',
266        dest='fuzzing_mode',
267        action='store_false',
268        help='Do not define FUZZING_BUILD_MORE_UNSAFE_FOR_PRODUCTION')
269    parser.add_argument(
270        '--enable-stateful-fuzzing',
271        dest='stateful_fuzzing',
272        action='store_true',
273        help='Reuse contexts between runs (makes reproduction impossible)')
274    parser.add_argument(
275        '--cc',
276        dest='cc',
277        type=str,
278        default=CC,
279        help="CC (default: $CC='{}')".format(CC))
280    parser.add_argument(
281        '--cxx',
282        dest='cxx',
283        type=str,
284        default=CXX,
285        help="CXX (default: $CXX='{}')".format(CXX))
286    parser.add_argument(
287        '--cppflags',
288        dest='cppflags',
289        type=str,
290        default=CPPFLAGS,
291        help="CPPFLAGS (default: $CPPFLAGS='{}')".format(CPPFLAGS))
292    parser.add_argument(
293        '--cflags',
294        dest='cflags',
295        type=str,
296        default=CFLAGS,
297        help="CFLAGS (default: $CFLAGS='{}')".format(CFLAGS))
298    parser.add_argument(
299        '--cxxflags',
300        dest='cxxflags',
301        type=str,
302        default=CXXFLAGS,
303        help="CXXFLAGS (default: $CXXFLAGS='{}')".format(CXXFLAGS))
304    parser.add_argument(
305        '--ldflags',
306        dest='ldflags',
307        type=str,
308        default=LDFLAGS,
309        help="LDFLAGS (default: $LDFLAGS='{}')".format(LDFLAGS))
310    parser.add_argument(
311        '--mflags',
312        dest='mflags',
313        type=str,
314        default=MFLAGS,
315        help="Extra Make flags (default: $MFLAGS='{}')".format(MFLAGS))
316    parser.add_argument(
317        'TARGET',
318        nargs='*',
319        type=str,
320        help='Fuzz target(s) to build {{{}}}'.format(', '.join(ALL_TARGETS))
321    )
322    args = parser.parse_args(args)
323    args = parse_env_flags(args, ' '.join(
324        [args.cppflags, args.cflags, args.cxxflags, args.ldflags]))
325
326    # Check option sanitiy
327    if args.msan and (args.asan or args.ubsan):
328        raise RuntimeError('MSAN may not be used with any other sanitizers')
329    if args.msan_track_origins and not args.msan:
330        raise RuntimeError('--enable-msan-track-origins requires MSAN')
331    if args.ubsan_pointer_overflow and not args.ubsan:
332        raise RuntimeError('--enable-ubsan-pointer-overlow requires UBSAN')
333    if args.sanitize_recover and not args.sanitize:
334        raise RuntimeError('--enable-sanitize-recover but no sanitizers used')
335
336    return args
337
338
339def build(args):
340    try:
341        args = build_parser(args)
342    except Exception as e:
343        print(e)
344        return 1
345    # The compilation flags we are setting
346    targets = args.TARGET
347    cc = args.cc
348    cxx = args.cxx
349    cppflags = [args.cppflags]
350    cflags = [args.cflags]
351    ldflags = [args.ldflags]
352    cxxflags = [args.cxxflags]
353    mflags = [args.mflags] if args.mflags else []
354    # Flags to be added to both cflags and cxxflags
355    common_flags = []
356
357    cppflags += [
358        '-DZSTD_DEBUG={}'.format(args.debug),
359        '-DMEM_FORCE_MEMORY_ACCESS={}'.format(args.memory_access),
360        '-DFUZZ_RNG_SEED_SIZE={}'.format(args.fuzz_rng_seed_size),
361    ]
362
363    mflags += ['LIB_FUZZING_ENGINE={}'.format(args.lib_fuzzing_engine)]
364
365    # Set flags for options
366    if args.coverage:
367        common_flags += [
368            '-fsanitize-coverage=trace-pc-guard,indirect-calls,trace-cmp'
369        ]
370
371    if args.sanitize_recover:
372        recover_flags = ['-fsanitize-recover=all']
373    else:
374        recover_flags = ['-fno-sanitize-recover=all']
375    if args.sanitize:
376        common_flags += recover_flags
377
378    if args.msan:
379        msan_flags = ['-fsanitize=memory']
380        if args.msan_track_origins:
381            msan_flags += ['-fsanitize-memory-track-origins']
382        common_flags += msan_flags
383        # Append extra MSAN flags (it might require special setup)
384        cppflags += [args.msan_extra_cppflags]
385        cflags += [args.msan_extra_cflags]
386        cxxflags += [args.msan_extra_cxxflags]
387        ldflags += [args.msan_extra_ldflags]
388
389    if args.asan:
390        common_flags += ['-fsanitize=address']
391
392    if args.ubsan:
393        ubsan_flags = ['-fsanitize=undefined']
394        if not args.ubsan_pointer_overflow:
395            ubsan_flags += overflow_ubsan_flags(cc, cxx)
396        common_flags += ubsan_flags
397
398    if args.stateful_fuzzing:
399        cppflags += ['-DSTATEFUL_FUZZING']
400
401    if args.fuzzing_mode:
402        cppflags += ['-DFUZZING_BUILD_MORE_UNSAFE_FOR_PRODUCTION']
403
404    if args.lib_fuzzing_engine == 'libregression.a':
405        targets = ['libregression.a'] + targets
406
407    # Append the common flags
408    cflags += common_flags
409    cxxflags += common_flags
410
411    # Prepare the flags for Make
412    cc_str = "CC={}".format(cc)
413    cxx_str = "CXX={}".format(cxx)
414    cppflags_str = "CPPFLAGS={}".format(' '.join(cppflags))
415    cflags_str = "CFLAGS={}".format(' '.join(cflags))
416    cxxflags_str = "CXXFLAGS={}".format(' '.join(cxxflags))
417    ldflags_str = "LDFLAGS={}".format(' '.join(ldflags))
418
419    # Print the flags
420    print('MFLAGS={}'.format(' '.join(mflags)))
421    print(cc_str)
422    print(cxx_str)
423    print(cppflags_str)
424    print(cflags_str)
425    print(cxxflags_str)
426    print(ldflags_str)
427
428    # Clean and build
429    clean_cmd = ['make', 'clean'] + mflags
430    print(' '.join(clean_cmd))
431    subprocess.check_call(clean_cmd)
432    build_cmd = [
433        'make',
434        cc_str,
435        cxx_str,
436        cppflags_str,
437        cflags_str,
438        cxxflags_str,
439        ldflags_str,
440    ] + mflags + targets
441    print(' '.join(build_cmd))
442    subprocess.check_call(build_cmd)
443    return 0
444
445
446def libfuzzer_parser(args):
447    description = """
448    Runs a libfuzzer binary.
449    Passes all extra arguments to libfuzzer.
450    The fuzzer should have been build with LIB_FUZZING_ENGINE pointing to
451    libFuzzer.a.
452    Generates output in the CORPORA directory, puts crashes in the ARTIFACT
453    directory, and takes extra input from the SEED directory.
454    To merge AFL's output pass the SEED as AFL's output directory and pass
455    '-merge=1'.
456    """
457    parser = argparse.ArgumentParser(prog=args.pop(0), description=description)
458    parser.add_argument(
459        '--corpora',
460        type=str,
461        help='Override the default corpora dir (default: {})'.format(
462            abs_join(CORPORA_DIR, 'TARGET')))
463    parser.add_argument(
464        '--artifact',
465        type=str,
466        help='Override the default artifact dir (default: {})'.format(
467            abs_join(CORPORA_DIR, 'TARGET-crash')))
468    parser.add_argument(
469        '--seed',
470        type=str,
471        help='Override the default seed dir (default: {})'.format(
472            abs_join(CORPORA_DIR, 'TARGET-seed')))
473    parser.add_argument(
474        'TARGET',
475        type=str,
476        help='Fuzz target(s) to build {{{}}}'.format(', '.join(TARGETS)))
477    args, extra = parser.parse_known_args(args)
478    args.extra = extra
479
480    if args.TARGET and args.TARGET not in TARGETS:
481        raise RuntimeError('{} is not a valid target'.format(args.TARGET))
482
483    return args
484
485
486def libfuzzer(target, corpora=None, artifact=None, seed=None, extra_args=None):
487    if corpora is None:
488        corpora = abs_join(CORPORA_DIR, target)
489    if artifact is None:
490        artifact = abs_join(CORPORA_DIR, '{}-crash'.format(target))
491    if seed is None:
492        seed = abs_join(CORPORA_DIR, '{}-seed'.format(target))
493    if extra_args is None:
494        extra_args = []
495
496    target = abs_join(FUZZ_DIR, target)
497
498    corpora = [create(corpora)]
499    artifact = create(artifact)
500    seed = check(seed)
501
502    corpora += [artifact]
503    if seed is not None:
504        corpora += [seed]
505
506    cmd = [target, '-artifact_prefix={}/'.format(artifact)]
507    cmd += corpora + extra_args
508    print(' '.join(cmd))
509    subprocess.check_call(cmd)
510
511
512def libfuzzer_cmd(args):
513    try:
514        args = libfuzzer_parser(args)
515    except Exception as e:
516        print(e)
517        return 1
518    libfuzzer(args.TARGET, args.corpora, args.artifact, args.seed, args.extra)
519    return 0
520
521
522def afl_parser(args):
523    description = """
524    Runs an afl-fuzz job.
525    Passes all extra arguments to afl-fuzz.
526    The fuzzer should have been built with CC/CXX set to the AFL compilers,
527    and with LIB_FUZZING_ENGINE='libregression.a'.
528    Takes input from CORPORA and writes output to OUTPUT.
529    Uses AFL_FUZZ as the binary (set from flag or environment variable).
530    """
531    parser = argparse.ArgumentParser(prog=args.pop(0), description=description)
532    parser.add_argument(
533        '--corpora',
534        type=str,
535        help='Override the default corpora dir (default: {})'.format(
536            abs_join(CORPORA_DIR, 'TARGET')))
537    parser.add_argument(
538        '--output',
539        type=str,
540        help='Override the default AFL output dir (default: {})'.format(
541            abs_join(CORPORA_DIR, 'TARGET-afl')))
542    parser.add_argument(
543        '--afl-fuzz',
544        type=str,
545        default=AFL_FUZZ,
546        help='AFL_FUZZ (default: $AFL_FUZZ={})'.format(AFL_FUZZ))
547    parser.add_argument(
548        'TARGET',
549        type=str,
550        help='Fuzz target(s) to build {{{}}}'.format(', '.join(TARGETS)))
551    args, extra = parser.parse_known_args(args)
552    args.extra = extra
553
554    if args.TARGET and args.TARGET not in TARGETS:
555        raise RuntimeError('{} is not a valid target'.format(args.TARGET))
556
557    if not args.corpora:
558        args.corpora = abs_join(CORPORA_DIR, args.TARGET)
559    if not args.output:
560        args.output = abs_join(CORPORA_DIR, '{}-afl'.format(args.TARGET))
561
562    return args
563
564
565def afl(args):
566    try:
567        args = afl_parser(args)
568    except Exception as e:
569        print(e)
570        return 1
571    target = abs_join(FUZZ_DIR, args.TARGET)
572
573    corpora = create(args.corpora)
574    output = create(args.output)
575
576    cmd = [args.afl_fuzz, '-i', corpora, '-o', output] + args.extra
577    cmd += [target, '@@']
578    print(' '.join(cmd))
579    subprocess.call(cmd)
580    return 0
581
582
583def regression(args):
584    try:
585        description = """
586        Runs one or more regression tests.
587        The fuzzer should have been built with with
588        LIB_FUZZING_ENGINE='libregression.a'.
589        Takes input from CORPORA.
590        """
591        args = targets_parser(args, description)
592    except Exception as e:
593        print(e)
594        return 1
595    for target in args.TARGET:
596        corpora = create(abs_join(CORPORA_DIR, target))
597        target = abs_join(FUZZ_DIR, target)
598        cmd = [target, corpora]
599        print(' '.join(cmd))
600        subprocess.check_call(cmd)
601    return 0
602
603
604def gen_parser(args):
605    description = """
606    Generate a seed corpus appropiate for TARGET with data generated with
607    decodecorpus.
608    The fuzz inputs are prepended with a seed before the zstd data, so the
609    output of decodecorpus shouldn't be used directly.
610    Generates NUMBER samples prepended with FUZZ_RNG_SEED_SIZE random bytes and
611    puts the output in SEED.
612    DECODECORPUS is the decodecorpus binary, and must already be built.
613    """
614    parser = argparse.ArgumentParser(prog=args.pop(0), description=description)
615    parser.add_argument(
616        '--number',
617        '-n',
618        type=int,
619        default=100,
620        help='Number of samples to generate')
621    parser.add_argument(
622        '--max-size-log',
623        type=int,
624        default=13,
625        help='Maximum sample size to generate')
626    parser.add_argument(
627        '--seed',
628        type=str,
629        help='Override the default seed dir (default: {})'.format(
630            abs_join(CORPORA_DIR, 'TARGET-seed')))
631    parser.add_argument(
632        '--decodecorpus',
633        type=str,
634        default=DECODECORPUS,
635        help="decodecorpus binary (default: $DECODECORPUS='{}')".format(
636            DECODECORPUS))
637    parser.add_argument(
638        '--fuzz-rng-seed-size',
639        type=int,
640        default=4,
641        help="FUZZ_RNG_SEED_SIZE used for generate the samples (must match)"
642    )
643    parser.add_argument(
644        'TARGET',
645        type=str,
646        help='Fuzz target(s) to build {{{}}}'.format(', '.join(TARGETS)))
647    args, extra = parser.parse_known_args(args)
648    args.extra = extra
649
650    if args.TARGET and args.TARGET not in TARGETS:
651        raise RuntimeError('{} is not a valid target'.format(args.TARGET))
652
653    if not args.seed:
654        args.seed = abs_join(CORPORA_DIR, '{}-seed'.format(args.TARGET))
655
656    if not os.path.isfile(args.decodecorpus):
657        raise RuntimeError("{} is not a file run 'make -C {} decodecorpus'".
658                           format(args.decodecorpus, abs_join(FUZZ_DIR, '..')))
659
660    return args
661
662
663def gen(args):
664    try:
665        args = gen_parser(args)
666    except Exception as e:
667        print(e)
668        return 1
669
670    seed = create(args.seed)
671    with tmpdir() as compressed:
672        with tmpdir() as decompressed:
673            cmd = [
674                args.decodecorpus,
675                '-n{}'.format(args.number),
676                '-p{}/'.format(compressed),
677                '-o{}'.format(decompressed),
678            ]
679
680            if 'block_' in args.TARGET:
681                cmd += [
682                    '--gen-blocks',
683                    '--max-block-size-log={}'.format(args.max_size_log)
684                ]
685            else:
686                cmd += ['--max-content-size-log={}'.format(args.max_size_log)]
687
688            print(' '.join(cmd))
689            subprocess.check_call(cmd)
690
691            if '_round_trip' in args.TARGET:
692                print('using decompressed data in {}'.format(decompressed))
693                samples = decompressed
694            elif '_decompress' in args.TARGET:
695                print('using compressed data in {}'.format(compressed))
696                samples = compressed
697
698            # Copy the samples over and prepend the RNG seeds
699            for name in os.listdir(samples):
700                samplename = abs_join(samples, name)
701                outname = abs_join(seed, name)
702                rng_seed = os.urandom(args.fuzz_rng_seed_size)
703                with open(samplename, 'rb') as sample:
704                    with open(outname, 'wb') as out:
705                        out.write(rng_seed)
706                        CHUNK_SIZE = 131072
707                        chunk = sample.read(CHUNK_SIZE)
708                        while len(chunk) > 0:
709                            out.write(chunk)
710                            chunk = sample.read(CHUNK_SIZE)
711    return 0
712
713
714def minimize(args):
715    try:
716        description = """
717        Runs a libfuzzer fuzzer with -merge=1 to build a minimal corpus in
718        TARGET_seed_corpus. All extra args are passed to libfuzzer.
719        """
720        args = targets_parser(args, description)
721    except Exception as e:
722        print(e)
723        return 1
724
725    for target in args.TARGET:
726        # Merge the corpus + anything else into the seed_corpus
727        corpus = abs_join(CORPORA_DIR, target)
728        seed_corpus = abs_join(CORPORA_DIR, "{}_seed_corpus".format(target))
729        extra_args = [corpus, "-merge=1"] + args.extra
730        libfuzzer(target, corpora=seed_corpus, extra_args=extra_args)
731        seeds = set(os.listdir(seed_corpus))
732        # Copy all crashes directly into the seed_corpus if not already present
733        crashes = abs_join(CORPORA_DIR, '{}-crash'.format(target))
734        for crash in os.listdir(crashes):
735            if crash not in seeds:
736                shutil.copy(abs_join(crashes, crash), seed_corpus)
737                seeds.add(crash)
738
739
740def zip_cmd(args):
741    try:
742        description = """
743        Zips up the seed corpus.
744        """
745        args = targets_parser(args, description)
746    except Exception as e:
747        print(e)
748        return 1
749
750    for target in args.TARGET:
751        # Zip the seed_corpus
752        seed_corpus = abs_join(CORPORA_DIR, "{}_seed_corpus".format(target))
753        seeds = [abs_join(seed_corpus, f) for f in os.listdir(seed_corpus)]
754        zip_file = "{}.zip".format(seed_corpus)
755        cmd = ["zip", "-q", "-j", "-9", zip_file]
756        print(' '.join(cmd + [abs_join(seed_corpus, '*')]))
757        subprocess.check_call(cmd + seeds)
758
759
760def list_cmd(args):
761    print("\n".join(TARGETS))
762
763
764def short_help(args):
765    name = args[0]
766    print("Usage: {} [OPTIONS] COMMAND [ARGS]...\n".format(name))
767
768
769def help(args):
770    short_help(args)
771    print("\tfuzzing helpers (select a command and pass -h for help)\n")
772    print("Options:")
773    print("\t-h, --help\tPrint this message")
774    print("")
775    print("Commands:")
776    print("\tbuild\t\tBuild a fuzzer")
777    print("\tlibfuzzer\tRun a libFuzzer fuzzer")
778    print("\tafl\t\tRun an AFL fuzzer")
779    print("\tregression\tRun a regression test")
780    print("\tgen\t\tGenerate a seed corpus for a fuzzer")
781    print("\tminimize\tMinimize the test corpora")
782    print("\tzip\t\tZip the minimized corpora up")
783    print("\tlist\t\tList the available targets")
784
785
786def main():
787    args = sys.argv
788    if len(args) < 2:
789        help(args)
790        return 1
791    if args[1] == '-h' or args[1] == '--help' or args[1] == '-H':
792        help(args)
793        return 1
794    command = args.pop(1)
795    args[0] = "{} {}".format(args[0], command)
796    if command == "build":
797        return build(args)
798    if command == "libfuzzer":
799        return libfuzzer_cmd(args)
800    if command == "regression":
801        return regression(args)
802    if command == "afl":
803        return afl(args)
804    if command == "gen":
805        return gen(args)
806    if command == "minimize":
807        return minimize(args)
808    if command == "zip":
809        return zip_cmd(args)
810    if command == "list":
811        return list_cmd(args)
812    short_help(args)
813    print("Error: No such command {} (pass -h for help)".format(command))
814    return 1
815
816
817if __name__ == "__main__":
818    sys.exit(main())
819