#!/usr/bin/env python3
# Author:  cirakg & claude ai
# Summary: Compute the UI-test fixtures' PIN wraps (fake PINs 1234 and legacy 123456): no PBKDF2 on the watch at start.
"""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.
The extra UI-test scenario stores LEGACY_WRAP instead, the same master key
under the 6-digit PIN 123456 and without the marker: a PIN from before the
4-digit rule, which unlocks with OK and is then replaced.

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

from __future__ import annotations

import hashlib
import hmac
import struct
import subprocess

PIN = b"1234"
LEGACY_PIN = b"123456"  # the extra scenario's PIN from before the 4-digit rule
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(pin: bytes = PIN) -> 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("FIXTURE_LEGACY_PIN =", LEGACY_PIN.decode())
    print("FIXTURE_LEGACY_WRAP =", fixture_wrap(LEGACY_PIN).hex())
    print(f"marker: {PIN_LENGTH_KEY} = {PIN_DIGITS}")


if __name__ == "__main__":
    main()