406 lines
15 KiB
Python
406 lines
15 KiB
Python
import time
|
|
from typing import Dict, Optional, List, Set, Any
|
|
import logging
|
|
from blspy import G1Element, G2Element, AugSchemeMPL, PrivateKey
|
|
from src.types.coin import Coin
|
|
from src.types.coin_solution import CoinSolution
|
|
from src.types.program import Program
|
|
from src.types.spend_bundle import SpendBundle
|
|
from src.types.sized_bytes import bytes32
|
|
from src.util.condition_tools import (
|
|
conditions_for_solution,
|
|
conditions_dict_for_solution,
|
|
conditions_by_opcode,
|
|
pkm_pairs_for_conditions_dict,
|
|
)
|
|
from src.util.ints import uint64, uint32
|
|
from src.wallet.abstract_wallet import AbstractWallet
|
|
from src.wallet.puzzles.p2_conditions import puzzle_for_conditions
|
|
from src.wallet.puzzles.p2_delegated_puzzle import puzzle_for_pk
|
|
from src.wallet.puzzles.puzzle_utils import (
|
|
make_assert_my_coin_id_condition,
|
|
make_assert_time_exceeds_condition,
|
|
make_assert_coin_consumed_condition,
|
|
make_create_coin_condition,
|
|
make_assert_fee_condition,
|
|
)
|
|
from src.wallet.transaction_record import TransactionRecord
|
|
from src.wallet.wallet_coin_record import WalletCoinRecord
|
|
from src.wallet.wallet_info import WalletInfo
|
|
|
|
|
|
class Wallet(AbstractWallet):
|
|
wallet_state_manager: Any
|
|
log: logging.Logger
|
|
wallet_info: WalletInfo
|
|
_pk2sk: Dict[G1Element, PrivateKey]
|
|
|
|
@staticmethod
|
|
async def create(
|
|
wallet_state_manager: Any,
|
|
info: WalletInfo,
|
|
name: str = None,
|
|
):
|
|
self = Wallet()
|
|
if name:
|
|
self.log = logging.getLogger(name)
|
|
else:
|
|
self.log = logging.getLogger(__name__)
|
|
self.wallet_state_manager = wallet_state_manager
|
|
self.wallet_info = info
|
|
|
|
# HACK
|
|
self._pk2sk = {}
|
|
|
|
return self
|
|
|
|
async def get_confirmed_balance(self) -> uint64:
|
|
return await self.wallet_state_manager.get_confirmed_balance_for_wallet(
|
|
self.wallet_info.id
|
|
)
|
|
|
|
async def get_unconfirmed_balance(self) -> uint64:
|
|
return await self.wallet_state_manager.get_unconfirmed_balance(
|
|
self.wallet_info.id
|
|
)
|
|
|
|
async def get_frozen_amount(self) -> uint64:
|
|
return await self.wallet_state_manager.get_frozen_balance(self.wallet_info.id)
|
|
|
|
async def get_spendable_balance(self) -> uint64:
|
|
spendable_am = (
|
|
await self.wallet_state_manager.get_confirmed_spendable_balance_for_wallet(
|
|
self.wallet_info.id
|
|
)
|
|
)
|
|
return spendable_am
|
|
|
|
async def get_pending_change_balance(self) -> uint64:
|
|
unconfirmed_tx = (
|
|
await self.wallet_state_manager.tx_store.get_unconfirmed_for_wallet(
|
|
self.wallet_info.id
|
|
)
|
|
)
|
|
addition_amount = 0
|
|
|
|
for record in unconfirmed_tx:
|
|
our_spend = False
|
|
for coin in record.removals:
|
|
if await self.wallet_state_manager.does_coin_belong_to_wallet(
|
|
coin, self.wallet_info.id
|
|
):
|
|
our_spend = True
|
|
break
|
|
|
|
if our_spend is not True:
|
|
continue
|
|
|
|
for coin in record.additions:
|
|
if await self.wallet_state_manager.does_coin_belong_to_wallet(
|
|
coin, self.wallet_info.id
|
|
):
|
|
addition_amount += coin.amount
|
|
|
|
return uint64(addition_amount)
|
|
|
|
def puzzle_for_pk(self, pubkey: bytes) -> Program:
|
|
return puzzle_for_pk(pubkey)
|
|
|
|
async def puzzle_for_puzzle_hash(self, puzzle_hash: bytes32) -> Program:
|
|
maybe = await self.wallet_state_manager.get_keys(puzzle_hash)
|
|
if maybe is None:
|
|
error_msg = f"Wallet couldn't find keys for puzzle_hash {puzzle_hash}"
|
|
self.log.error(error_msg)
|
|
raise ValueError(error_msg)
|
|
|
|
# Get puzzle for pubkey
|
|
public_key, secret_key = maybe
|
|
|
|
# HACK
|
|
self._pk2sk[bytes(public_key)] = secret_key
|
|
|
|
return puzzle_for_pk(bytes(public_key))
|
|
|
|
async def get_new_puzzle(self) -> Program:
|
|
dr = await self.wallet_state_manager.get_unused_derivation_record(
|
|
self.wallet_info.id
|
|
)
|
|
return puzzle_for_pk(bytes(dr.pubkey))
|
|
|
|
async def get_new_puzzlehash(self) -> bytes32:
|
|
return (
|
|
await self.wallet_state_manager.get_unused_derivation_record(
|
|
self.wallet_info.id
|
|
)
|
|
).puzzle_hash
|
|
|
|
def make_solution(
|
|
self, primaries=None, min_time=0, me=None, consumed=None, fee=0
|
|
):
|
|
assert fee >= 0
|
|
condition_list = []
|
|
if primaries:
|
|
for primary in primaries:
|
|
condition_list.append(
|
|
make_create_coin_condition(primary["puzzlehash"], primary["amount"])
|
|
)
|
|
if consumed:
|
|
for coin in consumed:
|
|
condition_list.append(make_assert_coin_consumed_condition(coin))
|
|
if min_time > 0:
|
|
condition_list.append(make_assert_time_exceeds_condition(min_time))
|
|
if me:
|
|
condition_list.append(make_assert_my_coin_id_condition(me["id"]))
|
|
if fee:
|
|
condition_list.append(make_assert_fee_condition(fee))
|
|
return Program.to([puzzle_for_conditions(condition_list), []])
|
|
|
|
async def select_coins(self, amount, exclude: List[Coin] = None) -> Set[Coin]:
|
|
""" Returns a set of coins that can be used for generating a new transaction. """
|
|
async with self.wallet_state_manager.lock:
|
|
if exclude is None:
|
|
exclude = []
|
|
|
|
spendable_amount = await self.get_spendable_balance()
|
|
|
|
if amount > spendable_amount:
|
|
error_msg = (
|
|
f"Can't select amount higher than our spendable balance. Amount: {amount}, spendable: "
|
|
f" {spendable_amount}"
|
|
)
|
|
self.log.warning(error_msg)
|
|
raise ValueError(error_msg)
|
|
|
|
self.log.info(f"About to select coins for amount {amount}")
|
|
unspent: List[WalletCoinRecord] = list(
|
|
await self.wallet_state_manager.get_spendable_coins_for_wallet(
|
|
self.wallet_info.id
|
|
)
|
|
)
|
|
sum_value = 0
|
|
used_coins: Set = set()
|
|
|
|
# Use older coins first
|
|
unspent.sort(key=lambda r: r.confirmed_block_index)
|
|
|
|
# Try to use coins from the store, if there isn't enough of "unused"
|
|
# coins use change coins that are not confirmed yet
|
|
unconfirmed_removals: Dict[
|
|
bytes32, Coin
|
|
] = await self.wallet_state_manager.unconfirmed_removals_for_wallet(
|
|
self.wallet_info.id
|
|
)
|
|
for coinrecord in unspent:
|
|
if sum_value >= amount and len(used_coins) > 0:
|
|
break
|
|
if coinrecord.coin.name() in unconfirmed_removals:
|
|
continue
|
|
if coinrecord.coin in exclude:
|
|
continue
|
|
sum_value += coinrecord.coin.amount
|
|
used_coins.add(coinrecord.coin)
|
|
self.log.info(
|
|
f"Selected coin: {coinrecord.coin.name()} at height {coinrecord.confirmed_block_index}!"
|
|
)
|
|
|
|
# This happens when we couldn't use one of the coins because it's already used
|
|
# but unconfirmed, and we are waiting for the change. (unconfirmed_additions)
|
|
if sum_value < amount:
|
|
raise ValueError(
|
|
"Can't make this transaction at the moment. Waiting for the change from the previous transaction."
|
|
)
|
|
# TODO(straya): remove this
|
|
# unconfirmed_additions = await self.wallet_state_manager.unconfirmed_additions_for_wallet(
|
|
# self.wallet_info.id
|
|
# )
|
|
# for coin in unconfirmed_additions.values():
|
|
# if sum_value > amount:
|
|
# break
|
|
# if coin.name() in unconfirmed_removals:
|
|
# continue
|
|
|
|
# sum_value += coin.amount
|
|
# used_coins.add(coin)
|
|
# self.log.info(f"Selected used coin: {coin.name()}")
|
|
|
|
self.log.info(f"Successfully selected coins: {used_coins}")
|
|
return used_coins
|
|
|
|
async def generate_unsigned_transaction(
|
|
self,
|
|
amount: uint64,
|
|
newpuzzlehash: bytes32,
|
|
fee: uint64 = uint64(0),
|
|
origin_id: bytes32 = None,
|
|
coins: Set[Coin] = None,
|
|
) -> List[CoinSolution]:
|
|
"""
|
|
Generates a unsigned transaction in form of List(Puzzle, Solutions)
|
|
"""
|
|
if coins is None:
|
|
coins = await self.select_coins(amount + fee)
|
|
assert len(coins) > 0
|
|
|
|
self.log.info(f"coins is not None {coins}")
|
|
spend_value = sum([coin.amount for coin in coins])
|
|
change = spend_value - amount - fee
|
|
|
|
spends: List[CoinSolution] = []
|
|
output_created = False
|
|
|
|
for coin in coins:
|
|
self.log.info(f"coin from coins {coin}")
|
|
puzzle: Program = await self.puzzle_for_puzzle_hash(coin.puzzle_hash)
|
|
|
|
# Only one coin creates outputs
|
|
if not output_created and origin_id in (None, coin.name()):
|
|
primaries = [{"puzzlehash": newpuzzlehash, "amount": amount}]
|
|
if change > 0:
|
|
changepuzzlehash = await self.get_new_puzzlehash()
|
|
primaries.append({"puzzlehash": changepuzzlehash, "amount": change})
|
|
|
|
solution = self.make_solution(primaries=primaries, fee=fee)
|
|
output_created = True
|
|
else:
|
|
solution = self.make_solution()
|
|
|
|
puzzle_solution_pair = Program.to([puzzle, solution])
|
|
spends.append(CoinSolution(coin, puzzle_solution_pair))
|
|
|
|
self.log.info(f"Spends is {spends}")
|
|
return spends
|
|
|
|
def secret_key_for_public_key(self, public_key: G1Element) -> Optional[PrivateKey]:
|
|
return self._pk2sk.get(bytes(public_key))
|
|
|
|
async def sign_transaction(
|
|
self, coin_solutions: List[CoinSolution]
|
|
) -> SpendBundle:
|
|
signatures = []
|
|
|
|
for coin_solution in coin_solutions:
|
|
# Get AGGSIG conditions
|
|
err, con, cost = conditions_for_solution(coin_solution.solution)
|
|
if err or not con:
|
|
error_msg = f"Sign transaction failed, con:{con}, error: {err}"
|
|
self.log.error(error_msg)
|
|
raise ValueError(error_msg)
|
|
|
|
conditions_dict = conditions_by_opcode(con)
|
|
|
|
# Create signature
|
|
for _, msg in pkm_pairs_for_conditions_dict(
|
|
conditions_dict, bytes(coin_solution.coin)
|
|
):
|
|
secret_key = self.secret_key_for_public_key(_)
|
|
if secret_key is None:
|
|
self.log.error(f"no secret key for for {_}")
|
|
return None
|
|
signature = AugSchemeMPL.sign(secret_key, msg)
|
|
signatures.append(signature)
|
|
|
|
# Aggregate signatures
|
|
aggsig = AugSchemeMPL.aggregate(signatures)
|
|
return SpendBundle(coin_solutions, aggsig)
|
|
|
|
async def generate_signed_transaction(
|
|
self,
|
|
amount: uint64,
|
|
puzzle_hash: bytes32,
|
|
fee: uint64 = uint64(0),
|
|
origin_id: bytes32 = None,
|
|
coins: Set[Coin] = None,
|
|
) -> TransactionRecord:
|
|
""" Use this to generate transaction. """
|
|
|
|
transaction = await self.generate_unsigned_transaction(
|
|
amount, puzzle_hash, fee, origin_id, coins
|
|
)
|
|
assert len(transaction) > 0
|
|
|
|
self.log.info("About to sign a transaction")
|
|
spend_bundle: SpendBundle = await self.sign_transaction(transaction)
|
|
|
|
now = uint64(int(time.time()))
|
|
add_list: List[Coin] = list(spend_bundle.additions())
|
|
rem_list: List[Coin] = list(spend_bundle.removals())
|
|
|
|
return TransactionRecord(
|
|
confirmed_at_index=uint32(0),
|
|
created_at_time=now,
|
|
to_puzzle_hash=puzzle_hash,
|
|
amount=uint64(amount),
|
|
fee_amount=uint64(fee),
|
|
incoming=False,
|
|
confirmed=False,
|
|
sent=uint32(0),
|
|
spend_bundle=spend_bundle,
|
|
additions=add_list,
|
|
removals=rem_list,
|
|
wallet_id=self.wallet_info.id,
|
|
sent_to=[],
|
|
trade_id=None,
|
|
)
|
|
|
|
async def push_transaction(self, tx: TransactionRecord) -> None:
|
|
""" Use this API to send transactions. """
|
|
await self.wallet_state_manager.add_pending_transaction(tx)
|
|
|
|
# This is also defined in CCWallet as get_sigs()
|
|
# I think this should be a the default way the wallet gets signatures in sign_transaction()
|
|
async def get_sigs_for_innerpuz_with_innersol(
|
|
self, innerpuz: Program, innersol: Program
|
|
) -> List[G2Element]:
|
|
puzzle_hash = innerpuz.get_tree_hash()
|
|
pubkey, private = await self.wallet_state_manager.get_keys(puzzle_hash)
|
|
sigs: List[G2Element] = []
|
|
code_ = [innerpuz, innersol]
|
|
sexp = Program.to(code_)
|
|
error, conditions, cost = conditions_dict_for_solution(sexp)
|
|
if conditions is not None:
|
|
for _, msg in pkm_pairs_for_conditions_dict(conditions):
|
|
signature = AugSchemeMPL.sign(private, msg)
|
|
sigs.append(signature)
|
|
return sigs
|
|
|
|
# Create an offer spend bundle for chia given an amount of relative change (i.e -400 or 1000)
|
|
|
|
# This is to be aggregated together with a coloured coin offer to ensure that the trade happens
|
|
async def create_spend_bundle_relative_chia(
|
|
self, chia_amount: int, exclude: List[Coin]
|
|
) -> SpendBundle:
|
|
list_of_solutions = []
|
|
utxos = None
|
|
|
|
# If we're losing value then get coins with at least that much value
|
|
# If we're gaining value then our amount doesn't matter
|
|
if chia_amount < 0:
|
|
utxos = await self.select_coins(abs(chia_amount), exclude)
|
|
else:
|
|
utxos = await self.select_coins(0, exclude)
|
|
|
|
assert len(utxos) > 0
|
|
|
|
# Calculate output amount given sum of utxos
|
|
spend_value = sum([coin.amount for coin in utxos])
|
|
chia_amount = spend_value + chia_amount
|
|
|
|
# Create coin solutions for each utxo
|
|
output_created = None
|
|
sigs: List[G2Element] = []
|
|
for coin in utxos:
|
|
puzzle = await self.puzzle_for_puzzle_hash(coin.puzzle_hash)
|
|
if output_created is None:
|
|
newpuzhash = await self.get_new_puzzlehash()
|
|
primaries = [{"puzzlehash": newpuzhash, "amount": chia_amount}]
|
|
solution = self.make_solution(primaries=primaries)
|
|
output_created = coin
|
|
else:
|
|
solution = self.make_solution(consumed=[output_created.name()])
|
|
list_of_solutions.append(CoinSolution(coin, Program.to([puzzle, solution])))
|
|
new_sigs = await self.get_sigs_for_innerpuz_with_innersol(puzzle, solution)
|
|
sigs = sigs + new_sigs
|
|
|
|
aggsig = AugSchemeMPL.aggregate(sigs)
|
|
return SpendBundle(list_of_solutions, aggsig)
|