chia-blockchain/chia/wallet/notification_manager.py
2023-04-04 00:08:10 -05:00

111 lines
4.6 KiB
Python

from __future__ import annotations
import dataclasses
import logging
from typing import Any, Dict, List, Optional, Set
from blspy import G2Element
from chia.protocols.wallet_protocol import CoinState
from chia.types.announcement import Announcement
from chia.types.blockchain_format.coin import Coin
from chia.types.blockchain_format.program import Program
from chia.types.blockchain_format.sized_bytes import bytes32
from chia.types.coin_spend import CoinSpend
from chia.types.spend_bundle import SpendBundle
from chia.util.db_wrapper import DBWrapper2
from chia.util.ints import uint32, uint64
from chia.wallet.notification_store import Notification, NotificationStore
from chia.wallet.transaction_record import TransactionRecord
from chia.wallet.util.compute_memos import compute_memos_for_spend
from chia.wallet.util.notifications import construct_notification
from chia.wallet.util.wallet_types import WalletType
class NotificationManager:
wallet_state_manager: Any
log: logging.Logger
notification_store: NotificationStore
@staticmethod
async def create(
wallet_state_manager: Any,
db_wrapper: DBWrapper2,
name: Optional[str] = None,
) -> NotificationManager:
self = NotificationManager()
if name:
self.log = logging.getLogger(name)
else:
self.log = logging.getLogger(__name__)
self.wallet_state_manager = wallet_state_manager
self.notification_store = await NotificationStore.create(db_wrapper)
return self
async def potentially_add_new_notification(self, coin_state: CoinState, parent_spend: CoinSpend) -> bool:
coin_name: bytes32 = coin_state.coin.name()
if (
coin_state.spent_height is None
or not self.wallet_state_manager.wallet_node.config.get("enable_notifications", True)
or self.wallet_state_manager.wallet_node.config.get("required_notification_amount", 100000000)
> coin_state.coin.amount
or await self.notification_store.notification_exists(coin_name)
):
return False
else:
memos: Dict[bytes32, List[bytes]] = compute_memos_for_spend(parent_spend)
coin_memos: List[bytes] = memos.get(coin_name, [])
if len(coin_memos) == 0 or len(coin_memos[0]) != 32:
return False
wallet_identifier = await self.wallet_state_manager.get_wallet_identifier_for_puzzle_hash(
bytes32(coin_memos[0])
)
if (
wallet_identifier is not None
and wallet_identifier.type == WalletType.STANDARD_WALLET
and len(coin_memos) == 2
and construct_notification(bytes32(coin_memos[0]), uint64(coin_state.coin.amount)).get_tree_hash()
== coin_state.coin.puzzle_hash
):
if len(coin_memos[1]) > 10000: # 10KB
return False
await self.notification_store.add_notification(
Notification(
coin_state.coin.name(),
coin_memos[1],
uint64(coin_state.coin.amount),
uint32(coin_state.spent_height),
)
)
self.wallet_state_manager.state_changed("new_on_chain_notification")
return True
async def send_new_notification(
self, target: bytes32, msg: bytes, amount: uint64, fee: uint64 = uint64(0)
) -> TransactionRecord:
coins: Set[Coin] = await self.wallet_state_manager.main_wallet.select_coins(uint64(amount + fee))
origin_coin: bytes32 = next(iter(coins)).name()
notification_puzzle: Program = construct_notification(target, amount)
notification_hash: bytes32 = notification_puzzle.get_tree_hash()
notification_coin: Coin = Coin(origin_coin, notification_hash, amount)
notification_spend = CoinSpend(
notification_coin,
notification_puzzle,
Program.to(None),
)
extra_spend_bundle = SpendBundle([notification_spend], G2Element())
chia_tx = await self.wallet_state_manager.main_wallet.generate_signed_transaction(
amount,
notification_hash,
fee,
coins=coins,
origin_id=origin_coin,
coin_announcements_to_consume={Announcement(notification_coin.name(), b"")},
memos=[target, msg],
)
full_tx: TransactionRecord = dataclasses.replace(
chia_tx, spend_bundle=SpendBundle.aggregate([chia_tx.spend_bundle, extra_spend_bundle])
)
return full_tx