ctl.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337
  1. #!/usr/bin/env python3
  2. #
  3. # MMGen Wallet, a terminal-based cryptocurrency wallet
  4. # Copyright (C)2013-2025 The MMGen Project <mmgen@tuta.io>
  5. #
  6. # This program is free software: you can redistribute it and/or modify
  7. # it under the terms of the GNU General Public License as published by
  8. # the Free Software Foundation, either version 3 of the License, or
  9. # (at your option) any later version.
  10. #
  11. # This program is distributed in the hope that it will be useful,
  12. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14. # GNU General Public License for more details.
  15. #
  16. # You should have received a copy of the GNU General Public License
  17. # along with this program. If not, see <http://www.gnu.org/licenses/>.
  18. """
  19. tw.ctl: Tracking wallet control class for the MMGen suite
  20. """
  21. import json
  22. from collections import namedtuple
  23. from pathlib import Path
  24. from ..util import msg, msg_r, suf, die
  25. from ..base_obj import AsyncInit
  26. from ..objmethods import MMGenObject
  27. from ..obj import TwComment, get_obj
  28. from ..addr import CoinAddr, is_mmgen_id, is_coin_addr
  29. from ..rpc import rpc_init
  30. from .shared import TwMMGenID, TwLabel
  31. twmmid_addr_pair = namedtuple('addr_info', ['twmmid', 'coinaddr'])
  32. label_addr_pair = namedtuple('label_addr_pair', ['label', 'coinaddr'])
  33. # decorator for TwCtl
  34. def write_mode(orig_func):
  35. def f(self, *args, **kwargs):
  36. if self.mode != 'w':
  37. die(1, '{} opened in read-only mode: cannot execute method {}()'.format(
  38. type(self).__name__,
  39. locals()['orig_func'].__name__
  40. ))
  41. return orig_func(self, *args, **kwargs)
  42. return f
  43. class TwCtl(MMGenObject, metaclass=AsyncInit):
  44. caps = ('rescan', 'batch')
  45. data_key = 'addresses'
  46. use_tw_file = False
  47. aggressive_sync = False
  48. importing = False
  49. tw_fn = 'tracking-wallet.json'
  50. def __new__(cls, cfg, proto, *args, **kwargs):
  51. return MMGenObject.__new__(
  52. proto.base_proto_subclass(cls, 'tw.ctl', is_token=kwargs.get('token_addr')))
  53. async def __init__(
  54. self,
  55. cfg,
  56. proto,
  57. mode = 'r',
  58. token_addr = None,
  59. no_rpc = False,
  60. no_wallet_init = False,
  61. rpc_ignore_wallet = False):
  62. assert mode in ('r', 'w', 'i'), f"{mode!r}: wallet mode must be 'r', 'w' or 'i'"
  63. if mode == 'i':
  64. self.importing = True
  65. mode = 'w'
  66. self.cfg = cfg
  67. self.proto = proto
  68. self.mode = mode
  69. self.desc = self.base_desc = f'{self.proto.name} tracking wallet'
  70. self.cur_balances = {} # cache balances to prevent repeated lookups per program invocation
  71. if not no_rpc:
  72. self.rpc = await rpc_init(cfg, proto, ignore_wallet=rpc_ignore_wallet)
  73. if self.use_tw_file:
  74. if self.proto.coin == 'BTC':
  75. self.tw_dir = Path(self.cfg.data_dir)
  76. else:
  77. self.tw_dir = Path(
  78. self.cfg.data_dir_root,
  79. 'altcoins',
  80. self.proto.coin.lower(),
  81. ('' if self.proto.network == 'mainnet' else self.proto.network)
  82. )
  83. self.tw_path = self.tw_dir / self.tw_fn
  84. if no_wallet_init:
  85. return
  86. if self.use_tw_file:
  87. self.init_from_wallet_file()
  88. else:
  89. self.init_empty()
  90. if self.data['coin'] != self.proto.coin: # TODO remove?
  91. die('WalletFileError',
  92. f'Tracking wallet coin ({self.data["coin"]}) does not match current coin ({self.proto.coin})!')
  93. self.conv_types(self.data[self.data_key])
  94. def init_from_wallet_file(self):
  95. from ..fileutil import check_or_create_dir, get_data_from_file
  96. check_or_create_dir(self.tw_dir)
  97. try:
  98. self.orig_data = get_data_from_file(self.cfg, self.tw_path, quiet=True)
  99. self.data = json.loads(self.orig_data)
  100. except:
  101. try:
  102. self.tw_path.stat()
  103. except:
  104. self.orig_data = ''
  105. self.init_empty()
  106. self.force_write()
  107. else:
  108. die('WalletFileError', f'File ‘{self.tw_path}’ exists but does not contain valid JSON data')
  109. else:
  110. self.upgrade_wallet_maybe()
  111. # ensure that wallet file is written when user exits via KeyboardInterrupt:
  112. if self.mode == 'w':
  113. import atexit
  114. def del_twctl(twctl):
  115. self.cfg._util.dmsg(f'Running exit handler del_twctl() for {twctl!r}')
  116. del twctl
  117. atexit.register(del_twctl, self)
  118. def __del__(self):
  119. """
  120. TwCtl instances opened in write or import mode must be explicitly destroyed with ‘del
  121. twuo.twctl’ and the like to ensure the instance is deleted and wallet is written before
  122. global vars are destroyed by the interpreter at shutdown.
  123. Not that this code can only be debugged by examining the program output, as exceptions
  124. are ignored within __del__():
  125. /usr/share/doc/python3.6-doc/html/reference/datamodel.html#object.__del__
  126. Since no exceptions are raised, errors will not be caught by the test suite.
  127. """
  128. if getattr(self, 'mode', None) == 'w': # mode attr might not exist in this state
  129. self.write()
  130. elif self.cfg.debug:
  131. msg('read-only wallet, doing nothing')
  132. def conv_types(self, ad):
  133. for k, v in ad.items():
  134. if k not in ('params', 'coin'):
  135. v['mmid'] = TwMMGenID(self.proto, v['mmid'])
  136. v['comment'] = TwComment(v['comment'])
  137. @property
  138. def data_root(self):
  139. return self.data[self.data_key]
  140. @property
  141. def data_root_desc(self):
  142. return self.data_key
  143. def cache_balance(self, addr, bal, session_cache, data_root, force=False):
  144. if force or addr not in session_cache:
  145. session_cache[addr] = str(bal)
  146. if addr in data_root:
  147. data_root[addr]['balance'] = str(bal)
  148. if self.aggressive_sync:
  149. self.write()
  150. def get_cached_balance(self, addr, session_cache, data_root):
  151. if addr in session_cache:
  152. return self.proto.coin_amt(session_cache[addr])
  153. if not self.cfg.cached_balances:
  154. return None
  155. if addr in data_root and 'balance' in data_root[addr]:
  156. return self.proto.coin_amt(data_root[addr]['balance'])
  157. async def get_balance(self, addr, force_rpc=False):
  158. ret = None if force_rpc else self.get_cached_balance(addr, self.cur_balances, self.data_root)
  159. if ret is None:
  160. ret = await self.rpc_get_balance(addr)
  161. self.cache_balance(addr, ret, self.cur_balances, self.data_root)
  162. return ret
  163. def force_write(self):
  164. mode_save = self.mode
  165. self.mode = 'w'
  166. self.write()
  167. self.mode = mode_save
  168. @write_mode
  169. def write_changed(self, data, quiet):
  170. from ..fileutil import write_data_to_file
  171. write_data_to_file(
  172. self.cfg,
  173. self.tw_path,
  174. data,
  175. desc = f'{self.base_desc} data',
  176. ask_overwrite = False,
  177. ignore_opt_outdir = True,
  178. quiet = quiet,
  179. check_data = True, # die if wallet has been altered by another program
  180. cmp_data = self.orig_data)
  181. self.orig_data = data
  182. def write(self, quiet=True):
  183. if not self.use_tw_file:
  184. self.cfg._util.dmsg("'use_tw_file' is False, doing nothing")
  185. return
  186. self.cfg._util.dmsg(f'write(): checking if {self.desc} data has changed')
  187. wdata = json.dumps(self.data)
  188. if self.orig_data != wdata:
  189. self.write_changed(wdata, quiet=quiet)
  190. elif self.cfg.debug:
  191. msg('Data is unchanged\n')
  192. async def resolve_address(self, addrspec):
  193. twmmid, coinaddr = (None, None)
  194. pairs = await self.get_label_addr_pairs()
  195. if is_coin_addr(self.proto, addrspec):
  196. coinaddr = get_obj(CoinAddr, proto=self.proto, addr=addrspec)
  197. pair_data = [e for e in pairs if e.coinaddr == coinaddr]
  198. elif is_mmgen_id(self.proto, addrspec):
  199. twmmid = TwMMGenID(self.proto, addrspec)
  200. pair_data = [e for e in pairs if e.label.mmid == twmmid]
  201. else:
  202. msg(f'{addrspec!r}: invalid address for this network')
  203. return None
  204. if not pair_data:
  205. msg('{a} address {b!r} not found in tracking wallet'.format(
  206. a = 'MMGen' if twmmid else 'Coin',
  207. b = twmmid or coinaddr))
  208. return None
  209. return twmmid_addr_pair(
  210. twmmid or pair_data[0].label.mmid,
  211. coinaddr or pair_data[0].coinaddr)
  212. # returns on failure
  213. @write_mode
  214. async def set_comment(
  215. self,
  216. addrspec,
  217. comment = '',
  218. trusted_pair = None,
  219. silent = False):
  220. res = twmmid_addr_pair(*trusted_pair) if trusted_pair else await self.resolve_address(addrspec)
  221. if not res:
  222. return False
  223. comment = get_obj(TwComment, s=comment)
  224. if comment is False:
  225. return False
  226. lbl = get_obj(
  227. TwLabel,
  228. proto = self.proto,
  229. text = res.twmmid + (' ' + comment if comment else ''))
  230. if lbl is False:
  231. return False
  232. if await self.set_label(res.coinaddr, lbl):
  233. if not silent:
  234. desc = '{t} address {a} in tracking wallet'.format(
  235. t = res.twmmid.type.replace('mmgen', 'MMGen'),
  236. a = res.twmmid.addr.hl() if res.twmmid.type == 'mmgen' else
  237. res.twmmid.addr.hl(res.twmmid.addr.view_pref))
  238. msg(
  239. 'Added label {} to {}'.format(comment.hl2(encl='‘’'), desc) if comment else
  240. 'Removed label from {}'.format(desc))
  241. return True
  242. else:
  243. if not silent:
  244. msg('Label could not be {}'.format('added' if comment else 'removed'))
  245. return False
  246. @write_mode
  247. async def remove_comment(self, mmaddr):
  248. await self.set_comment(mmaddr, '')
  249. async def import_address_common(self, data, batch=False, gather=False):
  250. async def do_import(address, comment, message):
  251. try:
  252. res = await self.import_address(address, comment)
  253. self.cfg._util.qmsg(message)
  254. return res
  255. except Exception as e:
  256. die(2, f'\nImport of address {address!r} failed: {e.args[0]!r}')
  257. _d = namedtuple('formatted_import_data', data[0]._fields + ('mmid_disp',))
  258. pfx = self.proto.base_coin.lower() + ':'
  259. fdata = [_d(*d, 'non-MMGen' if d.twmmid.startswith(pfx) else d.twmmid) for d in data]
  260. fs = '{:%s}: {:%s} {:%s} - OK' % (
  261. len(str(len(fdata))) * 2 + 1,
  262. max(len(d.addr) for d in fdata),
  263. max(len(d.mmid_disp) for d in fdata) + 2
  264. )
  265. nAddrs = len(data)
  266. out = [( # create list, not generator, so we know data is valid before starting import
  267. CoinAddr(self.proto, d.addr),
  268. TwLabel(self.proto, d.twmmid + (f' {d.comment}' if d.comment else '')),
  269. fs.format(f'{n}/{nAddrs}', d.addr, f'({d.mmid_disp})')
  270. ) for n, d in enumerate(fdata, 1)]
  271. if batch:
  272. msg_r(f'Batch importing {len(out)} address{suf(data, "es")}...')
  273. ret = await self.batch_import_address((a, b, False) for a, b, c in out)
  274. msg(f'done\n{len(ret)} addresses imported')
  275. else:
  276. if gather: # this seems to provide little performance benefit
  277. import asyncio
  278. await asyncio.gather(*(do_import(*d) for d in out))
  279. else:
  280. for d in out:
  281. await do_import(*d)
  282. msg('Address import completed OK')