diff --git a/.gitignore b/.gitignore index eab3fc3e..ab855ca4 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ src/*.egg-info build .idea +.vscode diff --git a/setup.cfg b/setup.cfg index a3e04f90..d29274f4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -11,7 +11,7 @@ package_dir = packages = find: python_requires = >=3.7 install_requires = - momento-wire-types==0.8.1 + momento-wire-types==0.12.0 build setuptools pyjwt[crypto] diff --git a/src/momento/_utilities/_data_validation.py b/src/momento/_utilities/_data_validation.py index e1431909..fde22a35 100644 --- a/src/momento/_utilities/_data_validation.py +++ b/src/momento/_utilities/_data_validation.py @@ -45,6 +45,11 @@ def _validate_ttl(ttl_seconds: int) -> None: raise errors.InvalidArgumentError("TTL Seconds must be a non-negative integer") +def _validate_ttl_minutes(ttl_minutes: int) -> None: + if not isinstance(ttl_minutes, int) or ttl_minutes < 0: + raise errors.InvalidArgumentError("TTL Minutes must be a non-negative integer") + + def _validate_request_timeout(request_timeout_ms: Optional[int]) -> None: if request_timeout_ms is None: return diff --git a/src/momento/aio/_scs_control_client.py b/src/momento/aio/_scs_control_client.py index 29b68b10..8b35fc46 100644 --- a/src/momento/aio/_scs_control_client.py +++ b/src/momento/aio/_scs_control_client.py @@ -1,13 +1,23 @@ +from time import time from typing import Optional from momento_wire_types.controlclient_pb2 import _CreateCacheRequest from momento_wire_types.controlclient_pb2 import _DeleteCacheRequest from momento_wire_types.controlclient_pb2 import _ListCachesRequest +from momento_wire_types.controlclient_pb2 import _CreateSigningKeyRequest +from momento_wire_types.controlclient_pb2 import _RevokeSigningKeyRequest +from momento_wire_types.controlclient_pb2 import _ListSigningKeysRequest from .._utilities._data_validation import _validate_cache_name -from ..cache_operation_types import CreateCacheResponse +from .._utilities._data_validation import _validate_ttl_minutes +from ..cache_operation_types import ( + CreateCacheResponse, + ListSigningKeysResponse, + RevokeSigningKeyResponse, +) from ..cache_operation_types import DeleteCacheResponse from ..cache_operation_types import ListCachesResponse +from ..cache_operation_types import CreateSigningKeyResponse from .. import _cache_service_errors_converter from .. import _momento_logger @@ -70,5 +80,58 @@ async def list_caches(self, next_token: Optional[str] = None) -> ListCachesRespo except Exception as e: raise _cache_service_errors_converter.convert(e) + async def create_signing_key( + self, ttl_minutes: int, endpoint: str + ) -> CreateSigningKeyResponse: + _validate_ttl_minutes(ttl_minutes) + try: + _momento_logger.debug( + f"Creating signing key with ttl (in minutes): {ttl_minutes}" + ) + create_signing_key_request = _CreateSigningKeyRequest() + create_signing_key_request.ttl_minutes = ttl_minutes + return CreateSigningKeyResponse( + await self._grpc_manager.async_stub().CreateSigningKey( + create_signing_key_request, timeout=_DEADLINE_SECONDS + ), + endpoint, + ) + except Exception as e: + _momento_logger.debug(f"Failed to create signing key with exception: {e}") + raise _cache_service_errors_converter.convert(e) + + async def revoke_signing_key(self, key_id: str) -> RevokeSigningKeyResponse: + try: + _momento_logger.debug(f"Revoking signing key with key_id {key_id}") + request = _RevokeSigningKeyRequest() + request.key_id = key_id + return RevokeSigningKeyResponse( + await self._grpc_manager.async_stub().RevokeSigningKey( + request, timeout=_DEADLINE_SECONDS + ) + ) + except Exception as e: + _momento_logger.debug( + f"Failed to revoke signing key with key_id {key_id} exception: {e}" + ) + raise _cache_service_errors_converter.convert(e) + + async def list_signing_keys( + self, endpoint: str, next_token: Optional[str] = None + ) -> ListSigningKeysResponse: + try: + list_signing_keys_request = _ListSigningKeysRequest() + list_signing_keys_request.next_token = ( + next_token if next_token is not None else "" + ) + return ListSigningKeysResponse( + await self._grpc_manager.async_stub().ListSigningKeys( + list_signing_keys_request, timeout=_DEADLINE_SECONDS + ), + endpoint, + ) + except Exception as e: + raise _cache_service_errors_converter.convert(e) + async def close(self) -> None: await self._grpc_manager.close() diff --git a/src/momento/aio/_scs_data_client.py b/src/momento/aio/_scs_data_client.py index 4d898c68..d115e5f8 100644 --- a/src/momento/aio/_scs_data_client.py +++ b/src/momento/aio/_scs_data_client.py @@ -37,6 +37,10 @@ def __init__( self._grpc_manager = _scs_grpc_manager._DataGrpcManager(auth_token, endpoint) _validate_ttl(default_ttl_seconds) self._default_ttlSeconds = default_ttl_seconds + self._endpoint = endpoint + + def get_endpoint(self) -> str: + return self._endpoint async def set( self, diff --git a/src/momento/aio/simple_cache_client.py b/src/momento/aio/simple_cache_client.py index 7d87e5a6..3cf18443 100644 --- a/src/momento/aio/simple_cache_client.py +++ b/src/momento/aio/simple_cache_client.py @@ -30,6 +30,9 @@ CreateCacheResponse, DeleteCacheResponse, ListCachesResponse, + CreateSigningKeyResponse, + RevokeSigningKeyResponse, + ListSigningKeysResponse, CacheSetResponse, CacheGetResponse, CacheMultiSetOperation, @@ -121,6 +124,59 @@ async def list_caches(self, next_token: Optional[str] = None) -> ListCachesRespo """ return await self._control_client.list_caches(next_token) + async def create_signing_key(self, ttl_minutes: int) -> CreateSigningKeyResponse: + """Creates a Momento signing key + + Args: + ttl_minutes: The key's time-to-live in minutes + + Returns: + CreateSigningKeyResponse + + Raises: + InvalidArgumentError: If provided ttl minutes is negative. + BadRequestError: If the ttl provided is not accepted + AuthenticationError: If the provided Momento Auth Token is invalid. + ClientSdkError: For any SDK checks that fail. + """ + return await self._control_client.create_signing_key( + ttl_minutes, self._data_client.get_endpoint() + ) + + async def revoke_signing_key(self, key_id: str) -> RevokeSigningKeyResponse: + """Revokes a Momento signing key, all tokens signed by which will be invalid + + Args: + key_id: The id of the Momento signing key to revoke + + Returns: + RevokeSigningKeyResponse + + Raises: + AuthenticationError: If the provided Momento Auth Token is invalid. + ClientSdkError: For any SDK checks that fail. + """ + return await self._control_client.revoke_signing_key(key_id) + + async def list_signing_keys( + self, next_token: Optional[str] = None + ) -> ListSigningKeysResponse: + """Lists all Momento signing keys for the provided auth token. + + Args: + next_token: Token to continue paginating through the list. It's used to handle large paginated lists. + + Returns: + ListSigningKeysResponse + + Raises: + AuthenticationError: If the provided Momento Auth Token is invalid. + ClientSdkError: For any SDK checks that fail. + """ + return await self._control_client.list_signing_keys( + self._data_client.get_endpoint(), next_token + ) + async def multi_set( self, cache_name: str, diff --git a/src/momento/cache_operation_types.py b/src/momento/cache_operation_types.py index 2d6a0f13..290f6019 100644 --- a/src/momento/cache_operation_types.py +++ b/src/momento/cache_operation_types.py @@ -1,3 +1,5 @@ +import json +from datetime import datetime from enum import Enum from typing import Any, Optional, List, Union from dataclasses import dataclass @@ -5,7 +7,6 @@ from momento_wire_types import cacheclient_pb2 as cache_client_types from . import _cache_service_errors_converter as error_converter from . import _momento_logger -from .errors import MomentoServiceError class CacheGetStatus(Enum): @@ -209,3 +210,91 @@ def next_token(self) -> Optional[str]: def caches(self) -> List[CacheInfo]: """Returns all caches.""" return self._caches + + +class CreateSigningKeyResponse: + def __init__(self, grpc_create_signing_key_response: Any, endpoint: str): # type: ignore[misc] + """Initializes CreateSigningKeyResponse to handle create signing key response. + + Args: + grpc_create_signing_key_response: Protobuf based response returned by Scs. + """ + self._key_id: str = json.loads(grpc_create_signing_key_response.key)["kid"] # type: ignore[misc] + self._endpoint: str = endpoint + self._key: str = grpc_create_signing_key_response.key # type: ignore[misc] + self._expires_at: datetime = datetime.fromtimestamp( + grpc_create_signing_key_response.expires_at # type: ignore[misc] + ) + + def key_id(self) -> str: + """Returns the id of the signing key""" + return self._key_id + + def endpoint(self) -> str: + """Returns the endpoint of the signing key""" + return self._endpoint + + def key(self) -> str: + """Returns the JSON string of the key itself""" + return self._key + + def expires_at(self) -> datetime: + """Returns the datetime representation of when the key expires""" + return self._expires_at + + +class RevokeSigningKeyResponse: + def __init__(self, grpc_revoke_signing_key_response: Any): # type: ignore[misc] + pass + + +class SigningKey: + def __init__(self, grpc_listed_signing_key: Any, endpoint: str): # type: ignore[misc] + """Initializes SigningKey to handle signing keys returned from list signing keys operation. + + Args: + grpc_listed_signing_key: Protobuf based response returned by Scs. + """ + self._key_id: str = grpc_listed_signing_key.key_id # type: ignore[misc] + self._expires_at: datetime = datetime.fromtimestamp( + grpc_listed_signing_key.expires_at # type: ignore[misc] + ) + self._endpoint: str = endpoint + + def key_id(self) -> str: + """Returns the id of the Momento signing key""" + return self._key_id + + def expires_at(self) -> datetime: + """Returns the time the key expires""" + return self._expires_at + + def endpoint(self) -> str: + """Returns the endpoint of the Momento signing key""" + return self._endpoint + + +class ListSigningKeysResponse: + def __init__(self, grpc_list_signing_keys_response: Any, endpoint: str): # type: ignore[misc] + """Initializes ListSigningKeysResponse to handle list signing keys response. + + Args: + grpc_list_signing_keys_response: Protobuf based response returned by Scs. + """ + self._next_token: Optional[str] = ( + grpc_list_signing_keys_response.next_token # type: ignore[misc] + if grpc_list_signing_keys_response.next_token != "" # type: ignore[misc] + else None + ) + self._signing_keys: List[SigningKey] = [ # type: ignore[misc] + SigningKey(signing_key, endpoint) # type: ignore[misc] + for signing_key in grpc_list_signing_keys_response.signing_key # type: ignore[misc] + ] + + def next_token(self) -> Optional[str]: + """Returns next token.""" + return self._next_token + + def signing_keys(self) -> List[SigningKey]: + """Returns all signing keys.""" + return self._signing_keys diff --git a/src/momento/simple_cache_client.py b/src/momento/simple_cache_client.py index e7206aa7..d71d13e0 100644 --- a/src/momento/simple_cache_client.py +++ b/src/momento/simple_cache_client.py @@ -9,6 +9,7 @@ CacheGetResponse, CacheSetResponse, CreateCacheResponse, + CreateSigningKeyResponse, DeleteCacheResponse, ListCachesResponse, CacheMultiGetResponse, @@ -17,6 +18,8 @@ CacheMultiSetResponse, CacheMultiSetFailureResponse, CacheMultiGetFailureResponse, + ListSigningKeysResponse, + RevokeSigningKeyResponse, ) from ._utilities._data_validation import _validate_request_timeout @@ -116,6 +119,58 @@ def list_caches(self, next_token: Optional[str] = None) -> ListCachesResponse: coroutine = self._momento_async_client.list_caches(next_token) return wait_for_coroutine(self._loop, coroutine) + def create_signing_key(self, ttl_minutes: int) -> CreateSigningKeyResponse: + """Creates a Momento signing key + + Args: + ttl_minutes: The key's time-to-live in minutes + + Returns: + CreateSigningKeyResponse + + Raises: + InvalidArgumentError: If provided ttl minutes is negative. + BadRequestError: If the ttl provided is not accepted + AuthenticationError: If the provided Momento Auth Token is invalid. + ClientSdkError: For any SDK checks that fail. + """ + coroutine = self._momento_async_client.create_signing_key(ttl_minutes) + return wait_for_coroutine(self._loop, coroutine) + + def revoke_signing_key(self, key_id: str) -> RevokeSigningKeyResponse: + """Revokes a Momento signing key, all tokens signed by which will be invalid + + Args: + key_id: The id of the Momento signing key to revoke + + Returns: + RevokeSigningKeyResponse + + Raises: + AuthenticationError: If the provided Momento Auth Token is invalid. + ClientSdkError: For any SDK checks that fail. + """ + coroutine = self._momento_async_client.revoke_signing_key(key_id) + return wait_for_coroutine(self._loop, coroutine) + + def list_signing_keys( + self, next_token: Optional[str] = None + ) -> ListSigningKeysResponse: + """Lists all Momento signing keys for the provided auth token. + + Args: + next_token: Token to continue paginating through the list. It's used to handle large paginated lists. + + Returns: + ListSigningKeysResponse + + Raises: + AuthenticationError: If the provided Momento Auth Token is invalid. + ClientSdkError: For any SDK checks that fail. + """ + coroutine = self._momento_async_client.list_signing_keys(next_token) + return wait_for_coroutine(self._loop, coroutine) + def set( self, cache_name: str, diff --git a/tests/test_momento_async.py b/tests/test_momento_async.py index 31293ec7..32a9d71f 100644 --- a/tests/test_momento_async.py +++ b/tests/test_momento_async.py @@ -178,6 +178,22 @@ async def test_list_caches_with_next_token_works(self): # skip until pagination is actually implemented, see # https://github.com/momentohq/control-plane-service/issues/83 self.skipTest("pagination not yet implemented") + + # signing keys + async def test_create_list_revoke_signing_keys(self): + try: + create_resp = await self.client.create_signing_key(30) + list_resp = await self.client.list_signing_keys() + self.assertTrue(0 < len(list_resp.signing_keys())) + signing_key = list_resp.signing_keys()[0] + self.assertEqual(create_resp.endpoint(), signing_key.endpoint()) + self.assertEqual(create_resp.key_id(), signing_key.key_id()) + finally: + list_resp = await self.client.list_signing_keys() + for signing_key in list_resp.signing_keys(): + await self.client.revoke_signing_key(signing_key.key_id()) + list_resp = await self.client.list_signing_keys() + self.assertEqual(0, len(list_resp.signing_keys())) # setting and getting