Implement key-based authentication (#56)

Marenz · web-flow · commit 434960be8b4f · 2024-06-20T09:41:50.000Z
- **Add key-based authentication support**
- **Update REPL client to work better with host/port/key parameters**
diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md
@@ -6,11 +6,12 @@
 
 ## Upgrading
 
-<!-- Here goes notes on how to upgrade from previous versions, including deprecations and what they should be replaced with -->
+* An API key for authorization must now be passed to the `DispatchClient`.
+* The client constructor now requires named parameters. If you are using the client directly, you will need to update your code to use named parameters.
 
 ## New Features
 
-<!-- Here goes the main new features and examples or instructions on how to use them -->
+* TLS is now enabled by default for the CLI client.
 
 ## Bug Fixes
 
diff --git a/src/frequenz/client/dispatch/__main__.py b/src/frequenz/client/dispatch/__main__.py
@@ -5,7 +5,6 @@
 
 import asyncio
 import os
-import sys
 from pprint import pformat
 from typing import Any, List
 
@@ -37,22 +36,66 @@
 DEFAULT_DISPATCH_API_PORT = 50051
 
 
-def get_client(host: str, port: int) -> Client:
+def ssl_channel_credentials_from_files(
+    root_cert_path: str | None = None,
+    client_cert_path: str | None = None,
+    client_key_path: str | None = None,
+) -> grpc.ChannelCredentials:
+    """Create credentials for use with an SSL-enabled Channel.
+
+    Using the provided certificate and key files.
+
+    Args:
+      root_cert_path: Path to the PEM-encoded root certificates file,
+                      or None to retrieve them from a default location chosen by gRPC runtime.
+      client_cert_path: Path to the PEM-encoded client certificate file.
+      client_key_path: Path to the PEM-encoded client private key file.
+
+    Returns:
+      A ChannelCredentials for use with an SSL-enabled Channel.
+    """
+    root_certificates = None
+    if root_cert_path is not None:
+        with open(root_cert_path, "rb") as f:
+            root_certificates = f.read()
+
+    certificate_chain = None
+    if client_cert_path is not None:
+        with open(client_cert_path, "rb") as f:
+            certificate_chain = f.read()
+
+    private_key = None
+    if client_key_path is not None:
+        with open(client_key_path, "rb") as f:
+            private_key = f.read()
+
+    return grpc.ssl_channel_credentials(
+        root_certificates=root_certificates,
+        private_key=private_key,
+        certificate_chain=certificate_chain,
+    )
+
+
+def get_client(*, host: str, port: int, key: str) -> Client:
     """Get a new client instance.
 
     Args:
         host: The host of the dispatch service.
         port: The port of the dispatch service.
+        key: The API key for authentication.
 
     Returns:
         Client: A new client instance.
     """
-    channel = grpc.aio.insecure_channel(f"{host}:{port}")
-    return Client(channel, f"{host}:{port}")
+    channel = grpc.aio.secure_channel(
+        f"{host}:{port}",
+        credentials=ssl_channel_credentials_from_files(),
+    )
+    return Client(grpc_channel=channel, svc_addr=f"{host}:{port}", key=key)
 
 
 # Click command groups
-@click.group()
+@click.group(invoke_without_command=True)
 @click.option(
     "--host",
     default=DEFAULT_DISPATCH_API_HOST,
@@ -69,11 +112,29 @@ def get_client(host: str, port: int) -> Client:
     show_envvar=True,
     show_default=True,
 )
+@click.option(
+    "--key",
+    help="API key for authentication",
+    envvar="DISPATCH_API_KEY",
+    show_envvar=True,
+    required=True,
+)
 @click.pass_context
-async def cli(ctx: click.Context, host: str, port: int) -> None:
+async def cli(ctx: click.Context, host: str, port: int, key: str) -> None:
     """Dispatch Service CLI."""
-    ctx.ensure_object(dict)
-    ctx.obj["client"] = get_client(host, port)
+    if ctx.obj is None:
+        ctx.obj = {}
+
+    ctx.obj["client"] = get_client(host=host, port=port, key=key)
+    ctx.obj["params"] = {
+        "host": host,
+        "port": port,
+        "key": key,
+    }
+
+    # Check if a subcommand was given
+    if ctx.invoked_subcommand is None:
+        await interactive_mode(host, port, key)
 
 
 @cli.command("list")
@@ -327,6 +388,18 @@ async def get(ctx: click.Context, dispatch_ids: List[int]) -> None:
         raise click.ClickException("Some gets failed.")
 
 
+@cli.command()
+@click.pass_obj
+async def repl(
+    obj: dict[str, Any],
+) -> None:
+    """Start an interactive interface."""
+    click.echo(f"Parameters: {obj}")
+    await interactive_mode(
+        obj["params"]["host"], obj["params"]["port"], obj["params"]["key"]
+    )
+
+
 @cli.command()
 @click.argument("dispatch_ids", type=FuzzyIntRange(), nargs=-1)  # Allow multiple IDs
 @click.pass_context
@@ -359,7 +432,7 @@ async def delete(ctx: click.Context, dispatch_ids: list[list[int]]) -> None:
         raise click.ClickException("Some deletions failed.")
 
 
-async def interactive_mode() -> None:
+async def interactive_mode(host: str, port: int, key: str) -> None:
     """Interactive mode for the CLI."""
     hist_file = os.path.expanduser("~/.dispatch_cli_history.txt")
     session: PromptSession[str] = PromptSession(history=FileHistory(filename=hist_file))
@@ -390,7 +463,15 @@ async def display_help() -> None:
             break
         else:
             # Split, but keep quoted strings together
-            params = click.parser.split_arg_string(user_input)
+            params = [
+                "--host",
+                host,
+                "--port",
+                str(port),
+                "--key",
+                key,
+            ] + click.parser.split_arg_string(user_input)
+
             try:
                 await cli.main(args=params, standalone_mode=False)
             except click.ClickException as e:
@@ -405,10 +486,7 @@ async def display_help() -> None:
 
 def main() -> None:
     """Entrypoint for the CLI."""
-    if len(sys.argv) > 1:
-        asyncio.run(cli.main())
-    else:
-        asyncio.run(interactive_mode())
+    asyncio.run(cli.main())
 
 
 if __name__ == "__main__":
diff --git a/src/frequenz/client/dispatch/_client.py b/src/frequenz/client/dispatch/_client.py
@@ -36,15 +36,19 @@
 class Client:
     """Dispatch API client."""
 
-    def __init__(self, grpc_channel: grpc.aio.Channel, svc_addr: str) -> None:
+    def __init__(
+        self, *, grpc_channel: grpc.aio.Channel, svc_addr: str, key: str
+    ) -> None:
         """Initialize the client.
 
         Args:
             grpc_channel: gRPC channel to use for communication with the API.
             svc_addr: Address of the service to connect to.
+            key: API key to use for authentication.
         """
         self._svc_addr = svc_addr
         self._stub = dispatch_pb2_grpc.MicrogridDispatchServiceStub(grpc_channel)
+        self._metadata = (("key", key),)
 
     # pylint: disable=too-many-arguments, too-many-locals
     async def list(
@@ -64,7 +68,7 @@ async def list(
 
         ```python
         grpc_channel = grpc.aio.insecure_channel("example")
-        client = Client(grpc_channel, "localhost:50051")
+        client = Client(grpc_channel=grpc_channel, svc_addr="localhost:50051", key="key")
         async for dispatch in client.list(microgrid_id=1):
             print(dispatch)
         ```
@@ -108,7 +112,9 @@ async def list(
         )
         request = DispatchListRequest(microgrid_id=microgrid_id, filter=filters)
 
-        response = await self._stub.ListMicrogridDispatches(request)  # type: ignore
+        response = await self._stub.ListMicrogridDispatches(
+            request, metadata=self._metadata
+        )  # type: ignore
         for dispatch in response.dispatches:
             yield Dispatch.from_protobuf(dispatch)
 
@@ -166,7 +172,9 @@ async def create(
             recurrence=recurrence or RecurrenceRule(),
         )
 
-        await self._stub.CreateMicrogridDispatch(request.to_protobuf())  # type: ignore
+        await self._stub.CreateMicrogridDispatch(
+            request.to_protobuf(), metadata=self._metadata
+        )  # type: ignore
 
         if dispatch := await self._try_fetch_created_dispatch(request):
             return dispatch
@@ -246,7 +254,9 @@ async def update(
 
             msg.update_mask.paths.append(key)
 
-        await self._stub.UpdateMicrogridDispatch(msg)  # type: ignore
+        await self._stub.UpdateMicrogridDispatch(
+            msg, metadata=self._metadata
+        )  # type: ignore
 
     async def get(self, dispatch_id: int) -> Dispatch:
         """Get a dispatch.
@@ -258,7 +268,9 @@ async def get(self, dispatch_id: int) -> Dispatch:
             Dispatch: The dispatch.
         """
         request = DispatchGetRequest(id=dispatch_id)
-        response = await self._stub.GetMicrogridDispatch(request)  # type: ignore
+        response = await self._stub.GetMicrogridDispatch(
+            request, metadata=self._metadata
+        )  # type: ignore
         return Dispatch.from_protobuf(response)
 
     async def delete(self, dispatch_id: int) -> None:
@@ -268,7 +280,9 @@ async def delete(self, dispatch_id: int) -> None:
             dispatch_id: The dispatch_id to delete.
         """
         request = DispatchDeleteRequest(id=dispatch_id)
-        await self._stub.DeleteMicrogridDispatch(request)  # type: ignore
+        await self._stub.DeleteMicrogridDispatch(
+            request, metadata=self._metadata
+        )  # type: ignore
 
     async def _try_fetch_created_dispatch(
         self,
diff --git a/src/frequenz/client/dispatch/test/_service.py b/src/frequenz/client/dispatch/test/_service.py
@@ -35,6 +35,12 @@
 from .._internal_types import DispatchCreateRequest
 from ..types import Dispatch
 
+ALL_KEY = "all"
+"""Key that has access to all resources in the FakeService."""
+
+NONE_KEY = "none"
+"""Key that has no access to any resources in the FakeService."""
+
 
 @dataclass
 class FakeService:
@@ -49,18 +55,59 @@ class FakeService:
     _shuffle_after_create: bool = False
     """Whether to shuffle the dispatches after creating them."""
 
+    def _check_access(self, metadata: grpc.aio.Metadata) -> None:
+        """Check if the access key is valid.
+
+        Args:
+            metadata: The metadata.
+
+        Raises:
+            grpc.RpcError: If the access key is invalid.
+        """
+        # metadata is a weird tuple of tuples, we don't like it
+        metadata_dict = dict(metadata)
+
+        if "key" not in metadata_dict:
+            raise grpc.RpcError(
+                grpc.StatusCode.UNAUTHENTICATED,
+                "No access key provided",
+            )
+
+        key = metadata_dict["key"]
+
+        if key is None:
+            raise grpc.RpcError(
+                grpc.StatusCode.UNAUTHENTICATED,
+                "No access key provided",
+            )
+
+        if key == NONE_KEY:
+            raise grpc.RpcError(
+                grpc.StatusCode.PERMISSION_DENIED,
+                "Permission denied",
+            )
+
+        if key != ALL_KEY:
+            raise grpc.RpcError(
+                grpc.StatusCode.UNAUTHENTICATED,
+                "Invalid access key",
+            )
+
     # pylint: disable=invalid-name
     async def ListMicrogridDispatches(
-        self, request: PBDispatchListRequest
+        self, request: PBDispatchListRequest, metadata: grpc.aio.Metadata
     ) -> DispatchList:
         """List microgrid dispatches.
 
         Args:
             request: The request.
+            metadata: The metadata.
 
         Returns:
             The dispatch list.
         """
+        self._check_access(metadata)
+
         return DispatchList(
             dispatches=map(
                 lambda d: d.to_protobuf(),
@@ -108,8 +155,10 @@ def _filter_dispatch(dispatch: Dispatch, request: PBDispatchListRequest) -> bool
     async def CreateMicrogridDispatch(
         self,
         request: PBDispatchCreateRequest,
+        metadata: grpc.aio.Metadata,
     ) -> Empty:
         """Create a new dispatch."""
+        self._check_access(metadata)
         self._last_id += 1
 
         self.dispatches.append(
@@ -128,8 +177,10 @@ async def CreateMicrogridDispatch(
     async def UpdateMicrogridDispatch(
         self,
         request: DispatchUpdateRequest,
+        metadata: grpc.aio.Metadata,
     ) -> Empty:
         """Update a dispatch."""
+        self._check_access(metadata)
         index = next(
             (i for i, d in enumerate(self.dispatches) if d.id == request.id),
             None,
@@ -194,8 +245,10 @@ async def UpdateMicrogridDispatch(
     async def GetMicrogridDispatch(
         self,
         request: DispatchGetRequest,
+        metadata: grpc.aio.Metadata,
     ) -> PBDispatch:
         """Get a single dispatch."""
+        self._check_access(metadata)
         dispatch = next((d for d in self.dispatches if d.id == request.id), None)
 
         if dispatch is None:
@@ -211,8 +264,10 @@ async def GetMicrogridDispatch(
     async def DeleteMicrogridDispatch(
         self,
         request: DispatchDeleteRequest,
+        metadata: grpc.aio.Metadata,
     ) -> Empty:
         """Delete a given dispatch."""
+        self._check_access(metadata)
         num_dispatches = len(self.dispatches)
         self.dispatches = [d for d in self.dispatches if d.id != request.id]
 
diff --git a/src/frequenz/client/dispatch/test/client.py b/src/frequenz/client/dispatch/test/client.py
diff --git a/tests/test_dispatch_cli.py b/tests/test_dispatch_cli.py
