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
async def set_user_info(
self, interval: uint64, limit: uint64, origin_id: str, admin_pubkey: str
):
if admin_pubkey.startswith("0x"):
admin_pubkey = admin_pubkey[2:]
admin_pubkey_bytes = bytes.fromhex(admin_pubkey)
assert self.rl_info.user_pubkey is not None
rl_puzzle = rl_puzzle_for_pk(
pubkey=self.rl_info.user_pubkey,
rate_amount=limit,
interval_time=interval,
origin_id=bytes.fromhex(origin_id),
clawback_pk=admin_pubkey_bytes,
)
rl_puzzle_hash = rl_puzzle.get_hash()
new_rl_info = RLInfo(
"admin",
admin_pubkey_bytes,
self.rl_info.user_pubkey,
limit,
interval,
None,
origin_id,
rl_puzzle_hash,
)
rl_puzzle_hash = rl_puzzle.get_hash()
index = await self.wallet_state_manager.puzzle_store.index_for_pubkey(
self.rl_info.user_pubkey.hex()
)
assert index is not None
record = DerivationRecord(
index,
rl_puzzle_hash,
self.rl_info.user_pubkey,
WalletType.RATE_LIMITED,
self.wallet_info.id,
)
await self.wallet_state_manager.puzzle_store.add_derivation_paths([record])
data_str = json.dumps(new_rl_info.to_json_dict())
new_wallet_info = WalletInfo(
self.wallet_info.id, self.wallet_info.name, self.wallet_info.type, data_str
)
await self.wallet_state_manager.user_store.update_wallet(new_wallet_info)
self.wallet_info = new_wallet_info
self.rl_info = new_rl_info
async def create_rl_user(
config: Dict,
key_config: Dict,
wallet_state_manager: Any,
wallet: Wallet,
name: str = None,
):
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
sk_hex = key_config["wallet_sk"]
private_key = PrivateKey.from_bytes(bytes.fromhex(sk_hex))
pubkey_bytes: bytes = bytes(
master_sk_to_wallet_sk(private_key, unused).get_g1()
)
rl_info = RLInfo("user", None, pubkey_bytes, None, None, None, None, None)
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(
config, key_config, wallet_state_manager, wallet_info, wallet, name
)
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)
return self
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_bytes: bytes = bytes(
master_sk_to_wallet_sk(private_key, unused).get_g1()
)
rl_info = RLInfo(
"user", None, pubkey_bytes, 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,
token_bytes(),
pubkey_bytes,
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 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.
"""
for wallet_id in self.wallets.keys():
target_wallet = self.wallets[wallet_id]
unused: Optional[
uint32] = await self.puzzle_store.get_unused_derivation_path()
last: Optional[
uint32] = await self.puzzle_store.get_last_derivation_path()
to_generate = 500
start_index = 0
derivation_paths: List[DerivationRecord] = []
if last is None:
assert unused is None
if unused is not None:
assert last is not None
start_index = last + 1
to_generate -= last - unused
# If the key was replaced (from_zero=True), we should generate the puzzle hashes for the new key
end = start_index + to_generate
if from_zero:
start_index = 0
for index in range(start_index, end):
pubkey: PublicKey = target_wallet.get_public_key(index)
puzzle: Program = target_wallet.puzzle_for_pk(bytes(pubkey))
puzzlehash: bytes32 = puzzle.get_hash()
self.log.info(
f"Generating public key at index {index} puzzle hash {puzzlehash.hex()}"
)
derivation_paths.append(
DerivationRecord(
uint32(index),
puzzlehash,
pubkey,
WalletType.STANDARD_WALLET,
uint32(target_wallet.wallet_info.id),
))
await self.puzzle_store.add_derivation_paths(derivation_paths)
if from_zero and unused is not None and unused > 0:
await self.puzzle_store.set_used_up_to(uint32(unused - 1))
async def create_rl_admin(
config: Dict,
key_config: Dict,
wallet_state_manager: Any,
wallet: Wallet,
name: str = None,
):
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
sk_hex = key_config["wallet_sk"]
private_key = ExtendedPrivateKey.from_bytes(bytes.fromhex(sk_hex))
pubkey_bytes: bytes = bytes(
private_key.public_child(unused).get_public_key())
rl_info = RLInfo("admin", pubkey_bytes, None, None, None, None, None,
None)
info_as_string = json.dumps(rl_info.to_json_dict())
await wallet_state_manager.user_store.create_wallet(
"RL Admin", WalletType.RATE_LIMITED, info_as_string)
wallet_info = await wallet_state_manager.user_store.get_last_wallet()
if wallet_info is None:
raise
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(config, key_config, wallet_state_manager,
wallet_info, wallet, name)
return self
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.value, 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.wallet_info.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_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 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."
)
index = await self.wallet_state_manager.puzzle_store.index_for_pubkey(
G1Element.from_bytes(self.rl_info.user_pubkey)
)
assert index is not None
record = DerivationRecord(
index,
rl_puzzle_hash,
self.rl_info.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,
self.rl_info.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)
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,
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)
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
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),
PrivateKey.from_seed(token_bytes(32)).get_g1(),
WalletType.STANDARD_WALLET,
uint32(1),
)
)
derivation_recs.append(
DerivationRecord(
uint32(i),
token_bytes(32),
PrivateKey.from_seed(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
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()
