"""
Extract the static IV string from DLL and find how key derivation works.

Key findings from disassembly:
1. Static 30-byte string at RVA 0x02725C60 used as IV (truncated to 16)
2. SHA256(combined) used as AES key material
3. Combined = some_function(key_string, iv_from_data, flag)
4. Function at 0x18006c3d0 combines key + iv_prefix

Need to:
a) Read the static IV string
b) Disassemble function 0x18006c3d0 to understand combination
c) Try decryption
"""
import struct, hashlib
from capstone import Cs, CS_ARCH_X86, CS_MODE_64

DLL_PATH = r"c:\Users\MattyMroz\Desktop\PROJECTS\ONEOCR\ocr_data\oneocr.dll"
MODEL_PATH = r"c:\Users\MattyMroz\Desktop\PROJECTS\ONEOCR\ocr_data\oneocr.onemodel"

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

with open(MODEL_PATH, "rb") as f:
    model = f.read()

# Parse PE sections for RVA → file offset mapping
e_lfanew = struct.unpack_from('<I', dll, 0x3c)[0]
num_sections = struct.unpack_from('<H', dll, e_lfanew + 6)[0]
opt_size = struct.unpack_from('<H', dll, e_lfanew + 20)[0]
sections_off = e_lfanew + 24 + opt_size

sections = []
for i in range(num_sections):
    so = sections_off + i * 40
    name = dll[so:so+8].rstrip(b'\x00').decode('ascii', errors='replace')
    vsize = struct.unpack_from('<I', dll, so + 8)[0]
    vrva = struct.unpack_from('<I', dll, so + 12)[0]
    rawsize = struct.unpack_from('<I', dll, so + 16)[0]
    rawoff = struct.unpack_from('<I', dll, so + 20)[0]
    sections.append((name, vrva, vsize, rawoff, rawsize))
    print(f"Section {name}: RVA=0x{vrva:08x} VSize=0x{vsize:08x} Raw=0x{rawoff:08x} RawSize=0x{rawsize:08x}")

def rva_to_foff(rva):
    for name, vrva, vsize, rawoff, rawsize in sections:
        if vrva <= rva < vrva + rawsize:
            return rawoff + (rva - vrva)
    return None

IMAGE_BASE = 0x180000000
TEXT_VA = 0x1000
TEXT_FILE_OFFSET = 0x400

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

# 1. Read the static 30-byte string at RVA 0x02725C60
# The LEA instruction was at RVA 0x0015baac:
# lea rdx, [rip + 0x25ca1ad]
# RIP = 0x0015baac + 7 = 0x0015bab3
# Target RVA = 0x0015bab3 + 0x25ca1ad = ?
target_rva = 0x0015bab3 + 0x25ca1ad
print(f"\nStatic IV string RVA: 0x{target_rva:08x}")
foff = rva_to_foff(target_rva)
print(f"File offset: 0x{foff:08x}" if foff else "NOT FOUND")

if foff:
    static_iv_30 = dll[foff:foff+30]
    print(f"Static IV (30 bytes): {static_iv_30.hex()}")
    print(f"Static IV (30 chars): {static_iv_30}")
    static_iv_16 = static_iv_30[:16]
    print(f"Static IV truncated to 16: {static_iv_16.hex()}")
    print(f"Static IV truncated (repr): {static_iv_16}")

# Also check nearby strings for context
if foff:
    print(f"\nContext around static IV string (foff-16 to foff+48):")
    for i in range(-16, 64):
        c = dll[foff+i]
        print(f"  +{i:3d}: 0x{c:02x} ({chr(c) if 32 <= c < 127 else '.'})")

# 2. Read the first 16 bytes of encrypted data (the "prefix" extracted before key derivation)
header_offset = struct.unpack_from('<Q', model, 0)[0]
prefix_16 = model[8:24]
print(f"\nData prefix (first 16 bytes after offset): {prefix_16.hex()}")
print(f"Data prefix repr: {prefix_16}")

# 3. Disassemble the key combination function at 0x18006c3d0
# RVA = 0x0006c3d0
combo_rva = 0x0006c3d0
combo_foff = rva_to_foff(combo_rva)
print(f"\nKey combination function RVA: 0x{combo_rva:08x}, file: 0x{combo_foff:08x}" if combo_foff else "NOT FOUND")

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

if combo_foff:
    code = dll[combo_foff:combo_foff + 0x200]
    print(f"\n{'='*100}")
    print(f"Key combination function at RVA 0x{combo_rva:08x}")
    print(f"{'='*100}")
    for insn in md.disasm(code, IMAGE_BASE + combo_rva):
        foff2 = rva_to_foff(insn.address - IMAGE_BASE)
        line = f"  {insn.address - IMAGE_BASE:08x} ({foff2:08x}): {insn.bytes.hex():<40s} {insn.mnemonic:<10s} {insn.op_str}"
        if insn.mnemonic == 'ret':
            print(line)
            break
        print(line)

# 4. Now let's also check what string is compared when key is empty
# At 0x0015b88d: lea rdx, [rip + 0x25b59db]
# RIP = 0x0015b88d + 7 = 0x0015b894
# Target = 0x0015b894 + 0x25b59db = ?
empty_key_str_rva = 0x0015b894 + 0x25b59db
empty_foff = rva_to_foff(empty_key_str_rva)
if empty_foff:
    s = dll[empty_foff:empty_foff+64].split(b'\x00')[0]
    print(f"\nDefault key comparison string: {s}")

# 5. Try decryption with various key derivation methods
print("\n" + "="*100)
print("DECRYPTION ATTEMPTS WITH STATIC IV AND DERIVED KEY")
print("="*100)

KEY_RAW = b'kj)TGtrK>f]b[Piow.gU+nC@s""""""4'

# Encrypted header: model[8 : header_offset]
enc_header = model[8:header_offset]
# The prefix/IV: first 16 bytes
data_prefix = enc_header[:16]
# The ciphertext: remaining bytes
ciphertext = enc_header[16:]

MAGIC = 0x252b081a4a

from Crypto.Cipher import AES

def try_dec(aes_key, iv, ct, label):
    """Try AES-256-CFB128 decryption."""
    try:
        cipher = AES.new(aes_key, AES.MODE_CFB, iv=iv, segment_size=128)
        pt = cipher.decrypt(ct[:256])
        if len(pt) >= 24:
            magic = struct.unpack_from('<Q', pt, 0x10)[0]
            if magic == MAGIC:
                print(f"  *** SUCCESS *** {label}")
                print(f"  First 64 bytes: {pt[:64].hex()}")
                return True
            else:
                unique = len(set(pt[:64]))
                # Only print if somewhat promising
                if unique < 45 or magic & 0xFF == 0x4a:
                    print(f"  {label}: magic=0x{magic:016x}, unique_64={unique}")
    except Exception as e:
        print(f"  {label}: ERROR {e}")
    return False

if foff:
    iv = static_iv_16
    
    # Try 1: SHA256(key + data_prefix)
    combined1 = KEY_RAW + data_prefix
    aes_key1 = hashlib.sha256(combined1).digest()
    try_dec(aes_key1, iv, ciphertext, "SHA256(key + data_prefix)")
    
    # Try 2: SHA256(data_prefix + key)  
    combined2 = data_prefix + KEY_RAW
    aes_key2 = hashlib.sha256(combined2).digest()
    try_dec(aes_key2, iv, ciphertext, "SHA256(data_prefix + key)")
    
    # Try 3: SHA256(key) with static IV
    aes_key3 = hashlib.sha256(KEY_RAW).digest()
    try_dec(aes_key3, iv, ciphertext, "SHA256(key) + static_iv")
    
    # Try 4: Raw key with static IV
    try_dec(KEY_RAW, iv, ciphertext, "raw_key + static_iv")
    
    # Try 5: SHA256(key + data_prefix) on full enc_header (no prefix removal)
    try_dec(aes_key1, iv, enc_header, "SHA256(key+prefix) + full_header")
    try_dec(aes_key2, iv, enc_header, "SHA256(prefix+key) + full_header")
    
    # Try 6: Maybe prefix is NOT stripped from ciphertext for BCrypt
    try_dec(aes_key3, iv, enc_header, "SHA256(key) + static_iv + full_header")
    try_dec(KEY_RAW, iv, enc_header, "raw_key + static_iv + full_header")
    
    # Also try the full static_iv_30 string as both key and IV source
    # Maybe the static string IS the key, and data_prefix IS the IV
    try_dec(hashlib.sha256(static_iv_30).digest(), data_prefix, ciphertext, "SHA256(static30) + data_prefix_iv")
    
    # What if key derivation involves the static string too?
    # SHA256(key + static_string)
    combined3 = KEY_RAW + static_iv_30
    aes_key6 = hashlib.sha256(combined3).digest()
    try_dec(aes_key6, data_prefix, ciphertext, "SHA256(key + static30) + prefix_iv")
    try_dec(aes_key6, iv, ciphertext, "SHA256(key + static30) + static_iv")
    
    # What if the function combines key with static string, and data_prefix is IV?
    # Try many concatenation variants
    variants = [
        (KEY_RAW + data_prefix, iv, ciphertext, "key||prefix"),
        (data_prefix + KEY_RAW, iv, ciphertext, "prefix||key"),
        (KEY_RAW + static_iv_16, iv, ciphertext, "key||static16"),
        (KEY_RAW + static_iv_30, iv, ciphertext, "key||static30"),
        (static_iv_16 + KEY_RAW, iv, ciphertext, "static16||key"),
        (static_iv_30 + KEY_RAW, iv, ciphertext, "static30||key"),
        (KEY_RAW + data_prefix, data_prefix, ciphertext, "key||prefix, iv=prefix"),
        (data_prefix + KEY_RAW, data_prefix, ciphertext, "prefix||key, iv=prefix"),
    ]
    
    for combo, iv_used, ct, desc in variants:
        aes_key = hashlib.sha256(combo).digest()
        try_dec(aes_key, iv_used, ct, f"SHA256({desc})")
    
    # Maybe the function at 0x06c3d0 does something more complex
    # Let's also try: the "combined" is just the key (no IV involvement),
    # and the function just copies/formats the key
    # With different IV sources
    
    # Try with BCrypt API directly
    print("\n--- BCrypt API tests with static IV ---")
    import ctypes
    bcrypt = ctypes.windll.bcrypt
    
    def bcrypt_dec(key_bytes, iv_bytes, ct_bytes, label):
        hAlg = ctypes.c_void_p()
        status = bcrypt.BCryptOpenAlgorithmProvider(ctypes.byref(hAlg), "AES", None, 0)
        if status != 0:
            print(f"  {label}: OpenAlg failed {status}")
            return None
        
        mode = "ChainingModeCFB".encode('utf-16-le') + b'\x00\x00'
        bcrypt.BCryptSetProperty(hAlg, "ChainingMode", mode, len(mode), 0)
        
        block_len = ctypes.c_ulong(16)
        bcrypt.BCryptSetProperty(hAlg, "MessageBlockLength", 
                                ctypes.byref(block_len), 4, 0)
        
        hKey = ctypes.c_void_p()
        kb = (ctypes.c_byte * len(key_bytes))(*key_bytes)
        status = bcrypt.BCryptGenerateSymmetricKey(
            hAlg, ctypes.byref(hKey), None, 0, kb, len(key_bytes), 0)
        if status != 0:
            bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
            print(f"  {label}: GenKey failed {status}")
            return None
        
        ct_buf = (ctypes.c_byte * len(ct_bytes))(*ct_bytes)
        iv_buf = (ctypes.c_byte * len(iv_bytes))(*iv_bytes)
        
        out_size = ctypes.c_ulong(0)
        status = bcrypt.BCryptDecrypt(
            hKey, ct_buf, len(ct_bytes), None,
            iv_buf, len(iv_bytes), None, 0,
            ctypes.byref(out_size), 0)
        
        if status != 0:
            bcrypt.BCryptDestroyKey(hKey)
            bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
            print(f"  {label}: Decrypt size query failed {status:#x}")
            return None
        
        pt_buf = (ctypes.c_byte * out_size.value)()
        iv_buf2 = (ctypes.c_byte * len(iv_bytes))(*iv_bytes)
        result = ctypes.c_ulong(0)
        status = bcrypt.BCryptDecrypt(
            hKey, ct_buf, len(ct_bytes), None,
            iv_buf2, len(iv_bytes), pt_buf, out_size.value,
            ctypes.byref(result), 0)
        
        bcrypt.BCryptDestroyKey(hKey)
        bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
        
        if status != 0:
            print(f"  {label}: Decrypt failed {status:#x}")
            return None
        
        pt = bytes(pt_buf[:result.value])
        if len(pt) >= 24:
            magic = struct.unpack_from('<Q', pt, 0x10)[0]
            if magic == MAGIC:
                print(f"  *** BCrypt SUCCESS *** {label}")
                print(f"  First 64: {pt[:64].hex()}")
                return pt
        return pt
    
    # BCrypt tests with various key derivations
    for combo_data, desc in [
        (KEY_RAW, "raw_key"),
        (hashlib.sha256(KEY_RAW).digest(), "SHA256(key)"),
        (hashlib.sha256(KEY_RAW + data_prefix).digest(), "SHA256(key+prefix)"),
        (hashlib.sha256(data_prefix + KEY_RAW).digest(), "SHA256(prefix+key)"),
    ]:
        for iv_data, iv_desc in [(iv, "static16"), (data_prefix, "data_prefix")]:
            for ct_data, ct_desc in [(ciphertext, "ct_no_prefix"), (enc_header, "full_header")]:
                result = bcrypt_dec(combo_data, iv_data, ct_data[:512], 
                                   f"key={desc}, iv={iv_desc}, ct={ct_desc}")
                if result:
                    magic = struct.unpack_from('<Q', result, 0x10)[0] if len(result) >= 24 else 0
                    if magic == MAGIC:
                        print("FOUND THE CORRECT PARAMETERS!")

print("\nDone.")