krux PR #866 — does pre-PR taproot BIP-137 signing round-trip through a lenient verifier? (pytest)

test_bip137_p2tr.py

"""test_bip137_p2tr.py — does the pre-PR krux taproot message-signing
path round-trip through a "lenient" BIP-137 verifier?

Context (selfcustody/krux PR #866):

  On `develop`, single-sig P2TR message signing flows through
  `_sign_at_address` -> `Key.sign_at`. The signature is over the
  standard Bitcoin-signed-message commitment, produced as an
  ECDSA-recoverable sig over the BIP-86 INTERNAL key (the raw
  m/86'/.../ private scalar). The header byte stays in 31-34
  (P2PKH-compressed range), because BIP-137 has no taproot header
  range.

  PR #866 routes signing through `src/krux/bip137.py:sign`. Its
  `build_header` raises `ValueError("p2tr legacy sign not supported ")`
  for "p2tr", which the menu surfaces as an error screen -> P2TR
  message signing breaks.

  odudex's review (the blocker on PR #866):
    > Lenient verifiers (Sparrow / Electrum-style) recover the pubkey
    > and reconstruct the address of the claimed type from it - for
    > taproot they apply the taproot tweak (same as script.p2tr) and
    > the recovered+tweaked key matches the bc1p... address.

  This file tests that round-trip mechanically. It mocks `Key.sign_at`
  with the develop-branch implementation (verbatim from
  src/krux/key.py on `origin/develop`) so the test does not depend on
  which branch is checked out, and exercises the recover-pubkey +
  apply-taproot-tweak + compare-address path that a lenient verifier
  would run.

Drop into `tests/` and run:

    poetry run pytest tests/test_bip137_p2tr.py -v -s
"""
import pytest


@pytest.fixture
def tdata():
    """Test mnemonic + derivation. Well-known BIP-39 vector — never
    use with real funds."""
    from collections import namedtuple

    return namedtuple("TData", ["MNEMONIC", "MESSAGE", "DERIVATION_TNET"])(
        MNEMONIC=(
            "abandon abandon abandon abandon abandon abandon "
            "abandon abandon abandon abandon abandon about"
        ),
        MESSAGE=b"a test message.",
        DERIVATION_TNET="m/86h/1h/0h/0/3",
    )


def _mock_embit(mocker):
    """Standard krux test-suite shim (mirrors tests/test_key.py)."""
    from embit import bip32, bip39

    mocker.patch("krux.key.bip32", new=mocker.MagicMock(wraps=bip32))
    mocker.patch("krux.key.bip39", new=mocker.MagicMock(wraps=bip39))


def _develop_sign_at(root, derivation, message_hash):
    """Pinned copy of `Key.sign_at` from `origin/develop` (src/krux/key.py).

    Used as a substitute for the actual method so this test is
    branch-independent: even if a future PR rewrites `Key.sign_at`,
    these tests still exercise the develop-era behavior that odudex's
    claim is about.
    """
    from embit import ec
    from embit.util import secp256k1

    prv = root.derive(derivation).key
    sig = secp256k1.ecdsa_sign_recoverable(message_hash, prv._secret)
    flag = bytes([27 + sig[64] + 4])
    ec_signature = ec.Signature(sig[:64])
    return flag + secp256k1.ecdsa_signature_serialize_compact(ec_signature._sig)


def _bip137_commitment(message):
    """BIP-137 double-SHA256 over the magic-prefixed payload."""
    from embit import compact
    import hashlib

    return hashlib.sha256(
        hashlib.sha256(
            b"\x18Bitcoin Signed Message:\n"
            + compact.to_bytes(len(message))
            + message
        ).digest()
    ).digest()


def test_develop_p2tr_recovered_pubkey_tweaked_matches_address(
    mocker, m5stickv, tdata
):
    """Reproduce the develop-branch P2TR signing path and verify the
    lenient-verifier round-trip.

    Steps:
      1. Build the BIP-137 commitment.
      2. Sign it via the pinned develop-era `sign_at` (mocks any
         post-PR wrapping that might intercept p2tr).
      3. Parse the 65-byte sig: header || compact ECDSA.
      4. ECDSA-recover the pubkey (this is the m/86' internal key).
      5. Apply the BIP-341 tweak via `script.p2tr(pub).address(...)`.
      6. Assert the recovered+tweaked address equals the address the
         wallet derives from the same xpub.

    Pass = odudex's claim is mechanically correct.
    Fail = claim refuted.
    """
    _mock_embit(mocker)

    from embit import bip32, ec, script
    from embit.networks import NETWORKS
    from embit.util import secp256k1
    from krux.key import Key, TYPE_SINGLESIG

    # Replace Key.sign_at with the develop-era implementation so this
    # test does not depend on whether the current branch wraps signing
    # through krux.bip137 or any other future module.
    mocker.patch.object(
        Key,
        "sign_at",
        autospec=True,
        side_effect=lambda self, deriv, h: _develop_sign_at(self.root, deriv, h),
    )

    key = Key(tdata.MNEMONIC, TYPE_SINGLESIG)
    derivation = bip32.parse_path(tdata.DERIVATION_TNET)
    commitment = _bip137_commitment(tdata.MESSAGE)

    sig = key.sign_at(derivation, commitment)

    assert len(sig) == 65, "BIP-137 sig is 1-byte header + 64-byte compact ECDSA"
    header = sig[0]
    assert 31 <= header <= 34, (
        "develop's Key.sign_at always emits 31-34 (P2PKH-compressed range), "
        "regardless of script type — this is what PR #866 noticed for segwit, "
        "but for taproot there is no defined BIP-137 header range anyway."
    )

    recid = header - 31
    parsed_recoverable = secp256k1.ecdsa_recoverable_signature_parse_compact(
        sig[1:], recid
    )
    raw_recovered = secp256k1.ecdsa_recover(parsed_recoverable, commitment)
    recovered_sec = secp256k1.ec_pubkey_serialize(raw_recovered)
    recovered_pub = ec.PublicKey.parse(recovered_sec)

    expected_internal_pub = key.root.derive(derivation).to_public().key
    assert recovered_pub.sec() == expected_internal_pub.sec(), (
        "Recovered pubkey must equal the m/86' INTERNAL key — develop signs "
        "the internal scalar directly, no taproot tweak at signing time."
    )

    derived_addr = script.p2tr(expected_internal_pub).address(
        network=NETWORKS["test"]
    )
    recovered_addr = script.p2tr(recovered_pub).address(network=NETWORKS["test"])

    import base64

    print()
    print("  derivation:    %s" % tdata.DERIVATION_TNET)
    print("  bc1p address:  %s" % derived_addr)
    print("  header byte:   0x%02x (recid=%d)" % (header, recid))
    print("  base64 sig:    %s" % base64.b64encode(sig).decode())
    print("  recover+tweak: %s" % recovered_addr)
    print("  matches:       %s" % (recovered_addr == derived_addr))

    assert recovered_addr == derived_addr, (
        "ODUDEX CLAIM REFUTED: the recover+tweak round-trip does not yield "
        "the wallet's bc1p... address."
    )


def test_fix_bip137_module_rejects_p2tr_currently(mocker, m5stickv, tdata):
    """Documents PR #866's introduced regression: `bip137.sign` raises
    `ValueError` for P2TR.

    Lazy import so the file is also runnable on `develop`, where
    `krux.bip137` does not exist (test will be skipped there).
    """
    _mock_embit(mocker)
    pytest.importorskip("krux.bip137")

    from krux.bip137 import sign
    from krux.key import Key, TYPE_SINGLESIG

    key = Key(tdata.MNEMONIC, TYPE_SINGLESIG)

    with pytest.raises(ValueError, match="legacy sign not supported"):
        sign(tdata.MESSAGE, key, tdata.DERIVATION_TNET, script_type="p2tr")