Convert from base58/pubkey/privkey to class method

libp2p/__init__.py

@@ -4,7 +4,7 @@ from Crypto.PublicKey import RSA
 
 from .security.insecure_security import InsecureTransport
 from .peer.peerstore import PeerStore
-from .peer.id import id_from_public_key
+from .peer.id import ID
 from .network.swarm import Swarm
 from .host.basic_host import BasicHost
 from .transport.upgrader import TransportUpgrader
@@ -29,7 +29,7 @@ async def cleanup_done_tasks():
 
 def generate_id():
     new_key = RSA.generate(2048, e=65537)
-    new_id = id_from_public_key(new_key.publickey())
+    new_id = ID.from_pubkey(new_key.publickey())
     # private_key = new_key.exportKey("PEM")
     return new_id
 

libp2p/network/swarm.py

@@ -3,7 +3,7 @@ from typing import Awaitable, Callable, Dict, List, Sequence
 
 from multiaddr import Multiaddr
 
-from libp2p.peer.id import ID, id_b58_decode
+from libp2p.peer.id import ID
 from libp2p.peer.peerstore import PeerStore
 from libp2p.protocol_muxer.multiselect import Multiselect
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
@@ -185,7 +185,7 @@ class Swarm(INetwork):
                 reader: asyncio.StreamReader, writer: asyncio.StreamWriter
             ) -> None:
                 # Read in first message (should be peer_id of initiator) and ack
-                peer_id = id_b58_decode((await reader.read(1024)).decode())
+                peer_id = ID.from_base58((await reader.read(1024)).decode())
 
                 writer.write("received peer id".encode())
                 await writer.drain()

libp2p/peer/id.py

@@ -53,6 +53,29 @@ class ID:
     def __hash__(self) -> int:
         return hash(self._bytes)
 
+    @classmethod
+    def from_base58(cls, b58_encoded_peer_id_str: str) -> 'ID':
+        peer_id_bytes = base58.b58decode(b58_encoded_peer_id_str)
+        pid = ID(peer_id_bytes)
+        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)
+        return cls(mh_digest.encode())
+
+    @classmethod
+    def from_privkey(cls, key: RsaKey) -> 'ID':
+        return cls.from_pubkey(key.publickey())
+
 
 def id_b58_encode(peer_id: ID) -> str:
     """

libp2p/peer/peerinfo.py

@@ -2,7 +2,7 @@ from typing import List
 
 import multiaddr
 
-from .id import ID, id_b58_decode
+from .id import ID
 from .peerdata import PeerData
 
 
@@ -39,7 +39,7 @@ def info_from_p2p_addr(addr: multiaddr.Multiaddr) -> PeerInfo:
 
     # make sure the /p2p value parses as a peer.ID
     peer_id_str: str = p2p_part.value_for_protocol(multiaddr.protocols.P_P2P)
-    peer_id: ID = id_b58_decode(peer_id_str)
+    peer_id: ID = ID.from_base58(peer_id_str)
 
     # we might have received just an / p2p part, which means there's no addr.
     if len(parts) > 1:

libp2p/pubsub/floodsub.py

@@ -1,6 +1,6 @@
 from typing import Iterable, List, Sequence
 
-from libp2p.peer.id import ID, id_b58_decode
+from libp2p.peer.id import ID
 
 from .pb import rpc_pb2
 from .pubsub import Pubsub
@@ -106,7 +106,7 @@ class FloodSub(IPubsubRouter):
             if topic not in self.pubsub.peer_topics:
                 continue
             for peer_id_str in self.pubsub.peer_topics[topic]:
-                peer_id = id_b58_decode(peer_id_str)
+                peer_id = ID.from_base58(peer_id_str)
                 if peer_id in (msg_forwarder, origin):
                     continue
                 # FIXME: Should change `self.pubsub.peers` to Dict[PeerID, ...]

libp2p/pubsub/gossipsub.py

@@ -1,9 +1,16 @@
 from ast import literal_eval
 import asyncio
 import random
-from typing import Any, Dict, Iterable, List, Set, Sequence
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Set,
+    Sequence,
+)
 
-from libp2p.peer.id import ID, id_b58_decode
+from libp2p.peer.id import ID
 
 from .mcache import MessageCache
 from .pb import rpc_pb2
@@ -201,7 +208,7 @@ class GossipSub(IPubsubRouter):
                 # FIXME: `gossipsub.peers_floodsub` can be changed to `gossipsub.peers` in go.
                 #   This will improve the efficiency when searching for a peer's protocol id.
                 if peer_id_str in self.peers_floodsub:
-                    peer_id = id_b58_decode(peer_id_str)
+                    peer_id = ID.from_base58(peer_id_str)
                     send_to.add(peer_id)
 
             # gossipsub peers
@@ -223,7 +230,7 @@ class GossipSub(IPubsubRouter):
                     )
                 in_topic_gossipsub_peers = self.fanout[topic]
             for peer_id_str in in_topic_gossipsub_peers:
-                send_to.add(id_b58_decode(peer_id_str))
+                send_to.add(ID.from_base58(peer_id_str))
         # Excludes `msg_forwarder` and `origin`
         yield from send_to.difference([msg_forwarder, origin])
 

tests/peer/test_peerid.py

@@ -6,9 +6,6 @@ from Crypto.PublicKey import RSA
 from libp2p.peer.id import (
     ID,
     id_b58_encode,
-    id_b58_decode,
-    id_from_public_key,
-    id_from_private_key,
 )
 
 
@@ -107,7 +104,7 @@ def test_id_b58_decode():
     for _ in range(10):
         random_id_string += random.SystemRandom().choice(ALPHABETS)
     expected = ID(base58.b58decode(random_id_string.encode()))
-    actual = id_b58_decode(random_id_string.encode())
+    actual = ID.from_base58(random_id_string.encode())
 
     assert actual == expected
 
@@ -119,13 +116,13 @@ 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_public_key(key)
+    actual = ID.from_pubkey(key)
 
     assert actual == expected
 
 
 def test_id_from_private_key():
     key = RSA.generate(2048, e=65537)
-    id_from_pub = id_from_public_key(key.publickey())
-    id_from_priv = id_from_private_key(key)
+    id_from_pub = ID.from_pubkey(key.publickey())
+    id_from_priv = ID.from_privkey(key)
     assert id_from_pub == id_from_priv