#include <stdio.h>

int main(int argc, char *argv[]) {

    char *message;

    if (argv[1][0] == 'k') {
        message = "Congrats!\n";
    } else {
        message = "Try again\n";
    }
    
    printf("%s", message);
    
    return 0;
}

gcc -O0 -g -o simple_branch ./simple_branch.c

# Script 1, Cell 1
import os
import claripy
import angr

# Initialize the symbolic execution
current_dir = os.getcwd()
binary_path = os.path.join(current_dir, "simple_branch")
proj = angr.Project(binary_path, auto_load_libs=False)
sym_var = claripy.BVS("input", 8)
state = proj.factory.entry_state(args=["./simple_branch", sym_var])
simgr = proj.factory.simulation_manager(state)


# Advance till the only state diverges into two
while len(simgr.active) > 0:
    simgr.step() # it will execute the whole basic block
    print("Number of active states:", len(simgr.active))
    for state in simgr.active:
        print("\tstate addr", hex(state.addr))
    if len(simgr.active) == 2:
        break

# Script 1, Cell 2
print("\n-- Divergence point reached, checking the constraints --")
# Let's check what the constraints are in the two states
for state in simgr.active:
    print("State addr:", hex(state.addr))
    print("\tConstraints on sym_var:")
    for constraint in state.solver.constraints:
        print(f"\t\t{constraint}") 

print("\n-- Keep going until the result is printed --")
# 1) One possible way of doing this
# simgr.explore(find=0x4011a1) # states reaching the address will be moved to the 'found' stash

# 2) An alternative way without hardcoding the address:
def until_something_is_printed(sigmr: angr.SimulationManager):
    for s in sigmr.active:
        if s.posix.dumps(1):
            return True
    for s in sigmr.deadended:
        if s.posix.dumps(1):
            return True
    return False

# advance states till something is printed in some state(s), then pick them out
# do this until we have no state active
while len(simgr.active) > 0:
    simgr.run(stash="active", until=until_something_is_printed)
    simgr.move(from_stash="active", to_stash="found", filter_func=lambda s: s.posix.dumps(1))

for state in simgr.found:
    print("State advanced to:", hex(state.addr))
    print("\tConstraints on sym_var:")
    for constraint in state.solver.constraints:
        print(f"\t\t{constraint}")
    print("\tExpected output:", state.posix.dumps(1).decode('utf-8'))

# Script 2, cell 1
import os
import claripy
import angr

# Initialize the symbolic execution
current_dir = os.getcwd()
binary_path = os.path.join(current_dir, "simple_branch")
proj = angr.Project(binary_path, auto_load_libs=False)
sym_var = claripy.BVS("input", 8)
state = proj.factory.entry_state(args=["./simple_branch", sym_var],
                                add_options={angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY,
                                            angr.options.ZERO_FILL_UNCONSTRAINED_REGISTERS})
simgr = proj.factory.simulation_manager(state)


# Advance till the first instruction after the divergence
simgr.explore(find=0x401186)

# Merge the two states
after_merge = simgr.found[0]
for state in simgr.found[1:]:
    after_merge,_ ,_ = after_merge.merge(state)
simgr = proj.factory.simulation_manager(after_merge)
assert len(simgr.active) == 1

# print the constraints
for c in simgr.active[0].solver.constraints:
    print(c)

# Script 2, cell 2
# Stop right before the call to printf
simgr.explore(find=0x40119c)
assert len(simgr.found) == 1
print("Address of current state:", hex(simgr.found[0].addr))
state = simgr.found[0]
rsi_val = state.solver.eval(state.regs.rsi)
print(f"RSI is symbolic: {state.regs.rsi.symbolic}")

# check all possible values of RSI
print(f"\n-- Possible messages to be printed --")
possible_rsi_vals = state.solver.eval_upto(state.regs.rsi, 128)
for val in possible_rsi_vals:
    print(f"Possible RSI: {hex(val)}")
    message = state.solver.eval(state.memory.load(val, 8), cast_to=bytes)
    print(f"Message: {message.decode('utf-8', errors='ignore')}")

print("\n-- Adding constraint on input --")
print(f"RSI is symbolic: {state.regs.rsi.symbolic}")
possible_rsi_vals = state.solver.eval_upto(state.regs.rsi, 128, extra_constraints=[sym_var == b'k'])
for val in possible_rsi_vals:
    print(f"Possible RSI: {hex(val)}")
    message = state.solver.eval(state.memory.load(val, 8),
                                cast_to=bytes)
    print(f"Message: {message.decode('utf-8', errors='ignore')}")

print("\n-- Running the simulation till the end --")
state.solver.add(sym_var == b'k')
simgr.run(stash='found')
state = simgr.deadended[0]
print(f"Output: {state.posix.dumps(1).decode('utf-8', errors='ignore')}")

// real_example.c
// compile with:
// gcc -O0 -g -o real_example ./real_example.c
#include <stdio.h>
#include <string.h>

int main(int argc, char *argv[]) {

    if (argc != 2) {
        printf("Usage: %s <input_string>\n", argv[0]);
        return 1;
    }


    char *input = argv[1];

    if (strlen(input) != 32) {
        printf("Input string must be exactly 32 characters long.\n");
        return 1;
    }

    int x_pos = 0;
    int y_pos = 0;

    for (int i = 0; i < 32; i++) {
        char byte = input[i];
        int x_delta = 0;
        int y_delta = 0;


        if (byte == 'U') {
            y_delta = 1;
        }

        if (byte == 'D') {
            y_delta = -1;
        }

        if (byte == 'L') {
            x_delta = -1;
        }

        if (byte == 'R') {
            x_delta = 1;
        }

        x_pos += x_delta;
        y_pos += y_delta;
    }

    // choosing 0725 bc it is the birthdate of my Oshimen
    if (x_pos == 7 && y_pos == 25) {
        printf("Success: Reached target position (7, 25)\n");
    } else {
        printf("Failed: Current position is (%d, %d), not (7, 25)\n", x_pos, y_pos);
    }

    return 0;
}

# Script 3, cell 1
import os
import claripy
import angr

# Initialize the symbolic execution
current_dir = os.getcwd()
binary_path = os.path.join(current_dir, "./real_example")
proj = angr.Project(binary_path, auto_load_libs=False)

# In some examples, the symbolic input is initialized as a whole byte chunk,
# but the nature of this input is a string of 32 bytes (null-terminated),
# so one symbolic variable per byte is created.
sym_vars = [claripy.BVS(f"input_{x}", 8) for x in range(33)]
sym_string = claripy.Concat(*sym_vars)
state = proj.factory.entry_state(args=["./real_example", sym_string],
                                add_options={angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY,
                                            angr.options.ZERO_FILL_UNCONSTRAINED_REGISTERS})

# Constraint: the len of the input must be 32 bytes
for idx, char in enumerate(sym_vars):
    if idx < 32:
        state.solver.add(char >= 0x20)
        state.solver.add(char <= 0x7e)
    else:
        state.solver.add(char == 0x00)

simgr = proj.factory.simulation_manager(state)
# the merge points are the first instruction
# where the two divergent paths resulted from a conditional branch meet

# For example, 0x401240 is the merge point after branching at 0x40123d
#
# 00401239  807de355           cmp     byte ptr [rbp-0x1d {var_25_1}], 0x55
# 0040123d  7507               jnz     0x401246
# 0040123f  c745f401000000     mov     dword ptr [rbp-0xc {var_14_1}], 0x1
# 00401246  807de344           cmp     byte ptr [rbp-0x1d {var_25_1}], 0x44

mergepoints = [
    0x401246,
    0x401253,
    0x401260,
    0x40126d,
]


for i in range(32):
    for mp in mergepoints:
        simgr.explore(find=mp)
        for s in simgr.found:
            assert s.addr == mp
        for s in simgr.active:
            assert s.addr == mp
        simgr.stash(filter_func=None,
                    from_stash='found', to_stash='active')
        first = simgr.active[0]
        the_rest = simgr.active[1:]
        for s in the_rest:
            first = first.merge(s)[0]
        new_simgr = proj.factory.simulation_manager(first)
        simgr = new_simgr
        state = simgr.active[0]

        if i == 0:
            rbp = state.regs.rbp
            rbp_value = state.solver.eval(rbp)
            partial_sum = state.memory.load(rbp_value - 0x10, 4)
            possible_values = state.solver.eval_upto(partial_sum, 1024, cast_to=bytes)
            # usually cast_to=int will work, but in this case since the offset could be -1
            # which will be interpreted as unsigned by default, we need to cast manually
            possible_values = [int.from_bytes(x, 'little', signed=True) for x in possible_values]
            print(f"Possible values: {sorted(possible_values)}")
            print(f'RIP: {hex(state.solver.eval(state.regs.rip))}')