Make a KeyPair dataclass for passing around key pairs

libp2p/__init__.py

@@ -34,7 +34,8 @@ async def cleanup_done_tasks() -> None:
 
 
 def generate_peer_id_from_rsa_identity() -> ID:
-    _, new_public_key = create_new_key_pair()
+    new_key_pair = create_new_key_pair()
+    new_public_key = new_key_pair.public_key
     new_id = ID.from_pubkey(new_public_key)
     return new_id
 

libp2p/crypto/keys.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from enum import Enum, unique
 
 from .pb import crypto_pb2 as protobuf
@@ -73,3 +74,9 @@ class PrivateKey(ABC, Key):
         protobuf_key.key_type = _type.value
         protobuf_key.data = data
         return protobuf_key
+
+
+@dataclass(frozen=True)
+class KeyPair:
+    private_key: PrivateKey
+    public_key: PublicKey

libp2p/crypto/rsa.py

@@ -1,9 +1,7 @@
-from typing import Tuple
-
 import Crypto.PublicKey.RSA as RSA
 from Crypto.PublicKey.RSA import RsaKey
 
-from libp2p.crypto.keys import KeyType, PrivateKey, PublicKey
+from libp2p.crypto.keys import KeyPair, KeyType, PrivateKey, PublicKey
 
 
 class RSAPublicKey(PublicKey):
@@ -42,13 +40,11 @@ class RSAPrivateKey(PrivateKey):
         return RSAPublicKey(self.impl.publickey())
 
 
-def create_new_key_pair(
+def create_new_key_pair(bits: int = 2048, e: int = 65537) -> KeyPair:
-    bits: int = 2048, e: int = 65537
-) -> Tuple[PrivateKey, PublicKey]:
     """
     Returns a new RSA keypair with the requested key size (``bits``) and the given public
     exponent ``e``. Sane defaults are provided for both values.
     """
     private_key = RSAPrivateKey.new(bits, e)
     public_key = private_key.get_public_key()
-    return private_key, public_key
+    return KeyPair(private_key, public_key)

libp2p/crypto/secp256k1.py

@@ -1,8 +1,6 @@
-from typing import Tuple
-
 import coincurve
 
-from libp2p.crypto.keys import KeyType, PrivateKey, PublicKey
+from libp2p.crypto.keys import KeyPair, KeyType, PrivateKey, PublicKey
 
 
 class Secp256k1PublicKey(PublicKey):
@@ -42,7 +40,7 @@ class Secp256k1PrivateKey(PrivateKey):
         return Secp256k1PublicKey(public_key_impl)
 
 
-def create_new_key_pair(secret: bytes = None) -> Tuple[PrivateKey, PublicKey]:
+def create_new_key_pair(secret: bytes = None) -> KeyPair:
     """
     Returns a new Secp256k1 keypair derived from the provided ``secret``,
     a sequence of bytes corresponding to some integer between 0 and the group order.
@@ -51,4 +49,4 @@ def create_new_key_pair(secret: bytes = None) -> Tuple[PrivateKey, PublicKey]:
     """
     private_key = Secp256k1PrivateKey.new()
     public_key = private_key.get_public_key()
-    return private_key, public_key
+    return KeyPair(private_key, public_key)

tests/peer/test_peerid.py

@@ -89,7 +89,8 @@ def test_id_from_base58():
 
 
 def test_id_from_public_key():
-    _, public_key = create_new_key_pair()
+    key_pair = create_new_key_pair()
+    public_key = key_pair.public_key
 
     key_bin = public_key.serialize_to_protobuf().SerializeToString()
     algo = multihash.Func.sha2_256