"""
OneOCR .onemodel decryption using Windows BCrypt CNG API directly.
Replicates the exact behavior of oneocr.dll's Crypto.cpp.

Known from DLL analysis:
- BCryptOpenAlgorithmProvider with L"AES"
- BCryptSetProperty L"ChainingMode" = L"ChainingModeCFB"
- BCryptGetProperty L"BlockLength" (→ 16)
- BCryptSetProperty L"MessageBlockLength" = 16 (→ CFB128)
- BCryptGenerateSymmetricKey with raw key bytes
- BCryptDecrypt
- SHA256Hash function exists (uses BCryptCreateHash/BCryptHashData/BCryptFinishHash)
"""

import ctypes
import ctypes.wintypes as wintypes
import struct
import hashlib
import zlib
from collections import Counter
import math
import os

# ═══════════════════════════════════════════════════════════════
# Windows BCrypt API via ctypes
# ═══════════════════════════════════════════════════════════════

bcrypt = ctypes.WinDLL("bcrypt")

BCRYPT_ALG_HANDLE = ctypes.c_void_p
BCRYPT_KEY_HANDLE = ctypes.c_void_p
NTSTATUS = ctypes.c_long

# Constants
BCRYPT_AES_ALGORITHM = "AES"
BCRYPT_SHA256_ALGORITHM = "SHA256"
BCRYPT_CHAINING_MODE = "ChainingMode"
BCRYPT_CHAIN_MODE_CFB = "ChainingModeCFB"
BCRYPT_BLOCK_LENGTH = "BlockLength"
BCRYPT_MESSAGE_BLOCK_LENGTH = "MessageBlockLength"

def check_status(status, msg=""):
    if status != 0:
        raise OSError(f"BCrypt error 0x{status & 0xFFFFFFFF:08x}: {msg}")


def bcrypt_sha256(data: bytes) -> bytes:
    """Compute SHA256 using Windows BCrypt API."""
    hAlg = BCRYPT_ALG_HANDLE()
    status = bcrypt.BCryptOpenAlgorithmProvider(
        ctypes.byref(hAlg),
        ctypes.c_wchar_p(BCRYPT_SHA256_ALGORITHM),
        None, 0)
    check_status(status, "SHA256 OpenAlgorithmProvider")

    hHash = ctypes.c_void_p()
    status = bcrypt.BCryptCreateHash(hAlg, ctypes.byref(hHash), None, 0, None, 0, 0)
    check_status(status, "CreateHash")

    status = bcrypt.BCryptHashData(hHash, data, len(data), 0)
    check_status(status, "HashData")

    hash_out = (ctypes.c_ubyte * 32)()
    status = bcrypt.BCryptFinishHash(hHash, hash_out, 32, 0)
    check_status(status, "FinishHash")

    bcrypt.BCryptDestroyHash(hHash)
    bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)

    return bytes(hash_out)


def bcrypt_aes_cfb_decrypt(ciphertext: bytes, key: bytes, iv: bytes,
                            message_block_length: int = 16) -> bytes:
    """Decrypt using AES-CFB via Windows BCrypt CNG API.
    
    message_block_length: 1 for CFB8, 16 for CFB128
    """
    hAlg = BCRYPT_ALG_HANDLE()
    status = bcrypt.BCryptOpenAlgorithmProvider(
        ctypes.byref(hAlg),
        ctypes.c_wchar_p(BCRYPT_AES_ALGORITHM),
        None, 0)
    check_status(status, "AES OpenAlgorithmProvider")

    # Set chaining mode to CFB
    mode_str = BCRYPT_CHAIN_MODE_CFB
    mode_buf = ctypes.create_unicode_buffer(mode_str)
    mode_size = (len(mode_str) + 1) * 2  # UTF-16 with null terminator
    status = bcrypt.BCryptSetProperty(
        hAlg,
        ctypes.c_wchar_p(BCRYPT_CHAINING_MODE),
        mode_buf, mode_size, 0)
    check_status(status, "SetProperty ChainingMode")

    # Set message block length (feedback size)
    mbl = ctypes.c_ulong(message_block_length)
    status = bcrypt.BCryptSetProperty(
        hAlg,
        ctypes.c_wchar_p(BCRYPT_MESSAGE_BLOCK_LENGTH),
        ctypes.byref(mbl), ctypes.sizeof(mbl), 0)
    check_status(status, f"SetProperty MessageBlockLength={message_block_length}")

    # Generate symmetric key
    hKey = BCRYPT_KEY_HANDLE()
    key_buf = (ctypes.c_ubyte * len(key))(*key)
    status = bcrypt.BCryptGenerateSymmetricKey(
        hAlg, ctypes.byref(hKey), None, 0, key_buf, len(key), 0)
    check_status(status, "GenerateSymmetricKey")

    # Prepare IV (BCrypt modifies it during decryption, so use a copy)
    iv_buf = (ctypes.c_ubyte * 16)(*iv)

    # Prepare input/output buffers
    ct_buf = (ctypes.c_ubyte * len(ciphertext))(*ciphertext)
    pt_buf = (ctypes.c_ubyte * len(ciphertext))()
    result_len = ctypes.c_ulong(0)

    # Decrypt
    status = bcrypt.BCryptDecrypt(
        hKey, ct_buf, len(ciphertext), None,
        iv_buf, 16,
        pt_buf, len(ciphertext),
        ctypes.byref(result_len), 0)
    check_status(status, "BCryptDecrypt")

    # Cleanup
    bcrypt.BCryptDestroyKey(hKey)
    bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)

    return bytes(pt_buf[:result_len.value])


def entropy(data: bytes) -> float:
    """Shannon entropy (bits per byte)."""
    if not data:
        return 0.0
    freq = Counter(data)
    total = len(data)
    return -sum((c / total) * math.log2(c / total) for c in freq.values())


def hex_dump(data: bytes, offset: int = 0, max_lines: int = 8) -> str:
    lines = []
    for i in range(0, min(len(data), max_lines * 16), 16):
        hex_part = " ".join(f"{b:02x}" for b in data[i:i+16])
        ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in data[i:i+16])
        lines.append(f"  {offset+i:08x}: {hex_part:<48s}  {ascii_part}")
    return "\n".join(lines)


def check_decrypted(data: bytes, label: str) -> bool:
    """Check if decrypted data looks valid. Return True if promising."""
    if not data or len(data) < 16:
        return False

    ent = entropy(data[:min(4096, len(data))])
    u32_le = struct.unpack_from("<I", data, 0)[0]

    # Check for magic_number = 1
    magic_match = (u32_le == 1)

    # Check for protobuf
    protobuf = data[0] == 0x08 or data[0] == 0x0a

    # Check for compression headers
    zlib_header = data[:2] in [b"\x78\x01", b"\x78\x5e", b"\x78\x9c", b"\x78\xda"]
    gzip_header = data[:2] == b"\x1f\x8b"
    lz4_header = data[:4] == b"\x04\x22\x4d\x18"

    is_promising = magic_match or (ent < 7.0) or zlib_header or gzip_header or lz4_header

    if is_promising or protobuf:
        print(f"\n  ★★★ {'MAGIC=1 !!!' if magic_match else 'Promising'}: {label}")
        print(f"    Entropy: {ent:.3f}, uint32_LE[0]={u32_le}, first_byte=0x{data[0]:02x}")
        print(f"    First 128 bytes:")
        print(hex_dump(data[:128]))
        if zlib_header:
            print(f"    → ZLIB header detected!")
        if gzip_header:
            print(f"    → GZIP header detected!")
        if lz4_header:
            print(f"    → LZ4 header detected!")
        if magic_match:
            print(f"    → MAGIC_NUMBER = 1 !! This is likely correct decryption!")
            # Try decompression after offset 4 or later
            for skip in [0, 4, 8, 16, 32, 64]:
                chunk = data[skip:skip+min(10000, len(data)-skip)]
                try:
                    dec = zlib.decompress(chunk)
                    print(f"    → ZLIB decompress SUCCESS at skip={skip}: {len(dec)} bytes!")
                    print(f"      First 64: {dec[:64].hex()}")
                    return True
                except:
                    pass
                try:
                    dec = zlib.decompress(chunk, -15)
                    print(f"    → Raw DEFLATE decompress SUCCESS at skip={skip}: {len(dec)} bytes!")
                    print(f"      First 64: {dec[:64].hex()}")
                    return True
                except:
                    pass
        return True
    return False


# ═══════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════

MODEL_PATH = r"c:\Users\MattyMroz\Desktop\PROJECTS\ONEOCR\ocr_data\oneocr.onemodel"
KEY_RAW = b'kj)TGtrK>f]b[Piow.gU+nC@s""""""4'
KEY_SHA256 = bcrypt_sha256(KEY_RAW)

print("=" * 80)
print("OneOCR Decryption via Windows BCrypt CNG API")
print("=" * 80)

print(f"\nKey (raw): {KEY_RAW.hex()}")
print(f"Key (SHA256): {KEY_SHA256.hex()}")
print(f"Python hashlib SHA256: {hashlib.sha256(KEY_RAW).digest().hex()}")
print(f"BCrypt SHA256 match: {KEY_SHA256 == hashlib.sha256(KEY_RAW).digest()}")

# Read file
with open(MODEL_PATH, "rb") as f:
    full_data = f.read()
filesize = len(full_data)

header_offset = struct.unpack_from("<I", full_data, 0)[0]  # 22636
payload_size = struct.unpack_from("<Q", full_data, header_offset + 8)[0]  # 58431147
payload_start = header_offset + 16  # 22652

print(f"\nFile size: {filesize:,}")
print(f"Header offset: {header_offset}")
print(f"Payload size: {payload_size:,}")
print(f"Payload start: {payload_start}")

# ═══════════════════════════════════════════════════════════════
# Test 1: Try standard combinations with BCrypt API
# ═══════════════════════════════════════════════════════════════

print("\n" + "=" * 80)
print("TEST 1: Standard combinations via BCrypt CFB128")
print("=" * 80)

iv_zero = b"\x00" * 16
iv_candidates = {
    "zeros": iv_zero,
    "file[8:24]": full_data[8:24],
    "file[4:20]": full_data[4:20],
    "file[0:16]": full_data[0:16],
    f"file[{header_offset}:{header_offset+16}]": full_data[header_offset:header_offset+16],
    f"file[{payload_start}:{payload_start+16}]": full_data[payload_start:payload_start+16],
    "SHA256(key)[:16]": KEY_SHA256[:16],
    "SHA256(key)[16:]": KEY_SHA256[16:],
    "key_raw[:16]": KEY_RAW[:16],
    "key_raw[16:]": KEY_RAW[16:],
}

key_candidates = {
    "raw": KEY_RAW,
    "SHA256": KEY_SHA256,
}

data_regions = {
    "header[8:]": full_data[8:8+4096],
    f"payload[{payload_start}:]": full_data[payload_start:payload_start+4096],
}

for mbl in [16, 1]:  # CFB128 first (most likely), then CFB8
    for key_name, key in key_candidates.items():
        for iv_name, iv in iv_candidates.items():
            for region_name, region_data in data_regions.items():
                label = f"CFB{'128' if mbl == 16 else '8'} key={key_name} iv={iv_name} data={region_name}"
                try:
                    dec = bcrypt_aes_cfb_decrypt(region_data, key, iv, mbl)
                    if check_decrypted(dec, label):
                        pass  # Already printed
                except Exception as e:
                    pass  # Silently skip errors

# ═══════════════════════════════════════════════════════════════
# Test 2: Known-plaintext IV search
# ═══════════════════════════════════════════════════════════════

print("\n" + "=" * 80)
print("TEST 2: Known-plaintext IV search (magic_number=1)")
print("=" * 80)
print("  Searching for IV that produces magic_number=1 (0x01000000) at start...")

# For AES-CFB128, first block:
# plaintext[0:16] = AES_ECB_encrypt(IV, key) XOR ciphertext[0:16]
# We want plaintext[0:4] = 01 00 00 00 (LE)
# So: AES_ECB_encrypt(IV, key)[0:4] = ciphertext[0:4] XOR 01 00 00 00

# We can't easily predict AES output, so we try each IV candidate
# Try every 4-byte aligned position in header as IV, with both key candidates

found = False
for key_name, key in key_candidates.items():
    for mbl in [16, 1]:
        # Try IV from file at every 4-byte step in the first 22700 bytes  
        for iv_offset in range(0, min(22700, filesize - 16), 4):
            iv = full_data[iv_offset:iv_offset + 16]
            
            # Try decrypting header encrypted data (byte 8+)
            ct = full_data[8:24]  # Just decrypt first 16 bytes
            try:
                dec = bcrypt_aes_cfb_decrypt(ct, key, iv, mbl)
                u32 = struct.unpack_from("<I", dec, 0)[0]
                if u32 == 1:
                    print(f"\n  ★★★ FOUND! magic_number=1 with iv_offset={iv_offset}, key={key_name}, CFB{'128' if mbl==16 else '8'}")
                    print(f"    IV: {iv.hex()}")
                    print(f"    Decrypted first 16 bytes: {dec[:16].hex()}")
                    # Decrypt more data
                    dec_full = bcrypt_aes_cfb_decrypt(full_data[8:8+4096], key, iv, mbl)
                    check_decrypted(dec_full, f"FULL header with iv_offset={iv_offset}")
                    found = True
            except:
                pass
            
            # Try decrypting payload data
            ct2 = full_data[payload_start:payload_start+16]
            try:
                dec2 = bcrypt_aes_cfb_decrypt(ct2, key, iv, mbl)
                u32_2 = struct.unpack_from("<I", dec2, 0)[0]
                if u32_2 == 1:
                    print(f"\n  ★★★ FOUND! magic_number=1 with iv_offset={iv_offset}, key={key_name}, CFB{'128' if mbl==16 else '8'}")
                    print(f"    IV: {iv.hex()}")
                    print(f"    Decrypted first 16 bytes: {dec2[:16].hex()}")
                    # Decrypt more data
                    dec_full2 = bcrypt_aes_cfb_decrypt(full_data[payload_start:payload_start+4096], key, iv, mbl)
                    check_decrypted(dec_full2, f"FULL payload with iv_offset={iv_offset}")
                    found = True
            except:
                pass
        
        if found:
            break
    if found:
        break

if not found:
    print("  No IV found in file that produces magic_number=1")

# ═══════════════════════════════════════════════════════════════
# Test 3: Try derived IVs not from file
# ═══════════════════════════════════════════════════════════════

print("\n" + "=" * 80)
print("TEST 3: Derived IV strategies via BCrypt")
print("=" * 80)

derived_ivs = {
    "zeros": b"\x00" * 16,
    "SHA256(key)[:16]": KEY_SHA256[:16],
    "SHA256(key)[16:]": KEY_SHA256[16:],
    "key[:16]": KEY_RAW[:16],
    "key[16:]": KEY_RAW[16:],
    "SHA256('')[:16]": hashlib.sha256(b"").digest()[:16],
    "SHA256('\\0')[:16]": hashlib.sha256(b"\x00").digest()[:16],
    "MD5(key)": hashlib.md5(KEY_RAW).digest(),
    "SHA256('oneocr')[:16]": hashlib.sha256(b"oneocr").digest()[:16],
    "SHA256(key+\\0)[:16]": hashlib.sha256(KEY_RAW + b"\x00").digest()[:16],
    "SHA256(key_reversed)[:16]": hashlib.sha256(KEY_RAW[::-1]).digest()[:16],
    "key XOR 0x36 [:16]": bytes(b ^ 0x36 for b in KEY_RAW[:16]),  # HMAC ipad
    "key XOR 0x5c [:16]": bytes(b ^ 0x5c for b in KEY_RAW[:16]),  # HMAC opad
}

for iv_name, iv in derived_ivs.items():
    for key_name, key in key_candidates.items():
        for mbl in [16, 1]:
            for region_name, ct in [("header[8:]", full_data[8:8+4096]),
                                     (f"payload", full_data[payload_start:payload_start+4096])]:
                try:
                    dec = bcrypt_aes_cfb_decrypt(ct, key, iv, mbl)
                    label = f"CFB{'128' if mbl==16 else '8'} key={key_name} iv={iv_name} data={region_name}"
                    check_decrypted(dec, label)
                except:
                    pass

# ═══════════════════════════════════════════════════════════════
# Test 4: What if entire file from byte 0 is encrypted?
# ═══════════════════════════════════════════════════════════════

print("\n" + "=" * 80)
print("TEST 4: Entire file encrypted from byte 0")
print("=" * 80)

for key_name, key in key_candidates.items():
    for mbl in [16, 1]:
        for iv_name, iv in [("zeros", iv_zero), ("SHA256(key)[:16]", KEY_SHA256[:16]),
                             ("key[:16]", KEY_RAW[:16])]:
            try:
                dec = bcrypt_aes_cfb_decrypt(full_data[:4096], key, iv, mbl)
                label = f"CFB{'128' if mbl==16 else '8'} key={key_name} iv={iv_name} data=file[0:]"
                check_decrypted(dec, label)
            except:
                pass

# ═══════════════════════════════════════════════════════════════
# Test 5: Decrypt with IV prepended to ciphertext in file
# ═══════════════════════════════════════════════════════════════

print("\n" + "=" * 80)
print("TEST 5: IV prepended to ciphertext at various offsets")
print("=" * 80)

for data_start in [0, 4, 8, 16, 24, header_offset, payload_start]:
    iv_test = full_data[data_start:data_start+16]
    ct_test = full_data[data_start+16:data_start+16+4096]
    for key_name, key in key_candidates.items():
        for mbl in [16, 1]:
            try:
                dec = bcrypt_aes_cfb_decrypt(ct_test, key, iv_test, mbl)
                label = f"CFB{'128' if mbl==16 else '8'} key={key_name} IV=file[{data_start}:{data_start+16}] CT=file[{data_start+16}:]"
                check_decrypted(dec, label)
            except:
                pass

print("\n" + "=" * 80)
print("DONE")
print("=" * 80)