<commit_before>import os
import angr
import shutil
import tempfile
import subprocess
import shellphish_afl
from .fuzzer import Fuzzer

import logging
l = logging.getLogger("fuzzer.Showmap")

class Showmap(object):
    """Show map"""

    def __init__(self, binary_path, testcase, timeout=None):
        """
        :param binary_path: path to the binary which the testcase applies to
        :param testcase: string representing the contents of the testcase
        :param timeout: millisecond timeout
        """

        self.binary_path = binary_path
        self.testcase = testcase
        self.timeout = None

        if timeout is not None:
            if isinstance(timeout, (int, long)):
                self.timeout = str(timeout)
            elif isinstance(timeout, (str)):
                self.timeout = timeout
            else:
                raise ValueError("timeout param must be of type int or str")

        # will be set by showmap's return code
        self.causes_crash = False

        Fuzzer._perform_env_checks()

        self.base = Fuzzer._get_base()
        l.debug("got base dir %s", self.base)

        # unfortunately here is some code reuse between Fuzzer and Minimizer
        p = angr.Project(self.binary_path)
        tracer_id = 'cgc' if p.loader.main_bin.os == 'cgc' else p.arch.qemu_name
        self.showmap_path = os.path.join(shellphish_afl.afl_dir(tracer_id), "afl-showmap")
        self.afl_path_var = shellphish_afl.afl_path_var(tracer_id)

        l.debug("showmap_path: %s", self.showmap_path)
        l.debug("afl_path_var: %s", self.afl_path_var)

        os.environ['AFL_PATH'] = self.afl_path_var

        # create temp
        self.work_dir = tempfile.mkdtemp(prefix='showmap-', dir='/tmp/')

        # flag for work directory removal
        self._removed = False

        self.input_testcase = os.path.join(self.work_dir, 'testcase')
        self.output = os.path.join(self.work_dir, 'out')

        l.debug("input_testcase: %s", self.input_testcase)
        l.debug("output: %s", self.output)

        # populate contents of input testcase
        with open(self.input_testcase, 'w') as f:
            f.write(testcase)

    def __del__(self):
        if not self._removed:
            shutil.rmtree(self.work_dir)

    def showmap(self):
        """Create the map"""

        if self._start_showmap().wait() == 2:
            self.causes_crash = True

        with open(self.output) as f: result = f.read()

        shutil.rmtree(self.work_dir)
        self._removed = True

        shownmap = dict()
        for line in result.split("\n")[:-1]:
            key, h_count = map(int, line.split(":"))
            shownmap[key] = h_count

        return shownmap

    def _start_showmap(self, memory="8G"):

        args = [self.showmap_path]

        args += ["-o", self.output]
        args += ["-m", memory]
        args += ["-Q"]

        if self.timeout:
            args += ["-t", self.timeout]

        args += ["--"]
        args += [self.binary_path]

        outfile = "minimizer.log"

        l.debug("execing: %s > %s", " ".join(args), outfile)

        outfile = os.path.join(self.work_dir, outfile)
        with open(outfile, "w") as fp, open(self.input_testcase, 'rb') as it, open("/dev/null", 'wb') as devnull:
            return subprocess.Popen(args, stdin=it, stdout=devnull, stderr=fp)<commit_msg>Add theoretical support for MultiCBs<commit_after>import os
import angr
import shutil
import tempfile
import subprocess
import shellphish_afl
from .fuzzer import Fuzzer

import logging
l = logging.getLogger("fuzzer.Showmap")

class Showmap(object):
    """Show map"""

    def __init__(self, binary_path, testcase, timeout=None):
        """
        :param binary_path: path to the binary which the testcase applies to
        :param testcase: string representing the contents of the testcase
        :param timeout: millisecond timeout
        """

        self.binary_path = binary_path
        self.testcase = testcase
        self.timeout = None

        if isinstance(binary_path, basestring):
            self.is_multicb = False
            self.binaries = [binary_path]
        elif isinstance(binary_path, (list,tuple)):
            self.is_multicb = True
            self.binaries = binary_path
        else:
            raise ValueError("Was expecting either a string or a list/tuple for binary_path! "
                "It's {} instead.".format(type(binary_path)))

        if timeout is not None:
            if isinstance(timeout, (int, long)):
                self.timeout = str(timeout)
            elif isinstance(timeout, (str)):
                self.timeout = timeout
            else:
                raise ValueError("timeout param must be of type int or str")

        # will be set by showmap's return code
        self.causes_crash = False

        Fuzzer._perform_env_checks()

        self.base = Fuzzer._get_base()
        l.debug("got base dir %s", self.base)

        # unfortunately here is some code reuse between Fuzzer and Minimizer (and Showmap!)
        p = angr.Project(self.binaries[0])
        tracer_id = 'cgc' if p.loader.main_bin.os == 'cgc' else p.arch.qemu_name
        if self.is_multicb:
            tracer_id = 'multi-{}'.format(tracer_id)

        self.showmap_path = os.path.join(shellphish_afl.afl_dir(tracer_id), "afl-showmap")
        self.afl_path_var = shellphish_afl.afl_path_var(tracer_id)

        l.debug("showmap_path: %s", self.showmap_path)
        l.debug("afl_path_var: %s", self.afl_path_var)

        os.environ['AFL_PATH'] = self.afl_path_var

        # create temp
        self.work_dir = tempfile.mkdtemp(prefix='showmap-', dir='/tmp/')

        # flag for work directory removal
        self._removed = False

        self.input_testcase = os.path.join(self.work_dir, 'testcase')
        self.output = os.path.join(self.work_dir, 'out')

        l.debug("input_testcase: %s", self.input_testcase)
        l.debug("output: %s", self.output)

        # populate contents of input testcase
        with open(self.input_testcase, 'w') as f:
            f.write(testcase)

    def __del__(self):
        if not self._removed:
            shutil.rmtree(self.work_dir)

    def showmap(self):
        """Create the map"""

        if self._start_showmap().wait() == 2:
            self.causes_crash = True

        with open(self.output) as f: result = f.read()

        shutil.rmtree(self.work_dir)
        self._removed = True

        shownmap = dict()
        for line in result.split("\n")[:-1]:
            key, h_count = map(int, line.split(":"))
            shownmap[key] = h_count

        return shownmap

    def _start_showmap(self, memory="8G"):

        args = [self.showmap_path]

        args += ["-o", self.output]
        args += ["-m", memory]
        args += ["-Q"]

        if self.timeout:
            args += ["-t", self.timeout]

        args += ["--"]
        args += self.binaries

        outfile = "minimizer.log"

        l.debug("execing: %s > %s", " ".join(args), outfile)

        outfile = os.path.join(self.work_dir, outfile)
        with open(outfile, "w") as fp, open(self.input_testcase, 'rb') as it, open("/dev/null", 'wb') as devnull:
            return subprocess.Popen(args, stdin=it, stdout=devnull, stderr=fp)
{'fraction_non_alphanumeric': 0.11933534743202417, 'fraction_numerical': 0.0009612743751716561, 'mean_word_length': 2.2659192825112107, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6566037', 'n_tokens_mistral': 2341, 'n_tokens_neox': 2274, 'n_words': 553}