"""Patch Rust HKDF key derivation to output 16 bytes (AES-128).

Element Call JS SDK derives 128-bit AES-GCM keys via:
  deriveKey({name:"AES-GCM", length:128}, ...)

The C++ FrameCryptor may allocate a larger derived_key buffer (32+ bytes).
This patch ensures only 16 bytes are filled, matching the JS output.
"""
import sys

path = sys.argv[1]
with open(path) as f:
    content = f.read()

old = 'hkdf.expand(&[0u8; 128], derived_key).is_ok()'
new = """{
                // MAT-258: Derive 16 bytes (AES-128) matching Element Call JS SDK
                let mut buf = [0u8; 16];
                let ok = hkdf.expand(&[0u8; 128], &mut buf).is_ok();
                if ok {
                    // Fill first 16 bytes of derived_key, zero-pad rest
                    let len = derived_key.len().min(16);
                    derived_key[..len].copy_from_slice(&buf[..len]);
                    for b in derived_key[len..].iter_mut() { *b = 0; }
                }
                ok
            }"""

if old not in content:
    print(f"WARNING: Could not find HKDF expand line in {path}")
    print("File may already be patched or format changed")
    sys.exit(0)

content = content.replace(old, new)
with open(path, 'w') as f:
    f.write(content)
print(f"Patched HKDF output to 16 bytes in {path}")