# в случае успеха в файл "example.txt" допечатывается строка "Full" 

import angr
import sys, time
import claripy

GREEN = "\033[92m"
PURPLE = "\033[35m"
RESET = "\033[0m"
RED = "\033[31m"

INPUT_SIZE = None
MAX_INPUT_SIZE = 20
# ключи - имена функций инпута, значения - словари вида {"input_data": BVS, "addr": int},
# где input_data - символический битовый вектор, а "addr" - адрес функции
INPUTS = {}

def CreateFileW(state: angr.SimState): 
    state.regs.rax = claripy.BVV(1, 64)

def ReadFile(state: angr.SimState):
    input_data = state.solver.BVS('input_data', 8 * INPUT_SIZE)

    state.memory.store(state.regs.rdx, input_data)

    state.memory.eax = claripy.BVV(1, 32)

    func_name = sys._getframe().f_code.co_name
    INPUTS[func_name] = {"input_data": input_data,
    "addr": int(state.regs.rip.concrete_value)}

def memchr(state: angr.SimState):
    buf, val, maxCount = state.regs.rcx, state.regs.dl, state.regs.r8
    bytes_list = INPUTS["ReadFile"]["input_data"].chop(8)
    constraint = claripy.Or(*[byte == val for byte in bytes_list])
    state.solver.add(constraint)

def WriteFile(state: angr.SimState):
    state.memory.eax = claripy.BVV(1, 32)

def CloseHandle(state: angr.SimState):
    state.memory.eax = claripy.BVV(1, 32)

if __name__ == "__main__":
    # Загрузка бинарного файла
    print(f"{GREEN}Uploading executable file to angr...{RESET}")
    proj = angr.Project('test6.exe', auto_load_libs=False)
    print(f"{GREEN}The executable has been uploaded to angr!{RESET}")
    ENTRY_POINT = 0x140001000  # Адрес функции main
    TARGET_ADDR = 0x1400010BA  # Целевой адрес, нахождение на нем означает что программа взломана

    # Хуки
    proj.hook(0x140001046, CreateFileW, length=6)
    proj.hook(0x140001071, ReadFile, length=6)
    proj.hook(0x140001089, memchr, length=5)
    proj.hook(0x1400010AD, WriteFile, length=6)
    proj.hook(0x1400010BA, CloseHandle, length=6)

    for INPUT_SIZE in range(1, MAX_INPUT_SIZE+1):
        
        INPUTS.clear()
        print(f"{GREEN}Trying input size: {INPUT_SIZE}{RESET}")

        # Создаем начальное состояние
        state = proj.factory.entry_state(
            addr=ENTRY_POINT,
            add_options={angr.options.ZERO_FILL_UNCONSTRAINED_REGISTERS, 
                            angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY}
        )

        # Инициализация стека
        state.regs.rbp = 0x7ffffffffffffff
        state.regs.rsp = 0x7ffffffffffffff

        # Запуск симуляции
        simgr = proj.factory.simulation_manager(state)
        start_time = time.time()
        simgr.explore(find=TARGET_ADDR, avoid=0x1400010C4)
        elapsed = time.time() - start_time

        # Проверка результата
        if simgr.found:
            found_state = simgr.found[0]
            
            for input in INPUTS:
                input_data = found_state.solver.eval(INPUTS[input]['input_data'], cast_to=bytes)
                print(f"{GREEN}Success cracked! {PURPLE}"
               f"input data = {input_data} "
               f"[0x{INPUTS[input]['addr']:X}]{PURPLE}{RESET}")

            print(f"{GREEN}Pwd cracking time = {elapsed}{RESET}")
            break
        else:
            print(f"{RED}Failed with input size {INPUT_SIZE}{RESET}")
    else:
        print(f"{RED}Failed to find a solution with input sizes from 1 to {MAX_INPUT_SIZE}{RESET}")