Move keypair into PeerData
This commit is contained in:
NIC619
parent
828ae69c66
commit
0bfbdf7fab
|
|
@ -2,11 +2,15 @@ from typing import Any, Dict, List, Sequence
|
|||
|
||||
from multiaddr import Multiaddr
|
||||
|
||||
from libp2p.crypto.keys import PrivateKey, PublicKey
|
||||
|
||||
from .peerdata_interface import IPeerData
|
||||
|
||||
|
||||
class PeerData(IPeerData):
|
||||
|
||||
pubkey: PublicKey
|
||||
privkey: PrivateKey
|
||||
metadata: Dict[Any, Any]
|
||||
protocols: List[str]
|
||||
addrs: List[Multiaddr]
|
||||
|
|
@ -17,31 +21,86 @@ class PeerData(IPeerData):
|
|||
self.addrs = []
|
||||
|
||||
def get_protocols(self) -> List[str]:
|
||||
"""
|
||||
:return: all protocols associated with given peer
|
||||
"""
|
||||
return self.protocols
|
||||
|
||||
def add_protocols(self, protocols: Sequence[str]) -> None:
|
||||
"""
|
||||
:param protocols: protocols to add
|
||||
"""
|
||||
self.protocols.extend(list(protocols))
|
||||
|
||||
def set_protocols(self, protocols: Sequence[str]) -> None:
|
||||
"""
|
||||
:param protocols: protocols to set
|
||||
"""
|
||||
self.protocols = list(protocols)
|
||||
|
||||
def add_addrs(self, addrs: Sequence[Multiaddr]) -> None:
|
||||
"""
|
||||
:param addrs: multiaddresses to add
|
||||
"""
|
||||
self.addrs.extend(addrs)
|
||||
|
||||
def get_addrs(self) -> List[Multiaddr]:
|
||||
"""
|
||||
:return: all multiaddresses
|
||||
"""
|
||||
return self.addrs
|
||||
|
||||
def clear_addrs(self) -> None:
|
||||
"""Clear all addresses."""
|
||||
self.addrs = []
|
||||
|
||||
def put_metadata(self, key: str, val: Any) -> None:
|
||||
"""
|
||||
:param key: key in KV pair
|
||||
:param val: val to associate with key
|
||||
"""
|
||||
self.metadata[key] = val
|
||||
|
||||
def get_metadata(self, key: str) -> Any:
|
||||
"""
|
||||
:param key: key in KV pair
|
||||
:return: val for key
|
||||
:raise PeerDataError: key not found
|
||||
"""
|
||||
if key in self.metadata:
|
||||
return self.metadata[key]
|
||||
raise PeerDataError("key not found")
|
||||
|
||||
def add_pubkey(self, pubkey: PublicKey) -> None:
|
||||
"""
|
||||
:param pubkey:
|
||||
"""
|
||||
self.pubkey = pubkey
|
||||
|
||||
def get_pubkey(self) -> PublicKey:
|
||||
"""
|
||||
:return: public key of the peer
|
||||
:raise PeerDataError: if public key not found
|
||||
"""
|
||||
if self.pubkey is None:
|
||||
raise PeerDataError("public key not found")
|
||||
return self.pubkey
|
||||
|
||||
def add_privkey(self, privkey: PrivateKey) -> None:
|
||||
"""
|
||||
:param privkey:
|
||||
"""
|
||||
self.privkey = privkey
|
||||
|
||||
def get_privkey(self) -> PrivateKey:
|
||||
"""
|
||||
:return: private key of the peer
|
||||
:raise PeerDataError: if private key not found
|
||||
"""
|
||||
if self.privkey is None:
|
||||
raise PeerDataError("private key not found")
|
||||
return self.privkey
|
||||
|
||||
|
||||
class PeerDataError(KeyError):
|
||||
"""Raised when a key is not found in peer metadata."""
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ from typing import Any, List, Sequence
|
|||
|
||||
from multiaddr import Multiaddr
|
||||
|
||||
from libp2p.crypto.keys import PrivateKey, PublicKey
|
||||
|
||||
from .peermetadata_interface import IPeerMetadata
|
||||
|
||||
|
||||
|
|
@ -22,7 +24,7 @@ class IPeerData(ABC):
|
|||
@abstractmethod
|
||||
def set_protocols(self, protocols: Sequence[str]) -> None:
|
||||
"""
|
||||
:param protocols: protocols to add
|
||||
:param protocols: protocols to set
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
|
|
@ -46,7 +48,6 @@ class IPeerData(ABC):
|
|||
"""
|
||||
:param key: key in KV pair
|
||||
:param val: val to associate with key
|
||||
:raise Exception: unsuccesful put
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
|
|
@ -54,5 +55,31 @@ class IPeerData(ABC):
|
|||
"""
|
||||
:param key: key in KV pair
|
||||
:return: val for key
|
||||
:raise Exception: key not found
|
||||
:raise PeerDataError: key not found
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def add_pubkey(self, pubkey: PublicKey) -> None:
|
||||
"""
|
||||
:param pubkey:
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_pubkey(self) -> PublicKey:
|
||||
"""
|
||||
:return: public key of the peer
|
||||
:raise PeerDataError: if public key not found
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def add_privkey(self, privkey: PrivateKey) -> None:
|
||||
"""
|
||||
:param privkey:
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_privkey(self) -> PrivateKey:
|
||||
"""
|
||||
:return: private key of the peer
|
||||
:raise PeerDataError: if private key not found
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import Any, Dict, List, Optional, Sequence, Set
|
||||
from typing import Any, Dict, List, Sequence
|
||||
|
||||
from multiaddr import Multiaddr
|
||||
|
||||
|
|
@ -12,16 +12,12 @@ from .peerstore_interface import IPeerStore
|
|||
|
||||
class PeerStore(IPeerStore):
|
||||
|
||||
peer_pubkey_map: Dict[ID, PublicKey]
|
||||
peer_privkey_map: Dict[ID, PrivateKey]
|
||||
peer_data_map: Dict[ID, PeerData]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.peer_pubkey_map = {}
|
||||
self.peer_privkey_map = {}
|
||||
self.peer_data_map = {}
|
||||
|
||||
def __create_or_get_peer(self, peer_id: ID) -> PeerData:
|
||||
def __create_or_get_peer_data(self, peer_id: ID) -> PeerData:
|
||||
"""
|
||||
Returns the peer data for peer_id or creates a new peer data (and
|
||||
stores it in peer_data_map) if peer data for peer_id does not yet
|
||||
|
|
@ -36,15 +32,15 @@ class PeerStore(IPeerStore):
|
|||
self.peer_data_map[peer_id] = data
|
||||
return data
|
||||
|
||||
def peer_info(self, peer_id: ID) -> Optional[PeerInfo]:
|
||||
def peer_info(self, peer_id: ID) -> PeerInfo:
|
||||
"""
|
||||
:param peer_id: peer ID to get info for
|
||||
:return: peer info object
|
||||
"""
|
||||
if peer_id in self.peer_data_map:
|
||||
peer_data = self.peer_data_map[peer_id]
|
||||
return PeerInfo(peer_id, peer_data.addrs)
|
||||
return None
|
||||
return PeerInfo(peer_id, peer_data.get_addrs())
|
||||
raise PeerStoreError("peer ID not found")
|
||||
|
||||
def get_protocols(self, peer_id: ID) -> List[str]:
|
||||
"""
|
||||
|
|
@ -61,16 +57,16 @@ class PeerStore(IPeerStore):
|
|||
:param peer_id: peer ID to add protocols for
|
||||
:param protocols: protocols to add
|
||||
"""
|
||||
peer = self.__create_or_get_peer(peer_id)
|
||||
peer.add_protocols(list(protocols))
|
||||
peer_data = self.__create_or_get_peer_data(peer_id)
|
||||
peer_data.add_protocols(list(protocols))
|
||||
|
||||
def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None:
|
||||
"""
|
||||
:param peer_id: peer ID to set protocols for
|
||||
:param protocols: protocols to set
|
||||
"""
|
||||
peer = self.__create_or_get_peer(peer_id)
|
||||
peer.set_protocols(list(protocols))
|
||||
peer_data = self.__create_or_get_peer_data(peer_id)
|
||||
peer_data.set_protocols(list(protocols))
|
||||
|
||||
def peer_ids(self) -> List[ID]:
|
||||
"""
|
||||
|
|
@ -101,8 +97,8 @@ class PeerStore(IPeerStore):
|
|||
"""
|
||||
# <<?>>
|
||||
# This can output an error, not sure what the possible errors are
|
||||
peer = self.__create_or_get_peer(peer_id)
|
||||
peer.put_metadata(key, val)
|
||||
peer_data = self.__create_or_get_peer_data(peer_id)
|
||||
peer_data.put_metadata(key, val)
|
||||
|
||||
def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None:
|
||||
"""
|
||||
|
|
@ -119,8 +115,8 @@ class PeerStore(IPeerStore):
|
|||
:param ttl: time-to-live for the this record
|
||||
"""
|
||||
# Ignore ttl for now
|
||||
peer = self.__create_or_get_peer(peer_id)
|
||||
peer.add_addrs(list(addrs))
|
||||
peer_data = self.__create_or_get_peer_data(peer_id)
|
||||
peer_data.add_addrs(list(addrs))
|
||||
|
||||
def addrs(self, peer_id: ID) -> List[Multiaddr]:
|
||||
"""
|
||||
|
|
@ -156,60 +152,57 @@ class PeerStore(IPeerStore):
|
|||
"""
|
||||
:param peer_id: peer ID to add public key for
|
||||
:param pubkey:
|
||||
:raise PeerStoreError: if peer ID already has pubkey set
|
||||
"""
|
||||
if peer_id in self.peer_pubkey_map:
|
||||
raise PeerStoreError(
|
||||
f"peer ID already has pubkey: {self.peer_pubkey_map[peer_id]}"
|
||||
)
|
||||
self.peer_pubkey_map[peer_id] = pubkey
|
||||
peer_data = self.__create_or_get_peer_data(peer_id)
|
||||
# TODO: Check if pubkey matches peer ID
|
||||
peer_data.add_pubkey(pubkey)
|
||||
|
||||
def pubkey(self, peer_id: ID) -> PublicKey:
|
||||
"""
|
||||
:param peer_id: peer ID to get public key for
|
||||
:return: public key of the peer
|
||||
:raise PeerStoreError: if peer ID not found
|
||||
:raise PeerStoreError: if peer ID or peer pubkey not found
|
||||
"""
|
||||
if peer_id in self.peer_pubkey_map:
|
||||
return self.peer_pubkey_map[peer_id]
|
||||
if peer_id in self.peer_data_map:
|
||||
peer_data = self.peer_data_map[peer_id]
|
||||
try:
|
||||
pubkey = peer_data.get_pubkey()
|
||||
except PeerDataError:
|
||||
raise PeerStoreError("peer pubkey not found")
|
||||
return pubkey
|
||||
raise PeerStoreError("peer ID not found")
|
||||
|
||||
def add_privkey(self, peer_id: ID, privkey: PrivateKey) -> None:
|
||||
"""
|
||||
:param peer_id: peer ID to add private key for
|
||||
:param privkey:
|
||||
:raise PeerStoreError: if peer ID already has privkey set
|
||||
"""
|
||||
if peer_id in self.peer_privkey_map:
|
||||
raise PeerStoreError(
|
||||
f"peer ID already has privkey: {self.peer_privkey_map[peer_id]}"
|
||||
)
|
||||
self.peer_privkey_map[peer_id] = privkey
|
||||
peer_data = self.__create_or_get_peer_data(peer_id)
|
||||
# TODO: Check if privkey matches peer ID
|
||||
peer_data.add_privkey(privkey)
|
||||
|
||||
def privkey(self, peer_id: ID) -> PrivateKey:
|
||||
"""
|
||||
:param peer_id: peer ID to get private key for
|
||||
:return: private key of the peer
|
||||
:raise PeerStoreError: if peer ID not found
|
||||
:raise PeerStoreError: if peer ID or peer privkey not found
|
||||
"""
|
||||
if peer_id in self.peer_pubkey_map:
|
||||
return self.peer_privkey_map[peer_id]
|
||||
if peer_id in self.peer_data_map:
|
||||
peer_data = self.peer_data_map[peer_id]
|
||||
try:
|
||||
privkey = peer_data.get_privkey()
|
||||
except PeerDataError:
|
||||
raise PeerStoreError("peer privkey not found")
|
||||
return privkey
|
||||
raise PeerStoreError("peer ID not found")
|
||||
|
||||
def add_key_pair(self, peer_id: ID, key_pair: KeyPair) -> None:
|
||||
"""
|
||||
:param peer_id: peer ID to add private key for
|
||||
:param key_pair:
|
||||
:raise PeerStoreError: if peer ID already has pubkey or privkey set
|
||||
"""
|
||||
self.add_pubkey(peer_id, key_pair.public_keypubkey)
|
||||
self.add_privkey(peer_id, key_pair.private_key)
|
||||
|
||||
def peers_with_keys(self) -> Set[ID]:
|
||||
"""
|
||||
:return: all of the peer IDs which has pubkey/privkey stored in peer store
|
||||
"""
|
||||
return set(self.peer_pubkey_map.keys()).union(self.peer_privkey_map.keys())
|
||||
self.add_pubkey(key_pair.public_key)
|
||||
self.add_privkey(key_pair.private_key)
|
||||
|
||||
|
||||
class PeerStoreError(KeyError):
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user