import os import stat import json import asyncio import logging from typing import Optional, Dict from lbry.db import Database from .wallet import Wallet from .account import SingleKey, HierarchicalDeterministic log = logging.getLogger(__name__) class WalletManager: def __init__(self, db: Database): self.db = db self.ledger = db.ledger self.wallets: Dict[str, Wallet] = {} if self.ledger.conf.wallet_storage == "file": self.storage = FileWallet(self.db, self.ledger.conf.wallet_dir) elif self.ledger.conf.wallet_storage == "database": self.storage = DatabaseWallet(self.db) else: raise Exception(f"Unknown wallet storage format: {self.ledger.conf.wallet_storage}") def __getitem__(self, wallet_id: str) -> Wallet: try: return self.wallets[wallet_id] except KeyError: raise ValueError(f"Couldn't find wallet: {wallet_id}.") async def generate_addresses(self): for wallet in self.wallets.values(): await wallet.generate_addresses() @property def default(self) -> Optional[Wallet]: for wallet in self.wallets.values(): return wallet def get_or_default(self, wallet_id: Optional[str]) -> Wallet: if wallet_id: return self[wallet_id] wallet = self.default if not wallet: raise ValueError("No wallets available.") return wallet def get_or_default_for_spending(self, wallet_id: Optional[str]) -> Wallet: wallet = self.get_or_default(wallet_id) if wallet.is_locked: raise ValueError("Cannot spend funds with locked wallet, unlock first.") return wallet async def initialize(self): for wallet_id in self.ledger.conf.wallets: if wallet_id in self.wallets: log.warning("Ignoring duplicate wallet_id in config: %s", wallet_id) continue await self.load(wallet_id) default_wallet = self.default if default_wallet is None: if self.ledger.conf.create_default_wallet: assert self.ledger.conf.wallets[0] == "default_wallet", ( "Requesting to generate the default wallet but the 'wallets' " "config setting does not include 'default_wallet' as the first wallet." ) await self.create( self.ledger.conf.wallets[0], 'Wallet', create_account=self.ledger.conf.create_default_account ) elif not default_wallet.has_accounts and self.ledger.conf.create_default_account: await default_wallet.accounts.generate() async def load(self, wallet_id: str) -> Optional[Wallet]: wallet = await self.storage.get(wallet_id) if wallet is not None: return self.add(wallet) async def create( self, wallet_id: str, name: str = "", create_account=False, language="en", single_key=False ) -> Wallet: if wallet_id in self.wallets: raise Exception(f"Wallet with id '{wallet_id}' is already loaded and cannot be created.") if await self.storage.exists(wallet_id): raise Exception(f"Wallet '{wallet_id}' already exists, use 'wallet_add' to load wallet.") wallet = Wallet(wallet_id, self.db, name) if create_account: await wallet.accounts.generate(language=language, address_generator={ 'name': SingleKey.name if single_key else HierarchicalDeterministic.name }) await self.storage.save(wallet) return self.add(wallet) def add(self, wallet: Wallet) -> Wallet: self.wallets[wallet.id] = wallet wallet.on_change.listen(lambda _: self.storage.save(wallet)) return wallet async def _report_state(self): try: for account in self.accounts: balance = dewies_to_lbc(await account.get_balance(include_claims=True)) _, channel_count = await account.get_channels(limit=1) claim_count = await account.get_claim_count() if isinstance(account.receiving, SingleKey): log.info("Loaded single key account %s with %s LBC. " "%d channels, %d certificates and %d claims", account.id, balance, channel_count, len(account.channel_keys), claim_count) else: total_receiving = len(await account.receiving.get_addresses()) total_change = len(await account.change.get_addresses()) log.info("Loaded account %s with %s LBC, %d receiving addresses (gap: %d), " "%d change addresses (gap: %d), %d channels, %d certificates and %d claims. ", account.id, balance, total_receiving, account.receiving.gap, total_change, account.change.gap, channel_count, len(account.channel_keys), claim_count) except Exception as err: if isinstance(err, asyncio.CancelledError): # TODO: remove when updated to 3.8 raise log.exception( 'Failed to display wallet state, please file issue ' 'for this bug along with the traceback you see below:' ) class WalletStorage: async def prepare(self): raise NotImplementedError async def exists(self, walllet_id: str) -> bool: raise NotImplementedError async def get(self, wallet_id: str) -> Wallet: raise NotImplementedError async def save(self, wallet: Wallet): raise NotImplementedError class FileWallet(WalletStorage): def __init__(self, db, wallet_dir): self.db = db self.wallet_dir = wallet_dir def get_wallet_path(self, wallet_id: str): return os.path.join(self.wallet_dir, wallet_id) async def prepare(self): await asyncio.get_running_loop().run_in_executor( None, self.sync_ensure_wallets_directory_exists ) def sync_ensure_wallets_directory_exists(self): if not os.path.exists(self.wallet_dir): os.mkdir(self.wallet_dir) async def exists(self, wallet_id: str) -> bool: return os.path.exists(self.get_wallet_path(wallet_id)) async def get(self, wallet_id: str) -> Wallet: wallet_dict = await asyncio.get_running_loop().run_in_executor( None, self.sync_read, wallet_id ) if wallet_dict is not None: return await Wallet.from_dict(wallet_id, wallet_dict, self.db) def sync_read(self, wallet_id): try: with open(self.get_wallet_path(wallet_id), 'r') as f: json_data = f.read() return json.loads(json_data) except FileNotFoundError: return None async def save(self, wallet: Wallet): return await asyncio.get_running_loop().run_in_executor( None, self.sync_write, wallet ) def sync_write(self, wallet: Wallet): temp_path = os.path.join(self.wallet_dir, f".tmp.{os.path.basename(wallet.id)}") with open(temp_path, "w") as f: f.write(wallet.to_serialized()) f.flush() os.fsync(f.fileno()) wallet_path = self.get_wallet_path(wallet.id) if os.path.exists(wallet_path): mode = os.stat(wallet_path).st_mode else: mode = stat.S_IREAD | stat.S_IWRITE try: os.rename(temp_path, wallet_path) except Exception: # pylint: disable=broad-except os.remove(wallet_path) os.rename(temp_path, wallet_path) os.chmod(wallet_path, mode) class DatabaseWallet(WalletStorage): def __init__(self, db: 'Database'): self.db = db async def prepare(self): pass async def exists(self, wallet_id: str) -> bool: return await self.db.has_wallet(wallet_id) async def get(self, wallet_id: str) -> Wallet: data = await self.db.get_wallet(wallet_id) if data: wallet_dict = json.loads(data['data']) if wallet_dict is not None: return await Wallet.from_dict(wallet_id, wallet_dict, self.db) async def save(self, wallet: Wallet): await self.db.add_wallet( wallet.id, wallet.to_serialized() )