From 144d93a02370679b44d20c2624a842bb4d96592b Mon Sep 17 00:00:00 2001
From: NIC619 <twedusuck@gmail.com>
Date: Sun, 24 Nov 2019 18:04:00 +0800
Subject: [PATCH] Add pubkey/privkey info to `PeerStore`

---
 libp2p/peer/peerstore.py | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/libp2p/peer/peerstore.py b/libp2p/peer/peerstore.py
index 6d6b1f1..8a19a3c 100644
--- a/libp2p/peer/peerstore.py
+++ b/libp2p/peer/peerstore.py
@@ -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."""