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