zepben.ewb 1.0.0__py3-none-any.whl

zepben/ewb/streaming/get/query_network_state_client.py

@@ -0,0 +1,64 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["QueryNetworkStateClient"]
+
+from datetime import datetime
+from typing import List, Callable, AsyncGenerator
+
+from zepben.protobuf.ns.network_state_pb2_grpc import QueryNetworkStateServiceStub
+from zepben.protobuf.ns.network_state_requests_pb2 import GetCurrentStatesRequest
+
+from zepben.ewb.streaming.data.current_state_event import CurrentStateEvent
+from zepben.ewb.streaming.data.current_state_event_batch import CurrentStateEventBatch
+from zepben.ewb.streaming.data.set_current_states_status import SetCurrentStatesStatus
+from zepben.ewb.streaming.grpc.grpc import GrpcClient
+from zepben.ewb.util import datetime_to_timestamp
+
+
+class QueryNetworkStateClient(GrpcClient):
+    """
+    A client class that provides functionality to interact with the gRPC service for querying network states.
+    A gRPC channel or stub must be provided.
+    """
+
+    _stub: QueryNetworkStateServiceStub = None
+
+    def __init__(self, channel=None, stub: QueryNetworkStateServiceStub = None, error_handlers: List[Callable[[Exception], bool]] = None, timeout: int = 60):
+        super().__init__(error_handlers=error_handlers, timeout=timeout)
+        if channel is None and stub is None:
+            raise ValueError("Must provide either a channel or a stub")
+        if stub is not None:
+            self._stub = stub
+        else:
+            self._stub = QueryNetworkStateServiceStub(channel)
+
+    async def get_current_states(self, query_id: int, from_datetime: datetime, to_datetime: datetime) -> AsyncGenerator[CurrentStateEventBatch, None]:
+        """
+        Asynchronously retrieves a collection containing CurrentStateEvent objects, representing the network states
+        within a specified time range.
+
+        Args:
+            query_id: A unique identifier for the query being processed.
+            from_datetime: The start time, as a datetime, for the query range.
+            to_datetime: The end time, as a datetime, for the query range.
+
+        Returns:
+            A stream of batched network state events in the specified time range.
+        """
+        request = GetCurrentStatesRequest(
+            messageId=query_id,
+            fromTimestamp=datetime_to_timestamp(from_datetime),
+            toTimestamp=datetime_to_timestamp(to_datetime)
+        )
+        async for response in self._stub.getCurrentStates(request):
+            yield CurrentStateEventBatch(response.messageId, [CurrentStateEvent.from_pb(event) for event in response.event])
+
+    def report_batch_status(self, status: SetCurrentStatesStatus):
+        """
+        Send a response to a previous `getCurrentStates` request to let the server know how we handled its response.
+        :param status: The batch status to report.
+        """
+        self._stub.reportBatchStatus(iter([status.to_pb()]))

zepben/ewb/streaming/get/query_network_state_service.py

@@ -0,0 +1,94 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["QueryNetworkStateService"]
+
+from datetime import datetime
+from typing import Callable, AsyncGenerator, Optional
+
+from google.protobuf import empty_pb2
+from zepben.protobuf.ns.network_state_pb2_grpc import QueryNetworkStateServiceServicer
+from zepben.protobuf.ns.network_state_requests_pb2 import GetCurrentStatesRequest
+from zepben.protobuf.ns.network_state_responses_pb2 import GetCurrentStatesResponse, SetCurrentStatesResponse
+
+from zepben.ewb.streaming.data.current_state_event_batch import CurrentStateEventBatch
+from zepben.ewb.streaming.data.set_current_states_status import SetCurrentStatesStatus
+
+
+class QueryNetworkStateService(QueryNetworkStateServiceServicer):
+    """
+    A service class that provides a simplified interface for retrieving current state events
+    via gRPC without exposing the underlying complexity of gRPC mechanisms.
+
+    This class serves as a wrapper around the gRPC-generated service implementation,
+    allowing users to interact with it using a more convenient function type.
+
+    Attributes:
+        on_get_current_states: An function that retrieves CurrentStateEvent objects within the specified time range.
+        on_current_states_status: A callback triggered when the response status of an event returned via `on_get_current_states` is received from the client.
+        on_processing_error: A function that takes a message and optional cause. Called when `on_current_states_status` raises an exception, or the
+          `SetCurrentStatesResponse` is for an unknown event status.
+    """
+
+    def __init__(
+        self,
+        on_get_current_states: Callable[[datetime, datetime], AsyncGenerator[CurrentStateEventBatch, None]],
+        on_current_states_status: Callable[[SetCurrentStatesStatus], None],
+        on_processing_error: Callable[[str, Optional[Exception]], None]
+    ):
+        self.on_get_current_states = on_get_current_states
+        self.on_current_states_status = on_current_states_status
+        self.on_processing_error = on_processing_error
+
+    async def getCurrentStates(self, request: GetCurrentStatesRequest, context) -> AsyncGenerator[GetCurrentStatesResponse, None]:
+        """
+        Handles the incoming request for retrieving current state events.
+
+        You shouldn't be calling this method directly, it will be invoked automatically via the gRPC engine. Each
+        `GetCurrentStatesRequest` retrieves the corresponding current state events using the `onGetCurrentStates`
+        callback function passed in the constructor, and asynchronously returns a `GetCurrentStatesResponse`.
+
+        It acts as a bridge between the gRPC request and the business logic that fetches the current state events.
+
+        Args:
+            request: The request object containing parameters for fetching current state events,
+              including the time range for the query.
+            context: The gRPC context.
+
+        Returns:
+           A stream of gRPC response messages
+        """
+        async for batch in self.on_get_current_states(request.fromTimestamp.ToDatetime(), request.toTimestamp.ToDatetime()):
+            yield GetCurrentStatesResponse(messageId=batch.batch_id, event=[event.to_pb() for event in batch.events])
+
+    # noinspection PyUnresolvedReferences
+    async def reportBatchStatus(self, status_responses: AsyncGenerator[SetCurrentStatesResponse, None], context) -> empty_pb2.Empty:
+        """
+        Handles incoming status reports in response to an event batch returned via `getCurrentStates`.
+
+        You shouldn't be calling this method directly, it will be invoked automatically via the gRPC engine. Each
+        `SetCurrentStatesResponse` will trigger the `on_current_states_status` callback function passed in the constructor.
+
+        It acts as a bridge between the gRPC request and the business logic that validates/logs the status of event
+        handling. Any errors in this handling, or unexpected status message will trigger the `on_processing_error`
+        callback function passed in the constructor.
+
+        Args:
+            status_responses: The request object stream (yes you read that correctly that the request is a response...) containing the response to our current state batch.
+            context: The gRPC context.
+        """
+        async for status_response in status_responses:
+            status = SetCurrentStatesStatus.from_pb(status_response)
+            if status is not None:
+                try:
+                    self.on_current_states_status(status)
+                except Exception as e:
+                    message = f"Exception thrown in status response handler for batch `{status_response.messageId}`: {str(e)}"
+                    self.on_processing_error(message, e)
+            else:
+                self.on_processing_error(f"Failed to decode status response for batch `{status_response.messageId}`: Unsupported type {status_response.WhichOneof('status')}", None)
+
+        # noinspection PyUnresolvedReferences
+        return empty_pb2.Empty()

zepben/ewb/streaming/grpc/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.

zepben/ewb/streaming/grpc/auth_token_plugin.py

@@ -0,0 +1,24 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["AuthTokenPlugin"]
+
+import grpc
+from zepben.ewb.auth import ZepbenTokenFetcher
+
+_AUTH_HEADER_KEY = 'authorization'
+
+
+class AuthTokenPlugin(grpc.AuthMetadataPlugin):
+    """
+    Metadata plugin that injects tokens into the metadata of each call.
+    Tokens are fetched using a provided :class:`ZepbenTokenFetcher`.
+    """
+
+    def __init__(self, token_fetcher: ZepbenTokenFetcher):
+        self.token_fetcher = token_fetcher
+
+    def __call__(self, context, callback):
+        callback(((_AUTH_HEADER_KEY, self.token_fetcher.fetch_token()),), None)

zepben/ewb/streaming/grpc/connect.py

@@ -0,0 +1,209 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = ["connect_tls", "connect_insecure", "connect_with_password", "connect_with_secret", "connect_with_identity", "connect_with_token"]
+
+from typing import Optional, Union
+
+import grpc
+from zepben.ewb.auth import ZepbenTokenFetcher, create_token_fetcher, create_token_fetcher_managed_identity
+
+from zepben.ewb import GrpcChannelBuilder
+
+GRPC_READY_TIMEOUT = 20  # seconds
+
+
+def connect_insecure(
+    host: str = "localhost",
+    rpc_port: int = 50051,
+    **kwargs
+) -> grpc.aio.Channel:
+    """
+    Create a :class:`grpc.aio.Channel` that communicates with the gRPC service over plaintext.
+
+    :param host: The hostname where the gRPC service is hosted
+    :param rpc_port: The port of the gRPC service
+    :return: A plaintext connection to the gRPC service
+    """
+    return GrpcChannelBuilder().for_address(host, rpc_port).build(**kwargs)
+
+
+def connect_tls(
+    host: str = "localhost",
+    rpc_port: int = 50051,
+    ca_filename: Optional[str] = None,
+    **kwargs
+) -> grpc.aio.Channel:
+    """
+    Create a :class:`grpc.aio.Channel` that communicates with the gRPC service using SSL/TLS transport security.
+
+    :param host: The hostname where the gRPC service is hosted
+    :param rpc_port: The port of the gRPC service
+    :param ca_filename: The filename of a truststore containing additional trusted root certificates. This parameter is optional
+                        and defaults to null, in which case only the system CAs are used to verify certificates.
+    :return:An encrypted connection to the gRPC service
+    """
+    return GrpcChannelBuilder().for_address(host, rpc_port).make_secure(root_certificates=ca_filename).build(**kwargs)
+
+
+def connect_with_identity(host: str,
+                          rpc_port: int,
+                          identity_url: str,
+                          ca_filename: Optional[str] = None,
+                          verify_auth: bool = True,
+                          **kwargs) -> grpc.aio.Channel:
+    """
+    Create a :class:`grpc.aio.Channel` that communicates with the gRPC service using SSL/TLS transport security and the OAuth client credentials flow,
+    utilising an Azure managed identity for fetching access tokens.
+
+    :param host: The hostname where the gRPC service is hosted
+    :param rpc_port: The port of the gRPC service
+    :param identity_url: The URL including resource identifier (API) to fetch token from. Typically something like:
+                         http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=<SOME_IDENTIFIER>
+    :param ca_filename: The filename of a truststore containing additional trusted root certificates. This parameter is optional
+                        and defaults to null, in which case only the system CAs are used to verify certificates.
+    :param verify_auth: Passed through to `requests.post()` when fetching access tokens to verify certificates of the auth service.
+    :return: An authenticated, encrypted connection to the gRPC service based on OAuth2 flows. If the authentication configuration specifies that no
+             authentication is required, a non-authenticated, encrypted connection is returned instead.
+    """
+    token_fetcher = create_token_fetcher_managed_identity(identity_url, verify_auth)
+    return GrpcChannelBuilder().for_address(host, rpc_port).make_secure(root_certificates=ca_filename).with_token_fetcher(token_fetcher).build(**kwargs)
+
+
+def connect_with_secret(
+    client_id: str,
+    client_secret: str,
+    host: str = "localhost",
+    rpc_port: int = 50051,
+    conf_address: Optional[str] = None,
+    verify_conf: Union[bool, str] = True,
+    verify_auth: Union[bool, str] = True,
+    ca_filename: Optional[str] = None,
+    **kwargs
+) -> grpc.aio.Channel:
+    """
+    Create a :class:`grpc.aio.Channel` that communicates with the gRPC service using SSL/TLS transport security and the OAuth client credentials flow.
+    The OAuth provider's domain and the "audience" parameter of the token request are fetched as JSON from a specified URL.
+
+    :param client_id: The client ID of the OAuth application to authenticate for
+    :param client_secret: The client secret of the OAuth application to authenticate for
+    :param host: The hostname where the gRPC service is hosted
+    :param rpc_port: The port of the gRPC service
+    :param conf_address: The address of the authentication configuration
+    :param verify_conf: Passed through to `requests.get()` when fetching the authentication configuration
+    :param verify_auth: Passed through to `requests.post()` when fetching access tokens to verify certificates of the auth service.
+    :param ca_filename: The filename of a truststore containing additional trusted root certificates. This parameter is optional
+                        and defaults to null, in which case only the system CAs are used to verify certificates.
+    :return: An authenticated, encrypted connection to the gRPC service based on OAuth2 flows. If the authentication configuration specifies that no
+             authentication is required, a non-authenticated, encrypted connection is returned instead.
+    """
+    token_fetcher = create_token_fetcher(
+        conf_address=conf_address or f"https://{host}/ewb/auth",
+        verify_conf=verify_conf,
+        verify_auth=verify_auth
+    )
+
+    if token_fetcher:
+        return _connect_with_secret_using_token_fetcher(token_fetcher, client_id, client_secret, host, rpc_port, ca_filename, **kwargs)
+    else:
+        return connect_tls(host, rpc_port, ca_filename, **kwargs)
+
+
+def connect_with_token(
+    access_token: str,
+    host: str = "localhost",
+    rpc_port: int = 50051,
+    ca_filename: Optional[str] = None,
+    **kwargs
+) -> grpc.aio.Channel:
+    """
+    Create a :class:`grpc.aio.Channel` that communicates with the gRPC service using SSL/TLS transport security and a personal access token generated from Evolve App Server or Evolve Web Client.
+
+    :param access_token: The token string of the client generated using Evolve App
+    :param host: The hostname where the gRPC service is hosted
+    :param rpc_port: The port of the gRPC service
+    :param ca_filename: The filename of a truststore containing additional trusted root certificates. This parameter is optional
+                        and defaults to null, in which case only the system CAs are used to verify certificates.
+    :return: An authenticated, encrypted connection to the gRPC service based on OAuth2 flows. If the authentication configuration specifies that no
+             authentication is required, a non-authenticated, encrypted connection is returned instead.
+    """
+
+    return GrpcChannelBuilder().for_address(host, rpc_port).make_secure(root_certificates=ca_filename).with_client_token(access_token).build(**kwargs)
+
+
+def connect_with_password(
+    client_id: str,
+    username: str,
+    password: str,
+    host: str = "localhost",
+    rpc_port: int = 50051,
+    conf_address: Optional[str] = None,
+    verify_conf: Union[bool, str] = True,
+    verify_auth: Union[bool, str] = True,
+    ca_filename: Optional[str] = None,
+    **kwargs
+) -> grpc.aio.Channel:
+    """
+    Create a :class:`grpc.aio.Channel` that communicates with the gRPC service using SSL/TLS transport security and the OAuth password grant flow.
+    The OAuth provider's domain and the "audience" parameter of the token request are fetched as JSON from a specified URL.
+
+    :param client_id: The client ID of the OAuth application to authenticate for
+    :param username: The username of the user to authenticate with
+    :param password: The password of the user to authenticate with
+    :param host: The hostname where the gRPC service is hosted
+    :param rpc_port: The port of the gRPC service
+    :param conf_address: The address of the authentication configuration
+    :param verify_conf: Passed through to `requests.get()` when fetching the authentication configuration
+    :param verify_auth: Passed through to `requests.post()` when fetching access tokens to verify certificates of the auth service.
+    :param ca_filename: The filename of a truststore containing additional trusted root certificates. This parameter is optional
+                        and defaults to null, in which case only the system CAs are used to verify certificates.
+    :return: An authenticated, encrypted connection to the gRPC service based on OAuth2 flows. If the authentication configuration specifies that no
+             authentication is required, a non-authenticated, encrypted connection is returned instead.
+    """
+    token_fetcher = create_token_fetcher(
+        conf_address=conf_address or f"https://{host}/ewb/auth",
+        verify_conf=verify_conf,
+        verify_auth=verify_auth
+    )
+
+    if token_fetcher:
+        return _connect_with_password_using_token_fetcher(token_fetcher, client_id, username, password, host, rpc_port, ca_filename, **kwargs)
+    else:
+        return connect_tls(host, rpc_port, ca_filename, **kwargs)
+
+
+def _connect_with_secret_using_token_fetcher(
+    token_fetcher: ZepbenTokenFetcher,
+    client_id: str,
+    client_secret: str,
+    host: str,
+    rpc_port: int,
+    ca_filename: Optional[str],
+    **kwargs
+) -> grpc.aio.Channel:
+    token_fetcher.token_request_data["client_id"] = client_id
+    token_fetcher.token_request_data["client_secret"] = client_secret
+    token_fetcher.token_request_data["grant_type"] = "client_credentials"
+
+    return GrpcChannelBuilder().for_address(host, rpc_port).make_secure(root_certificates=ca_filename).with_token_fetcher(token_fetcher).build(**kwargs)
+
+
+def _connect_with_password_using_token_fetcher(
+    token_fetcher: ZepbenTokenFetcher,
+    client_id: str,
+    username: str,
+    password: str,
+    host: str,
+    rpc_port: int,
+    ca_filename: Optional[str],
+    **kwargs
+) -> grpc.aio.Channel:
+    token_fetcher.token_request_data["client_id"] = client_id
+    token_fetcher.token_request_data["username"] = username
+    token_fetcher.token_request_data["password"] = password
+    token_fetcher.token_request_data["grant_type"] = "password"
+    token_fetcher.token_request_data["scope"] = "offline_access"
+
+    return GrpcChannelBuilder().for_address(host, rpc_port).make_secure(root_certificates=ca_filename).with_token_fetcher(token_fetcher).build(**kwargs)

zepben/ewb/streaming/grpc/grpc.py

@@ -0,0 +1,107 @@
+#  Copyright 2024 Zeppelin Bend Pty Ltd
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+from __future__ import annotations
+
+__all__ = ["GrpcResult", "GrpcClient"]
+
+from typing import TypeVar, Generic, Callable, List, Union, Coroutine
+
+from zepben.ewb.dataclassy import dataclass
+
+T = TypeVar("T")
+
+
+@dataclass(slots=True)
+class GrpcResult(Generic[T]):
+    result: Union[T, Exception]
+    was_error_handled: bool = False
+
+    @property
+    def value(self) -> T:
+        if self.result is not None and isinstance(self.result, Exception):
+            raise TypeError("You can only call value on a successful result.")
+        return self.result
+
+    @property
+    def thrown(self) -> Exception:
+        if self.result is None or not isinstance(self.result, Exception):
+            raise TypeError("You can only call thrown on an unsuccessful result.")
+        return self.result
+
+    @property
+    def was_successful(self):
+        return not self.was_failure
+
+    @property
+    def was_failure(self):
+        return isinstance(self.result, Exception)
+
+    @property
+    def was_error_unhandled(self):
+        return not self.was_error_handled
+
+    def on_success(self, handler: Callable[[T], None]) -> GrpcResult[T]:
+        """Calls `handler` with the `result` if this `was_successful`"""
+        if self.was_successful:
+            handler(self.result)
+        return self
+
+    def on_error(self, handler: Callable[[Exception, bool], GrpcResult[T]]) -> GrpcResult[T]:
+        """Calls `handler` with the `thrown` exception and `self.was_error_handled` if `self.was_failure`."""
+        if self.was_failure and self.was_error_handled:
+            return handler(self.result, self.was_error_handled)
+        return self
+
+    def on_handled_error(self, handler: Callable[[Exception], None]) -> GrpcResult[T]:
+        """Calls `handler` with the `thrown` exception if `self.was_failure` only if `self.was_error_handled`."""
+        if self.was_failure and self.was_error_handled:
+            handler(self.result)
+        return self
+
+    def on_unhandled_error(self, handler: Callable[[Exception], None]) -> GrpcResult[T]:
+        """Calls `handler` with the `thrown` exception if `self.was_failure` only if not `self.was_error_handled`."""
+        if self.was_failure and not self.was_error_handled:
+            handler(self.result)
+        return self
+
+    def throw_on_error(self) -> GrpcResult[T]:
+        """Throws `self.result` if `self.was_failure`"""
+        if self.was_failure:
+            raise self.result
+        return self
+
+    def throw_on_unhandled_error(self) -> GrpcResult[T]:
+        """Throws `self.result` only if `self.was_failure` and not `self.was_error_handled`"""
+        if self.was_failure and not self.was_error_handled:
+            raise self.result
+        return self
+
+
+@dataclass(init=False, slots=True)
+class GrpcClient(object):
+    error_handlers: List[Callable[[Exception], bool]] = []
+
+    timeout: int = 0
+    '''Timeout for client gRPC requests'''
+
+    def __init__(self, error_handlers: List[Callable[[Exception], bool]] = None, timeout: int = 60):
+        if error_handlers:
+            self.error_handlers = error_handlers.copy()
+        else:
+            self.error_handlers = list()
+        self.timeout = timeout
+
+    def try_handle_error(self, e: Exception) -> bool:
+        for handler in self.error_handlers:
+            if handler(e):
+                return True
+        return False
+
+    async def try_rpc(self, rpc: Callable[[], Coroutine[None, None, T]]) -> GrpcResult[T]:
+        try:
+            return GrpcResult(await rpc())
+        except Exception as e:
+            return GrpcResult(e, self.try_handle_error(e))