import angr_mbf
from collections import defaultdict
import angr
import logging
import claripy
import sys

BENEVOLENT_ENDING  = b'he'
DESTRUCTIVE_ENDING = b'sa'
CHOSEN_ENDING = BENEVOLENT_ENDING

logging.getLogger('angr').setLevel('WARNING')

log = logging.getLogger('main')

project = angr.Project('./hellas.mbf')

bytes_list = [claripy.BVS('byte_%d' % i, 8) for i in range(33)]
bytes_ast = claripy.Concat(claripy.BVV(CHOSEN_ENDING), *bytes_list)
simfile = angr.SimFile('/dev/stdin', content=bytes_ast, has_end=True)
state = project.factory.entry_state(stdin=simfile)

sim = project.factory.simgr(state)

known_printable = 2


def improvement_found(sm):
    """ Greedy strategy: discard paths that don't lead to more characters immediately. """
    global known_printable
    better = []
    best = None
    new_kp = known_printable

    for state in sm.active:
        num_printable = len([n for n in state.posix.dumps(0)+state.posix.dumps(1) if n >= 0x20 and n <= 0xf0])
        if num_printable > known_printable:
            if num_printable > new_kp:
                new_kp = num_printable
                best = state
            better.append(state)
    if better:
        sm.move(from_stash='active', to_stash='stagnant', filter_func=lambda x: x not in better)
        known_printable = new_kp
        log.warning(f"Best state solves to {known_printable} printable characters; input is {best.posix.dumps(0)!r}")
        return True
    else:
        return False

while len(sim.active):
    sim.explore(until=improvement_found)

assert len(sim.deadended) == 1

inp = sim.one_deadended.posix.dumps(0)
out = sim.one_deadended.posix.dumps(1).decode('utf-8')

print(f"Input: {inp!r}\n\nOutput:\n{out}")
