gazprea-fuzzer-python/gazprea_fuzzer.py

import os
import random
from contextlib import redirect_stdout

import yaml

import fuzzer as fz
import argparse

import xml.dom.minidom
import xml.etree.ElementTree as ET

class Fuzzer():
    def __init__(self, config: str, batch: int, seed: str, file_name: str = "fuzz"):
        with open(config) as yaml_file:
            settings: dict = yaml.safe_load(yaml_file)
        self.settings = settings
        self.batch = batch
        random.seed(seed)
        self.file_name = file_name

        self.fuzzer = fz.GazpreaFuzzer(config)

    def fuzz(self):
        os.system("rm -rf fuzzer")
        os.mkdir("fuzzer")
        os.mkdir("fuzzer/input")
        os.mkdir("fuzzer/debug")
        os.mkdir("fuzzer/instream")
        os.mkdir("fuzzer/outputs")
        os.mkdir("fuzzer/ground_truth")
        for i in range(self.batch):
            self.fuzzer.fuzz()
            dom = xml.dom.minidom.parseString(ET.tostring(self.fuzzer.ast).decode('utf-8'))
            pretty: str = dom.toprettyxml()
            with open("fuzzer/input/{}_{}.in".format(self.file_name, i), 'w') as f:
                f.write(self.fuzzer.source)
            with open("fuzzer/debug/{}_{}.out".format(self.file_name, i), 'w') as f:
                f.write(pretty)
            with open("fuzzer/ground_truth/{}_{}.py".format(self.file_name, i), 'w') as f:
                f.write(self.fuzzer.ground_truth)
            with open("fuzzer/ground_truth/{}_{}.py".format(self.file_name, i), 'r') as f:
                with open("fuzzer/outputs/{}_{}.out".format(self.file_name, i), 'w') as y:
                    with redirect_stdout(y): # Workaround for fuzzer.py:49
                        exec(f.read(), globals(), locals())
                    # y.write(self.fuzzer.out)
            # with open("fuzzer/instream/{}.in".format(i), 'w') as f:
            #     f.write(self.fuzzer.source)
            # with open("fuzzer/outputs/{}.out".format(i), 'w') as f:
            #     f.write(self.fuzzer.out)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Procedurally generate a test case for Gazprea'
    )
    parser.add_argument('-b', '--batch_size', type=int, required=False, default=1,
                        help="generate BATCH cases (fuzzer/source/nameX.in, /instream/..., /outputs/...)")
    parser.add_argument('--seed', type=int, required=False, action="store",
                        help="rng seed")
    parser.add_argument('config_file', type=str, action="store",
                        help="path to your configuration file")
    parser.add_argument('file_name', type=str, action="store",
                        help="name for the generated files")

    args = parser.parse_args()
    fuzzer = Fuzzer(config=args.config_file, batch=args.batch_size, seed=args.seed)
    fuzzer.fuzz()