From bd8d45fbc161962d187c986daf473471103aea20 Mon Sep 17 00:00:00 2001
From: Alex Stokes <r.alex.stokes@gmail.com>
Date: Fri, 2 Aug 2019 16:03:40 -0700
Subject: [PATCH] Refactor ID to not use third-party type for cryptographic
 keys

Remove `ID.from_privkey` which would require specific knowledge per cryptosystem
---
 libp2p/__init__.py        |  5 ++---
 libp2p/peer/id.py         | 27 +++------------------------
 tests/peer/test_peerid.py |  9 +--------
 3 files changed, 6 insertions(+), 35 deletions(-)

diff --git a/libp2p/__init__.py b/libp2p/__init__.py
index 110aba49..bebd120e 100644
--- a/libp2p/__init__.py
+++ b/libp2p/__init__.py
@@ -28,9 +28,8 @@ async def cleanup_done_tasks():
 
 
 def generate_id():
-    new_key = RSA.generate(2048, e=65537)
-    new_id = ID.from_pubkey(new_key.publickey())
-    # private_key = new_key.exportKey("PEM")
+    new_key = RSA.generate(2048, e=65537).publickey().export_key("DER")
+    new_id = ID.from_pubkey(new_key)
     return new_id
 
 
diff --git a/libp2p/peer/id.py b/libp2p/peer/id.py
index 7a3f9190..91df71a0 100644
--- a/libp2p/peer/id.py
+++ b/libp2p/peer/id.py
@@ -5,16 +5,6 @@ import base58
 
 import multihash
 
-from Crypto.PublicKey.RSA import RsaKey
-
-# MaxInlineKeyLength is the maximum length a key can be for it to be inlined in
-# the peer ID.
-# * When `len(pubKey.Bytes()) <= MaxInlineKeyLength`, the peer ID is the
-#   identity multihash hash of the public key.
-# * When `len(pubKey.Bytes()) > MaxInlineKeyLength`, the peer ID is the
-#   sha2-256 multihash of the public key.
-MAX_INLINE_KEY_LENGTH = 42
-
 
 class ID:
 
@@ -64,22 +54,11 @@ class ID:
         return pid
 
     @classmethod
-    def from_pubkey(cls, key: RsaKey) -> "ID":
-        # export into binary format
-        key_bin = key.exportKey("DER")
-
-        algo: int = multihash.Func.sha2_256
-        # TODO: seems identity is not yet supported in pymultihash
-        # if len(b) <= MAX_INLINE_KEY_LENGTH:
-        #     algo multihash.func.identity
-
-        mh_digest: multihash.Multihash = multihash.digest(key_bin, algo)
+    def from_pubkey(cls, key: bytes) -> "ID":
+        algo = multihash.Func.sha2_256
+        mh_digest = multihash.digest(key, algo)
         return cls(mh_digest.encode())
 
-    @classmethod
-    def from_privkey(cls, key: RsaKey) -> "ID":
-        return cls.from_pubkey(key.publickey())
-
 
 def digest(data: Union[str, bytes]) -> bytes:
     if isinstance(data, str):
diff --git a/tests/peer/test_peerid.py b/tests/peer/test_peerid.py
index 692a8e71..4c034e3f 100644
--- a/tests/peer/test_peerid.py
+++ b/tests/peer/test_peerid.py
@@ -111,13 +111,6 @@ def test_id_from_public_key():
     algo = multihash.Func.sha2_256
     mh_digest = multihash.digest(key_bin, algo)
     expected = ID(mh_digest.encode())
-    actual = ID.from_pubkey(key)
+    actual = ID.from_pubkey(key_bin)
 
     assert actual == expected
-
-
-def test_id_from_private_key():
-    key = RSA.generate(2048, e=65537)
-    id_from_pub = ID.from_pubkey(key.publickey())
-    id_from_priv = ID.from_privkey(key)
-    assert id_from_pub == id_from_priv