#!/usr/bin/env python3

import angr
import claripy
import os

base_addr = 0x400000
main_fun = base_addr + 0x1080

binarypath = "./lineq2"

# Load the binary
p = angr.Project(binarypath, auto_load_libs=False)

# If you want debug output
import logging
logging.getLogger('angr').setLevel('INFO')



class ConcretizationHandler(object):
    def __init__(self, dbg_addr=0xDEADAFFE0000):
        self.DEBUG_ADDR = dbg_addr

    def before_address_concretization(self, state):
        self.index_expr = state.inspect.address_concretization_expr
        state.inspect.address_concretization_expr = self.DEBUG_ADDR

    """
    Angr API:


    """
    def after_address_concretization(self, state):
        ## TODO
        # Set self.DEBUG_ADDR to a value modeling the array access / mathematical operation that is actually performed
        # memory can be set via state.mem[addr].uint8_t = ..., state.mem[addr].uint16_t = ..., etc.

        # Do not add new path constraints
        state.inspect.address_concretization_add_constraints = False
        # Specify our DEBUG_ADDR as the target of the concretization
        state.inspect.address_concretization_result = [self.DEBUG_ADDR]







proto = "int check_flag(void)" # Example for main

## Initialize call state with useful options
#   ZERO_FILL_UNCONSTRAINED_MEMORY: Memory that is used but not initialized is set to 0
#   ZERO_FILL_UNCONSTRAINED_REGISTERS: Registers that are used but are not initialized are set to 0, this prevents warnings whencallee-saved registers are pushed onto the stack
#   LAZY_SOLVES: This is somewhat of a magic "hack": Do not query the constrained solver unless really necessary. If you're trying to go through a specific code path, this can save tons of time! Experiment with and without it!
state = p.factory.call_state(main_fun, prototype=proto,
                             add_options=({angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY,
                                           angr.options.ZERO_FILL_UNCONSTRAINED_REGISTERS,
                                           angr.options.LAZY_SOLVES}),
                             )

# Register concretization hooks
conc_handler = ConcretizationHandler()
state.inspect.b('address_concretization', when=angr.BP_BEFORE, action=conc_handler.before_address_concretization)
state.inspect.b('address_concretization', when=angr.BP_AFTER, action=conc_handler.after_address_concretization)



# Initialize a simulation manager with the state
sm = p.factory.simulation_manager(state)


# Perform explore
print("Starting explore...")

# Try to find a state that fulfills a certain condition via exploration
# Useful: s.posix.dumps(1) to access stdout and s.posix.dumps(2) to access stderr
# Note: find and avoid can also have single integers (addresses) or lists of integers
sm.explore(find=lambda s: False, avoid=lambda s: False) # TODO

# dump stdin and stdout of found states
for s in sm.found:
    print(s.posix.dumps(1))
    print(s.posix.dumps(0).strip(b'\x00').decode())