wallet tools cache hack

almando · Feb 18, 2020 · abbb85c · abbb85c
1 parent 1e3de32
commit abbb85c
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/tests/wallet_tools.py b/tests/wallet_tools.py
@@ -43,6 +43,7 @@ def __init__(self):
         self.extended_secret_key = ExtendedPrivateKey.from_seed(self.seed)
         self.generator_lookups: Dict = {}
         self.name = "MyChiaWallet"
+        self.puzzle_pk_cache = {}
 
     def get_next_public_key(self):
         pubkey = self.extended_secret_key.public_child(
@@ -69,13 +70,19 @@ def can_generate_puzzle_hash(self, hash):
         )
 
     def get_keys(self, puzzle_hash):
-        for child in range(self.next_address):
+        if puzzle_hash in self.puzzle_pk_cache:
+            child = self.puzzle_pk_cache[puzzle_hash]
             pubkey = self.extended_secret_key.public_child(child).get_public_key()
-            if puzzle_hash == puzzle_for_pk(pubkey.serialize()).get_hash():
-                return (
-                    pubkey,
-                    self.extended_secret_key.private_child(child).get_private_key(),
-                )
+            private = self.extended_secret_key.private_child(child).get_private_key()
+            return pubkey, private
+        else:
+            for child in range(self.next_address):
+                pubkey = self.extended_secret_key.public_child(child).get_public_key()
+                if puzzle_hash == puzzle_for_pk(pubkey.serialize()).get_hash():
+                    return (
+                        pubkey,
+                        self.extended_secret_key.private_child(child).get_private_key(),
+                    )
 
     def puzzle_for_pk(self, pubkey):
         return puzzle_for_pk(pubkey)
@@ -84,6 +91,7 @@ def get_new_puzzle(self):
         pubkey_a = self.get_next_public_key()
         pubkey = pubkey_a.serialize()
         puzzle = puzzle_for_pk(pubkey)
+        self.puzzle_pk_cache[puzzle.get_hash()] = self.next_address - 1
         return puzzle
 
     def get_new_puzzlehash(self):