diff --git a/libp2p/__init__.py b/libp2p/__init__.py index f6c6894..9510e23 100644 --- a/libp2p/__init__.py +++ b/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 diff --git a/libp2p/network/swarm.py b/libp2p/network/swarm.py index 76141a0..2b9c9e1 100644 --- a/libp2p/network/swarm.py +++ b/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() diff --git a/libp2p/peer/id.py b/libp2p/peer/id.py index a61b8c8..62374d4 100644 --- a/libp2p/peer/id.py +++ b/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: """ diff --git a/libp2p/peer/peerinfo.py b/libp2p/peer/peerinfo.py index 89bd307..ed4eed0 100644 --- a/libp2p/peer/peerinfo.py +++ b/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: diff --git a/libp2p/pubsub/floodsub.py b/libp2p/pubsub/floodsub.py index 0fcefe0..d09a7de 100644 --- a/libp2p/pubsub/floodsub.py +++ b/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, ...] diff --git a/libp2p/pubsub/gossipsub.py b/libp2p/pubsub/gossipsub.py index d8e8268..d7db309 100644 --- a/libp2p/pubsub/gossipsub.py +++ b/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]) diff --git a/tests/peer/test_peerid.py b/tests/peer/test_peerid.py index d6e0dbc..e7d9569 100644 --- a/tests/peer/test_peerid.py +++ b/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