#!/usr/bin/env python3
# Author:  cirakg & claude ai
# Summary: Compute the UI-test fixture's PIN wrap (fake PIN 1234) so the watch needn't run PBKDF2 at start-up.
"""Print the fake PIN wrap that source/app/Fixtures.mc stores in the UI-test build.

The wrap follows docs/SECURITY.md ("PIN wrap, version 1") with fixed, fake
inputs: PIN 1234, salt 10..1f, IV 20..2f, master key 40..5f, 2000
PBKDF2 iterations. Deriving the PIN key on the watch in one event trips the
watchdog, and the UI tests need the store ready before the first screen, so
the fixture stores this precomputed value instead. The app's own unlock code
then checks it like any other wrap, which also cross-checks the format.
The fixture also stores the 4-digit PIN marker (KeyManager.PIN_LENGTH_KEY =
PIN_DIGITS), so the PIN isn't treated as a legacy 4-to-10-digit one.

Usage: gen_fixture_wrap.py   (prints FIXTURE_PIN, FIXTURE_WRAP and FIXTURE_MASTER_KEY in hex,
                              and the marker the fixture stores)
"""

from __future__ import annotations

import hashlib
import hmac
import struct
import subprocess

PIN = b"1234"
PIN_DIGITS = 4  # source/storage/KeyManager.mc: every new PIN, and the marker value
PIN_LENGTH_KEY = "wl"  # source/storage/KeyManager.mc
ITERATIONS = 2000
SALT = bytes(range(0x10, 0x20))
IV = bytes(range(0x20, 0x30))
MASTER_KEY = bytes(range(0x40, 0x60))


def hkdf(ikm: bytes, info: bytes, length: int = 32) -> bytes:
    """HKDF-SHA256 with an empty salt (RFC 5869), as in source/crypto/Kdf.mc."""
    prk = hmac.new(bytes(32), ikm, hashlib.sha256).digest()
    out, block, counter = b"", b"", 1
    while len(out) < length:
        block = hmac.new(prk, block + info + bytes([counter]), hashlib.sha256).digest()
        out += block
        counter += 1
    return out[:length]


def aes_cbc_encrypt(key: bytes, iv: bytes, plaintext: bytes) -> bytes:
    """AES-256-CBC with PKCS#7 padding, through openssl (not in the standard library)."""
    result = subprocess.run(
        ["openssl", "enc", "-aes-256-cbc", "-K", key.hex(), "-iv", iv.hex()],
        input=plaintext,
        capture_output=True,
        check=True,
    )
    return result.stdout


def sealed_box(enc_key: bytes, mac_key: bytes, aad: bytes, plaintext: bytes) -> bytes:
    """source/crypto/SealedBox.mc: 0x01 | IV | ciphertext | HMAC tag."""
    body = b"\x01" + IV + aes_cbc_encrypt(enc_key, IV, plaintext)
    tag = hmac.new(mac_key, b"OTPW-REC-1" + struct.pack(">I", len(aad)) + aad + body, hashlib.sha256).digest()
    return body + tag


def fixture_wrap() -> bytes:
    pin_key = hashlib.pbkdf2_hmac("sha256", PIN, SALT, ITERATIONS, 32)
    enc_key, mac_key = hkdf(pin_key, b"OTPW enc v1"), hkdf(pin_key, b"OTPW mac v1")
    count = struct.pack(">I", ITERATIONS)
    aad = b"OTPW-WRAP-1" + SALT + count
    return SALT + count + sealed_box(enc_key, mac_key, aad, MASTER_KEY)


def main() -> None:
    assert len(PIN) == PIN_DIGITS
    print("FIXTURE_PIN =", PIN.decode())
    print("FIXTURE_WRAP =", fixture_wrap().hex())
    print("FIXTURE_MASTER_KEY =", MASTER_KEY.hex())
    print(f"marker: {PIN_LENGTH_KEY} = {PIN_DIGITS}")


if __name__ == "__main__":
    main()