129 lines
4.6 KiB
Python
129 lines
4.6 KiB
Python
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
|
|
|
from multiaddr import Multiaddr
|
|
|
|
from libp2p.crypto.keys import PrivateKey, PublicKey
|
|
|
|
from .id import ID
|
|
from .peerdata import PeerData, PeerDataError
|
|
from .peerinfo import PeerInfo
|
|
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:
|
|
IPeerStore.__init__(self)
|
|
self.peer_pubkey_map = {}
|
|
self.peer_privkey_map = {}
|
|
self.peer_data_map = {}
|
|
|
|
def __create_or_get_peer(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 exist.
|
|
|
|
:param peer_id: peer ID
|
|
:return: peer data
|
|
"""
|
|
if peer_id in self.peer_data_map:
|
|
return self.peer_data_map[peer_id]
|
|
data = PeerData()
|
|
self.peer_data_map[peer_id] = data
|
|
return self.peer_data_map[peer_id]
|
|
|
|
def peer_info(self, peer_id: ID) -> Optional[PeerInfo]:
|
|
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
|
|
|
|
def get_protocols(self, peer_id: ID) -> List[str]:
|
|
if peer_id in self.peer_data_map:
|
|
return self.peer_data_map[peer_id].get_protocols()
|
|
raise PeerStoreError("peer ID not found")
|
|
|
|
def add_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None:
|
|
peer = self.__create_or_get_peer(peer_id)
|
|
peer.add_protocols(list(protocols))
|
|
|
|
def set_protocols(self, peer_id: ID, protocols: Sequence[str]) -> None:
|
|
peer = self.__create_or_get_peer(peer_id)
|
|
peer.set_protocols(list(protocols))
|
|
|
|
def peer_ids(self) -> List[ID]:
|
|
return list(self.peer_data_map.keys())
|
|
|
|
def get(self, peer_id: ID, key: str) -> Any:
|
|
if peer_id in self.peer_data_map:
|
|
try:
|
|
val = self.peer_data_map[peer_id].get_metadata(key)
|
|
except PeerDataError as error:
|
|
raise PeerStoreError(error)
|
|
return val
|
|
raise PeerStoreError("peer ID not found")
|
|
|
|
def put(self, peer_id: ID, key: str, val: Any) -> None:
|
|
# <<?>>
|
|
# 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)
|
|
|
|
def add_addr(self, peer_id: ID, addr: Multiaddr, ttl: int) -> None:
|
|
self.add_addrs(peer_id, [addr], ttl)
|
|
|
|
def add_addrs(self, peer_id: ID, addrs: Sequence[Multiaddr], ttl: int) -> None:
|
|
# Ignore ttl for now
|
|
peer = self.__create_or_get_peer(peer_id)
|
|
peer.add_addrs(list(addrs))
|
|
|
|
def addrs(self, peer_id: ID) -> List[Multiaddr]:
|
|
if peer_id in self.peer_data_map:
|
|
return self.peer_data_map[peer_id].get_addrs()
|
|
raise PeerStoreError("peer ID not found")
|
|
|
|
def clear_addrs(self, peer_id: ID) -> None:
|
|
# Only clear addresses if the peer is in peer map
|
|
if peer_id in self.peer_data_map:
|
|
self.peer_data_map[peer_id].clear_addrs()
|
|
|
|
def peers_with_addrs(self) -> List[ID]:
|
|
# Add all peers with addrs at least 1 to output
|
|
output: List[ID] = []
|
|
|
|
for peer_id in self.peer_data_map:
|
|
if len(self.peer_data_map[peer_id].get_addrs()) >= 1:
|
|
output.append(peer_id)
|
|
return output
|
|
|
|
def add_pubkey(self, peer_id: ID, pubkey: PublicKey) -> None:
|
|
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
|
|
|
|
def pubkey(self, peer_id: ID) -> PublicKey:
|
|
if peer_id in self.peer_pubkey_map:
|
|
return self.peer_pubkey_map[peer_id]
|
|
raise PeerStoreError("peer ID not found")
|
|
|
|
def add_privkey(self, peer_id: ID, privkey: PrivateKey) -> None:
|
|
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
|
|
|
|
def privkey(self, peer_id: ID) -> PrivateKey:
|
|
if peer_id in self.peer_pubkey_map:
|
|
return self.peer_privkey_map[peer_id]
|
|
raise PeerStoreError("peer ID not found")
|
|
|
|
def peers_with_keys(self) -> Tuple[ID]:
|
|
return set(self.peer_pubkey_map.keys()).union(self.peer_privkey_map.keys())
|
|
|
|
|
|
class PeerStoreError(KeyError):
|
|
"""Raised when peer ID is not found in peer store."""
|