pwntools bruteforce progress logging

pwntools-progress.py

import random
from pwn import *

SEED = 0xdeadbeef
RANGE_MIN = 1
RANGE_MAX = 5
ROUNDS = 20
SLEEP_TIME = 0.100 # 100ms

random.seed(SEED)

rounds = [random.randint(RANGE_MIN, RANGE_MAX) for _ in range(ROUNDS)]
total = 0

log.info('Need to guess %d numbers, each within the range [%d, %d]' % \
        (ROUNDS, RANGE_MIN, RANGE_MAX))

time.sleep(2)
log.info('Starting blind guesses')
time.sleep(2)

# We use a local context to suppress the internal loggign that
# pwntools produces when starting processes. We also use 
# progress loggers with their levels set to WARN to
# bypass the log filter we set on the local context
# so we can still print update messages.

with context.local(log_level='warn'), \
    log.progress('Numbers remaining', level=logging.WARN) as prog_remaining, \
    log.progress('Correct guesses', level=logging.WARN) as prog_correct, \
    log.progress('Total guesses', level=logging.WARN) as prog_total, \
    log.progress('Current guess', level=logging.WARN) as prog_guess:

    log.info('This wont print because we suppress everything < WARN')

    correct_guesses = []
    for answer in rounds:
        prog_remaining.status(str(ROUNDS - len(correct_guesses)))
        prog_correct.status(repr(correct_guesses))
        guess = RANGE_MIN - 1 # intentionally wrong guess
        
        while guess is not answer:
            # Using shuf to demonstrate log suppression for process creation
            io = process('shuf -i %s-%s -n 1' % (RANGE_MIN, RANGE_MAX), shell=True)
            guess = int(io.recvall())
            
            # Update progress
            # status() expects a string
            prog_guess.status(str(guess))
            time.sleep(SLEEP_TIME)

            total += 1
            prog_total.status(str(total))

        # We found the right guess!
        correct_guesses.append(guess)

    # Set final messages, otherwise 'Done' is displayed
    prog_correct.success(repr(correct_guesses))
    prog_total.success(str(total))

log.info("Wasn't that fun?")

Demo at https://twitter.com/byte_how/status/1245499798284722176