324 lines
11 KiB
Python
324 lines
11 KiB
Python
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import time
|
|
from typing import cast
|
|
|
|
import pytest
|
|
from coincurve import PrivateKey, PublicKey
|
|
from Cryptodome import Random
|
|
from Cryptodome.Cipher import AES
|
|
from Cryptodome.Util.Padding import pad, unpad
|
|
from websockets import ServerConnection
|
|
from websockets import serve as ws_serve
|
|
|
|
from lnbits.wallets.nwc import NWCConnection, NWCWallet
|
|
from tests.wallets.helpers import (
|
|
WalletTest,
|
|
build_test_id,
|
|
check_assertions,
|
|
load_funding_source,
|
|
wallet_fixtures_from_json,
|
|
)
|
|
|
|
|
|
def encrypt_content(priv_key, dest_pub_key, content):
|
|
p = PublicKey(bytes.fromhex("02" + dest_pub_key))
|
|
shared = p.multiply(bytes.fromhex(priv_key)).format()[1:]
|
|
iv = Random.new().read(AES.block_size)
|
|
aes = AES.new(shared, AES.MODE_CBC, iv)
|
|
|
|
content_bytes = content.encode("utf-8")
|
|
content_bytes = pad(content_bytes, AES.block_size)
|
|
|
|
encrypted_b64 = base64.b64encode(aes.encrypt(content_bytes)).decode("ascii")
|
|
iv_b64 = base64.b64encode(iv).decode("ascii")
|
|
encrypted_content = encrypted_b64 + "?iv=" + iv_b64
|
|
return encrypted_content
|
|
|
|
|
|
def decrypt_content(priv_key, source_pub_key, content):
|
|
p = PublicKey(bytes.fromhex("02" + source_pub_key))
|
|
shared = p.multiply(bytes.fromhex(priv_key)).format()[1:]
|
|
encrypted_content_b64, iv_b64 = content.split("?iv=")
|
|
encrypted_content = base64.b64decode(encrypted_content_b64.encode("ascii"))
|
|
iv = base64.b64decode(iv_b64.encode("ascii"))
|
|
aes = AES.new(shared, AES.MODE_CBC, iv)
|
|
decrypted_bytes = aes.decrypt(encrypted_content)
|
|
decrypted_bytes = unpad(decrypted_bytes, AES.block_size)
|
|
return decrypted_bytes.decode("utf-8")
|
|
|
|
|
|
def json_dumps(data):
|
|
if isinstance(data, dict):
|
|
data = {k: v for k, v in data.items() if v is not None}
|
|
return json.dumps(data, separators=(",", ":"), ensure_ascii=False)
|
|
|
|
|
|
def sign_event(pub_key, priv_key, event):
|
|
signature_data = json_dumps(
|
|
[
|
|
0,
|
|
pub_key,
|
|
event["created_at"],
|
|
event["kind"],
|
|
event["tags"],
|
|
event["content"],
|
|
]
|
|
)
|
|
event_id = hashlib.sha256(signature_data.encode()).hexdigest()
|
|
event["id"] = event_id
|
|
event["pubkey"] = pub_key
|
|
s = PrivateKey(bytes.fromhex(priv_key))
|
|
signature = s.sign_schnorr(bytes.fromhex(event_id)).hex()
|
|
event["sig"] = signature
|
|
return event
|
|
|
|
|
|
async def handle( # noqa: C901
|
|
wallet, mock_settings, data, websocket: ServerConnection
|
|
):
|
|
async for message in websocket:
|
|
if not wallet:
|
|
continue
|
|
msg = json.loads(message)
|
|
if msg[0] == "REQ":
|
|
sub_id = msg[1]
|
|
sub_filter = msg[2]
|
|
kinds = sub_filter["kinds"]
|
|
if 13194 in kinds: # Send info event
|
|
event = {
|
|
"kind": 13194,
|
|
"content": " ".join(mock_settings["supported_methods"]),
|
|
"created_at": int(time.time()),
|
|
"tags": [],
|
|
}
|
|
sign_event(
|
|
mock_settings["service_public_key"],
|
|
mock_settings["service_private_key"],
|
|
event,
|
|
)
|
|
await websocket.send(json.dumps(["EVENT", sub_id, event]))
|
|
elif 23195 in kinds:
|
|
assert sub_filter["authors"] == [mock_settings["service_public_key"]]
|
|
elif msg[0] == "EVENT":
|
|
event = msg[1]
|
|
decrypted_content = decrypt_content(
|
|
mock_settings["service_private_key"],
|
|
mock_settings["user_public_key"],
|
|
event["content"],
|
|
)
|
|
content = json.loads(decrypted_content)
|
|
mock = None
|
|
for m in data.mocks:
|
|
rb = m.request_body
|
|
if rb and rb["method"] == content["method"]:
|
|
p1 = rb["params"]
|
|
p2 = content["params"]
|
|
p1 = json_dumps({k: v for k, v in p1.items() if v is not None})
|
|
p2 = json_dumps({k: v for k, v in p2.items() if v is not None})
|
|
if p1 == p2:
|
|
mock = m
|
|
break
|
|
if mock:
|
|
sub_id = None
|
|
nwcwallet = cast(NWCWallet, wallet)
|
|
for subscription in nwcwallet.conn.subscriptions.values():
|
|
if subscription["event_id"] == event["id"]:
|
|
sub_id = subscription["sub_id"]
|
|
break
|
|
if sub_id:
|
|
response = mock.response
|
|
encrypted_content = encrypt_content(
|
|
mock_settings["service_private_key"],
|
|
mock_settings["user_public_key"],
|
|
json_dumps(response),
|
|
)
|
|
response_event = {
|
|
"kind": 23195,
|
|
"content": encrypted_content,
|
|
"created_at": int(time.time()),
|
|
"tags": [
|
|
["e", event["id"]],
|
|
["p", mock_settings["user_public_key"]],
|
|
],
|
|
}
|
|
sign_event(
|
|
mock_settings["service_public_key"],
|
|
mock_settings["service_private_key"],
|
|
response_event,
|
|
)
|
|
await websocket.send(json.dumps(["EVENT", sub_id, response_event]))
|
|
else:
|
|
raise Exception(
|
|
"No mock found for "
|
|
+ content["method"]
|
|
+ " "
|
|
+ json_dumps(content["params"])
|
|
)
|
|
|
|
|
|
async def run(data: WalletTest):
|
|
if data.skip:
|
|
pytest.skip()
|
|
|
|
wallet = None
|
|
mock_settings = data.funding_source.mock_settings
|
|
if mock_settings is None:
|
|
return
|
|
|
|
def handler(websocket):
|
|
return handle(wallet, mock_settings, data, websocket)
|
|
|
|
if mock_settings is not None:
|
|
async with ws_serve(handler, "localhost", mock_settings["port"]) as server:
|
|
await server.start_serving()
|
|
wallet = load_funding_source(data.funding_source)
|
|
await check_assertions(wallet, data)
|
|
nwcwallet = cast(NWCWallet, wallet)
|
|
await nwcwallet.cleanup()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_nwc_rejects_event_from_unexpected_pubkey(mocker):
|
|
async def _noop(*args, **kwargs):
|
|
return None
|
|
|
|
mocker.patch("lnbits.wallets.nwc.NWCConnection._connect_to_relay", new=_noop)
|
|
mocker.patch("lnbits.wallets.nwc.NWCConnection._handle_timeouts", new=_noop)
|
|
|
|
service_private_key = PrivateKey()
|
|
service_public_key = service_private_key.public_key.format().hex()[2:]
|
|
attacker_private_key = PrivateKey()
|
|
attacker_public_key = attacker_private_key.public_key.format().hex()[2:]
|
|
account_private_key = PrivateKey()
|
|
|
|
conn = NWCConnection(
|
|
service_public_key,
|
|
account_private_key.secret.hex(),
|
|
"ws://127.0.0.1:8555",
|
|
)
|
|
try:
|
|
event = {
|
|
"kind": 23195,
|
|
"content": "{}",
|
|
"created_at": int(time.time()),
|
|
"tags": [["e", "request-event-id"]],
|
|
}
|
|
sign_event(attacker_public_key, attacker_private_key.secret.hex(), event)
|
|
|
|
with pytest.raises(Exception, match="Invalid event signature"):
|
|
await conn._on_event_message(["EVENT", "subid", event])
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_nwc_marks_pending_invoice_settled_only_once():
|
|
wallet = NWCWallet.__new__(NWCWallet)
|
|
wallet.pending_invoice_details = {"checking-id": {"checking_id": "checking-id"}}
|
|
wallet.pending_invoices = ["checking-id"]
|
|
wallet.paid_invoices_queue = asyncio.Queue(0)
|
|
|
|
wallet._mark_invoice_settled("checking-id", source="notification")
|
|
wallet._mark_invoice_settled("checking-id", source="notification")
|
|
|
|
assert wallet.paid_invoices_queue.qsize() == 1
|
|
assert await wallet.paid_invoices_queue.get() == "checking-id"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_nwc_registers_notification_subscriptions(mocker):
|
|
async def _noop(*args, **kwargs):
|
|
return None
|
|
|
|
mocker.patch("lnbits.wallets.nwc.NWCConnection._connect_to_relay", new=_noop)
|
|
mocker.patch("lnbits.wallets.nwc.NWCConnection._handle_timeouts", new=_noop)
|
|
|
|
service_private_key = PrivateKey()
|
|
service_public_key = service_private_key.public_key.format().hex()[2:]
|
|
account_private_key = PrivateKey()
|
|
|
|
conn = NWCConnection(
|
|
service_public_key,
|
|
account_private_key.secret.hex(),
|
|
"ws://127.0.0.1:8555",
|
|
)
|
|
send_mock = mocker.patch.object(conn, "_send", mocker.AsyncMock())
|
|
|
|
try:
|
|
await conn._subscribe_to_notifications()
|
|
|
|
assert len(conn.notification_subscription_ids) == 2
|
|
assert len(conn.subscriptions) == 2
|
|
assert set(conn.subscriptions.keys()) == conn.notification_subscription_ids
|
|
assert all(
|
|
subscription["method"] == "notification_sub"
|
|
and subscription["event_id"] == subscription["sub_id"]
|
|
for subscription in conn.subscriptions.values()
|
|
)
|
|
assert send_mock.await_count == 2
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_nwc_spreads_fallback_lookups_with_cooldown(mocker):
|
|
def _schedule_next_lookup(
|
|
invoice: dict[str, object], now: float | None = None
|
|
) -> None:
|
|
assert now is not None
|
|
invoice["next_lookup_at"] = now + 1
|
|
|
|
wallet = NWCWallet.__new__(NWCWallet)
|
|
wallet.shutdown = False
|
|
wallet.pending_invoices = ["checking-1", "checking-2"]
|
|
wallet.pending_invoice_details = {
|
|
"checking-1": {
|
|
"checking_id": "checking-1",
|
|
"next_lookup_at": 0.0,
|
|
"lookup_attempts": 0,
|
|
},
|
|
"checking-2": {
|
|
"checking_id": "checking-2",
|
|
"next_lookup_at": 0.0,
|
|
"lookup_attempts": 0,
|
|
},
|
|
}
|
|
wallet.pending_invoices_lookup_cooldown = 1.0
|
|
wallet._is_shutting_down = lambda: False
|
|
wallet._payment_data_is_settled = lambda payment_data: False
|
|
wallet._cache_payment_data = lambda *args, **kwargs: None
|
|
wallet._schedule_next_lookup = _schedule_next_lookup
|
|
wallet.conn = mocker.Mock()
|
|
wallet.conn.get_info = mocker.AsyncMock()
|
|
wallet.conn.supports_method = mocker.Mock(return_value=True)
|
|
wallet.conn.call = mocker.AsyncMock(return_value={"settled_at": None})
|
|
sleep_mock = mocker.patch("lnbits.wallets.nwc.asyncio.sleep", mocker.AsyncMock())
|
|
|
|
await wallet._run_fallback_lookups(100.0)
|
|
|
|
assert wallet.conn.call.await_count == 2
|
|
sleep_mock.assert_awaited_once_with(1.0)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize(
|
|
"test_data",
|
|
wallet_fixtures_from_json("tests/wallets/fixtures/json/fixtures_nwc.json"),
|
|
ids=build_test_id,
|
|
)
|
|
async def test_nwc_wallet(test_data: WalletTest):
|
|
await run(test_data)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize(
|
|
"test_data",
|
|
wallet_fixtures_from_json("tests/wallets/fixtures/json/fixtures_nwc_bad.json"),
|
|
ids=build_test_id,
|
|
)
|
|
async def test_nwc_wallet_bad(test_data: WalletTest):
|
|
await run(test_data)
|