diff --git a/itests/tests/conftest.py b/itests/tests/conftest.py index 82d68fff..e791ad29 100644 --- a/itests/tests/conftest.py +++ b/itests/tests/conftest.py @@ -13,6 +13,8 @@ from proto import ( squeak_server_pb2_grpc, ) +from tests.util import generate_signing_key +from tests.util import get_address from tests.util import load_lightning_client @@ -38,3 +40,38 @@ def admin_stub(): @pytest.fixture def lightning_client(): return load_lightning_client() + +@pytest.fixture +def whitelisted_signing_key(server_stub, admin_stub): + # Create a signing key + signing_key = generate_signing_key() + + # Create a new contact profile + profile_name = "whitelisted_contact" + profile_address = get_address(signing_key) + create_contact_profile_response = admin_stub.CreateContactProfile( + squeak_admin_pb2.CreateContactProfileRequest( + profile_name=profile_name, + address=profile_address, + ) + ) + contact_profile_id = create_contact_profile_response.profile_id + + # Set the profile to be whitelisted + admin_stub.SetSqueakProfileWhitelisted( + squeak_admin_pb2.SetSqueakProfileWhitelistedRequest( + profile_id=contact_profile_id, + whitelisted=True, + ) + ) + + # Yield the signing key + yield signing_key + +@pytest.fixture +def nonwhitelisted_signing_key(server_stub, admin_stub): + # Create a signing key + signing_key = generate_signing_key() + + # Yield the signing key + yield signing_key diff --git a/itests/tests/test_squeak_node.py b/itests/tests/test_squeak_node.py index 925d0903..464be707 100644 --- a/itests/tests/test_squeak_node.py +++ b/itests/tests/test_squeak_node.py @@ -31,15 +31,36 @@ from tests.util import bxor from tests.util import string_to_hex -def test_buy_squeak(server_stub, admin_stub, lightning_client): +def test_buy_squeak(server_stub, admin_stub, lightning_client, whitelisted_signing_key): balance_from_client = lightning_client.get_wallet_balance() print("Balance from direct client: %s" % balance_from_client) assert balance_from_client.total_balance >= 1505000000000 + # # Create a signing key + # signing_key = generate_signing_key() + + # # Create a new contact profile + # profile_name = "whitelisted_contact" + # profile_address = get_address(signing_key) + # create_contact_profile_response = admin_stub.CreateContactProfile( + # squeak_admin_pb2.CreateContactProfileRequest( + # profile_name=profile_name, + # address=profile_address, + # ) + # ) + # contact_profile_id = create_contact_profile_response.profile_id + + # # Set the profile to be whitelisted + # admin_stub.SetSqueakProfileWhitelisted( + # squeak_admin_pb2.SetSqueakProfileWhitelistedRequest( + # profile_id=contact_profile_id, + # whitelisted=True, + # ) + # ) + # Post a squeak with a direct request to the server - signing_key = generate_signing_key() block_height, block_hash = get_latest_block_info(lightning_client) - squeak = make_squeak(signing_key, "hello from itest!", block_hash, block_height) + squeak = make_squeak(whitelisted_signing_key, "hello from itest!", block_hash, block_height) squeak_hash = get_hash(squeak) squeak_msg = build_squeak_msg(squeak) @@ -61,7 +82,7 @@ def test_buy_squeak(server_stub, admin_stub, lightning_client): assert get_hash(get_response_squeak) == get_hash(squeak) # Lookup squeaks based on address - signing_address = get_address(signing_key) + signing_address = get_address(whitelisted_signing_key) addresses = [ signing_address, get_address(generate_signing_key()), @@ -266,7 +287,7 @@ def test_make_squeak(server_stub, admin_stub): ) print("Get followed squeak displays response: " + str(get_followed_squeak_display_response)) assert ( - len(get_followed_squeak_display_response.squeak_display_entries) == 2 + len(get_followed_squeak_display_response.squeak_display_entries) >= 2 ) # Get all squeak displays for the known address diff --git a/squeakserver/node/squeak_node.py b/squeakserver/node/squeak_node.py index 2c8673da..a17af99b 100644 --- a/squeakserver/node/squeak_node.py +++ b/squeakserver/node/squeak_node.py @@ -53,6 +53,9 @@ class SqueakNode: self.squeak_block_queue_worker.start_running() def save_squeak(self, squeak): + if not self.squeak_whitelist.should_allow_squeak(squeak): + raise Exception("Squeak upload not allowed by whitelist.") + if not self.squeak_rate_limiter.should_rate_limit_allow(squeak): raise Exception("Excedeed allowed number of squeaks per block.") @@ -165,6 +168,7 @@ class SqueakNode: def set_squeak_profile_whitelisted(self, profile_id, whitelisted): self.postgres_db.set_profile_whitelisted(profile_id, whitelisted) + self.squeak_whitelist.refresh() def make_squeak(self, profile_id, content_str, replyto_hash): squeak_profile = self.postgres_db.get_profile(profile_id) diff --git a/squeakserver/node/squeak_whitelist.py b/squeakserver/node/squeak_whitelist.py index 9847e6ff..0fe0a458 100644 --- a/squeakserver/node/squeak_whitelist.py +++ b/squeakserver/node/squeak_whitelist.py @@ -1,6 +1,8 @@ import logging import queue +from squeak.core.signing import CSqueakAddress + from squeakserver.server.util import get_hash from squeakserver.node.block_info import BlockInfo @@ -14,10 +16,11 @@ class SqueakWhitelist: self.refresh() def should_allow_squeak(self, squeak): - squeak_hash = get_hash(squeak) - squeak_address = CSqueakAddress.from_verifying_key(verifying_key) - squeak_address_str = str(squeak_address_str) - logger.info("Checking whitelist for squeak hash: {}, squeak address".format(squeak_hash, squeak_address_str)) + squeak_hash = get_hash(squeak).hex() + squeak_address = squeak.GetAddress() + squeak_address_str = str(squeak_address) + logger.info("Checking whitelist for squeak hash: {}, squeak address: {}".format(squeak_hash, squeak_address_str)) + logger.info("Allowed addresses: {}".format(self.allowed_addresses)) is_allowed = squeak_address_str in self.allowed_addresses logger.info("Is squeak in whitelist: {}".format(is_allowed)) return is_allowed