Simplify wallet_tools.

This commit is contained in:
Richard Kiss 2020-09-08 15:07:47 -07:00 committed by Gene Hoffman
parent 68bffd7158
commit 06d5bef204

View File

@ -1,6 +1,5 @@
from typing import List, Optional, Dict, Tuple from typing import List, Optional, Dict, Tuple
from secrets import token_bytes
from blspy import PrivateKey, AugSchemeMPL from blspy import PrivateKey, AugSchemeMPL
from src.types.condition_var_pair import ConditionVarPair from src.types.condition_var_pair import ConditionVarPair
@ -31,8 +30,11 @@ from src.wallet.puzzles.puzzle_utils import (
from src.wallet.derive_keys import master_sk_to_wallet_sk from src.wallet.derive_keys import master_sk_to_wallet_sk
DEFAULT_SEED = b"seed" * 8
assert len(DEFAULT_SEED) == 32
class WalletTool: class WalletTool:
seed = b"seed"
next_address = 0 next_address = 0
pubkey_num_lookup: Dict[str, int] = {} pubkey_num_lookup: Dict[str, int] = {}
@ -42,33 +44,14 @@ class WalletTool:
if sk is not None: if sk is not None:
self.private_key = sk self.private_key = sk
else: else:
self.private_key = AugSchemeMPL.key_gen(token_bytes(32)) self.private_key = AugSchemeMPL.key_gen(DEFAULT_SEED)
self.generator_lookups: Dict = {} self.generator_lookups: Dict = {}
self.name = "MyChiaWallet"
self.puzzle_pk_cache: Dict = {} self.puzzle_pk_cache: Dict = {}
self.get_new_puzzle() self.get_new_puzzle()
def get_next_public_key(self): def get_next_address_index(self):
pubkey = master_sk_to_wallet_sk(self.private_key, self.next_address).get_g1()
self.pubkey_num_lookup[bytes(pubkey)] = self.next_address
self.next_address = self.next_address + 1 self.next_address = self.next_address + 1
return pubkey return self.next_address
def set_name(self, name):
self.name = name
def can_generate_puzzle_hash(self, hash):
return any(
map(
lambda child: hash
== puzzle_for_pk(
bytes(
master_sk_to_wallet_sk(self.private_key, uint32(child)).get_g1()
)
).get_tree_hash(),
reversed(range(self.next_address)),
)
)
def get_keys(self, puzzle_hash): def get_keys(self, puzzle_hash):
if puzzle_hash in self.puzzle_pk_cache: if puzzle_hash in self.puzzle_pk_cache:
@ -91,17 +74,14 @@ class WalletTool:
def puzzle_for_pk(self, pubkey: bytes): def puzzle_for_pk(self, pubkey: bytes):
return puzzle_for_pk(pubkey) return puzzle_for_pk(pubkey)
def get_first_puzzle(self):
pubkey = master_sk_to_wallet_sk(self.private_key, self.next_address).get_g1()
puzzle = puzzle_for_pk(bytes(pubkey))
self.puzzle_pk_cache[puzzle.get_tree_hash()] = 0
return puzzle
def get_new_puzzle(self): def get_new_puzzle(self):
pubkey_a = self.get_next_public_key() next_address_index = self.get_next_address_index()
pubkey = bytes(pubkey_a) pubkey = master_sk_to_wallet_sk(self.private_key, next_address_index).get_g1()
puzzle = puzzle_for_pk(pubkey) self.pubkey_num_lookup[bytes(pubkey)] = next_address_index
self.puzzle_pk_cache[puzzle.get_tree_hash()] = self.next_address - 1
puzzle = puzzle_for_pk(bytes(pubkey))
self.puzzle_pk_cache[puzzle.get_tree_hash()] = next_address_index
return puzzle return puzzle
def get_new_puzzlehash(self): def get_new_puzzlehash(self):