flwr 1.22.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

flwr/cli/utils.py

@@ -26,16 +26,18 @@ from typing import Any, Callable, Optional, Union, cast
 import grpc
 import typer
 
-from flwr.cli.cli_user_auth_interceptor import CliUserAuthInterceptor
-from flwr.common.auth_plugin import CliAuthPlugin
 from flwr.common.constant import (
-    AUTH_TYPE_JSON_KEY,
+    AUTHN_TYPE_JSON_KEY,
     CREDENTIALS_DIR,
     FLWR_DIR,
+    NO_ACCOUNT_AUTH_MESSAGE,
     NO_ARTIFACT_PROVIDER_MESSAGE,
-    NO_USER_AUTH_MESSAGE,
+    NODE_NOT_FOUND_MESSAGE,
+    PUBLIC_KEY_ALREADY_IN_USE_MESSAGE,
+    PUBLIC_KEY_NOT_VALID,
     PULL_UNFINISHED_RUN_MESSAGE,
     RUN_ID_NOT_FOUND_MESSAGE,
+    AuthnType,
 )
 from flwr.common.grpc import (
     GRPC_MAX_MESSAGE_LENGTH,
@@ -43,7 +45,8 @@ from flwr.common.grpc import (
     on_channel_state_change,
 )
 
-from .auth_plugin import get_cli_auth_plugins
+from .auth_plugin import CliAuthPlugin, get_cli_plugin_class
+from .cli_account_auth_interceptor import CliAccountAuthInterceptor
 from .config_utils import validate_certificate_in_federation_config
 
 
@@ -166,8 +169,8 @@ def get_sha256_hash(file_path_or_int: Union[Path, int]) -> str:
     return sha256.hexdigest()
 
 
-def get_user_auth_config_path(root_dir: Path, federation: str) -> Path:
-    """Return the path to the user auth config file.
+def get_account_auth_config_path(root_dir: Path, federation: str) -> Path:
+    """Return the path to the account auth config file.
 
     Additionally, a `.gitignore` file will be created in the Flower directory to
     include the `.credentials` folder to be excluded from git. If the `.gitignore`
@@ -217,71 +220,68 @@ def get_user_auth_config_path(root_dir: Path, federation: str) -> Path:
     return credentials_dir / f"{federation}.json"
 
 
-def try_obtain_cli_auth_plugin(
+def account_auth_enabled(federation_config: dict[str, Any]) -> bool:
+    """Check if account authentication is enabled in the federation config."""
+    enabled: bool = federation_config.get("enable-user-auth", False)
+    enabled |= federation_config.get("enable-account-auth", False)
+    if "enable-user-auth" in federation_config:
+        typer.secho(
+            "`enable-user-auth` is deprecated and will be removed in a future "
+            "release. Please use `enable-account-auth` instead.",
+            fg=typer.colors.YELLOW,
+            bold=True,
+        )
+    return enabled
+
+
+def retrieve_authn_type(config_path: Path) -> str:
+    """Retrieve the auth type from the config file or return NOOP if not found."""
+    try:
+        with config_path.open("r", encoding="utf-8") as file:
+            json_file = json.load(file)
+        authn_type: str = json_file[AUTHN_TYPE_JSON_KEY]
+        return authn_type
+    except (FileNotFoundError, KeyError):
+        return AuthnType.NOOP
+
+
+def load_cli_auth_plugin(
     root_dir: Path,
     federation: str,
     federation_config: dict[str, Any],
-    auth_type: Optional[str] = None,
-) -> Optional[CliAuthPlugin]:
-    """Load the CLI-side user auth plugin for the given auth type."""
-    # Check if user auth is enabled
-    if not federation_config.get("enable-user-auth", False):
-        return None
-
-    config_path = get_user_auth_config_path(root_dir, federation)
-
-    # Get the auth type from the config if not provided
-    # auth_type will be None for all CLI commands except login
-    if auth_type is None:
-        try:
-            with config_path.open("r", encoding="utf-8") as file:
-                json_file = json.load(file)
-            auth_type = json_file[AUTH_TYPE_JSON_KEY]
-        except (FileNotFoundError, KeyError):
-            typer.secho(
-                "❌ Missing or invalid credentials for user authentication. "
-                "Please run `flwr login` to authenticate.",
-                fg=typer.colors.RED,
-                bold=True,
-            )
-            raise typer.Exit(code=1) from None
+    authn_type: Optional[str] = None,
+) -> CliAuthPlugin:
+    """Load the CLI-side account auth plugin for the given authn type."""
+    # Find the path to the account auth config file
+    config_path = get_account_auth_config_path(root_dir, federation)
+
+    # Determine the auth type if not provided
+    # Only `flwr login` command can provide `authn_type` explicitly, as it can query the
+    # SuperLink for the auth type.
+    if authn_type is None:
+        authn_type = AuthnType.NOOP
+        if account_auth_enabled(federation_config):
+            authn_type = retrieve_authn_type(config_path)
 
     # Retrieve auth plugin class and instantiate it
     try:
-        all_plugins: dict[str, type[CliAuthPlugin]] = get_cli_auth_plugins()
-        auth_plugin_class = all_plugins[auth_type]
+        auth_plugin_class = get_cli_plugin_class(authn_type)
         return auth_plugin_class(config_path)
-    except KeyError:
-        typer.echo(f"❌ Unknown user authentication type: {auth_type}")
-        raise typer.Exit(code=1) from None
-    except ImportError:
-        typer.echo("❌ No authentication plugins are currently supported.")
+    except ValueError:
+        typer.echo(f"❌ Unknown account authentication type: {authn_type}")
         raise typer.Exit(code=1) from None
 
 
 def init_channel(
-    app: Path, federation_config: dict[str, Any], auth_plugin: Optional[CliAuthPlugin]
+    app: Path, federation_config: dict[str, Any], auth_plugin: CliAuthPlugin
 ) -> grpc.Channel:
     """Initialize gRPC channel to the Control API."""
     insecure, root_certificates_bytes = validate_certificate_in_federation_config(
         app, federation_config
     )
 
-    # Initialize the CLI-side user auth interceptor
-    interceptors: list[grpc.UnaryUnaryClientInterceptor] = []
-    if auth_plugin is not None:
-        # Check if TLS is enabled. If not, raise an error
-        if insecure:
-            typer.secho(
-                "❌ User authentication requires TLS to be enabled. "
-                "Remove `insecure = true` from the federation configuration.",
-                fg=typer.colors.RED,
-                bold=True,
-            )
-            raise typer.Exit(code=1)
-
-        auth_plugin.load_tokens()
-        interceptors.append(CliUserAuthInterceptor(auth_plugin))
+    # Load tokens
+    auth_plugin.load_tokens()
 
     # Create the gRPC channel
     channel = create_channel(
@@ -289,14 +289,14 @@ def init_channel(
         insecure=insecure,
         root_certificates=root_certificates_bytes,
         max_message_length=GRPC_MAX_MESSAGE_LENGTH,
-        interceptors=interceptors or None,
+        interceptors=[CliAccountAuthInterceptor(auth_plugin)],
     )
     channel.subscribe(on_channel_state_change)
     return channel
 
 
 @contextmanager
-def flwr_cli_grpc_exc_handler() -> Iterator[None]:
+def flwr_cli_grpc_exc_handler() -> Iterator[None]:  # pylint: disable=too-many-branches
     """Context manager to handle specific gRPC errors.
 
     It catches grpc.RpcError exceptions with UNAUTHENTICATED, UNIMPLEMENTED,
@@ -315,9 +315,9 @@ def flwr_cli_grpc_exc_handler() -> Iterator[None]:
             )
             raise typer.Exit(code=1) from None
         if e.code() == grpc.StatusCode.UNIMPLEMENTED:
-            if e.details() == NO_USER_AUTH_MESSAGE:  # pylint: disable=E1101
+            if e.details() == NO_ACCOUNT_AUTH_MESSAGE:  # pylint: disable=E1101
                 typer.secho(
-                    "❌ User authentication is not enabled on this SuperLink.",
+                    "❌ Account authentication is not enabled on this SuperLink.",
                     fg=typer.colors.RED,
                     bold=True,
                 )
@@ -354,16 +354,21 @@ def flwr_cli_grpc_exc_handler() -> Iterator[None]:
                 bold=True,
             )
             raise typer.Exit(code=1) from None
-        if (
-            e.code() == grpc.StatusCode.NOT_FOUND
-            and e.details() == RUN_ID_NOT_FOUND_MESSAGE  # pylint: disable=E1101
-        ):
-            typer.secho(
-                "❌ Run ID not found.",
-                fg=typer.colors.RED,
-                bold=True,
-            )
-            raise typer.Exit(code=1) from None
+        if e.code() == grpc.StatusCode.NOT_FOUND:
+            if e.details() == RUN_ID_NOT_FOUND_MESSAGE:  # pylint: disable=E1101
+                typer.secho(
+                    "❌ Run ID not found.",
+                    fg=typer.colors.RED,
+                    bold=True,
+                )
+                raise typer.Exit(code=1) from None
+            if e.details() == NODE_NOT_FOUND_MESSAGE:  # pylint: disable=E1101
+                typer.secho(
+                    "❌ Node ID not found for this account.",
+                    fg=typer.colors.RED,
+                    bold=True,
+                )
+                raise typer.Exit(code=1) from None
         if e.code() == grpc.StatusCode.FAILED_PRECONDITION:
             if e.details() == PULL_UNFINISHED_RUN_MESSAGE:  # pylint: disable=E1101
                 typer.secho(
@@ -373,4 +378,22 @@ def flwr_cli_grpc_exc_handler() -> Iterator[None]:
                     bold=True,
                 )
                 raise typer.Exit(code=1) from None
+            if (
+                e.details() == PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
+            ):  # pylint: disable=E1101
+                typer.secho(
+                    "❌ The provided public key is already in use by another "
+                    "SuperNode.",
+                    fg=typer.colors.RED,
+                    bold=True,
+                )
+                raise typer.Exit(code=1) from None
+            if e.details() == PUBLIC_KEY_NOT_VALID:  # pylint: disable=E1101
+                typer.secho(
+                    "❌ The provided public key is invalid. Please provide a valid "
+                    "NIST EC public key.",
+                    fg=typer.colors.RED,
+                    bold=True,
+                )
+                raise typer.Exit(code=1) from None
         raise

flwr/client/__init__.py

@@ -15,10 +15,11 @@
 """Flower client."""
 
 
+from flwr.clientapp import ClientApp
+
 from ..compat.client.app import start_client as start_client  # Deprecated
 from ..compat.client.app import start_numpy_client as start_numpy_client  # Deprecated
 from .client import Client as Client
-from .client_app import ClientApp as ClientApp
 from .numpy_client import NumPyClient as NumPyClient
 from .typing import ClientFn as ClientFn
 from .typing import ClientFnExt as ClientFnExt

flwr/client/grpc_adapter_client/connection.py

@@ -44,10 +44,9 @@ def grpc_adapter(  # pylint: disable=R0913,too-many-positional-arguments
     ] = None,
 ) -> Iterator[
     tuple[
+        int,
         Callable[[], Optional[tuple[Message, ObjectTree]]],
         Callable[[Message, ObjectTree], set[str]],
-        Callable[[], Optional[int]],
-        Callable[[], None],
         Callable[[int], Run],
         Callable[[str, int], Fab],
         Callable[[int, str], bytes],
@@ -77,22 +76,21 @@ def grpc_adapter(  # pylint: disable=R0913,too-many-positional-arguments
         connection using the certificates will be established to an SSL-enabled
         Flower server. Bytes won't work for the REST API.
     authentication_keys : Optional[Tuple[PrivateKey, PublicKey]] (default: None)
-        Client authentication is not supported for this transport type.
+        SuperNode authentication is not supported for this transport type.
 
     Returns
     -------
+    node_id : int
     receive : Callable[[], Optional[tuple[Message, ObjectTree]]]
     send : Callable[[Message, ObjectTree], set[str]]
-    create_node : Optional[Callable]
-    delete_node : Optional[Callable]
-    get_run : Optional[Callable]
-    get_fab : Optional[Callable]
+    get_run : Callable[[int], Run]
+    get_fab : Callable[[str, int], Fab]
     pull_object : Callable[[str], bytes]
     push_object : Callable[[str, bytes], None]
     confirm_message_received : Callable[[str], None]
     """
     if authentication_keys is not None:
-        log(ERROR, "Client authentication is not supported for this transport type.")
+        log(ERROR, "SuperNode authentication is not supported for this transport type.")
     with grpc_request_response(
         server_address=server_address,
         insecure=insecure,

flwr/client/grpc_rere_client/connection.py

@@ -19,7 +19,7 @@ from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
 from logging import ERROR
 from pathlib import Path
-from typing import Callable, Optional, Union, cast
+from typing import Callable, Optional, Union
 
 import grpc
 from cryptography.hazmat.primitives.asymmetric import ec
@@ -36,19 +36,24 @@ from flwr.common.inflatable_protobuf_utils import (
 from flwr.common.logger import log
 from flwr.common.message import Message, remove_content_from_message
 from flwr.common.retry_invoker import RetryInvoker, _wrap_stub
-from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    generate_key_pairs,
+from flwr.common.serde import (
+    fab_from_proto,
+    message_from_proto,
+    message_to_proto,
+    run_from_proto,
 )
-from flwr.common.serde import message_from_proto, message_to_proto, run_from_proto
 from flwr.common.typing import Fab, Run
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    DeleteNodeRequest,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
+    RegisterNodeFleetRequest,
+    UnregisterNodeFleetRequest,
 )
 from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
 from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
@@ -58,9 +63,10 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
 from flwr.proto.message_pb2 import ObjectTree  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
+from flwr.supercore.primitives.asymmetric import generate_key_pairs, public_key_to_bytes
 
-from .client_interceptor import AuthenticateClientInterceptor
 from .grpc_adapter import GrpcAdapter
+from .node_auth_client_interceptor import NodeAuthClientInterceptor
 
 
 @contextmanager
@@ -76,10 +82,9 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     adapter_cls: Optional[Union[type[FleetStub], type[GrpcAdapter]]] = None,
 ) -> Iterator[
     tuple[
+        int,
         Callable[[], Optional[tuple[Message, ObjectTree]]],
         Callable[[Message, ObjectTree], set[str]],
-        Callable[[], Optional[int]],
-        Callable[[], None],
         Callable[[int], Run],
         Callable[[str, int], Fab],
         Callable[[int, str], bytes],
@@ -122,11 +127,11 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
 
     Returns
     -------
-    receive : Callable
-    send : Callable
-    create_node : Optional[Callable]
-    delete_node : Optional[Callable]
-    get_run : Optional[Callable]
+    node_id : int
+    receive : Callable[[], Optional[tuple[Message, ObjectTree]]]
+    send : Callable[[Message, ObjectTree], set[str]]
+    get_run : Callable[[int], Run]
+    get_fab : Callable[[str, int], Fab]
     pull_object : Callable[[str], bytes]
     push_object : Callable[[str, bytes], None]
     confirm_message_received : Callable[[str], None]
@@ -135,13 +140,16 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         root_certificates = Path(root_certificates).read_bytes()
 
     # Automatic node auth: generate keys if user didn't provide any
+    self_registered = False
     if authentication_keys is None:
+        self_registered = True
         authentication_keys = generate_key_pairs()
 
     # Always configure auth interceptor, with either user-provided or generated keys
     interceptors: Sequence[grpc.UnaryUnaryClientInterceptor] = [
-        AuthenticateClientInterceptor(*authentication_keys),
+        NodeAuthClientInterceptor(*authentication_keys),
     ]
+    node_pk = public_key_to_bytes(authentication_keys[1])
     channel = create_channel(
         server_address=server_address,
         insecure=insecure,
@@ -160,7 +168,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     # Wrap stub
     _wrap_stub(stub, retry_invoker)
     ###########################################################################
-    # send_node_heartbeat/create_node/delete_node/receive/send/get_run functions
+    # SuperNode functions
     ###########################################################################
 
     def send_node_heartbeat() -> bool:
@@ -197,22 +205,26 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
 
     heartbeat_sender = HeartbeatSender(send_node_heartbeat)
 
-    def create_node() -> Optional[int]:
-        """Set create_node."""
-        # Call FleetAPI
-        create_node_request = CreateNodeRequest(
-            heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL
+    def register_node() -> None:
+        """Register node with SuperLink."""
+        stub.RegisterNode(RegisterNodeFleetRequest(public_key=node_pk))
+
+    def activate_node() -> int:
+        """Activate node and start heartbeat."""
+        req = ActivateNodeRequest(
+            public_key=node_pk,
+            heartbeat_interval=HEARTBEAT_DEFAULT_INTERVAL,
         )
-        create_node_response = stub.CreateNode(request=create_node_request)
+        res: ActivateNodeResponse = stub.ActivateNode(req)
 
         # Remember the node and start the heartbeat sender
         nonlocal node
-        node = cast(Node, create_node_response.node)
+        node = Node(node_id=res.node_id)
         heartbeat_sender.start()
         return node.node_id
 
-    def delete_node() -> None:
-        """Set delete_node."""
+    def deactivate_node() -> None:
+        """Deactivate node and stop heartbeat."""
         # Get Node
         nonlocal node
         if node is None:
@@ -223,8 +235,20 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         heartbeat_sender.stop()
 
         # Call FleetAPI
-        delete_node_request = DeleteNodeRequest(node=node)
-        stub.DeleteNode(request=delete_node_request)
+        req = DeactivateNodeRequest(node_id=node.node_id)
+        stub.DeactivateNode(req)
+
+    def unregister_node() -> None:
+        """Unregister node from SuperLink."""
+        # Get Node
+        nonlocal node
+        if node is None:
+            log(ERROR, "Node instance missing")
+            return
+
+        # Call FleetAPI
+        req = UnregisterNodeFleetRequest(node_id=node.node_id)
+        stub.UnregisterNode(req)
 
         # Cleanup
         node = None
@@ -289,7 +313,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         get_fab_request = GetFabRequest(node=node, hash_str=fab_hash, run_id=run_id)
         get_fab_response: GetFabResponse = stub.GetFab(request=get_fab_request)
 
-        return Fab(get_fab_response.fab.hash_str, get_fab_response.fab.content)
+        return fab_from_proto(get_fab_response.fab)
 
     def pull_object(run_id: int, object_id: str) -> bytes:
         """Pull the object from the SuperLink."""
@@ -331,12 +355,14 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
         fn(object_id)
 
     try:
+        if self_registered:
+            register_node()
+        node_id = activate_node()
         # Yield methods
         yield (
+            node_id,
             receive,
             send,
-            create_node,
-            delete_node,
             get_run,
             get_fab,
             pull_object,
@@ -351,7 +377,9 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
             if node is not None:
                 # Disable retrying
                 retry_invoker.max_tries = 1
-                delete_node()
+                deactivate_node()
+                if self_registered:
+                    unregister_node()
         except grpc.RpcError:
             pass
         channel.close()

flwr/client/grpc_rere_client/grpc_adapter.py

@@ -34,14 +34,18 @@ from flwr.common.constant import (
 from flwr.common.version import package_name, package_version
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
-    CreateNodeRequest,
-    CreateNodeResponse,
-    DeleteNodeRequest,
-    DeleteNodeResponse,
+    ActivateNodeRequest,
+    ActivateNodeResponse,
+    DeactivateNodeRequest,
+    DeactivateNodeResponse,
     PullMessagesRequest,
     PullMessagesResponse,
     PushMessagesRequest,
     PushMessagesResponse,
+    RegisterNodeFleetRequest,
+    RegisterNodeFleetResponse,
+    UnregisterNodeFleetRequest,
+    UnregisterNodeFleetResponse,
 )
 from flwr.proto.grpcadapter_pb2 import MessageContainer  # pylint: disable=E0611
 from flwr.proto.grpcadapter_pb2_grpc import GrpcAdapterStub
@@ -118,17 +122,29 @@ class GrpcAdapter:
         response.ParseFromString(container_res.grpc_message_content)
         return response
 
-    def CreateNode(  # pylint: disable=C0103
-        self, request: CreateNodeRequest, **kwargs: Any
-    ) -> CreateNodeResponse:
+    def RegisterNode(  # pylint: disable=C0103
+        self, request: RegisterNodeFleetRequest, **kwargs: Any
+    ) -> RegisterNodeFleetResponse:
         """."""
-        return self._send_and_receive(request, CreateNodeResponse, **kwargs)
+        return self._send_and_receive(request, RegisterNodeFleetResponse, **kwargs)
 
-    def DeleteNode(  # pylint: disable=C0103
-        self, request: DeleteNodeRequest, **kwargs: Any
-    ) -> DeleteNodeResponse:
+    def ActivateNode(  # pylint: disable=C0103
+        self, request: ActivateNodeRequest, **kwargs: Any
+    ) -> ActivateNodeResponse:
         """."""
-        return self._send_and_receive(request, DeleteNodeResponse, **kwargs)
+        return self._send_and_receive(request, ActivateNodeResponse, **kwargs)
+
+    def DeactivateNode(  # pylint: disable=C0103
+        self, request: DeactivateNodeRequest, **kwargs: Any
+    ) -> DeactivateNodeResponse:
+        """."""
+        return self._send_and_receive(request, DeactivateNodeResponse, **kwargs)
+
+    def UnregisterNode(  # pylint: disable=C0103
+        self, request: UnregisterNodeFleetRequest, **kwargs: Any
+    ) -> UnregisterNodeFleetResponse:
+        """."""
+        return self._send_and_receive(request, UnregisterNodeFleetResponse, **kwargs)
 
     def SendNodeHeartbeat(  # pylint: disable=C0103
         self, request: SendNodeHeartbeatRequest, **kwargs: Any

flwr/client/grpc_rere_client/{client_interceptor.py → node_auth_client_interceptor.py}

@@ -23,14 +23,11 @@ from google.protobuf.message import Message as GrpcMessage
 
 from flwr.common import now
 from flwr.common.constant import PUBLIC_KEY_HEADER, SIGNATURE_HEADER, TIMESTAMP_HEADER
-from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    public_key_to_bytes,
-    sign_message,
-)
+from flwr.supercore.primitives.asymmetric import public_key_to_bytes, sign_message
 
 
-class AuthenticateClientInterceptor(grpc.UnaryUnaryClientInterceptor):  # type: ignore
-    """Client interceptor for client authentication."""
+class NodeAuthClientInterceptor(grpc.UnaryUnaryClientInterceptor):  # type: ignore
+    """Client interceptor for node authentication."""
 
     def __init__(
         self,

flwr/client/mod/secure_aggregation/secaggplus_mod.py

@@ -35,14 +35,9 @@ from flwr.common.constant import MessageType
 from flwr.common.logger import log
 from flwr.common.secure_aggregation.crypto.shamir import create_shares
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (
-    bytes_to_private_key,
-    bytes_to_public_key,
     decrypt,
     encrypt,
-    generate_key_pairs,
     generate_shared_key,
-    private_key_to_bytes,
-    public_key_to_bytes,
 )
 from flwr.common.secure_aggregation.ndarrays_arithmetic import (
     factor_combine,
@@ -64,6 +59,13 @@ from flwr.common.secure_aggregation.secaggplus_utils import (
     share_keys_plaintext_separate,
 )
 from flwr.common.typing import ConfigRecordValues
+from flwr.supercore.primitives.asymmetric import (
+    bytes_to_private_key,
+    bytes_to_public_key,
+    generate_key_pairs,
+    private_key_to_bytes,
+    public_key_to_bytes,
+)
 
 
 @dataclass