from binaryninja.architecture import Architecture, ArchitectureHook
from dataclasses import dataclass
import struct


@dataclass
class ExceptionStruct:
    rva: int
    case: int
    loc1: int
    loc2: int

# Build breakpoint EIP adjustment table
#######################################
def get_struct_data(addr) -> ExceptionStruct:
    rva, case, loc1, loc2 = struct.unpack('<LBll', bv.read(addr, 0xd))
    return ExceptionStruct(rva, case, loc1, loc2)

def create_breakpoint_lut():
    breakpoint_lut = {}
    STRUCT_RVA = 0x1280a0
    BASE_ADDRESS = 0x60000
    offset = 0
    while True:
        x = get_struct_data(BASE_ADDRESS + STRUCT_RVA + offset)
        if x.rva == 0:  # Continue scanning for structs until we encounter null data
            break
        breakpoint_lut[BASE_ADDRESS + x.rva] = x
        offset += 0xd
    return breakpoint_lut


# Helper functions for our case functions.
##########################################
def get_label(il, address):
    label = il.get_label_for_address(Architecture['x86'], address)
    if label == None:
        il.add_label_for_address(Architecture['x86'], address)
        label = il.get_label_for_address(Architecture['x86'], address)
        assert label != None
    return label

def create_branch(il, condition, t_addr, f_addr):
    t = get_label(il, t_addr)
    f = get_label(il, f_addr)
    il.append(il.if_expr(condition, t, f))


# Code responsible for creating the appropriate branch type for each int3
#########################################################################
def case0(il, t_addr, f_addr):
    flag = il.flag_group('ne')
    create_branch(il, flag, t_addr, f_addr)

def case1(il, t_addr, f_addr):
    flag = il.flag_group('e')   
    create_branch(il, flag, t_addr, f_addr)

def case2(il, t_addr, f_addr):
    flag = il.flag_group('ge')   
    create_branch(il, flag, t_addr, f_addr)

def case3(il, t_addr, f_addr):
    flag = il.flag_group('lt')   
    create_branch(il, flag, t_addr, f_addr)

def case4(il, t_addr, f_addr):
    flag = il.flag_group('gt')   
    create_branch(il, flag, t_addr, f_addr)

def case5(il, t_addr, f_addr):
    flag = il.flag_group('le')   
    create_branch(il, flag, t_addr, f_addr)

def case6(il, t_addr, f_addr):
    flag = il.flag_condition(LowLevelILFlagCondition.LLFC_SGE)  
    create_branch(il, flag, t_addr, f_addr)

def case7(il, t_addr, f_addr):
    flag = il.flag_condition(LowLevelILFlagCondition.LLFC_SLE)  
    create_branch(il, flag, f_addr, t_addr) # Reserved order of branch addresses

def case8(il, t_addr, f_addr):
    flag = il.flag_condition(LowLevelILFlagCondition.LLFC_SGT)  
    create_branch(il, flag, t_addr, f_addr)

def case9(il, t_addr, f_addr):
    flag = il.flag_condition(LowLevelILFlagCondition.LLFC_SLE)  
    create_branch(il, flag, t_addr, f_addr)

def casea(il, t_addr, f_addr):
    flag = il.flag_group('pe')   
    create_branch(il, flag, t_addr, f_addr)

def caseb(il, t_addr, f_addr):
    flag = il.flag_group('po')   
    create_branch(il, flag, t_addr, f_addr)

def cased(il, t_addr, f_addr):
    flag = il.flag_condition(LowLevelILFlagCondition.LLFC_POS)  
    create_branch(il, flag, t_addr, f_addr)

def casee(il, t_addr, f_addr):
    flag = il.flag_condition(LowLevelILFlagCondition.LLFC_NEG)  
    create_branch(il, flag, t_addr, f_addr)

def case11(il, t_addr, f_addr):
    target = get_label(il, f_addr)
    il.append(il.goto(target))

exception_handler_lut = {
    0x0: case0,
    0x1: case1,
    0x2: case2,
    0x3: case3,
    0x4: case4,
    0x5: case5,
    0x6: case6,
    0x7: case7,
    0x8: case8,
    0x9: case9,
    0xa: casea,
    0xb: caseb,
    0xd: cased,
    0xe: casee,
    0x11: case11
}

# Create and insert our IL hook
###############################
class ILLifterHook(ArchitectureHook):
    lut: dict[int, ExceptionStruct]
    
    def __init__(self, base_arch, lut):
        super(ILLifterHook, self).__init__(base_arch)
        self.lut = lut
    
    def get_instruction_info(self, data, addr):
        # This callback is responsible for informing BN on what branches an instruction has.
        # We need to hook this so that we can inform BN that some int3s have a true/false branch.
        # If we don't hook this, binary ninja will not include the branch locations within the functions' address space.
        # If the address isn't within the function, we wont be able to acquire address labels for if_expr IL's.
        if addr in self.lut and bv.read(addr, 1) == b'\xcc':
            cf_data = self.lut[addr]
            info = InstructionInfo()
            info.length = 1
            if cf_data.case == 0x11:
                info.add_branch(enums.BranchType.UnconditionalBranch, addr + cf_data.loc2)
            else:
                info.add_branch(enums.BranchType.TrueBranch, addr + cf_data.loc1)
                info.add_branch(enums.BranchType.FalseBranch, addr + cf_data.loc2)
            return info
        else:
            return super(ILLifterHook, self).get_instruction_info(data, addr)
    
    def get_instruction_low_level_il(self, data, addr, il):
        # This callback is responsible for inserting LLIL into the IL tree/graph.
        # We'll process the int3s as if_expr's.
        if data[0] == 0xCC and addr in self.lut:
            control_flow = self.lut[addr]
            if control_flow.case in exception_handler_lut:
                exception_handler_lut[control_flow.case](il, addr + control_flow.loc1, addr + control_flow.loc2)
            else:
                print(f'Unhandled @ {hex(addr)} for case: {control_flow.case}')
                il.append(il.unimplemented())
            return 1
        else:
            return super(ILLifterHook, self).get_instruction_low_level_il(data, addr, il)

hook = ILLifterHook(Architecture['x86'], create_breakpoint_lut())
hook.register()