From ad20d8cb008fb56d3c3936070936cf34373787dd Mon Sep 17 00:00:00 2001
From: Alex Stokes <r.alex.stokes@gmail.com>
Date: Tue, 13 Aug 2019 20:23:07 -0700
Subject: [PATCH] Make a `KeyPair` dataclass for passing around key pairs

---
 libp2p/__init__.py         |  3 ++-
 libp2p/crypto/keys.py      |  7 +++++++
 libp2p/crypto/rsa.py       | 10 +++-------
 libp2p/crypto/secp256k1.py |  8 +++-----
 tests/peer/test_peerid.py  |  3 ++-
 5 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/libp2p/__init__.py b/libp2p/__init__.py
index 41f0a554..af72cb81 100644
--- a/libp2p/__init__.py
+++ b/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
 
diff --git a/libp2p/crypto/keys.py b/libp2p/crypto/keys.py
index 58ceb125..97517058 100644
--- a/libp2p/crypto/keys.py
+++ b/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
diff --git a/libp2p/crypto/rsa.py b/libp2p/crypto/rsa.py
index df202af5..ceea0088 100644
--- a/libp2p/crypto/rsa.py
+++ b/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(
-    bits: int = 2048, e: int = 65537
-) -> Tuple[PrivateKey, PublicKey]:
+def create_new_key_pair(bits: int = 2048, e: int = 65537) -> KeyPair:
     """
     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)
diff --git a/libp2p/crypto/secp256k1.py b/libp2p/crypto/secp256k1.py
index 4ad2364b..4cebe54a 100644
--- a/libp2p/crypto/secp256k1.py
+++ b/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)
diff --git a/tests/peer/test_peerid.py b/tests/peer/test_peerid.py
index 7d7bbfe9..30b74c17 100644
--- a/tests/peer/test_peerid.py
+++ b/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