Adding generate_tenant_token method

meilisearch/client.py

@@ -1,10 +1,15 @@
 from typing import Any, Dict, List, Optional
 
+import base64
+import hashlib
+import hmac
+import json
 from meilisearch.index import Index
 from meilisearch.config import Config
 from meilisearch.task import get_task, get_tasks, wait_for_task
 from meilisearch._httprequests import HttpRequests
 from meilisearch.errors import MeiliSearchError
+from typing import Any, Dict, List, Optional, Union
 
 class Client():
     """
@@ -464,3 +469,67 @@ def wait_for_task(
             An error containing details about why Meilisearch can't process your request. Meilisearch error codes are described here: https://docs.meilisearch.com/errors/#meilisearch-errors
         """
         return wait_for_task(self.config, uid, timeout_in_ms, interval_in_ms)
+
+    def generate_tenant_token(
+        self,
+        search_rules: Union[Dict[str, Any], List[str]],
+        expires_at: Optional[int] = None,
+        api_key: Optional[str] = None
+    ) -> str:
+        """Generate a JWT token for the use of multitenancy.
+
+        Parameters
+        ----------
+        search_rules:
+            A Dictionary or list of string which contains the rules to be enforced at search time for all or specific
+            accessible indexes for the signing API Key.
+            In the specific case where you do not want to have any restrictions you can also use a list ["*"].
+        expires_at (optional):
+            Date and time when the key will expire.
+            Note that if an expires_at value is included it should a `timestamp`.
+        api_key (optional):
+            The API key parent of the token. If you leave it empty the client API Key will be used.
+
+        Returns
+        -------
+        jwt_token:
+           A string containing the jwt tenant token.
+           Note: If your token does not work remember that the search_rules is mandatory and should be well formatted.
+           `exp` must be a timestamp in the future.
+        """
+        # Standard JWT header for encryption with SHA256/HS256 algorithm
+        header = {
+            "typ": "JWT",
+            "alg": "HS256"
+        }
+
+        api_key = str(self.config.api_key) if api_key is None else str(api_key)
+
+        # Add the required fields to the payload
+        payload = {
+            'apiKeyPrefix': api_key[0:8],
+            'searchRules': search_rules,
+            'exp': expires_at
+        }
+
+        # Serialize the header and the payload
+        json_header = json.dumps(header, separators=(",",":")).encode()
+        json_payload = json.dumps(payload, separators=(",",":")).encode()
+
+        # Encode the header and the payload to Base64Url String
+        header_encode = self._base64url_encode(json_header)
+        header_payload = self._base64url_encode(json_payload)
+
+        secret_encoded = api_key.encode()
+        # Create Signature Hash
+        signature = hmac.new(secret_encoded, (header_encode + "." + header_payload).encode(), hashlib.sha256).digest()
+        # Create JWT
+        jwt_token = header_encode + '.' + header_payload + '.' + self._base64url_encode(signature)
+
+        return jwt_token
+
+    @staticmethod
+    def _base64url_encode(
+        data: bytes
+    ) -> str:
+        return base64.urlsafe_b64encode(data).decode('utf-8').replace('=','')

tests/client/test_client_tenant_token.py

@@ -0,0 +1,92 @@
+# pylint: disable=invalid-name
+
+from re import search
+import pytest
+import meilisearch
+from tests import BASE_URL, MASTER_KEY
+from meilisearch.errors import MeiliSearchApiError
+import datetime
+
+def test_generate_tenant_token_with_search_rules(get_private_key, index_with_documents):
+    """Tests create a tenant token with only search rules."""
+    index_with_documents()
+    key = get_private_key
+
+    client = meilisearch.Client(BASE_URL, key['key'])
+
+    token = client.generate_tenant_token(search_rules=["*"])
+
+    token_client = meilisearch.Client(BASE_URL, token)
+    response = token_client.index('indexUID').search('')
+    assert isinstance(response, dict)
+    assert len(response['hits']) == 20
+    assert response['query'] == ''
+
+def test_generate_tenant_token_with_api_key(client, get_private_key, index_with_documents):
+    """Tests create a tenant token with only search rules."""
+    index_with_documents()
+    key = get_private_key
+
+    token = client.generate_tenant_token(search_rules=["*"], api_key=key['key'])
+
+    token_client = meilisearch.Client(BASE_URL, token)
+    response = token_client.index('indexUID').search('')
+    assert isinstance(response, dict)
+    assert len(response['hits']) == 20
+    assert response['query'] == ''
+
+def test_generate_tenant_token_with_expires_at(client, get_private_key, index_with_documents):
+    """Tests create a tenant token with search rules and expiration date."""
+    index_with_documents()
+    key = get_private_key
+
+    client = meilisearch.Client(BASE_URL, key['key'])
+
+    tomorrow = datetime.datetime.now() + datetime.timedelta(days=1)
+    timestamp = datetime.datetime.timestamp(tomorrow)
+
+    token = client.generate_tenant_token(search_rules=["*"], expires_at=int(timestamp))
+
+    token_client = meilisearch.Client(BASE_URL, token)
+    response = token_client.index('indexUID').search('')
+    assert isinstance(response, dict)
+    assert len(response['hits']) == 20
+    assert response['query'] == ''
+
+def test_generate_tenant_token_without_search_rules(get_private_key, index_with_documents):
+    """Tests create a tenant token without search rules."""
+    index_with_documents()
+    key = get_private_key
+
+    client = meilisearch.Client(BASE_URL, key['key'])
+
+    token = client.generate_tenant_token(search_rules='')
+
+    token_client = meilisearch.Client(BASE_URL, token)
+    with pytest.raises(MeiliSearchApiError):
+        token_client.index('indexUID').search('')
+
+def test_generate_tenant_token_with_master_key(client, get_private_key, index_with_documents):
+    """Tests create a tenant token with master key."""
+    index_with_documents()
+    token = client.generate_tenant_token(search_rules=['*'])
+
+    token_client = meilisearch.Client(BASE_URL, token)
+    with pytest.raises(MeiliSearchApiError):
+        token_client.index('indexUID').search('')
+
+def test_generate_tenant_token_with_bad_expires_at(client, get_private_key, index_with_documents):
+    """Tests create a tenant token with only search rules."""
+    index_with_documents()
+    key = get_private_key
+
+    client = meilisearch.Client(BASE_URL, key['key'])
+
+    yesterday = datetime.datetime.now() + datetime.timedelta(days=-1)
+    timestamp = datetime.datetime.timestamp(yesterday)
+
+    token = client.generate_tenant_token(search_rules=["*"], expires_at=int(timestamp))
+
+    token_client = meilisearch.Client(BASE_URL, token)
+    with pytest.raises(MeiliSearchApiError):
+        token_client.index('indexUID').search('')

tests/conftest.py

@@ -109,3 +109,9 @@ def test_key_info(client):
         client.delete_key(key['key'])
     except MeiliSearchApiError:
         pass
+
+@fixture(scope='function')
+def get_private_key(client):
+    keys = client.get_keys()['results']
+    key = next(x for x in keys if 'Default Admin API' in x['description'])
+    return key