byte_search_scaffolding.py

class ByteSearch:
    def __init__(self, oracle, confidence_threshold=0.9, quiet=True):
        self._counter = 0
        self.oracle = oracle
        self.queries = [[] for _ in range(256)]
        self.confidences = [1/256]*256
        self.confidence_threshold = confidence_threshold
        self.quiet = quiet

    def update_confidences(self, index, result):
        """Given an oracle result for a given byte, update the confidences for each byte."""
        self.confidences = self.get_updated_confidences(self.confidences, index, result)

    def pick_exhaustive(self):
        return self._counter % 256

    def pick_by_confidence(self):
        """Pick a byte to test based on the current confidences."""
        return max(range(256), key=lambda i: self.confidences[i])

    def pick_by_entropy(self):
        """Pick a byte to test based on expected reduction in entropy."""
        # NOTE: VERY SLOW - for demo, try replacing 256 with 16 here and in randrange
        entropies = []
        for i in range(256):
            e_if_t = self.get_entropy(self.get_updated_confidences(self.confidences, i, True))
            e_if_f = self.get_entropy(self.get_updated_confidences(self.confidences, i, False))
            p_t = self.confidences[i]
            p_f = 1 - p_t
            entropies.append(p_t * e_if_t + p_f * e_if_f)
        return min(range(256), key=lambda i: entropies[i])

    def query_byte(self, index):
        """Query the oracle for a given byte."""
        self._counter += 1
        result = self.oracle(index)
        self.queries[index].append(result)
        self.update_confidences(index, result)
        if not self.quiet and self._counter & 0xFF == 0:
            print(end=".", flush=True)
        return result

    def search(self, strategy):
        """Search for the plaintext byte by querying the oracle."""
        threshold = self.confidence_threshold
        while max(self.confidences) < threshold:
            self.query_byte(strategy())
        num_queries.append(sum(len(l) for l in self.queries))
        return max(range(256), key=lambda i: self.confidences[i])