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