forked from LBRYCommunity/lbry-sdk
137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
import os
|
|
import stat
|
|
import json
|
|
import zlib
|
|
import typing
|
|
from typing import Sequence, MutableSequence
|
|
from hashlib import sha256
|
|
from operator import attrgetter
|
|
from torba.client.hash import better_aes_encrypt, better_aes_decrypt
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from torba.client import basemanager, baseaccount, baseledger
|
|
|
|
|
|
class Wallet:
|
|
""" The primary role of Wallet is to encapsulate a collection
|
|
of accounts (seed/private keys) and the spending rules / settings
|
|
for the coins attached to those accounts. Wallets are represented
|
|
by physical files on the filesystem.
|
|
"""
|
|
|
|
def __init__(self, name: str = 'Wallet', accounts: MutableSequence['baseaccount.BaseAccount'] = None,
|
|
storage: 'WalletStorage' = None) -> None:
|
|
self.name = name
|
|
self.accounts = accounts or []
|
|
self.storage = storage or WalletStorage()
|
|
|
|
def add_account(self, account):
|
|
self.accounts.append(account)
|
|
|
|
def generate_account(self, ledger: 'baseledger.BaseLedger') -> 'baseaccount.BaseAccount':
|
|
return ledger.account_class.generate(ledger, self)
|
|
|
|
@classmethod
|
|
def from_storage(cls, storage: 'WalletStorage', manager: 'basemanager.BaseWalletManager') -> 'Wallet':
|
|
json_dict = storage.read()
|
|
wallet = cls(
|
|
name=json_dict.get('name', 'Wallet'),
|
|
storage=storage
|
|
)
|
|
account_dicts: Sequence[dict] = json_dict.get('accounts', [])
|
|
for account_dict in account_dicts:
|
|
ledger = manager.get_or_create_ledger(account_dict['ledger'])
|
|
ledger.account_class.from_dict(ledger, wallet, account_dict)
|
|
return wallet
|
|
|
|
def to_dict(self):
|
|
return {
|
|
'version': WalletStorage.LATEST_VERSION,
|
|
'name': self.name,
|
|
'accounts': [a.to_dict() for a in self.accounts]
|
|
}
|
|
|
|
def save(self):
|
|
self.storage.write(self.to_dict())
|
|
|
|
@property
|
|
def default_account(self):
|
|
for account in self.accounts:
|
|
return account
|
|
|
|
@property
|
|
def hash(self) -> str:
|
|
h = sha256()
|
|
for account in sorted(self.accounts, key=attrgetter('id')):
|
|
h.update(account.hash)
|
|
return h.digest()
|
|
|
|
def pack(self, password):
|
|
new_data = json.dumps(self.to_dict())
|
|
new_data_compressed = zlib.compress(new_data.encode())
|
|
return better_aes_encrypt(password, new_data_compressed)
|
|
|
|
@classmethod
|
|
def unpack(cls, password, encrypted):
|
|
decrypted = better_aes_decrypt(password, encrypted)
|
|
decompressed = zlib.decompress(decrypted)
|
|
return json.loads(decompressed)
|
|
|
|
|
|
class WalletStorage:
|
|
|
|
LATEST_VERSION = 1
|
|
|
|
def __init__(self, path=None, default=None):
|
|
self.path = path
|
|
self._default = default or {
|
|
'version': self.LATEST_VERSION,
|
|
'name': 'My Wallet',
|
|
'accounts': []
|
|
}
|
|
|
|
def read(self):
|
|
if self.path and os.path.exists(self.path):
|
|
with open(self.path, 'r') as f:
|
|
json_data = f.read()
|
|
json_dict = json.loads(json_data)
|
|
if json_dict.get('version') == self.LATEST_VERSION and \
|
|
set(json_dict) == set(self._default):
|
|
return json_dict
|
|
else:
|
|
return self.upgrade(json_dict)
|
|
else:
|
|
return self._default.copy()
|
|
|
|
def upgrade(self, json_dict):
|
|
json_dict = json_dict.copy()
|
|
version = json_dict.pop('version', -1)
|
|
if version == -1:
|
|
pass
|
|
upgraded = self._default.copy()
|
|
upgraded.update(json_dict)
|
|
return json_dict
|
|
|
|
def write(self, json_dict):
|
|
|
|
json_data = json.dumps(json_dict, indent=4, sort_keys=True)
|
|
if self.path is None:
|
|
return json_data
|
|
|
|
temp_path = "%s.tmp.%s" % (self.path, os.getpid())
|
|
with open(temp_path, "w") as f:
|
|
f.write(json_data)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
|
|
if os.path.exists(self.path):
|
|
mode = os.stat(self.path).st_mode
|
|
else:
|
|
mode = stat.S_IREAD | stat.S_IWRITE
|
|
try:
|
|
os.rename(temp_path, self.path)
|
|
except Exception: # pylint: disable=broad-except
|
|
os.remove(self.path)
|
|
os.rename(temp_path, self.path)
|
|
os.chmod(self.path, mode)
|