feat: add specific caveat for partners (#136)

synapse/handlers/auth.py

@@ -827,7 +827,13 @@ class AuthHandler(BaseHandler):
         ):
             await self.auth.check_auth_blocking(user_id)
 
+        """ watcha!
         access_token = self.macaroon_gen.generate_access_token(user_id)
+        !watcha """
+        # watcha+
+        extra_caveats = ["partner = true"] if await self.is_partner(user_id) else None
+        access_token = self.macaroon_gen.generate_access_token(user_id, extra_caveats)
+        # +watcha
         await self.store.add_access_token_to_user(
             user_id=user_id,
             token=access_token,
@@ -1589,7 +1595,7 @@ class AuthHandler(BaseHandler):
 
     # watcha+
     async def is_partner(self, user_id):
-        ret = await self.store.is_user_partner(
+        ret = await self.store.is_partner(
             user_id,
         )
         return ret

synapse/storage/databases/main/registration.py

@@ -1768,16 +1768,17 @@ class RegistrationStore(StatsStore, RegistrationBackgroundUpdateStore):
             start_or_continue_validation_session_txn,
         )
 
-    # watcha+ op318
-    async def is_user_partner(self, user_id):
+    # watcha+
+    async def is_partner(self, user_id):
 
         is_partner = await self.db_pool.simple_select_one_onecol(
             "users",
             keyvalues={"name": user_id},
             retcol="is_partner",
-            desc="is_user_partner",
+            allow_none=True,
+            desc="is_partner",
         )
-        return is_partner
+        return bool(is_partner)
 
     async def is_user_admin(self, user_id):
 

tests/api/test_auth.py

@@ -258,6 +258,7 @@ class AuthTestCase(unittest.HomeserverTestCase):
         USER_ID = "@percy:matrix.org"
         self.store.add_access_token_to_user = simple_async_mock(None)
         self.store.get_device = simple_async_mock(None)
+        self.hs.get_auth_handler().is_partner = simple_async_mock(return_value=False) # watcha+
 
         token = self.get_success(
             self.hs.get_auth_handler().get_access_token_for_user_id(

tests/handlers/test_auth.py

@@ -17,12 +17,20 @@ from mock import Mock
 import pymacaroons
 
 from synapse.api.errors import AuthError, ResourceLimitError
+from synapse.rest import admin # watcha+
 
 from tests import unittest
 from tests.test_utils import make_awaitable
 
 
 class AuthTestCase(unittest.HomeserverTestCase):
+
+    # watcha+
+    servlets = [
+        admin.register_servlets,
+    ]
+    # +watcha
+
     def prepare(self, reactor, clock, hs):
         self.auth_handler = hs.get_auth_handler()
         self.macaroon_generator = hs.get_macaroon_generator()
@@ -44,6 +52,18 @@ class AuthTestCase(unittest.HomeserverTestCase):
         if "some_user" not in macaroon.inspect():
             self.fail("some_user was not in %s" % macaroon.inspect())
 
+    # watcha+
+    def test_token_is_a_partner_macaroon(self):
+        partner_id = self.register_user("partner", "pass", is_partner=True)
+        access_token = self.get_success(
+            self.auth_handler.get_access_token_for_user_id(
+                partner_id, device_id=None, valid_until_ms=None
+            )
+        )
+        macaroon = pymacaroons.Macaroon.deserialize(access_token)
+        self.assertIn("partner = true", macaroon.inspect())
+    # +watcha
+
     def test_macaroon_caveats(self):
         token = self.macaroon_generator.generate_access_token("a_user")
         macaroon = pymacaroons.Macaroon.deserialize(token)

tests/storage/test_watcha_administration.py

@@ -102,7 +102,7 @@ class WatchaAdminTestCase(unittest.TestCase):
 
         for element in expected_values:
             yield self.store.watcha_update_user_role(user_id, element["role"])
-            is_partner = yield self.store.is_user_partner(user_id)
+            is_partner = yield self.store.is_partner(user_id)
             is_admin = yield self.store.is_user_admin(user_id)
             self.assertEquals(is_partner, element["values"]["is_partner"])
             self.assertEquals(is_admin, element["values"]["is_admin"])