"""Verify BCrypt CNG setup - test raw key + different CFB segment sizes."""
import ctypes
from ctypes import c_void_p, c_ulong, byref
from pathlib import Path
import struct

KEY = b'kj)TGtrK>f]b[Piow.gU+nC@s""""""4'
IV  = b"Copyright @ OneO"

bcrypt = ctypes.windll.bcrypt

# Known plaintext (DX header from hook dump)
dx_plain = bytes.fromhex("44580000000000005c58000000000000")
# Known ciphertext (from file at offset 24, first 16 bytes)
file_ct = bytes.fromhex("2e0c10c7c967f66b6d03821271115ad6")

# Full file data
file_data = Path("ocr_data/oneocr.onemodel").read_bytes()
hook_dx = Path("frida_dump/decrypt_1_in22624_out22624.bin").read_bytes()

print("=" * 70)
print("BCrypt CNG CFB Segment Size Test")
print("=" * 70)
print(f"KEY:         {KEY}")
print(f"IV:          {IV}")
print(f"Expected PT: {dx_plain.hex()}")
print(f"Expected CT: {file_ct.hex()}")
print()


def test_cfb(msg_block_length, use_blob=False):
    """Test BCrypt AES-CFB with given MessageBlockLength."""
    tag = "MBL={}".format("default" if msg_block_length is None else msg_block_length)
    if use_blob:
        tag += "+blob"

    hAlg = c_void_p()
    status = bcrypt.BCryptOpenAlgorithmProvider(
        byref(hAlg), "AES\0".encode("utf-16-le"), None, 0
    )
    if status != 0:
        print("  [{}] OpenAlgorithm failed: {:#010x}".format(tag, status))
        return None

    mode = "ChainingModeCFB\0".encode("utf-16-le")
    status = bcrypt.BCryptSetProperty(
        hAlg, "ChainingMode\0".encode("utf-16-le"), mode, len(mode), 0
    )
    if status != 0:
        print("  [{}] SetProperty ChainingMode failed: {:#010x}".format(tag, status))
        bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
        return None

    if msg_block_length is not None:
        mbl = c_ulong(msg_block_length)
        status = bcrypt.BCryptSetProperty(
            hAlg, "MessageBlockLength\0".encode("utf-16-le"),
            byref(mbl), 4, 0
        )
        if status != 0:
            print("  [{}] SetProperty MBL={} failed: {:#010x}".format(tag, msg_block_length, status))
            bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
            return None

    hKey = c_void_p()
    if use_blob:
        blob = struct.pack('<III', 0x4d42444b, 1, len(KEY)) + KEY
        status = bcrypt.BCryptGenerateSymmetricKey(
            hAlg, byref(hKey), None, 0, blob, len(blob), 0
        )
    else:
        status = bcrypt.BCryptGenerateSymmetricKey(
            hAlg, byref(hKey), None, 0, KEY, len(KEY), 0
        )

    if status != 0:
        print("  [{}] GenerateSymmetricKey failed: {:#010x}".format(tag, status))
        bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
        return None

    # Encrypt test
    iv_enc = bytearray(IV)
    result_size = c_ulong(0)
    bcrypt.BCryptEncrypt(hKey, dx_plain, len(dx_plain), None, None, 0,
                         None, 0, byref(result_size), 0)
    output = (ctypes.c_ubyte * result_size.value)()
    actual = c_ulong(0)
    iv_buf = (ctypes.c_ubyte * 16)(*iv_enc)
    bcrypt.BCryptEncrypt(hKey, dx_plain, len(dx_plain), None,
                         iv_buf, 16, output, result_size.value, byref(actual), 0)
    our_ct = bytes(output[:actual.value])
    ct_match = our_ct[:16] == file_ct

    # Decrypt test (fresh key)
    hKey2 = c_void_p()
    if use_blob:
        blob = struct.pack('<III', 0x4d42444b, 1, len(KEY)) + KEY
        bcrypt.BCryptGenerateSymmetricKey(hAlg, byref(hKey2), None, 0, blob, len(blob), 0)
    else:
        bcrypt.BCryptGenerateSymmetricKey(hAlg, byref(hKey2), None, 0, KEY, len(KEY), 0)

    iv_dec = bytearray(IV)
    encrypted_chunk = file_data[24:24 + 32]
    result_size = c_ulong(0)
    bcrypt.BCryptDecrypt(hKey2, encrypted_chunk, len(encrypted_chunk), None, None, 0,
                         None, 0, byref(result_size), 0)
    output2 = (ctypes.c_ubyte * result_size.value)()
    iv_buf2 = (ctypes.c_ubyte * 16)(*iv_dec)
    status = bcrypt.BCryptDecrypt(hKey2, encrypted_chunk, len(encrypted_chunk), None,
                                  iv_buf2, 16, output2, result_size.value, byref(actual), 0)
    our_pt = bytes(output2[:actual.value])
    pt_match = our_pt[:2] == b"DX"

    mark = "*** MATCH! ***" if ct_match else ""
    print("  [{}] Enc->CT: {}  {} {}".format(tag, our_ct[:16].hex(), "OK" if ct_match else "FAIL", mark))
    print("  [{}] Dec->PT: {}  {}".format(tag, our_pt[:16].hex(), "OK DX" if pt_match else "FAIL"))

    bcrypt.BCryptDestroyKey(hKey)
    bcrypt.BCryptDestroyKey(hKey2)
    bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
    return ct_match


print("--- Raw key (correct for BCryptGenerateSymmetricKey) ---")
test_cfb(None)
test_cfb(1)
test_cfb(16)
print()

print("--- Blob key (has 12-byte header prepended - wrong) ---")
test_cfb(None, use_blob=True)
test_cfb(1, use_blob=True)
test_cfb(16, use_blob=True)
print()

print("--- BCryptImportKey with BCRYPT_KEY_DATA_BLOB ---")
for mbl in [None, 1, 16]:
    tag = "Import+MBL={}".format("default" if mbl is None else mbl)
    hAlg = c_void_p()
    bcrypt.BCryptOpenAlgorithmProvider(byref(hAlg), "AES\0".encode("utf-16-le"), None, 0)
    mode = "ChainingModeCFB\0".encode("utf-16-le")
    bcrypt.BCryptSetProperty(hAlg, "ChainingMode\0".encode("utf-16-le"), mode, len(mode), 0)

    if mbl is not None:
        mbl_val = c_ulong(mbl)
        bcrypt.BCryptSetProperty(hAlg, "MessageBlockLength\0".encode("utf-16-le"),
                                 byref(mbl_val), 4, 0)

    blob = struct.pack('<III', 0x4d42444b, 1, len(KEY)) + KEY
    hKey = c_void_p()
    status = bcrypt.BCryptImportKey(
        hAlg, None, "KeyDataBlob\0".encode("utf-16-le"),
        byref(hKey), None, 0, blob, len(blob), 0
    )
    if status != 0:
        print("  [{}] ImportKey failed: {:#010x}".format(tag, status))
        bcrypt.BCryptCloseAlgorithmProvider(hAlg, 0)
        continue

    iv_dec = bytearray(IV)
    encrypted_chunk = file_data[24:24 + 32]
    result_size = c_ulong(0)
    bcrypt.BCryptDecrypt(hKey, encrypted_chunk, 32, None, None, 0,
                         None, 0, byref(result_size), 0)
    output = (ctypes.c_ubyte * result_size.value)()
    actual = c_ulong(0)
    iv_buf = (ctypes.c_ubyte * 16)(*iv_dec)
    status = bcrypt.BCryptDecrypt(hKey, encrypted_chunk, 32, None,
                                  iv_buf, 16, output, result_size.value, byref(actual), 0)
    dec = bytes(output[:actual.value])
    match = dec[:2] == b"DX"
    mark = "*** MATCH! ***" if match else ""
    print("  [{}] Decrypt: {}  {} {}".format(tag, dec[:16].hex(), "OK DX" if match else "FAIL", mark))

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

print()
print("=" * 70)
print("If no method matched, need to hook BCryptSetProperty in the DLL")
print("to discover ALL properties set before BCryptDecrypt is called.")