Add pubkey/privkey info to PeerStore
This commit is contained in:
parent
9837f30698
commit
144d93a023
|
@ -1,7 +1,9 @@
|
||||||
from typing import Any, Dict, List, Optional, Sequence
|
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||||
|
|
||||||
from multiaddr import Multiaddr
|
from multiaddr import Multiaddr
|
||||||
|
|
||||||
|
from libp2p.crypto.keys import PrivateKey, PublicKey
|
||||||
|
|
||||||
from .id import ID
|
from .id import ID
|
||||||
from .peerdata import PeerData, PeerDataError
|
from .peerdata import PeerData, PeerDataError
|
||||||
from .peerinfo import PeerInfo
|
from .peerinfo import PeerInfo
|
||||||
|
@ -10,10 +12,14 @@ from .peerstore_interface import IPeerStore
|
||||||
|
|
||||||
class PeerStore(IPeerStore):
|
class PeerStore(IPeerStore):
|
||||||
|
|
||||||
|
peer_pubkey_map: Dict[ID, PublicKey]
|
||||||
|
peer_privkey_map: Dict[ID, PrivateKey]
|
||||||
peer_data_map: Dict[ID, PeerData]
|
peer_data_map: Dict[ID, PeerData]
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
IPeerStore.__init__(self)
|
IPeerStore.__init__(self)
|
||||||
|
self.peer_pubkey_map = {}
|
||||||
|
self.peer_privkey_map = {}
|
||||||
self.peer_data_map = {}
|
self.peer_data_map = {}
|
||||||
|
|
||||||
def __create_or_get_peer(self, peer_id: ID) -> PeerData:
|
def __create_or_get_peer(self, peer_id: ID) -> PeerData:
|
||||||
|
@ -94,6 +100,29 @@ class PeerStore(IPeerStore):
|
||||||
output.append(peer_id)
|
output.append(peer_id)
|
||||||
return output
|
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):
|
class PeerStoreError(KeyError):
|
||||||
"""Raised when peer ID is not found in peer store."""
|
"""Raised when peer ID is not found in peer store."""
|
||||||
|
|
Loading…
Reference in New Issue
Block a user