diff --git a/itests/tests/test_squeak_node.py b/itests/tests/test_squeak_node.py index 8d245f95..925d0903 100644 --- a/itests/tests/test_squeak_node.py +++ b/itests/tests/test_squeak_node.py @@ -260,9 +260,6 @@ def test_make_squeak(server_stub, admin_stub): ) ) - # Wait a few seconds for the squeak to be verified on the server. - time.sleep(5) - # Get all followed squeak display items get_followed_squeak_display_response = admin_stub.GetFollowedSqueakDisplays( squeak_admin_pb2.GetFollowedSqueakDisplaysRequest() @@ -469,3 +466,38 @@ def test_make_contact_profile(server_stub, admin_stub): for profile in get_contact_profiles_response.squeak_profiles ] assert contact_name in contact_profile_names + +def test_set_profile_whitelisted(server_stub, admin_stub): + # Create a new contact profile + contact_name = "whitelisted_contact" + contact_signing_key = generate_signing_key() + contact_address = get_address(contact_signing_key) + create_contact_profile_response = admin_stub.CreateContactProfile( + squeak_admin_pb2.CreateContactProfileRequest( + profile_name=contact_name, + address=contact_address, + ) + ) + contact_profile_id = create_contact_profile_response.profile_id + + # Get the new squeak profile + get_squeak_profile_response = admin_stub.GetSqueakProfile( + squeak_admin_pb2.GetSqueakProfileRequest(profile_id=contact_profile_id,) + ) + assert get_squeak_profile_response.squeak_profile.profile_name == contact_name + assert get_squeak_profile_response.squeak_profile.whitelisted == False + + # Set the profile to be whitelisted + admin_stub.SetSqueakProfileWhitelisted( + squeak_admin_pb2.SetSqueakProfileWhitelistedRequest( + profile_id=contact_profile_id, + whitelisted=True, + ) + ) + + # Get the squeak profile again + get_squeak_profile_response = admin_stub.GetSqueakProfile( + squeak_admin_pb2.GetSqueakProfileRequest(profile_id=contact_profile_id,) + ) + assert get_squeak_profile_response.squeak_profile.profile_name == contact_name + assert get_squeak_profile_response.squeak_profile.whitelisted == True diff --git a/proto/squeak_admin.proto b/proto/squeak_admin.proto index 60c75ac3..f5c59163 100644 --- a/proto/squeak_admin.proto +++ b/proto/squeak_admin.proto @@ -46,6 +46,10 @@ service SqueakAdmin { */ rpc GetSqueakProfileByAddress (GetSqueakProfileByAddressRequest) returns (GetSqueakProfileByAddressReply) {} + /** sqkadmin: `getsqueakprofilebyaddress` + */ + rpc SetSqueakProfileWhitelisted (SetSqueakProfileWhitelistedRequest) returns (SetSqueakProfileWhitelistedReply) {} + /** sqkadmin: `makesqueak` */ rpc MakeSqueak (MakeSqueakRequest) returns (MakeSqueakReply) {} @@ -135,6 +139,17 @@ message GetSqueakProfileByAddressReply { SqueakProfile squeak_profile = 1; } +message SetSqueakProfileWhitelistedRequest { + /// The profile id + int32 profile_id = 1; + + /// Whitelisted + bool whitelisted = 2; +} + +message SetSqueakProfileWhitelistedReply { +} + message SqueakProfile { /// The profile id int32 profile_id = 1; diff --git a/squeakserver/admin/squeak_admin_server_handler.py b/squeakserver/admin/squeak_admin_server_handler.py index 1114bad9..b2329ab8 100644 --- a/squeakserver/admin/squeak_admin_server_handler.py +++ b/squeakserver/admin/squeak_admin_server_handler.py @@ -59,6 +59,13 @@ class SqueakAdminServerHandler(object): logger.info("Got squeak profile by address: {}".format(squeak_profile)) return squeak_profile + def handle_set_squeak_profile_whitelisted(self, profile_id, whitelisted): + logger.info("Handle set squeak profile whitelisted with profile id: {}, whitelisted: {}".format( + profile_id, + whitelisted, + )) + self.squeak_node.set_squeak_profile_whitelisted(profile_id, whitelisted) + def handle_make_squeak(self, profile_id, content_str, replyto_hash): logger.info("Handle make squeak profile with id: {}".format(profile_id)) inserted_squeak_hash = self.squeak_node.make_squeak( diff --git a/squeakserver/admin/squeak_admin_server_servicer.py b/squeakserver/admin/squeak_admin_server_servicer.py index cc292565..72a639bf 100644 --- a/squeakserver/admin/squeak_admin_server_servicer.py +++ b/squeakserver/admin/squeak_admin_server_servicer.py @@ -81,6 +81,12 @@ class SqueakAdminServerServicer(squeak_admin_pb2_grpc.SqueakAdminServicer): squeak_profile=squeak_profile_msg ) + def SetSqueakProfileWhitelisted(self, request, context): + profile_id = request.profile_id + whitelisted = request.whitelisted + self.handler.handle_set_squeak_profile_whitelisted(profile_id, whitelisted) + return squeak_admin_pb2.SetSqueakProfileWhitelistedReply() + def MakeSqueak(self, request, context): profile_id = request.profile_id content_str = request.content diff --git a/squeakserver/node/squeak_node.py b/squeakserver/node/squeak_node.py index 14cf07f2..2c8673da 100644 --- a/squeakserver/node/squeak_node.py +++ b/squeakserver/node/squeak_node.py @@ -10,6 +10,7 @@ from squeakserver.node.squeak_block_queue_worker import SqueakBlockQueueWorker from squeakserver.node.squeak_block_verifier import SqueakBlockVerifier from squeakserver.node.squeak_maker import SqueakMaker from squeakserver.node.squeak_rate_limiter import SqueakRateLimiter +from squeakserver.node.squeak_whitelist import SqueakWhitelist from squeakserver.server.buy_offer import BuyOffer from squeakserver.server.squeak_profile import SqueakProfile from squeakserver.server.util import generate_offer_preimage @@ -43,6 +44,9 @@ class SqueakNode: lightning_client, max_squeaks_per_block_per_address, ) + self.squeak_whitelist = SqueakWhitelist( + postgres_db, + ) def start_running(self): # self.squeak_block_periodic_worker.start_running() @@ -159,6 +163,9 @@ class SqueakNode: def get_squeak_profile_by_address(self, address): return self.postgres_db.get_profile_by_address(address) + def set_squeak_profile_whitelisted(self, profile_id, whitelisted): + self.postgres_db.set_profile_whitelisted(profile_id, whitelisted) + def make_squeak(self, profile_id, content_str, replyto_hash): squeak_profile = self.postgres_db.get_profile(profile_id) squeak_maker = SqueakMaker(self.lightning_client) diff --git a/squeakserver/node/squeak_whitelist.py b/squeakserver/node/squeak_whitelist.py new file mode 100644 index 00000000..9847e6ff --- /dev/null +++ b/squeakserver/node/squeak_whitelist.py @@ -0,0 +1,34 @@ +import logging +import queue + +from squeakserver.server.util import get_hash +from squeakserver.node.block_info import BlockInfo + +logger = logging.getLogger(__name__) + + +class SqueakWhitelist: + def __init__(self, postgres_db): + self.postgres_db = postgres_db + self.allowed_addresses = [] + 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)) + is_allowed = squeak_address_str in self.allowed_addresses + logger.info("Is squeak in whitelist: {}".format(is_allowed)) + return is_allowed + + def refresh(self): + whitelisted_addresses = self._get_whitelisted_addresses() + self.allowed_addresses = whitelisted_addresses + + def _get_whitelisted_addresses(self): + whitelisted_profiles = self.postgres_db.get_whitelisted_profiles() + return [ + profile.address + for profile in whitelisted_profiles + ] diff --git a/squeakserver/server/postgres_db.py b/squeakserver/server/postgres_db.py index e8b6222b..3ca2041d 100644 --- a/squeakserver/server/postgres_db.py +++ b/squeakserver/server/postgres_db.py @@ -235,6 +235,18 @@ class PostgresDb: profiles = [self._parse_squeak_profile(row) for row in rows] return profiles + def get_whitelisted_profiles(self): + """ Get all whitelisted profiles. """ + sql = """ + SELECT * FROM profile + WHERE whitelisted; + """ + with self.get_cursor() as curs: + curs.execute(sql) + rows = curs.fetchall() + profiles = [self._parse_squeak_profile(row) for row in rows] + return profiles + def get_profile(self, profile_id): """ Get a profile. """ sql = """ @@ -256,6 +268,16 @@ class PostgresDb: row = curs.fetchone() return self._parse_squeak_profile(row) + def set_profile_whitelisted(self, profile_id, whitelisted): + """ Set a profile is whitelisted. """ + sql = """ + UPDATE profile + SET whitelisted=%s + WHERE profile_id=%s; + """ + with self.get_cursor() as curs: + curs.execute(sql, (whitelisted, profile_id,)) + def get_unverified_block_squeaks(self): """ Get all squeaks without block header. """ sql = """