twctl.py 10.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327
  1. #!/usr/bin/env python3
  2. #
  3. # mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
  4. # Copyright (C)2013-2022 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. twctl: Tracking wallet control class for the MMGen suite
  20. """
  21. from .globalvars import g
  22. from .exception import WalletFileError
  23. from .util import (
  24. msg,
  25. dmsg,
  26. check_or_create_dir,
  27. write_data_to_file,
  28. get_data_from_file,
  29. write_mode,
  30. altcoin_subclass
  31. )
  32. from .base_obj import AsyncInit
  33. from .objmethods import MMGenObject
  34. from .obj import TwComment,get_obj
  35. from .addr import CoinAddr,is_mmgen_id,is_coin_addr
  36. from .rpc import rpc_init
  37. from .tw import TwMMGenID,TwLabel
  38. class TrackingWallet(MMGenObject,metaclass=AsyncInit):
  39. caps = ('rescan','batch')
  40. data_key = 'addresses'
  41. use_tw_file = False
  42. aggressive_sync = False
  43. importing = False
  44. def __new__(cls,proto,*args,**kwargs):
  45. return MMGenObject.__new__(altcoin_subclass(cls,proto,'twctl'))
  46. async def __init__(self,proto,mode='r',token_addr=None):
  47. assert mode in ('r','w','i'), f"{mode!r}: wallet mode must be 'r','w' or 'i'"
  48. if mode == 'i':
  49. self.importing = True
  50. mode = 'w'
  51. if g.debug:
  52. print_stack_trace(f'TW INIT {mode!r} {self!r}')
  53. self.rpc = await rpc_init(proto) # TODO: create on demand - only certain ops require RPC
  54. self.proto = proto
  55. self.mode = mode
  56. self.desc = self.base_desc = f'{self.proto.name} tracking wallet'
  57. if self.use_tw_file:
  58. self.init_from_wallet_file()
  59. else:
  60. self.init_empty()
  61. if self.data['coin'] != self.proto.coin: # TODO remove?
  62. raise WalletFileError(
  63. 'Tracking wallet coin ({}) does not match current coin ({})!'.format(
  64. self.data['coin'],
  65. self.proto.coin ))
  66. self.conv_types(self.data[self.data_key])
  67. self.cur_balances = {} # cache balances to prevent repeated lookups per program invocation
  68. def init_empty(self):
  69. self.data = { 'coin': self.proto.coin, 'addresses': {} }
  70. def init_from_wallet_file(self):
  71. import os,json
  72. tw_dir = (
  73. os.path.join(g.data_dir) if self.proto.coin == 'BTC' else
  74. os.path.join(
  75. g.data_dir_root,
  76. 'altcoins',
  77. self.proto.coin.lower(),
  78. ('' if self.proto.network == 'mainnet' else 'testnet')
  79. ))
  80. self.tw_fn = os.path.join(tw_dir,'tracking-wallet.json')
  81. check_or_create_dir(tw_dir)
  82. try:
  83. self.orig_data = get_data_from_file(self.tw_fn,quiet=True)
  84. self.data = json.loads(self.orig_data)
  85. except:
  86. try: os.stat(self.tw_fn)
  87. except:
  88. self.orig_data = ''
  89. self.init_empty()
  90. self.force_write()
  91. else:
  92. raise WalletFileError(f'File {self.tw_fn!r} exists but does not contain valid json data')
  93. else:
  94. self.upgrade_wallet_maybe()
  95. # ensure that wallet file is written when user exits via KeyboardInterrupt:
  96. if self.mode == 'w':
  97. import atexit
  98. def del_tw(tw):
  99. dmsg(f'Running exit handler del_tw() for {tw!r}')
  100. del tw
  101. atexit.register(del_tw,self)
  102. def __del__(self):
  103. """
  104. TrackingWallet instances opened in write or import mode must be explicitly destroyed
  105. with 'del twctl', 'del twuo.wallet' and the like to ensure the instance is deleted and
  106. wallet is written before global vars are destroyed by the interpreter at shutdown.
  107. Not that this code can only be debugged by examining the program output, as exceptions
  108. are ignored within __del__():
  109. /usr/share/doc/python3.6-doc/html/reference/datamodel.html#object.__del__
  110. Since no exceptions are raised, errors will not be caught by the test suite.
  111. """
  112. if g.debug:
  113. print_stack_trace(f'TW DEL {self!r}')
  114. if getattr(self,'mode',None) == 'w': # mode attr might not exist in this state
  115. self.write()
  116. elif g.debug:
  117. msg('read-only wallet, doing nothing')
  118. def upgrade_wallet_maybe(self):
  119. pass
  120. def conv_types(self,ad):
  121. for k,v in ad.items():
  122. if k not in ('params','coin'):
  123. v['mmid'] = TwMMGenID(self.proto,v['mmid'])
  124. v['comment'] = TwComment(v['comment'])
  125. @property
  126. def data_root(self):
  127. return self.data[self.data_key]
  128. @property
  129. def data_root_desc(self):
  130. return self.data_key
  131. def cache_balance(self,addr,bal,session_cache,data_root,force=False):
  132. if force or addr not in session_cache:
  133. session_cache[addr] = str(bal)
  134. if addr in data_root:
  135. data_root[addr]['balance'] = str(bal)
  136. if self.aggressive_sync:
  137. self.write()
  138. def get_cached_balance(self,addr,session_cache,data_root):
  139. if addr in session_cache:
  140. return self.proto.coin_amt(session_cache[addr])
  141. if not g.cached_balances:
  142. return None
  143. if addr in data_root and 'balance' in data_root[addr]:
  144. return self.proto.coin_amt(data_root[addr]['balance'])
  145. async def get_balance(self,addr,force_rpc=False):
  146. ret = None if force_rpc else self.get_cached_balance(addr,self.cur_balances,self.data_root)
  147. if ret == None:
  148. ret = await self.rpc_get_balance(addr)
  149. self.cache_balance(addr,ret,self.cur_balances,self.data_root)
  150. return ret
  151. async def rpc_get_balance(self,addr):
  152. raise NotImplementedError('not implemented')
  153. @property
  154. def sorted_list(self):
  155. return sorted(
  156. [ { 'addr':x[0],
  157. 'mmid':x[1]['mmid'],
  158. 'comment':x[1]['comment'] }
  159. for x in self.data_root.items() if x[0] not in ('params','coin') ],
  160. key=lambda x: x['mmid'].sort_key+x['addr'] )
  161. @property
  162. def mmid_ordered_dict(self):
  163. return dict((x['mmid'],{'addr':x['addr'],'comment':x['comment']}) for x in self.sorted_list)
  164. @write_mode
  165. async def import_address(self,addr,label,rescan):
  166. return await self.rpc.call('importaddress',addr,label,rescan,timeout=(False,3600)[rescan])
  167. @write_mode
  168. def batch_import_address(self,arg_list):
  169. return self.rpc.batch_call('importaddress',arg_list)
  170. def force_write(self):
  171. mode_save = self.mode
  172. self.mode = 'w'
  173. self.write()
  174. self.mode = mode_save
  175. @write_mode
  176. def write_changed(self,data):
  177. write_data_to_file(
  178. self.tw_fn,
  179. data,
  180. desc = f'{self.base_desc} data',
  181. ask_overwrite = False,
  182. ignore_opt_outdir = True,
  183. quiet = True,
  184. check_data = True,
  185. cmp_data = self.orig_data )
  186. self.orig_data = data
  187. def write(self): # use 'check_data' to check wallet hasn't been altered by another program
  188. if not self.use_tw_file:
  189. dmsg("'use_tw_file' is False, doing nothing")
  190. return
  191. dmsg(f'write(): checking if {self.desc} data has changed')
  192. import json
  193. wdata = json.dumps(self.data)
  194. if self.orig_data != wdata:
  195. if g.debug:
  196. print_stack_trace(f'TW DATA CHANGED {self!r}')
  197. print_diff(self.orig_data,wdata,from_json=True)
  198. self.write_changed(wdata)
  199. elif g.debug:
  200. msg('Data is unchanged\n')
  201. async def is_in_wallet(self,addr):
  202. from .twaddrs import TwAddrList
  203. return addr in (await TwAddrList(self.proto,[],0,True,True,True,wallet=self)).coinaddr_list()
  204. @write_mode
  205. async def set_label(self,coinaddr,lbl):
  206. # bitcoin-{abc,bchn} 'setlabel' RPC is broken, so use old 'importaddress' method to set label
  207. # broken behavior: new label is set OK, but old label gets attached to another address
  208. if 'label_api' in self.rpc.caps and self.proto.coin != 'BCH':
  209. args = ('setlabel',coinaddr,lbl)
  210. else:
  211. # NOTE: this works because importaddress() removes the old account before
  212. # associating the new account with the address.
  213. # RPC args: addr,label,rescan[=true],p2sh[=none]
  214. args = ('importaddress',coinaddr,lbl,False)
  215. try:
  216. return await self.rpc.call(*args)
  217. except Exception as e:
  218. rmsg(e.args[0])
  219. return False
  220. # returns on failure
  221. @write_mode
  222. async def add_label(self,arg1,label='',addr=None,silent=False,on_fail='return'):
  223. assert on_fail in ('return','raise'), 'add_label_chk1'
  224. mmaddr,coinaddr = None,None
  225. if is_coin_addr(self.proto,addr or arg1):
  226. coinaddr = get_obj(CoinAddr,proto=self.proto,addr=addr or arg1)
  227. if is_mmgen_id(self.proto,arg1):
  228. mmaddr = TwMMGenID(self.proto,arg1)
  229. if mmaddr and not coinaddr:
  230. from .addrdata import TwAddrData
  231. coinaddr = (await TwAddrData(self.proto)).mmaddr2coinaddr(mmaddr)
  232. try:
  233. if not is_mmgen_id(self.proto,arg1):
  234. assert coinaddr, f'Invalid coin address for this chain: {arg1}'
  235. assert coinaddr, f'{g.proj_name} address {mmaddr!r} not found in tracking wallet'
  236. assert await self.is_in_wallet(coinaddr), f'Address {coinaddr!r} not found in tracking wallet'
  237. except Exception as e:
  238. msg(str(e))
  239. return False
  240. # Allow for the possibility that BTC addr of MMGen addr was entered.
  241. # Do reverse lookup, so that MMGen addr will not be marked as non-MMGen.
  242. if not mmaddr:
  243. from .addrdata import TwAddrData
  244. mmaddr = (await TwAddrData(proto=self.proto)).coinaddr2mmaddr(coinaddr)
  245. if not mmaddr:
  246. mmaddr = f'{self.proto.base_coin.lower()}:{coinaddr}'
  247. mmaddr = TwMMGenID(self.proto,mmaddr)
  248. cmt = TwComment(label) if on_fail=='raise' else get_obj(TwComment,s=label)
  249. if cmt in (False,None):
  250. return False
  251. lbl_txt = mmaddr + (' ' + cmt if cmt else '')
  252. lbl = (
  253. TwLabel(self.proto,lbl_txt) if on_fail == 'raise' else
  254. get_obj(TwLabel,proto=self.proto,text=lbl_txt) )
  255. if await self.set_label(coinaddr,lbl) == False:
  256. if not silent:
  257. msg( 'Label could not be {}'.format('added' if label else 'removed') )
  258. return False
  259. else:
  260. desc = '{} address {} in tracking wallet'.format(
  261. mmaddr.type.replace('mmg','MMG'),
  262. mmaddr.replace(self.proto.base_coin.lower()+':','') )
  263. if label:
  264. msg(f'Added label {label!r} to {desc}')
  265. else:
  266. msg(f'Removed label from {desc}')
  267. return True
  268. @write_mode
  269. async def remove_label(self,mmaddr):
  270. await self.add_label(mmaddr,'')
  271. @write_mode
  272. async def remove_address(self,addr):
  273. raise NotImplementedError(f'address removal not implemented for coin {self.proto.coin}')