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 libp2p.crypto.keys import PrivateKey, PublicKey
|
||||
|
||||
from .id import ID
|
||||
from .peerdata import PeerData, PeerDataError
|
||||
from .peerinfo import PeerInfo
|
||||
|
@ -10,10 +12,14 @@ 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:
|
||||
|
@ -94,6 +100,29 @@ class PeerStore(IPeerStore):
|
|||
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."""
|
||||
|
|
Loading…
Reference in New Issue
Block a user