1
2
3
4

# Configure hookds on workspace
self.workspace.add_creation_hook(self.workspace.corpus_dir, self.__send_seed)
self.workspace.add_creation_hook(self.workspace.crash_dir, self.__send_crash)
self.workspace.add_file_modification_hook(self.workspace.stats_dir, self.__send_telemetry)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

def __send_seed(self, filename: Path):
    self.__send(filename, SeedType.INPUT)

def __send_crash(self, filename: Path):
    # Skip README file that AFL adds to the crash folder.
    if filename.name != 'README.txt':
        self.__send(filename, SeedType.CRASH)

def __send(self, filename: Path, typ: SeedType):
    self._tot_seeds += 1
    file = Path(filename)
    raw = file.read_bytes()
    h = self.hash_seed(raw)
    logging.debug(f'[{typ.name}] Sending new: {h} [{self._tot_seeds}]')
    if h not in self._seed_recvs:
        self._agent.send_seed(typ, raw)
    else:
        logging.info("seed (previously sent) do not send it back")
    self._queue_to_send.append((filename, True if typ == SeedType.CRASH else False))

1
2
3
4

def __setup_agent(self):
    # Register callbacks.
    self._agent.register_seed_callback(self.__seed_received)
    self._agent.register_stop_callback(self.__stop_received)

1
2
3
4
5

def __seed_received(self, typ: SeedType, seed: bytes):
    h = self.hash_seed(seed)
    logging.info(f"[SEED] received  {h} ({typ.name})")
    self._seed_recvs.add(h)
    self.add_seed(seed)

1
2
3

def add_seed(self, seed: bytes):
    seed_path = self.workspace.dynamic_input_dir / f"seed-{hashlib.md5(seed).hexdigest()}"
    seed_path.write_bytes(seed)

1
2
3
4
5
6
7
8

def register_seed_callback(self, cb: Callable) -> None:
    """
    Register a callback called when an input seed is received from the
    broker. The callback function take 2 parameters seed type and content.

    :param cb: callback function
    """
    self.register_callback(MessageType.INPUT_SEED, cb)

1
2
3
4
5
6
7

def _register_all(self):
    self.register_seed_callback(self.seed_received)
    self.register_hello_callback(self.hello_received)
    self.register_log_callback(self.log_received)
    self.register_telemetry_callback(self.telemetry_received)
    self.register_stop_coverage_callback(self.stop_coverage_received)
    self.register_data_callback(self.data_received)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

def seed_received(self, cli_id: bytes, typ: SeedType, seed: bytes):
    cli = self.get_client(cli_id)
    if not cli:
        return
    is_new = seed not in self._seed_pool
    h = md5(seed).hexdigest()

    # Show log message and save seed to file
    self.statmanager.update_seed_stat(cli, typ)  # Add info only if new
    cli.log(LogLevel.INFO, f"seed {h} [{self._colored_seed_type(typ)}][{self._colored_seed_newness(is_new)}]")
    cli.add_own_seed(seed)  # Add seed in client's seed
    self.write_seed(typ, cli.strid, seed) # Write seed to file

    if is_new:
        self._seed_pool[seed] = typ  # Save it in the local pool
    else:
        pass
        # logging.warning(f"receive duplicate seed {h} by {cli.strid}")

    # Iterate on all clients and send it to whomever never received it
    if self.broker_mode == BrokingMode.FULL:
        self.send_seed_to_all_others(cli.netid, typ, seed)

    if self.is_proxied:  # Forward it to the proxy
        self._proxy.send_seed(typ, seed)