neuron-tai/tests/test_wallet_binding_proof.py

"""C6: wallet binding requires proof of ownership, and rebinding is safe.

Any Bearer key could previously bind any wallet string with no proof of
ownership, and gossip `bind` events overwrote an existing binding directly.
These tests cover: signature-gated binding, rebind protection (with an
admin-override escape hatch), and gossip-safe conflict handling.
"""

import json
import urllib.error
import urllib.request

import pytest
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from meshnet_node.wallet import _b58encode
from meshnet_tracker.accounts import AccountStore
from meshnet_tracker.billing import BillingLedger
from meshnet_tracker.server import TrackerServer
from meshnet_tracker.wallet_proof import binding_message, verify_wallet_signature

HIVE_SECRET = "test-hive-secret"


def _keypair():
    priv = Ed25519PrivateKey.generate()
    return priv, _b58encode(priv.public_key().public_bytes_raw())


def _sign(priv, api_key: str, wallet: str) -> str:
    return priv.sign(binding_message(api_key, wallet)).hex()


# ---------------------------------------------------------------- unit tests


def test_verify_wallet_signature_accepts_valid_and_rejects_forged():
    priv, wallet = _keypair()
    other_priv, _ = _keypair()
    message = binding_message("client-key-1", wallet)

    assert verify_wallet_signature(wallet, message, priv.sign(message).hex())
    # signed by a different key entirely
    assert not verify_wallet_signature(wallet, message, other_priv.sign(message).hex())
    # signature for a different api_key can't be replayed
    other_message = binding_message("client-key-2", wallet)
    assert not verify_wallet_signature(wallet, other_message, priv.sign(message).hex())
    assert not verify_wallet_signature(wallet, message, "not-hex")


def test_bind_wallet_rejects_conflicting_rebind_without_admin_override():
    ledger = BillingLedger(starting_credit=0.0)
    ledger.bind_wallet("key-1", "WalletA")
    assert ledger.api_key_for_wallet("WalletA") == "key-1"

    event = ledger.bind_wallet("key-2", "WalletA")
    assert event.get("rejected") is True
    assert ledger.api_key_for_wallet("WalletA") == "key-1"  # unchanged

    override = ledger.bind_wallet("key-2", "WalletA", admin_override=True)
    assert not override.get("rejected")
    assert ledger.api_key_for_wallet("WalletA") == "key-2"


def test_gossip_bind_event_cannot_overwrite_existing_binding():
    """A conflicting `bind` event applied via gossip must not clobber."""
    leader = BillingLedger(starting_credit=0.0)
    leader.bind_wallet("key-1", "WalletA")

    follower = BillingLedger(starting_credit=0.0)
    follower.bind_wallet("key-9", "WalletA")  # follower already has its own binding

    events, _ = leader.events_since(0)
    applied = follower.apply_events(events)
    assert applied == 1  # event was processed (and rejected), not dropped/ignored
    # the follower's pre-existing binding is untouched by the conflicting gossip event
    assert follower.api_key_for_wallet("WalletA") == "key-9"


# ---------------------------------------------------------- HTTP integration


def _post_json(url: str, payload: dict, headers: dict | None = None) -> dict:
    req = urllib.request.Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json", **(headers or {})},
        method="POST",
    )
    with urllib.request.urlopen(req) as r:
        return json.loads(r.read())


@pytest.fixture
def tracker():
    ledger = BillingLedger(starting_credit=0.0)
    server = TrackerServer(billing=ledger, accounts=AccountStore(), hive_secret=HIVE_SECRET)
    port = server.start()
    yield f"http://127.0.0.1:{port}", ledger
    server.stop()


def test_bind_without_signature_is_rejected(tracker):
    url, ledger = tracker
    _, wallet = _keypair()
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _post_json(
            f"{url}/v1/wallet/register",
            {"wallet": wallet},
            headers={"Authorization": "Bearer client-key-1"},
        )
    assert exc_info.value.code == 400
    assert ledger.api_key_for_wallet(wallet) is None


def test_bind_with_forged_signature_is_rejected(tracker):
    url, ledger = tracker
    priv, wallet = _keypair()
    forger_priv, _ = _keypair()
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _post_json(
            f"{url}/v1/wallet/register",
            {"wallet": wallet, "signature": _sign(forger_priv, "client-key-1", wallet)},
            headers={"Authorization": "Bearer client-key-1"},
        )
    assert exc_info.value.code == 401
    assert ledger.api_key_for_wallet(wallet) is None


def test_valid_signature_binds_wallet(tracker):
    url, ledger = tracker
    priv, wallet = _keypair()
    reply = _post_json(
        f"{url}/v1/wallet/register",
        {"wallet": wallet, "signature": _sign(priv, "client-key-1", wallet)},
        headers={"Authorization": "Bearer client-key-1"},
    )
    assert reply == {"wallet": wallet, "bound": True}
    assert ledger.api_key_for_wallet(wallet) == "client-key-1"


def test_second_key_cannot_steal_bound_wallet(tracker):
    url, ledger = tracker
    priv, wallet = _keypair()
    _post_json(
        f"{url}/v1/wallet/register",
        {"wallet": wallet, "signature": _sign(priv, "client-key-1", wallet)},
        headers={"Authorization": "Bearer client-key-1"},
    )

    # key-2 gets the wallet owner's genuine signature over *their* api_key,
    # proving they control the private key — still must not steal the binding.
    with pytest.raises(urllib.error.HTTPError) as exc_info:
        _post_json(
            f"{url}/v1/wallet/register",
            {"wallet": wallet, "signature": _sign(priv, "client-key-2", wallet)},
            headers={"Authorization": "Bearer client-key-2"},
        )
    assert exc_info.value.code == 409
    assert ledger.api_key_for_wallet(wallet) == "client-key-1"


def test_admin_session_can_force_rebind(tracker):
    url, ledger = tracker
    priv, wallet = _keypair()
    _post_json(
        f"{url}/v1/wallet/register",
        {"wallet": wallet, "signature": _sign(priv, "client-key-1", wallet)},
        headers={"Authorization": "Bearer client-key-1"},
    )

    admin = _post_json(
        f"{url}/v1/auth/register",
        {"email": "admin@example.com", "password": "secret-123"},
    )
    reply = _post_json(
        f"{url}/v1/wallet/register",
        {"wallet": wallet, "api_key": "client-key-2"},
        headers={"Authorization": f"Bearer {admin['session_token']}"},
    )
    assert reply == {"wallet": wallet, "bound": True}
    assert ledger.api_key_for_wallet(wallet) == "client-key-2"