async def update_wallet_puzzle_hashes(self, wallet_id): derivation_paths: List[DerivationRecord] = [] target_wallet = self.wallets[wallet_id] last: Optional[ uint32] = await self.puzzle_store.get_last_derivation_path_for_wallet( wallet_id) unused: Optional[ uint32] = await self.puzzle_store.get_unused_derivation_path() if unused is None: # This handles the case where the database has entries but they have all been used unused = await self.puzzle_store.get_last_derivation_path() if unused is None: # This handles the case where the database is empty unused = uint32(0) for index in range(unused, last): pubkey: G1Element = self.get_public_key(uint32(index)) puzzle: Program = target_wallet.puzzle_for_pk(bytes(pubkey)) puzzlehash: bytes32 = puzzle.get_tree_hash() self.log.info( f"Generating public key at index {index} puzzle hash {puzzlehash.hex()}" ) derivation_paths.append( DerivationRecord( uint32(index), puzzlehash, pubkey, target_wallet.wallet_info.type, uint32(target_wallet.wallet_info.id), )) await self.puzzle_store.add_derivation_paths(derivation_paths)
async def get_derivation_record( self, index: uint32, wallet_id: uint32, hardened: bool ) -> Optional[DerivationRecord]: """ Returns the derivation record by index and wallet id. """ if hardened: hard = 1 else: hard = 0 cursor = await self.db_connection.execute( "SELECT * FROM derivation_paths WHERE derivation_index=? and wallet_id=? and hardened=?;", (index, wallet_id, hard), ) row = await cursor.fetchone() await cursor.close() if row is not None and row[0] is not None: return DerivationRecord( uint32(row[0]), bytes32.fromhex(row[2]), G1Element.from_bytes(bytes.fromhex(row[1])), WalletType(row[3]), uint32(row[4]), bool(row[5]), ) return None
async def get_derivation_record( self, index: uint32, wallet_id: uint32) -> Optional[DerivationRecord]: """ Returns the derivation record by index and wallet id. """ cursor = await self.db_connection.execute( "SELECT * FROM derivation_paths WHERE derivation_index=? and wallet_id=?;", ( index, wallet_id, ), ) row = await cursor.fetchone() await cursor.close() if row is not None and row[0] is not None: return DerivationRecord( row[0], bytes.fromhex(row[2]), G1Element.from_bytes(bytes.fromhex(row[1])), row[3], row[4], ) return None
def row_to_record(self, row) -> DerivationRecord: return DerivationRecord( uint32(row[0]), bytes32.fromhex(row[2]), G1Element.from_bytes(bytes.fromhex(row[1])), WalletType(row[3]), uint32(row[4]), bool(row[6]), )
async def create_rl_user(wallet_state_manager: Any, ): async with wallet_state_manager.puzzle_store.lock: unused: Optional[ uint32] = await wallet_state_manager.puzzle_store.get_unused_derivation_path( ) if unused is None: await wallet_state_manager.create_more_puzzle_hashes() unused = await wallet_state_manager.puzzle_store.get_unused_derivation_path( ) assert unused is not None private_key = wallet_state_manager.private_key pubkey: G1Element = master_sk_to_wallet_sk(private_key, unused).get_g1() rl_info = RLInfo("user", None, bytes(pubkey), None, None, None, None, None, False) info_as_string = json.dumps(rl_info.to_json_dict()) await wallet_state_manager.user_store.create_wallet( "RL User", WalletType.RATE_LIMITED, info_as_string) wallet_info = await wallet_state_manager.user_store.get_last_wallet( ) if wallet_info is None: raise Exception("wallet_info is None") self = await RLWallet.create(wallet_state_manager, wallet_info) await wallet_state_manager.puzzle_store.add_derivation_paths([ DerivationRecord( unused, bytes32(token_bytes(32)), pubkey, WalletType.RATE_LIMITED, wallet_info.id, ) ]) await wallet_state_manager.puzzle_store.set_used_up_to(unused) await wallet_state_manager.add_new_wallet(self, self.id()) return self
async def get_derivation_record_for_puzzle_hash( self, puzzle_hash: str) -> Optional[DerivationRecord]: """ Returns the derivation record by index and wallet id. """ cursor = await self.db_connection.execute( "SELECT * FROM derivation_paths WHERE puzzle_hash=?;", (puzzle_hash, ), ) row = await cursor.fetchone() await cursor.close() if row is not None and row[0] is not None: return DerivationRecord( row[0], bytes.fromhex(row[2]), G1Element.from_bytes(bytes.fromhex(row[1])), row[3], row[4], ) return None
async def create_rl_admin(wallet_state_manager: Any, ): unused: Optional[ uint32] = await wallet_state_manager.puzzle_store.get_unused_derivation_path( ) if unused is None: await wallet_state_manager.create_more_puzzle_hashes() unused = await wallet_state_manager.puzzle_store.get_unused_derivation_path( ) assert unused is not None private_key = master_sk_to_wallet_sk(wallet_state_manager.private_key, unused) pubkey_bytes: bytes = bytes(private_key.get_g1()) rl_info = RLInfo("admin", pubkey_bytes, None, None, None, None, None, None, False) info_as_string = json.dumps(rl_info.to_json_dict()) wallet_info: Optional[ WalletInfo] = await wallet_state_manager.user_store.create_wallet( "RL Admin", WalletType.RATE_LIMITED, info_as_string) if wallet_info is None: raise Exception("wallet_info is None") await wallet_state_manager.puzzle_store.add_derivation_paths([ DerivationRecord( unused, token_bytes(), pubkey_bytes, WalletType.RATE_LIMITED, wallet_info.id, ) ]) await wallet_state_manager.puzzle_store.set_used_up_to(unused) self = await RLWallet.create(wallet_state_manager, wallet_info) await wallet_state_manager.add_new_wallet(self, self.id()) return self
async def create_more_puzzle_hashes(self, from_zero: bool = False): """ For all wallets in the user store, generates the first few puzzle hashes so that we can restore the wallet from only the private keys. """ targets = list(self.wallets.keys()) unused: Optional[ uint32] = await self.puzzle_store.get_unused_derivation_path() if unused is None: # This handles the case where the database has entries but they have all been used unused = await self.puzzle_store.get_last_derivation_path() if unused is None: # This handles the case where the database is empty unused = uint32(0) if self.new_wallet: to_generate = self.config["initial_num_public_keys_new_wallet"] else: to_generate = self.config["initial_num_public_keys"] for wallet_id in targets: target_wallet = self.wallets[wallet_id] last: Optional[ uint32] = await self.puzzle_store.get_last_derivation_path_for_wallet( wallet_id) start_index = 0 derivation_paths: List[DerivationRecord] = [] if last is not None: start_index = last + 1 # If the key was replaced (from_zero=True), we should generate the puzzle hashes for the new key if from_zero: start_index = 0 for index in range(start_index, unused + to_generate): if WalletType(target_wallet.type()) == WalletType.RATE_LIMITED: if target_wallet.rl_info.initialized is False: break wallet_type = target_wallet.rl_info.type if wallet_type == "user": rl_pubkey = G1Element.from_bytes( target_wallet.rl_info.user_pubkey) else: rl_pubkey = G1Element.from_bytes( target_wallet.rl_info.admin_pubkey) rl_puzzle: Program = target_wallet.puzzle_for_pk(rl_pubkey) puzzle_hash: bytes32 = rl_puzzle.get_tree_hash() rl_index = self.get_derivation_index(rl_pubkey) if rl_index == -1: break derivation_paths.append( DerivationRecord( uint32(rl_index), puzzle_hash, rl_pubkey, target_wallet.type(), uint32(target_wallet.id()), )) break pubkey: G1Element = self.get_public_key(uint32(index)) puzzle: Program = target_wallet.puzzle_for_pk(bytes(pubkey)) if puzzle is None: self.log.warning( f"Unable to create puzzles with wallet {target_wallet}" ) break puzzlehash: bytes32 = puzzle.get_tree_hash() self.log.info( f"Puzzle at index {index} wallet ID {wallet_id} puzzle hash {puzzlehash.hex()}" ) derivation_paths.append( DerivationRecord( uint32(index), puzzlehash, pubkey, target_wallet.type(), uint32(target_wallet.id()), )) await self.puzzle_store.add_derivation_paths(derivation_paths) if unused > 0: await self.puzzle_store.set_used_up_to(uint32(unused - 1))
async def test_puzzle_store(self): db_filename = Path("puzzle_store_test.db") if db_filename.exists(): db_filename.unlink() con = await aiosqlite.connect(db_filename) db = await WalletPuzzleStore.create(con) try: derivation_recs = [] # wallet_types = [t for t in WalletType] [t for t in WalletType] for i in range(1000): derivation_recs.append( DerivationRecord( uint32(i), token_bytes(32), AugSchemeMPL.key_gen(token_bytes(32)).get_g1(), WalletType.STANDARD_WALLET, uint32(1), ) ) derivation_recs.append( DerivationRecord( uint32(i), token_bytes(32), AugSchemeMPL.key_gen(token_bytes(32)).get_g1(), WalletType.RATE_LIMITED, uint32(2), ) ) assert await db.puzzle_hash_exists(derivation_recs[0].puzzle_hash) is False assert await db.index_for_pubkey(derivation_recs[0].pubkey) is None assert await db.index_for_puzzle_hash(derivation_recs[2].puzzle_hash) is None assert await db.wallet_info_for_puzzle_hash(derivation_recs[2].puzzle_hash) is None assert len((await db.get_all_puzzle_hashes())) == 0 assert await db.get_last_derivation_path() is None assert await db.get_unused_derivation_path() is None assert await db.get_derivation_record(0, 2) is None await db.add_derivation_paths(derivation_recs) assert await db.puzzle_hash_exists(derivation_recs[0].puzzle_hash) is True phs_1 = [derivation_recs[0].puzzle_hash] phs_2 = [32 * bytes([1]), derivation_recs[0].puzzle_hash] phs_3 = [derivation_recs[0].puzzle_hash, 32 * bytes([1])] phs_4 = [32 * bytes([1]), 32 * bytes([2])] phs_5 = [] assert await db.one_of_puzzle_hashes_exists(phs_1) is True assert await db.one_of_puzzle_hashes_exists(phs_2) is True assert await db.one_of_puzzle_hashes_exists(phs_3) is True assert await db.one_of_puzzle_hashes_exists(phs_4) is False assert await db.one_of_puzzle_hashes_exists(phs_5) is False assert await db.index_for_pubkey(derivation_recs[4].pubkey) == 2 assert await db.index_for_puzzle_hash(derivation_recs[2].puzzle_hash) == 1 assert await db.wallet_info_for_puzzle_hash(derivation_recs[2].puzzle_hash) == ( derivation_recs[2].wallet_id, derivation_recs[2].wallet_type, ) assert len((await db.get_all_puzzle_hashes())) == 2000 assert await db.get_last_derivation_path() == 999 assert await db.get_unused_derivation_path() == 0 assert await db.get_derivation_record(0, 2) == derivation_recs[1] # Indeces up to 250 await db.set_used_up_to(249) assert await db.get_unused_derivation_path() == 250 except Exception as e: print(e, type(e)) await db._clear_database() await db.close() db_filename.unlink() raise e await db._clear_database() await db.close() db_filename.unlink()
async def test_make_double_output(self, two_wallet_nodes): num_blocks = 5 full_nodes, wallets = two_wallet_nodes full_node_1 = full_nodes[0] server_1 = full_node_1.server wallet_node, server_2 = wallets[0] wallet_node_2, server_3 = wallets[1] wallet = wallet_node.wallet_state_manager.main_wallet wallet2 = wallet_node_2.wallet_state_manager.main_wallet ph = await wallet.get_new_puzzlehash() await server_2.start_client( PeerInfo("localhost", uint16(server_1._port)), None) await server_3.start_client( PeerInfo("localhost", uint16(server_1._port)), None) for i in range(1, num_blocks): await full_node_1.farm_new_transaction_block( FarmNewBlockProtocol(ph)) funds = sum([ calculate_pool_reward(uint32(i)) + calculate_base_farmer_reward(uint32(i)) for i in range(1, num_blocks - 1) ]) await time_out_assert(15, wallet.get_confirmed_balance, funds) did_wallet: DIDWallet = await DIDWallet.create_new_did_wallet( wallet_node.wallet_state_manager, wallet, uint64(101)) ph2 = await wallet2.get_new_puzzlehash() for i in range(1, num_blocks): await full_node_1.farm_new_transaction_block( FarmNewBlockProtocol(ph2)) await time_out_assert(15, did_wallet.get_confirmed_balance, 101) await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101) await time_out_assert(15, did_wallet.get_spendable_balance, 101) # Lock up with non DID innerpuz so that we can create two outputs # Innerpuz will output the innersol, so we just pass in ((51 0xMyPuz 49) (51 0xMyPuz 51)) innerpuz = Program.to(binutils.assemble("1")) innerpuzhash = innerpuz.get_tree_hash() puz = did_wallet_puzzles.create_fullpuz( innerpuzhash, did_wallet.did_info.my_did, ) # Add the hacked puzzle to the puzzle store so that it is recognised as "our" puzzle old_devrec = await did_wallet.wallet_state_manager.get_unused_derivation_record( did_wallet.wallet_info.id) devrec = DerivationRecord( old_devrec.index, puz.get_tree_hash(), old_devrec.pubkey, old_devrec.wallet_type, old_devrec.wallet_id, ) await did_wallet.wallet_state_manager.puzzle_store.add_derivation_paths( [devrec]) await did_wallet.create_spend(puz.get_tree_hash()) for i in range(1, num_blocks): await full_node_1.farm_new_transaction_block( FarmNewBlockProtocol(ph2)) await time_out_assert(15, did_wallet.get_confirmed_balance, 101) await time_out_assert(15, did_wallet.get_unconfirmed_balance, 101) await time_out_assert(15, did_wallet.get_spendable_balance, 101) # Create spend by hand so that we can use the weird innersol coins = await did_wallet.select_coins(1) coin = coins.pop() # innerpuz is our desired output innersol = Program.to([[51, coin.puzzle_hash, 45], [51, coin.puzzle_hash, 56]]) # full solution is (corehash parent_info my_amount innerpuz_reveal solution) parent_info = await did_wallet.get_parent_for_coin(coin) fullsol = Program.to([ [ parent_info.parent_name, parent_info.inner_puzzle_hash, parent_info.amount, ], coin.amount, innersol, ]) try: cost, result = puz.run_with_cost(fullsol) except Exception as e: assert e.args == ("path into atom", ) else: assert False
async def set_user_info( self, interval: uint64, limit: uint64, origin_parent_id: str, origin_puzzle_hash: str, origin_amount: uint64, admin_pubkey: str, ) -> None: admin_pubkey_bytes = hexstr_to_bytes(admin_pubkey) assert self.rl_info.user_pubkey is not None origin = Coin( hexstr_to_bytes(origin_parent_id), hexstr_to_bytes(origin_puzzle_hash), origin_amount, ) rl_puzzle = rl_puzzle_for_pk( pubkey=self.rl_info.user_pubkey, rate_amount=limit, interval_time=interval, origin_id=origin.name(), clawback_pk=admin_pubkey_bytes, ) rl_puzzle_hash = rl_puzzle.get_tree_hash() new_rl_info = RLInfo( "user", admin_pubkey_bytes, self.rl_info.user_pubkey, limit, interval, origin, origin.name(), rl_puzzle_hash, True, ) rl_puzzle_hash = rl_puzzle.get_tree_hash() if await self.wallet_state_manager.puzzle_store.puzzle_hash_exists( rl_puzzle_hash): raise ValueError( "Cannot create multiple Rate Limited wallets under the same keys. This will change in a future release." ) user_pubkey: G1Element = G1Element.from_bytes(self.rl_info.user_pubkey) index = await self.wallet_state_manager.puzzle_store.index_for_pubkey( user_pubkey) assert index is not None record = DerivationRecord( index, rl_puzzle_hash, user_pubkey, WalletType.RATE_LIMITED, self.id(), ) aggregation_puzzlehash = self.rl_get_aggregation_puzzlehash( new_rl_info.rl_puzzle_hash) record2 = DerivationRecord( index + 1, aggregation_puzzlehash, user_pubkey, WalletType.RATE_LIMITED, self.id(), ) await self.wallet_state_manager.puzzle_store.add_derivation_paths( [record, record2]) self.wallet_state_manager.set_coin_with_puzzlehash_created_callback( aggregation_puzzlehash, self.aggregate_this_coin) data_str = json.dumps(new_rl_info.to_json_dict()) new_wallet_info = WalletInfo(self.id(), self.wallet_info.name, self.type(), data_str) await self.wallet_state_manager.user_store.update_wallet( new_wallet_info, False) await self.wallet_state_manager.add_new_wallet(self, self.id()) self.wallet_info = new_wallet_info self.rl_info = new_rl_info
async def admin_create_coin( self, interval: uint64, limit: uint64, user_pubkey: str, amount: uint64, fee: uint64, ) -> bool: coins = await self.wallet_state_manager.main_wallet.select_coins(amount ) if coins is None: return False origin = coins.copy().pop() origin_id = origin.name() user_pubkey_bytes = hexstr_to_bytes(user_pubkey) assert self.rl_info.admin_pubkey is not None rl_puzzle = rl_puzzle_for_pk( pubkey=user_pubkey_bytes, rate_amount=limit, interval_time=interval, origin_id=origin_id, clawback_pk=self.rl_info.admin_pubkey, ) rl_puzzle_hash = rl_puzzle.get_tree_hash() index = await self.wallet_state_manager.puzzle_store.index_for_pubkey( G1Element.from_bytes(self.rl_info.admin_pubkey)) assert index is not None record = DerivationRecord( index, rl_puzzle_hash, G1Element.from_bytes(self.rl_info.admin_pubkey), WalletType.RATE_LIMITED, self.id(), ) await self.wallet_state_manager.puzzle_store.add_derivation_paths( [record]) spend_bundle = await self.main_wallet.generate_signed_transaction( amount, rl_puzzle_hash, fee, origin_id, coins) if spend_bundle is None: return False await self.main_wallet.push_transaction(spend_bundle) new_rl_info = RLInfo( "admin", self.rl_info.admin_pubkey, user_pubkey_bytes, limit, interval, origin, origin.name(), rl_puzzle_hash, True, ) data_str = json.dumps(new_rl_info.to_json_dict()) new_wallet_info = WalletInfo(self.id(), self.wallet_info.name, self.type(), data_str) await self.wallet_state_manager.user_store.update_wallet( new_wallet_info, False) await self.wallet_state_manager.add_new_wallet(self, self.id()) self.wallet_info = new_wallet_info self.rl_info = new_rl_info return True