From cd4425ddf2868a22b7cd489a458c5cfe5bac792f Mon Sep 17 00:00:00 2001
From: MMGen <mmgen@tuta.io>
Date: Mon, 23 Jul 2018 19:07:41 +0000
Subject: [PATCH] TwGetBalance: reimplement

---
 mmgen/altcoins/eth/tw.py | 22 ++++++++++++++--------
 mmgen/tw.py              | 35 +++++++++++++++++++++--------------
 2 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/mmgen/altcoins/eth/tw.py b/mmgen/altcoins/eth/tw.py
index 63fbcb66..c2d95b35 100755
--- a/mmgen/altcoins/eth/tw.py
+++ b/mmgen/altcoins/eth/tw.py
@@ -199,15 +199,21 @@ class EthereumTwGetBalance(TwGetBalance):
 	def create_data(self):
 		data = TrackingWallet().mmid_ordered_dict()
 		for d in data:
-			keys = ['TOTAL']
-			keys += [str(d.obj.sid)] if d.type == 'mmgen' else ['Non-MMGen']
-			confs = 9999 # TODO
-			i = (1,2)[confs >= self.minconf]
+			if d.type == 'mmgen':
+				key = d.obj.sid
+				if key not in self.data:
+					self.data[key] = [g.proto.coin_amt('0')] * 4
+			else:
+				key = 'Non-MMGen'
 
-			for key in keys:
-				if key not in self.data: self.data[key] = [g.proto.coin_amt('0')] * 3
-				for j in ([],[0])[confs==0] + [i]:
-					self.data[key][j] += ETHAmt(int(g.rpch.eth_getBalance('0x'+data[d]['addr']),16),'wei')
+			conf_level = 2 # TODO
+			amt = self.get_addr_balance(data[d]['addr'])
+
+			self.data['TOTAL'][conf_level] += amt
+			self.data[key][conf_level] += amt
+
+	def get_addr_balance(self,addr):
+		return ETHAmt(int(g.rpch.eth_getBalance('0x'+addr),16),'wei')
 
 class EthereumAddrData(AddrData):
 
diff --git a/mmgen/tw.py b/mmgen/tw.py
index 02fceab6..22a22a3c 100755
--- a/mmgen/tw.py
+++ b/mmgen/tw.py
@@ -577,24 +577,28 @@ class TwGetBalance(MMGenObject):
 		rpc_init()
 		self.minconf = minconf
 		self.quiet = quiet
-		self.data = {}
+		self.data = dict([(k,[g.proto.coin_amt('0')] * 4) for k in 'TOTAL','Non-MMGen','Non-wallet'])
 		self.create_data()
 
 	def create_data(self):
+		# 0: unconfirmed, 1: below minconf, 2: confirmed, 3: spendable
 		for d in g.rpch.listunspent(0):
 			try:    lbl = TwLabel(d['account'],on_fail='silent')
 			except: lbl = None
-			keys = ['TOTAL']
-			if lbl and lbl.mmid.type == 'mmgen':
-				keys += [lbl.mmid.obj.sid]
-			if d['spendable']: keys += ['SPENDABLE']
-			confs = d['confirmations']
-			i = (1,2)[confs >= self.minconf]
+			if lbl:
+				if lbl.mmid.type == 'mmgen':
+					key = lbl.mmid.obj.sid
+					if key not in self.data:
+						self.data[key] = [g.proto.coin_amt('0')] * 4
+				else: key = 'Non-MMGen'
+			else: key = 'Non-wallet'
 
-			for key in keys:
-				if key not in self.data: self.data[key] = [g.proto.coin_amt('0')] * 3
-				for j in ([],[0])[confs==0] + [i]:
-					self.data[key][j] += d['amount']
+			conf_level = 0 if not d['confirmations'] else 1 if d['confirmations'] < self.minconf else 2
+
+			self.data['TOTAL'][conf_level] += d['amount']
+			self.data[key][conf_level] += d['amount']
+			if d['spendable']:
+				self.data[key][3] += d['amount']
 
 	def format(self):
 		if self.quiet:
@@ -605,11 +609,14 @@ class TwGetBalance(MMGenObject):
 								p=' <{} confirms'.format(self.minconf),
 								c=' >={} confirms'.format(self.minconf))
 			for key in sorted(self.data):
+				if not any(self.data[key]): continue
 				o += self.fs.format(**dict(zip(
 							('w','u','p','c'),
-							[key+':'] + [a.fmt(color=True,suf=' '+g.coin) for a in self.data[key]]
+							[key+':'] + [a.fmt(color=True,suf=' '+g.dcoin) for a in self.data[key]]
 							)))
 
-		if 'SPENDABLE' in self.data:
-			o += red('Warning: this wallet contains PRIVATE KEYS for the SPENDABLE balance!\n')
+		for key,vals in self.data.items():
+			if key == 'TOTAL': continue
+			if vals[3]:
+				o += red('Warning: this wallet contains PRIVATE KEYS for {} outputs!\n'.format(key))
 		return o