Add whitelist (#164)

* Add whitelist class

* Add itest for setting whitelisted field of a profile

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

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;

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(

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

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)

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
+        ]

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 = """