main.py

from __future__ import absolute_import, division, print_function

import argparse
import datetime
import json
import os
import pickle
import random
import shutil
import sys
import time

import inquirer
import toml
import torch

from fuzzer.construct_initial_seeds import construct_seeds, get_test_dataloader
from fuzzer.construct_profile import construct_profile
from fuzzer.coverage import Coverage
from fuzzer.dry_run import dry_run
from fuzzer.fuzzer import Fuzzer
from fuzzer.image_queue import ImageInputCorpus, TensorInputCorpus
from fuzzer.model_structure import metrics_param, model_structure
from fuzzer.util import profile_dict_to_gpu
from models.network_type import ClassificationNetworks
from utils.util import obj


def arguments():
    parser = argparse.ArgumentParser(description="coverage guided fuzzing for DNN")
    parser.add_argument(
        "-i",
        help="input seed directory",
        default="Insert path here",
    )  # Check for model weights path.
    parser.add_argument(
        "-o", help="output directory", default="Insert path here"
    )
    parser.add_argument("-config_file", help="choose configuration with which to run", default="config/svhn.toml")
    parser.add_argument("-construct_profile", help="0: No, 1: Yes", default="0")
    parser.add_argument("-construct_seeds", help="0: No, 1: Yes", default="0")
    return parser.parse_args()


if __name__ == "__main__":
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
    print("Torch using device:", device)
    start_time = datetime.datetime.now()
    random.seed(time.time())

    args = arguments()
    config = json.loads(json.dumps(toml.load(args.config_file)), object_hook=obj)
    config.device = device
    torch.set_grad_enabled(False)
    if config.model == "mnist-lenet5":
        model = model_structure[config.model]()
    elif config.model == "gtsrb-new":
        model = model_structure[config.model]()
    elif config.model == "gtsrb-lenet":
        model = model_structure[config.model]()
    elif config.model == "svhn-mixed":
        model = model_structure[config.model]()
    else:
        raise NotImplementedError("Model specified has not been implemented.")
    model.load_state_dict(torch.load(model.model_path, map_location=device), strict= False)
    model.to(device)
    model.eval()

    if config.network_type == "classification" and (
        config.data == "mnist" or config.data == "gtsrb" or config.data == "gtsrb-gray" or config.data == "svhn"
    ):
        network = ClassificationNetworks(config.init_seed_selection, num_classes=config.num_classes)
    else:
        raise NotImplementedError("Network type specified has not been implemented")

    # Important: until intermediate_outputs() of model class is called, construct profile and coverage class wont work.

    if (
        not hasattr(config, "input_channels")
        or not hasattr(config, "input_width")
        or not hasattr(config, "input_height")
    ):
        config.input_size = (3, 32, 32)
    else:
        config.input_size = (config.input_channels, config.input_width, config.input_height)

    if args.construct_seeds == "1":
        construct_seeds(model, network, config, args.i, partial_dataset=True)
    else:
        # For reproducibility with same test seeds, use this to save layers and num neurons in model object
        if config.data == "mnist":
            test_loader = get_test_dataloader(config)
            img, _ = next(iter(test_loader))
            print("test image shape:", img.shape)
            _ = model.intermediate_outputs(img.to(device))
        else:
            vimg = torch.randn(2, *config.input_size)
            vimg = vimg.to(device)
            _ = model.intermediate_outputs(vimg)

    # Load the profiling information which is needed by the coverage metrics.
    if args.construct_profile == "1":
        vimg = torch.randn(2, *config.input_size)
        vimg = vimg.to(device)
        _ = model.intermediate_outputs(vimg)  # Initializes intermediate layers attribute.
        construct_profile(model, config)

    profile_dict = pickle.load(open(config.model_profile_path, "rb"))
    print("Model Profile loaded from:", config.model_profile_path)

    if device.type == "cuda":
        profile_dict = profile_dict_to_gpu(profile_dict, device)

    # Create the output directory including seed queue and crash dir, replace if already exists
    if os.path.exists(args.o):
        continue_ = inquirer.prompt(
            [
                inquirer.Confirm(
                    "confirm",
                    default=False,
                    message=f"This will delete the existing output directory at {args.o}! Continue?",
                )
            ]
        )["confirm"]
        if not continue_:
            sys.exit()

        shutil.rmtree(args.o)

    os.makedirs(os.path.join(args.o, "queue"))
    os.makedirs(os.path.join(args.o, "crashes_rec"))
    os.makedirs(os.path.join(args.o, "crashes_f1"))
    os.makedirs(os.path.join(args.o, "crashes_hyb"))

    # The log file which records the plot data after each iteration of the fuzzing
    plot_file = open(os.path.join(args.o, "plot.log"), "a+")

    # Load the configuration for the selected metrics.
    cri = config.fuzz_criteria
    if cri == "lscd":
        if os.path.exists(config.model_profile_path):
            param = metrics_param[cri]  # replace param with centroid
            profile_centroids = profile_dict["centroids"]

        else:
            raise NotImplementedError("Please first generate centroid data to calculate LSCD metric.")
    else:
        param = metrics_param[cri]

    if (
        config.model == "gtsrb-new"
        or config.model == "gtsrb-lenet"
        or config.model == "svhn-mixed"
        or config.model == "mnist-lenet5"
    ):
        coverage_handler = Coverage(model, profile_dict, config, k=param)

    else:
        raise NotImplementedError

    # The seed queue
    if config.fuzz_criteria == "flann":
        queue = TensorInputCorpus(args.o, config.random, config.seed_selection, param, "kdtree")
    else:
        queue = ImageInputCorpus(
            args.o, config.random, config.seed_selection, coverage_handler.total_size, config.fuzz_criteria
        )

    # Perform the dry_run process from the initial seeds
    if (
        config.model == "gtsrb-new"
        or config.model == "gtsrb-lenet"
        or config.model == "svhn-mixed"
        or config.model == "mnist-lenet5"
    ):
        if config.fuzz_criteria == "lscd":
            dry_run(args.i, coverage_handler.predict, queue, config, profile_centroids)
        else:
            dry_run(args.i, coverage_handler.predict, queue, config)
    else:
        raise NotImplementedError

    # The fuzzing process
    fuzzer = Fuzzer(queue, network, model, coverage_handler.predict, config.num_mutants, config.seed_selection)

    if config.fuzz_criteria == "lscd":
        fuzzer.loop(config, profile_centroids)
    else:
        fuzzer.loop(config)

    print("Results stored at:", args.o)
    print("\nFinished! Took", datetime.datetime.now() - start_time)