204 lines
7.4 KiB
Python
204 lines
7.4 KiB
Python
from pathlib import Path
|
|
from typing import Dict, Optional, List
|
|
from blspy import ExtendedPrivateKey
|
|
import logging
|
|
import src.protocols.wallet_protocol
|
|
from src.full_node.full_node import OutboundMessageGenerator
|
|
from src.protocols.wallet_protocol import ProofHash
|
|
from src.server.outbound_message import OutboundMessage, NodeType, Message, Delivery
|
|
from src.server.server import ChiaServer
|
|
from src.types.full_block import additions_for_npc
|
|
from src.types.hashable.coin import Coin
|
|
from src.types.hashable.spend_bundle import SpendBundle
|
|
from src.types.name_puzzle_condition import NPC
|
|
from src.types.sized_bytes import bytes32
|
|
from src.util.hash import std_hash
|
|
from src.util.api_decorators import api_request
|
|
from src.util.ints import uint32
|
|
from src.util.mempool_check_conditions import get_name_puzzle_conditions
|
|
from src.wallet.wallet import Wallet
|
|
from src.wallet.wallet_state_manager import WalletStateManager
|
|
from src.wallet.wallet_store import WalletStore
|
|
from src.wallet.wallet_transaction_store import WalletTransactionStore
|
|
|
|
|
|
class WalletNode:
|
|
private_key: ExtendedPrivateKey
|
|
key_config: Dict
|
|
config: Dict
|
|
server: Optional[ChiaServer]
|
|
wallet_store: WalletStore
|
|
wallet_state_manager: WalletStateManager
|
|
header_hash: List[bytes32]
|
|
start_index: int
|
|
log: logging.Logger
|
|
wallet: Wallet
|
|
tx_store: WalletTransactionStore
|
|
|
|
@staticmethod
|
|
async def create(config: Dict, key_config: Dict, name: str = None):
|
|
self = WalletNode()
|
|
self.config = config
|
|
self.key_config = key_config
|
|
sk_hex = self.key_config["wallet_sk"]
|
|
self.private_key = ExtendedPrivateKey.from_bytes(bytes.fromhex(sk_hex))
|
|
if name:
|
|
self.log = logging.getLogger(name)
|
|
else:
|
|
self.log = logging.getLogger(__name__)
|
|
|
|
pub_hex = self.private_key.get_public_key().serialize().hex()
|
|
path = Path(f"wallet_db_{pub_hex}.db")
|
|
self.wallet_store = await WalletStore.create(path)
|
|
self.tx_store = await WalletTransactionStore.create(path)
|
|
|
|
self.wallet_state_manager = await WalletStateManager.create(
|
|
config, key_config, self.wallet_store, self.tx_store
|
|
)
|
|
self.wallet = await Wallet.create(config, key_config, self.wallet_state_manager)
|
|
|
|
self.server = None
|
|
|
|
return self
|
|
|
|
def set_server(self, server: ChiaServer):
|
|
self.server = server
|
|
self.wallet.set_server(server)
|
|
|
|
async def _on_connect(self) -> OutboundMessageGenerator:
|
|
"""
|
|
Whenever we connect to a FullNode we request new proof_hashes by sending last proof hash we have
|
|
"""
|
|
self.log.info(f"Requesting proof hashes")
|
|
request = ProofHash(std_hash(b"deadbeef"))
|
|
yield OutboundMessage(
|
|
NodeType.FULL_NODE,
|
|
Message("request_proof_hashes", request),
|
|
Delivery.BROADCAST,
|
|
)
|
|
|
|
@api_request
|
|
async def proof_hash(
|
|
self, request: src.protocols.wallet_protocol.ProofHash
|
|
) -> OutboundMessageGenerator:
|
|
"""
|
|
Received a proof hash from the FullNode
|
|
"""
|
|
self.log.info(f"Received a new proof hash: {request}")
|
|
reply_request = ProofHash(std_hash(b"a"))
|
|
# TODO Store and decide if we want full proof for this proof hash
|
|
yield OutboundMessage(
|
|
NodeType.FULL_NODE,
|
|
Message("request_full_proof_for_hash", reply_request),
|
|
Delivery.RESPOND,
|
|
)
|
|
|
|
@api_request
|
|
async def full_proof_for_hash(
|
|
self, request: src.protocols.wallet_protocol.FullProofForHash
|
|
):
|
|
"""
|
|
We've received a full proof for hash we requested
|
|
"""
|
|
# TODO Validate full proof
|
|
self.log.info(f"Received new proof: {request}")
|
|
|
|
@api_request
|
|
async def received_body(self, response: src.protocols.wallet_protocol.RespondBody):
|
|
"""
|
|
Called when body is received from the FullNode
|
|
"""
|
|
|
|
# Retry sending queued up transactions
|
|
await self.retry_send_queue()
|
|
|
|
additions: List[Coin] = []
|
|
|
|
if self.wallet.can_generate_puzzle_hash(response.body.coinbase.puzzle_hash):
|
|
await self.wallet_state_manager.coin_added(
|
|
response.body.coinbase, response.height, True
|
|
)
|
|
if self.wallet.can_generate_puzzle_hash(response.body.fees_coin.puzzle_hash):
|
|
await self.wallet_state_manager.coin_added(
|
|
response.body.fees_coin, response.height, True
|
|
)
|
|
|
|
npc_list: List[NPC]
|
|
if response.body.transactions:
|
|
error, npc_list, cost = get_name_puzzle_conditions(
|
|
response.body.transactions
|
|
)
|
|
|
|
additions.extend(additions_for_npc(npc_list))
|
|
|
|
for added_coin in additions:
|
|
if self.wallet.can_generate_puzzle_hash(added_coin.puzzle_hash):
|
|
await self.wallet_state_manager.coin_added(
|
|
added_coin, response.height, False
|
|
)
|
|
|
|
for npc in npc_list:
|
|
if self.wallet.can_generate_puzzle_hash(npc.puzzle_hash):
|
|
await self.wallet_state_manager.coin_removed(
|
|
npc.coin_name, response.height
|
|
)
|
|
|
|
@api_request
|
|
async def new_lca(self, header: src.protocols.wallet_protocol.Header):
|
|
self.log.info("new tip received")
|
|
|
|
async def retry_send_queue(self):
|
|
records = await self.wallet_state_manager.get_send_queue()
|
|
for record in records:
|
|
if record.spend_bundle:
|
|
await self._send_transaction(record.spend_bundle)
|
|
|
|
async def _send_transaction(self, spend_bundle: SpendBundle):
|
|
""" Sends spendbundle to connected full Nodes."""
|
|
await self.wallet_state_manager.add_pending_transaction(spend_bundle)
|
|
|
|
msg = OutboundMessage(
|
|
NodeType.FULL_NODE,
|
|
Message("wallet_transaction", spend_bundle),
|
|
Delivery.BROADCAST,
|
|
)
|
|
if self.server:
|
|
async for reply in self.server.push_message(msg):
|
|
self.log.info(reply)
|
|
|
|
async def _request_add_list(self, height: uint32, header_hash: bytes32):
|
|
obj = src.protocols.wallet_protocol.RequestAdditions(height, header_hash)
|
|
msg = OutboundMessage(
|
|
NodeType.FULL_NODE, Message("request_additions", obj), Delivery.BROADCAST,
|
|
)
|
|
if self.server:
|
|
async for reply in self.server.push_message(msg):
|
|
self.log.info(reply)
|
|
|
|
@api_request
|
|
async def response_additions(
|
|
self, response: src.protocols.wallet_protocol.Additions
|
|
):
|
|
print(response)
|
|
|
|
@api_request
|
|
async def response_additions_rejected(
|
|
self, response: src.protocols.wallet_protocol.RequestAdditions
|
|
):
|
|
print(f"request rejected {response}")
|
|
|
|
@api_request
|
|
async def transaction_ack(self, ack: src.protocols.wallet_protocol.TransactionAck):
|
|
if ack.status:
|
|
await self.wallet_state_manager.remove_from_queue(ack.txid)
|
|
self.log.info(f"SpendBundle has been received by the FullNode. id: {id}")
|
|
else:
|
|
self.log.info(f"SpendBundle has been rejected by the FullNode. id: {id}")
|
|
|
|
async def requestLCA(self):
|
|
msg = OutboundMessage(
|
|
NodeType.FULL_NODE, Message("request_lca", None), Delivery.BROADCAST,
|
|
)
|
|
async for reply in self.server.push_message(msg):
|
|
self.log.info(reply)
|