diff --git a/libp2p/__init__.py b/libp2p/__init__.py index 659c0eea..a69f6f61 100644 --- a/libp2p/__init__.py +++ b/libp2p/__init__.py @@ -1,7 +1,7 @@ import asyncio from typing import Mapping, Sequence -from libp2p.crypto.keys import KeyPair, PrivateKey +from libp2p.crypto.keys import KeyPair from libp2p.crypto.rsa import create_new_key_pair from libp2p.host.basic_host import BasicHost from libp2p.kademlia.network import KademliaServer @@ -69,7 +69,7 @@ def initialize_default_kademlia_router( def initialize_default_swarm( - private_key: PrivateKey, + key_pair: KeyPair, id_opt: ID = None, transport_opt: Sequence[str] = None, muxer_opt: Sequence[str] = None, @@ -89,7 +89,7 @@ def initialize_default_swarm( """ if not id_opt: - id_opt = generate_peer_id_from_rsa_identity() + id_opt = generate_peer_id_from_rsa_identity(key_pair) # TODO parse transport_opt to determine transport transport_opt = transport_opt or ["/ip4/127.0.0.1/tcp/8001"] @@ -99,9 +99,7 @@ def initialize_default_swarm( # TODO parse muxer and sec to pass into TransportUpgrader muxer = muxer_opt or ["mplex/6.7.0"] security_transports_by_protocol = sec_opt or { - TProtocol("insecure/1.0.0"): InsecureTransport( - private_key, private_key.get_public_key() - ) + TProtocol("insecure/1.0.0"): InsecureTransport(key_pair) } upgrader = TransportUpgrader(security_transports_by_protocol, muxer) @@ -113,7 +111,7 @@ def initialize_default_swarm( async def new_node( - private_key: PrivateKey = None, + key_pair: KeyPair = None, swarm_opt: INetwork = None, id_opt: ID = None, transport_opt: Sequence[str] = None, @@ -124,6 +122,7 @@ async def new_node( ) -> BasicHost: """ create new libp2p node + :param key_pair: key pair for deriving an identity :param swarm_opt: optional swarm :param id_opt: optional id for host :param transport_opt: optional choice of transport upgrade @@ -134,16 +133,15 @@ async def new_node( :return: return a host instance """ - if not private_key: + if not key_pair: key_pair = generate_new_rsa_identity() - private_key = key_pair.private_key if not id_opt: id_opt = generate_peer_id_from_rsa_identity(key_pair) if not swarm_opt: swarm_opt = initialize_default_swarm( - private_key=private_key, + key_pair=key_pair, id_opt=id_opt, transport_opt=transport_opt, muxer_opt=muxer_opt, diff --git a/libp2p/security/base_transport.py b/libp2p/security/base_transport.py index 3c1e1542..0f096bfc 100644 --- a/libp2p/security/base_transport.py +++ b/libp2p/security/base_transport.py @@ -1,4 +1,4 @@ -from libp2p.crypto.keys import PrivateKey, PublicKey +from libp2p.crypto.keys import KeyPair from libp2p.peer.id import ID from libp2p.security.secure_transport_interface import ISecureTransport @@ -9,8 +9,6 @@ class BaseSecureTransport(ISecureTransport): is only meant to be used in clases that derive from it. """ - def __init__( - self, local_private_key: PrivateKey, local_public_key: PublicKey - ) -> None: - self.local_private_key = local_private_key - self.local_peer = ID.from_pubkey(local_public_key) + def __init__(self, local_key_pair: KeyPair) -> None: + self.local_private_key = local_key_pair.private_key + self.local_peer = ID.from_pubkey(local_key_pair.public_key) diff --git a/libp2p/security/simple/transport.py b/libp2p/security/simple/transport.py index ac9c01e4..5ff07cdd 100644 --- a/libp2p/security/simple/transport.py +++ b/libp2p/security/simple/transport.py @@ -1,5 +1,6 @@ import asyncio +from libp2p.crypto.keys import KeyPair from libp2p.network.connection.raw_connection_interface import IRawConnection from libp2p.peer.id import ID from libp2p.security.base_transport import BaseSecureTransport @@ -10,10 +11,8 @@ from libp2p.security.secure_conn_interface import ISecureConn class SimpleSecurityTransport(BaseSecureTransport): key_phrase: str - def __init__( - self, local_private_key: bytes, local_public_key: bytes, key_phrase: str - ) -> None: - super().__init__(local_private_key, local_public_key) + def __init__(self, local_key_pair: KeyPair, key_phrase: str) -> None: + super().__init__(local_key_pair) self.key_phrase = key_phrase async def secure_inbound(self, conn: IRawConnection) -> ISecureConn: diff --git a/tests/libp2p/test_notify.py b/tests/libp2p/test_notify.py index f1262a1f..f4bd2efc 100644 --- a/tests/libp2p/test_notify.py +++ b/tests/libp2p/test_notify.py @@ -13,12 +13,12 @@ import multiaddr import pytest from libp2p import initialize_default_swarm, new_node +from libp2p.crypto.rsa import create_new_key_pair from libp2p.host.basic_host import BasicHost from libp2p.network.notifee_interface import INotifee from tests.utils import ( cleanup, echo_stream_handler, - generate_new_private_key, perform_two_host_set_up_custom_handler, ) @@ -173,17 +173,17 @@ async def test_one_notifier_on_two_nodes(): async def test_one_notifier_on_two_nodes_with_listen(): events_b = [] - node_a_key = generate_new_private_key() + node_a_key_pair = create_new_key_pair() node_a_transport_opt = ["/ip4/127.0.0.1/tcp/0"] - node_a = await new_node(node_a_key, transport_opt=node_a_transport_opt) + node_a = await new_node(node_a_key_pair, transport_opt=node_a_transport_opt) await node_a.get_network().listen(multiaddr.Multiaddr(node_a_transport_opt[0])) # Set up node_b swarm to pass into host - node_b_key = generate_new_private_key() + node_b_key_pair = create_new_key_pair() node_b_transport_opt = ["/ip4/127.0.0.1/tcp/0"] node_b_multiaddr = multiaddr.Multiaddr(node_b_transport_opt[0]) node_b_swarm = initialize_default_swarm( - node_b_key, transport_opt=node_b_transport_opt + node_b_key_pair, transport_opt=node_b_transport_opt ) node_b = BasicHost(node_b_swarm) diff --git a/tests/pubsub/factories.py b/tests/pubsub/factories.py index d25c2784..b57c29bc 100644 --- a/tests/pubsub/factories.py +++ b/tests/pubsub/factories.py @@ -1,6 +1,7 @@ import factory from libp2p import initialize_default_swarm +from libp2p.crypto.rsa import create_new_key_pair from libp2p.host.basic_host import BasicHost from libp2p.pubsub.floodsub import FloodSub from libp2p.pubsub.gossipsub import GossipSub @@ -11,11 +12,10 @@ from tests.pubsub.configs import ( GOSSIPSUB_PARAMS, GOSSIPSUB_PROTOCOL_ID, ) -from tests.utils import generate_new_private_key def swarm_factory(): - private_key = generate_new_private_key() + private_key = create_new_key_pair() return initialize_default_swarm(private_key, transport_opt=[str(LISTEN_MADDR)]) diff --git a/tests/security/test_security_multistream.py b/tests/security/test_security_multistream.py index bf4929ba..ddbae8e3 100644 --- a/tests/security/test_security_multistream.py +++ b/tests/security/test_security_multistream.py @@ -4,11 +4,12 @@ import multiaddr import pytest from libp2p import new_node +from libp2p.crypto.rsa import create_new_key_pair from libp2p.peer.peerinfo import info_from_p2p_addr from libp2p.protocol_muxer.multiselect_client import MultiselectClientError from libp2p.security.insecure.transport import InsecureSession, InsecureTransport from libp2p.security.simple.transport import SimpleSecurityTransport -from tests.utils import cleanup, connect, generate_new_private_key +from tests.utils import cleanup, connect # TODO: Add tests for multiple streams being opened on different # protocols through the same connection @@ -20,13 +21,9 @@ def peer_id_for_node(node): return info.peer_id -initiator_private_key = generate_new_private_key() -initiator_private_key_bytes = initiator_private_key.export_key("DER") -initiator_public_key_bytes = initiator_private_key.publickey().export_key("DER") +initiator_key_pair = create_new_key_pair() -noninitiator_private_key = generate_new_private_key() -noninitiator_private_key_bytes = noninitiator_private_key.export_key("DER") -noninitiator_public_key_bytes = noninitiator_private_key.publickey().export_key("DER") +noninitiator_key_pair = create_new_key_pair() async def perform_simple_test( @@ -68,16 +65,8 @@ async def perform_simple_test( @pytest.mark.asyncio async def test_single_insecure_security_transport_succeeds(): - transports_for_initiator = { - "foo": InsecureTransport( - initiator_private_key_bytes, initiator_public_key_bytes - ) - } - transports_for_noninitiator = { - "foo": InsecureTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes - ) - } + transports_for_initiator = {"foo": InsecureTransport(initiator_key_pair)} + transports_for_noninitiator = {"foo": InsecureTransport(noninitiator_key_pair)} def assertion_func(conn): assert isinstance(conn, InsecureSession) @@ -90,14 +79,10 @@ async def test_single_insecure_security_transport_succeeds(): @pytest.mark.asyncio async def test_single_simple_test_security_transport_succeeds(): transports_for_initiator = { - "tacos": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "tacos" - ) + "tacos": SimpleSecurityTransport(initiator_key_pair, "tacos") } transports_for_noninitiator = { - "tacos": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "tacos" - ) + "tacos": SimpleSecurityTransport(noninitiator_key_pair, "tacos") } def assertion_func(conn): @@ -111,17 +96,11 @@ async def test_single_simple_test_security_transport_succeeds(): @pytest.mark.asyncio async def test_two_simple_test_security_transport_for_initiator_succeeds(): transports_for_initiator = { - "tacos": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "tacos" - ), - "shleep": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "shleep" - ), + "tacos": SimpleSecurityTransport(initiator_key_pair, "tacos"), + "shleep": SimpleSecurityTransport(initiator_key_pair, "shleep"), } transports_for_noninitiator = { - "shleep": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "shleep" - ) + "shleep": SimpleSecurityTransport(noninitiator_key_pair, "shleep") } def assertion_func(conn): @@ -135,17 +114,11 @@ async def test_two_simple_test_security_transport_for_initiator_succeeds(): @pytest.mark.asyncio async def test_two_simple_test_security_transport_for_noninitiator_succeeds(): transports_for_initiator = { - "tacos": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "tacos" - ) + "tacos": SimpleSecurityTransport(initiator_key_pair, "tacos") } transports_for_noninitiator = { - "shleep": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "shleep" - ), - "tacos": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "tacos" - ), + "shleep": SimpleSecurityTransport(noninitiator_key_pair, "shleep"), + "tacos": SimpleSecurityTransport(noninitiator_key_pair, "tacos"), } def assertion_func(conn): @@ -159,20 +132,12 @@ async def test_two_simple_test_security_transport_for_noninitiator_succeeds(): @pytest.mark.asyncio async def test_two_simple_test_security_transport_for_both_succeeds(): transports_for_initiator = { - "a": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "a" - ), - "b": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "b" - ), + "a": SimpleSecurityTransport(initiator_key_pair, "a"), + "b": SimpleSecurityTransport(initiator_key_pair, "b"), } transports_for_noninitiator = { - "b": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "b" - ), - "c": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "c" - ), + "b": SimpleSecurityTransport(noninitiator_key_pair, "b"), + "c": SimpleSecurityTransport(noninitiator_key_pair, "c"), } def assertion_func(conn): @@ -186,20 +151,12 @@ async def test_two_simple_test_security_transport_for_both_succeeds(): @pytest.mark.asyncio async def test_multiple_security_none_the_same_fails(): transports_for_initiator = { - "a": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "a" - ), - "b": SimpleSecurityTransport( - initiator_private_key_bytes, initiator_public_key_bytes, "b" - ), + "a": SimpleSecurityTransport(initiator_key_pair, "a"), + "b": SimpleSecurityTransport(initiator_key_pair, "b"), } transports_for_noninitiator = { - "d": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "d" - ), - "c": SimpleSecurityTransport( - noninitiator_private_key_bytes, noninitiator_public_key_bytes, "c" - ), + "d": SimpleSecurityTransport(noninitiator_key_pair, "d"), + "c": SimpleSecurityTransport(noninitiator_key_pair, "c"), } def assertion_func(_): diff --git a/tests/utils.py b/tests/utils.py index 39368694..58a08807 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,7 +1,6 @@ import asyncio from contextlib import suppress -from Crypto.PublicKey import RSA import multiaddr from libp2p import new_node @@ -54,8 +53,3 @@ async def perform_two_host_set_up_custom_handler(handler): # Associate the peer with local ip address (see default parameters of Libp2p()) node_a.get_peerstore().add_addrs(node_b.get_id(), node_b.get_addrs(), 10) return node_a, node_b - - -def generate_new_private_key(): - new_key = RSA.generate(2048, e=65537) - return new_key