"""
Disassemble the actual BCrypt crypto operations at 0x18015ba45+ 
and map all indirect calls to IAT entries.
"""
import struct
from capstone import Cs, CS_ARCH_X86, CS_MODE_64

DLL_PATH = r"c:\Users\MattyMroz\Desktop\PROJECTS\ONEOCR\ocr_data\oneocr.dll"
IMAGE_BASE = 0x180000000
TEXT_VA = 0x1000
TEXT_FILE_OFFSET = 0x400

def rva_to_file(rva):
    return rva - TEXT_VA + TEXT_FILE_OFFSET

def file_to_rva(foff):
    return foff - TEXT_FILE_OFFSET + TEXT_VA

with open(DLL_PATH, "rb") as f:
    dll_data = f.read()

md = Cs(CS_ARCH_X86, CS_MODE_64)
md.detail = False

def disasm_region(name, file_start, file_end):
    rva_start = file_to_rva(file_start)
    va_start = IMAGE_BASE + rva_start
    code = dll_data[file_start:file_end]
    print(f"\n{'='*100}")
    print(f"{name}")
    print(f"File: 0x{file_start:08x}-0x{file_end:08x}, RVA: 0x{rva_start:08x}")
    print(f"{'='*100}")
    for insn in md.disasm(code, va_start):
        foff = rva_to_file(insn.address - IMAGE_BASE)
        line = f"  {insn.address - IMAGE_BASE:08x} ({foff:08x}): {insn.bytes.hex():<40s} {insn.mnemonic:<14s} {insn.op_str}"
        # Annotate indirect calls
        if insn.mnemonic == 'call' and insn.bytes[0] == 0xFF and insn.bytes[1] == 0x15:
            disp = struct.unpack_from('<i', bytes(insn.bytes), 2)[0]
            target_rva = (insn.address - IMAGE_BASE) + insn.size + disp
            line += f"  ; IAT@0x{target_rva:08x}"
        print(line)

# First, let's identify ALL BCrypt IAT entries
# From previous analysis:
# BCryptOpenAlgorithmProvider → IAT 0x0081a5e0
# BCryptGetProperty → IAT 0x0081a5d0 
# BCryptSetProperty → IAT 0x0081a608
# Let's find the rest by looking at the import section

# Parse PE to find BCrypt imports
print("="*100)
print("FINDING ALL BCRYPT IAT ENTRIES")
print("="*100)

# Parse PE headers
e_lfanew = struct.unpack_from('<I', dll_data, 0x3c)[0]
opt_hdr_off = e_lfanew + 24
import_dir_rva = struct.unpack_from('<I', dll_data, opt_hdr_off + 120)[0]  # Import RVA
import_dir_size = struct.unpack_from('<I', dll_data, opt_hdr_off + 124)[0]

# Find sections for RVA to file offset mapping
num_sections = struct.unpack_from('<H', dll_data, e_lfanew + 6)[0]
sections_off = e_lfanew + 24 + struct.unpack_from('<H', dll_data, e_lfanew + 20)[0]

sections = []
for i in range(num_sections):
    sec_off = sections_off + i * 40
    name = dll_data[sec_off:sec_off+8].rstrip(b'\x00').decode('ascii', errors='replace')
    vsize = struct.unpack_from('<I', dll_data, sec_off + 8)[0]
    vrva = struct.unpack_from('<I', dll_data, sec_off + 12)[0]
    rawsize = struct.unpack_from('<I', dll_data, sec_off + 16)[0]
    rawoff = struct.unpack_from('<I', dll_data, sec_off + 20)[0]
    sections.append((name, vrva, vsize, rawoff, rawsize))
    
def rva_to_foff(rva):
    for name, vrva, vsize, rawoff, rawsize in sections:
        if vrva <= rva < vrva + vsize:
            return rawoff + (rva - vrva)
    return None

if import_dir_rva:
    ioff = rva_to_foff(import_dir_rva)
    if ioff:
        idx = 0
        while True:
            desc_off = ioff + idx * 20
            ilt_rva = struct.unpack_from('<I', dll_data, desc_off)[0]
            name_rva = struct.unpack_from('<I', dll_data, desc_off + 12)[0]
            iat_rva = struct.unpack_from('<I', dll_data, desc_off + 16)[0]
            if ilt_rva == 0 and name_rva == 0:
                break
            name_off = rva_to_foff(name_rva)
            if name_off:
                dname = dll_data[name_off:name_off+64].split(b'\x00')[0].decode('ascii', errors='replace')
                if 'bcrypt' in dname.lower():
                    print(f"\nDLL: {dname}, ILT RVA: 0x{ilt_rva:08x}, IAT RVA: 0x{iat_rva:08x}")
                    # Walk the ILT to find function names
                    ilt_off = rva_to_foff(ilt_rva)
                    iat_entry_rva = iat_rva
                    j = 0
                    while ilt_off:
                        entry = struct.unpack_from('<Q', dll_data, ilt_off + j * 8)[0]
                        if entry == 0:
                            break
                        if entry & (1 << 63):
                            ordinal = entry & 0xFFFF
                            print(f"  IAT 0x{iat_entry_rva:08x}: Ordinal {ordinal}")
                        else:
                            hint_rva = entry & 0x7FFFFFFF
                            hint_off = rva_to_foff(hint_rva)
                            if hint_off:
                                hint = struct.unpack_from('<H', dll_data, hint_off)[0]
                                fname = dll_data[hint_off+2:hint_off+66].split(b'\x00')[0].decode('ascii', errors='replace')
                                print(f"  IAT 0x{iat_entry_rva:08x}: {fname} (hint {hint})")
                        iat_entry_rva += 8
                        j += 1
            idx += 1

# Now disassemble the actual BCrypt crypto operation block
# This is at RVA 0x0015ba45 = file 0x0015ae45
disasm_region(
    "BCrypt crypto operations (key gen + decrypt)",
    0x0015ae45, 0x0015b200
)

# Find all indirect calls in extended range
print("\n" + "="*100)
print("ALL INDIRECT CALLS (ff 15) in range 0x0015ae00-0x0015c200")
print("="*100)

for i in range(0x0015ae00, 0x0015c200):
    if dll_data[i] == 0xFF and dll_data[i+1] == 0x15:
        rva = file_to_rva(i)
        disp = struct.unpack_from('<i', dll_data, i + 2)[0]
        target_rva = rva + 6 + disp
        print(f"  File 0x{i:08x} (RVA 0x{rva:08x}): call [rip+0x{disp:x}] -> IAT@0x{target_rva:08x}")

# Also disassemble the function at 0x18015abd0 (called to process data when r14b=true)
# RVA 0x0015abd0, file 0x00159fd0
disasm_region(
    "Function at 0x18015abd0 (called on data when r14b=true)",
    0x00159fd0, 0x0015a0c0
)