iqm-station-control-client 11.3.1__py3-none-any.whl → 12.0.0__py3-none-any.whl

iqm/station_control/client/authentication.py

@@ -0,0 +1,239 @@
+# Copyright 2024 IQM client developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This module contains user authentication related classes and functions required by IQMClient."""
+
+from abc import ABC, abstractmethod
+from base64 import b64decode
+from collections.abc import Callable
+import json
+import os
+import time
+from typing import Any, TypeAlias
+
+from iqm.iqm_server_client.errors import ClientAuthenticationError, ClientConfigurationError
+
+REFRESH_MARGIN_SECONDS = 60
+
+
+AuthHeaderCallback: TypeAlias = Callable[[], str]
+"""Function that returns an authorization header containing a bearer token."""
+
+
+class TokenManager:
+    """TokenManager manages the access token required for user authentication.
+
+    Args:
+        token: Long-lived IQM token in plain text format.
+        tokens_file: Path to a tokens file used for authentication.
+        auth_header_callback: Callback function that returns an Authorization header containing a bearer token.
+        use_env_vars: Iff True, the first two parameters can also be read from the environment variables
+            :envvar:`IQM_TOKEN` and :envvar:`IQM_TOKENS_FILE`, respectively.
+            Environment variables can not be mixed with initialisation arguments.
+
+    At most one auth parameter should be given.
+
+    """
+
+    @staticmethod
+    def time_left_seconds(token: Any) -> int:
+        """Check how much time is left until the token expires.
+
+        Returns:
+            Time left on token in seconds.
+
+        """
+        if not token or not isinstance(token, str):
+            return 0
+        parts = token.split(".", 2)
+        if len(parts) != 3:
+            return 0
+        # Add padding to adjust body length to a multiple of 4 chars as required by base64 decoding
+        try:
+            body = parts[1] + ("=" * (-len(parts[1]) % 4))
+            exp_time = int(json.loads(b64decode(body)).get("exp", "0"))
+            return max(0, exp_time - int(time.time()))
+        except (UnicodeDecodeError, json.decoder.JSONDecodeError, ValueError, TypeError):
+            return 0
+
+    def __init__(
+        self,
+        token: str | None = None,
+        tokens_file: str | None = None,
+        auth_header_callback: AuthHeaderCallback | None = None,
+        *,
+        use_env_vars: bool = True,
+    ):
+        def _format_names(variable_names: list[str]) -> str:
+            """Format a list of variable names"""
+            return ", ".join(f'"{name}"' for name in variable_names)
+
+        auth_parameters: dict[str, str] = {}
+
+        init_parameters = {"token": token, "tokens_file": tokens_file}
+        init_params_given = [key for key, value in init_parameters.items() if value]
+        if auth_header_callback:
+            init_params_given.append("auth_header_callback")
+
+        if use_env_vars:
+            env_variables = {"token": "IQM_TOKEN", "tokens_file": "IQM_TOKENS_FILE"}
+            env_vars_given = [name for name in env_variables.values() if os.environ.get(name)]
+        else:
+            env_variables = {}
+            env_vars_given = []
+
+        if init_params_given and env_vars_given:
+            raise ClientConfigurationError(
+                "Authentication parameters given both as initialisation args and as environment variables: "
+                + f"initialisation args {_format_names(init_params_given)}, "
+                + f"environment variables {_format_names(env_vars_given)}."
+                + " Parameter sources must not be mixed."
+            )
+        auth_params_given = init_params_given + env_vars_given
+        if len(auth_params_given) > 1:
+            raise ClientConfigurationError(
+                f"No more than one authentication parameter may be given, received {auth_params_given}"
+            )
+
+        self._token_provider: TokenProviderInterface | None = None
+        self._auth_header_callback: AuthHeaderCallback | None = None
+        self._access_token: str | None = None
+
+        if auth_header_callback:
+            self._auth_header_callback = auth_header_callback
+            return
+
+        if env_vars_given:
+            auth_parameters = {key: value for key, name in env_variables.items() if (value := os.environ.get(name))}
+        else:
+            auth_parameters = {key: value for key, value in init_parameters.items() if value}
+
+        if not auth_parameters:
+            return  # no authentication
+        if set(auth_parameters) == {"token"}:
+            # This is not necessarily a JWT token
+            self._token_provider = ExternalToken(auth_parameters["token"])
+        elif set(auth_parameters) == {"tokens_file"}:
+            self._token_provider = TokensFileReader(auth_parameters["tokens_file"])
+        else:
+            raise ClientConfigurationError("Not possible.")
+
+    def get_auth_header_callback(self) -> AuthHeaderCallback | None:
+        """Return a callback providing an Authorization header, or None if we cannot provide it."""
+        return self._get_bearer_token if self._token_provider else self._auth_header_callback
+
+    def _get_bearer_token(self, retries: int = 1) -> str:
+        """Return a valid bearer token.
+
+        Raises:
+            ClientAuthenticationError: getting the token failed
+            RuntimeError: There is no token provider (authentication disabled)
+
+        """
+        if self._token_provider is None:
+            raise RuntimeError("There is no token provider.")
+
+        # Use the existing access token if it is still valid
+        if TokenManager.time_left_seconds(self._access_token) > REFRESH_MARGIN_SECONDS:
+            return f"Bearer {self._access_token}"
+
+        # Otherwise, get a new access token from token provider
+        try:
+            self._access_token = self._token_provider.get_token()
+            return f"Bearer {self._access_token}"
+        except ClientAuthenticationError:
+            if retries < 1:
+                raise
+
+        # Try again
+        return self._get_bearer_token(retries - 1)
+
+    def close(self) -> bool:
+        """Close the configured token provider.
+
+        Returns:
+            True if closing was successful
+
+        Raises:
+            ClientAuthenticationError: closing failed
+
+        """
+        if self._token_provider is None:
+            return False
+
+        self._token_provider.close()
+        self._token_provider = None
+        return True
+
+
+class TokenProviderInterface(ABC):
+    """Interface to token provider"""
+
+    @abstractmethod
+    def get_token(self) -> str:
+        """Return a valid access token.
+
+        Raises:
+            ClientAuthenticationError: acquiring the token failed
+
+        """
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the authentication session.
+
+        Raises:
+            ClientAuthenticationError: closing the session failed
+
+        """
+
+
+class ExternalToken(TokenProviderInterface):
+    """Holds an external token"""
+
+    def __init__(self, token: str):
+        self._token: str | None = token
+
+    def get_token(self) -> str:
+        if self._token is None:
+            raise ClientAuthenticationError("No external token available")
+        return self._token
+
+    def close(self) -> None:
+        self._token = None
+        raise ClientAuthenticationError("Can not close externally managed auth session")
+
+
+class TokensFileReader(TokenProviderInterface):
+    """Reads token from a file"""
+
+    def __init__(self, tokens_file: str):
+        self._path: str | None = tokens_file
+
+    def get_token(self) -> str:
+        try:
+            if self._path is None:
+                raise ClientAuthenticationError("No tokens file available")
+            with open(self._path, encoding="utf-8") as file:
+                raw_data = file.read()
+            json_data = json.loads(raw_data)
+            token = json_data.get("access_token")
+            if TokenManager.time_left_seconds(token) <= 0:
+                raise ClientAuthenticationError("Access token in file has expired or is not valid")
+        except (FileNotFoundError, IsADirectoryError, json.decoder.JSONDecodeError) as e:
+            raise ClientAuthenticationError(rf"Failed to read access token from file '{self._path}': {e}") from e
+        return token
+
+    def close(self) -> None:
+        self._path = None
+        raise ClientAuthenticationError("Can not close externally managed auth session")

iqm/station_control/client/iqm_server/error.py

@@ -1,30 +0,0 @@
-# Copyright 2025 IQM
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from exa.common.errors.station_control_errors import StationControlError
-
-
-class IqmServerError(StationControlError):
-    def __init__(self, message: str, status_code: str, error_code: str | None = None, details: dict | None = None):
-        super().__init__(message)
-        self.status_code = status_code
-        self.error_code = error_code
-        self.details = details
-
-    def __str__(self):
-        s = f"{self.message} (status_code = {self.status_code}, error_code = {self.error_code}"
-        if details := self.details:
-            s += f", details = {details}"
-        s += ")"
-        return s

iqm/station_control/client/iqm_server/grpc_utils.py

@@ -1,156 +0,0 @@
-# Copyright 2025 IQM
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Internal utility functions used by IqmServerClient."""
-
-from collections.abc import Callable, Iterable
-from dataclasses import dataclass
-from datetime import datetime
-import uuid
-
-from google.protobuf import json_format, struct_pb2, timestamp_pb2
-import grpc
-from grpc import Compression
-from pydantic import HttpUrl
-
-from iqm.station_control.client.iqm_server import proto
-from iqm.station_control.client.iqm_server.error import IqmServerError
-from iqm.station_control.interface.models.type_aliases import StrUUID
-
-
-class ClientCallDetails(grpc.ClientCallDetails):
-    def __init__(self, details):  # noqa: ANN001
-        self.method = details.method
-        self.metadata = list(details.metadata or [])
-        self.timeout = details.timeout
-        self.credentials = details.credentials
-        self.wait_for_ready = details.wait_for_ready
-        self.compression = details.compression
-
-
-class ApiTokenAuth(grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor):
-    def __init__(self, get_token_callback: Callable[[], str]):
-        self.get_token_callback = get_token_callback
-
-    def _add_auth_header(self, client_call_details) -> ClientCallDetails:  # noqa: ANN001
-        details = ClientCallDetails(client_call_details)
-        details.metadata.append(("authorization", self.get_token_callback()))
-        return details
-
-    def intercept_unary_stream(self, continuation, client_call_details, request):  # noqa: ANN001, ANN201
-        return continuation(self._add_auth_header(client_call_details), request)
-
-    def intercept_unary_unary(self, continuation, client_call_details, request):  # noqa: ANN001, ANN201
-        return continuation(self._add_auth_header(client_call_details), request)
-
-
-@dataclass(frozen=True, kw_only=True)
-class ConnectionParameters:
-    server_address: str
-    is_secure: bool
-    quantum_computer: str
-    use_timeslot: bool
-
-
-def parse_connection_params(qc_url: str) -> ConnectionParameters:
-    # Security measure: mitigate UTF-8 read order control character
-    # exploits by allowing only ASCII urls
-    if not qc_url.isascii():
-        raise ValueError("Invalid quantum computer URL")
-
-    # IQM Server QC urls are now form "https://cocos.<server_base_url>/<qc_name>[:timeslot]"
-    # In the future, "cocos." subdomain will be dropped. The parsing logic should work with
-    # the both url formats
-    url = HttpUrl(qc_url)
-    qc_name = (url.path or "").split("/")[-1].removesuffix(":timeslot")
-    use_timeslot = qc_url.endswith(":timeslot")
-    if not qc_name:
-        raise ValueError("Invalid quantum computer URL: device name is missing")
-
-    is_secure = url.scheme == "https"
-    hostname = (url.host or "").removeprefix("cocos.")
-    port = url.port or (443 if is_secure else 80)
-
-    return ConnectionParameters(
-        server_address=f"{hostname}:{port}",
-        is_secure=is_secure,
-        quantum_computer=qc_name,
-        use_timeslot=use_timeslot,
-    )
-
-
-def create_channel(
-    connection_params: ConnectionParameters,
-    get_token_callback: Callable[[], str] | None = None,
-    enable_compression: bool = True,
-) -> grpc.Channel:
-    compression = Compression.Gzip if enable_compression else None
-    options = [
-        # Let's try to parametrize this at least when we're merging station-control-client and iqm-client
-        ("grpc.keepalive_time_ms", 5000),
-        ("grpc.keepalive_permit_without_calls", 1),
-        ("grpc.http2.max_pings_without_data", 0),
-        ("grpc.keepalive_timeout_ms", 1000),
-    ]
-    address = connection_params.server_address
-    channel = (
-        grpc.secure_channel(
-            address, credentials=grpc.ssl_channel_credentials(), options=options, compression=compression
-        )
-        if connection_params.is_secure
-        else grpc.insecure_channel(address, options=options, compression=compression)
-    )
-    if get_token_callback is not None:
-        channel = grpc.intercept_channel(channel, ApiTokenAuth(get_token_callback))
-    return channel
-
-
-def to_proto_uuid(value: StrUUID) -> proto.Uuid:
-    if isinstance(value, str):
-        value = uuid.UUID(value)
-    return proto.Uuid(raw=value.bytes)
-
-
-def from_proto_uuid(value: proto.Uuid) -> uuid.UUID:
-    if value.WhichOneof("data") == "str":
-        return uuid.UUID(hex=value.str)
-    return uuid.UUID(bytes=value.raw)
-
-
-def to_datetime(timestamp: timestamp_pb2.Timestamp) -> datetime:
-    return timestamp.ToDatetime()
-
-
-def load_all(chunks: Iterable[proto.DataChunk]) -> bytes:
-    result = bytearray()
-    for chunk in chunks:
-        result.extend(chunk.data)
-    return bytes(result)
-
-
-def extract_error(error: grpc.RpcError, title: str | None = None) -> IqmServerError:
-    message = error.details()
-    status_code = str(error.code().name)
-    metadata = {k: v for k, v in list(error.initial_metadata()) + list(error.trailing_metadata())}
-    error_code = str(metadata.get("error_code")) if "error_code" in metadata else None
-    details = None
-    if details_bin := metadata.get("grpc-status-details-bin"):
-        value_proto = struct_pb2.Value()
-        value_proto.ParseFromString(details_bin)
-        details = json_format.MessageToJson(value_proto)
-    return IqmServerError(
-        message=f"{title}: {message}" if title else message,
-        status_code=status_code,
-        error_code=error_code,
-        details=details,  # type: ignore[arg-type]
-    )