import idautils
import idc

def my_nop(addr, endaddr):
    while addr < endaddr:
        patch_byte(addr, 0x90)
        addr += 1

pattern = "2E C4 E2 71 96 84 9A 0C 80 0E 08"
cur_addr = 0x80491FF
end_addr = 0x80E69F5

while cur_addr < end_addr:
    cur_addr = idc.find_binary(cur_addr,SEARCH_DOWN,pattern)
    print("patch address: " + str(cur_addr)) # 打印提示信息
    if cur_addr == idc.BADADDR:
        break
    else:
        my_nop(cur_addr,cur_addr + 11)
    cur_addr = idc.next_head(cur_addr)

from unicorn import *
from capstone import *
from unicorn.x86_const import *

import struct

import sys
import hashlib

# xor add sub ror rol bswap

def bswap(value):
    return ((value & 0xFF) << 24) | (((value >> 8) & 0xFF) << 16) | (((value >> 16) & 0xFF) << 8) | ((value >> 24) & 0xFF)

def rol(value, shift):
    shift %= 32 
    return ((value << shift) | (value >> (32 - shift))) & 0xFFFFFFFF

def ror(value, shift):
    shift %= 32
    return ((value >> shift) | (value << (32 - shift))) & 0xFFFFFFFF


binary = open('./masterpiece', 'rb').read()
#file = open("./print.txt",'w')

dump_insns = []
md = Cs(CS_ARCH_X86, CS_MODE_32)

BASE_ADDRESS = 0x8048000
TEXT_BEG = 0x8049000 - BASE_ADDRESS
TEXT_END = 0x80E6A0D - BASE_ADDRESS
SEG_TEXT = binary[TEXT_BEG:TEXT_END]

DATA_BEG = 0x80E8000 - BASE_ADDRESS
DATA_END = 0x80E82A8 - BASE_ADDRESS
SEG_DATA = binary[DATA_BEG:DATA_END]

ARR_ADDR = 0x80E8018
CMP_ADDR = 0x80E8118



def get_r_eax(index):
    cmp_data_offset = CMP_ADDR - BASE_ADDRESS - DATA_BEG + index * 4
    ret ,= struct.unpack('<I',SEG_DATA[cmp_data_offset:cmp_data_offset + 4])
    return ret

def get_byte(index):
    return SEG_DATA[ARR_ADDR - DATA_BEG - BASE_ADDRESS + index]

arr_list = []
for i in range(256):
    arr_list.append(get_byte(i))

cmp_list = []
for i in range(100):
    cmp_list.append(get_r_eax(i))


# code to be emulated

EMU_ADDRESS_BEG = 0x80491FF
EMU_ADDRESS_END = 0x80E69F5

def start(round):
    global dump_insns
    dump_insns = []

    try:

        # Initialize emulator

        mu = Uc(UC_ARCH_X86, UC_MODE_32)

        # map 8KB memory for this emulation

        mu.mem_map(BASE_ADDRESS, 2 * 1024 * 1024 )

        # write machine code to be emulated to memory

        mu.mem_write(TEXT_BEG + BASE_ADDRESS, SEG_TEXT)
        mu.mem_write(DATA_BEG + BASE_ADDRESS, SEG_DATA)

        # initialize machine registers

        mu.reg_write(UC_X86_REG_EAX, 0x0)
        mu.reg_write(UC_X86_REG_ECX, round)

        #hook code (capstone)
        mu.hook_add(UC_HOOK_CODE, hook_code)

        # emulate machine code in infinite time

        mu.emu_start(EMU_ADDRESS_BEG, EMU_ADDRESS_END)

    except UcError as e:
        print("ERROR: %s" % e)    

# callback for tracing instructions

def hook_code(uc, address, size, user_data):
    global md, dump_insns
    offset = address - TEXT_BEG - BASE_ADDRESS
    code = SEG_TEXT[offset:offset + size]
    insn = next(md.disasm(code, address))
    
    if insn.mnemonic not in ['nop', 'jmp']:
        dump_insns.append(insn)
        #file.write("0x%x:\t%s\t%s\n" %(insn.address, insn.mnemonic, insn.op_str))

flag = b''

def solve(r_eax):
    r_ecx = 0
    #逆序
    for insn in reversed(dump_insns):
        if 'eax' not in insn.op_str and 'al' not in insn.op_str and 'ah' not in insn.op_str:
            print(insn.mnemonic, insn.op_str)
            continue
        match (insn.mnemonic, insn.op_str):
            case ('inc', op_str):
                if('eax' == op_str):
                    r_eax -= 1
            case ('dec', op_str):
                if('eax' == op_str):
                    r_eax += 1
            case ('ror', op_str):
                reg, val = op_str.split(', ')
                if('eax' == reg):
                    r_eax = rol(r_eax, int(val, 16))
            case ('rol', op_str):
                reg, val = op_str.split(', ')
                if('eax' == reg):
                    r_eax = ror(r_eax, int(val, 16))
            case ('add', op_str):
                reg, val = op_str.split(', ')
                if('eax' == reg):
                    r_eax -= int(val, 16)
            case ('sub', op_str):
                reg, val = op_str.split(', ')
                if('eax' == reg):
                    r_eax += int(val, 16)
            case ('not', op_str):
                if('eax' == op_str):
                    r_eax ^= 0xffffffff
            case ('xor', op_str):
                reg, val = op_str.split(', ')
                if('eax' == reg):
                    r_eax ^= int(val, 16)
            case ('bswap', op_str):
                if('eax' == op_str):
                    r_eax = bswap(r_eax)
            case ('mov', op_str):
                #mov     cl, ah
                #mov     ah, byte_80E8018[ecx]
                #mov     cl, al
                #mov     al, byte_80E8018[ecx]
                op1, op2 = op_str.split(', ')
                if('cl' == op1):
                    if('al' == op2):
                        r_eax = (r_eax & 0xffffff00) | r_ecx
                    elif('ah' == op2):
                        r_eax = (r_eax & 0xffff00ff) | (r_ecx << 8)
                elif ('al' == op1) and ('byte ptr [ecx + 0x80e8018]' == op2):
                    r_ecx = arr_list.index(r_eax & 0xff)
                elif ('ah' == op1) and ('byte ptr [ecx + 0x80e8018]' == op2):
                    r_ecx = arr_list.index((r_eax >> 8) & 0xff)
        r_eax &= 0xffffffff
    
    return struct.pack('<I', r_eax)

for i, cmp_dword in enumerate(cmp_list):
    start(i)
    flag += solve(cmp_dword)

print("WACON2023{" + hashlib.sha256(flag).hexdigest() + "}")

#WACON2023{583f7c0189b232225e07c755f734c82029f58113c60f4eab75ab31159b568f49}

repiet ./chall.bmp -o baby_artist.python -b python -O 3

flag = [125, 63, 121, 114, 116, 95, 52, 78, 110, 52, 119, 95, 95, 33, 103, 110, 49, 87, 52, 114, 68, 95, 110, 85, 102, 123, 51, 50, 48, 50, 78, 79, 67, 65, 87]

right_flag = b''
for i in reversed(flag):
    right_flag += bytes([i])

print(right_flag.decode())
#WACON2023{fUn_Dr4W1ng!__w4nN4_try?}