Browse Source

new get_keccak() function to select internal keccak_256 module

The MMGen Project 3 years ago
parent
commit
8d957a15dd
6 changed files with 41 additions and 34 deletions
  1. 4 12
      mmgen/addr.py
  2. 10 11
      mmgen/altcoins/eth/contract.py
  3. 3 3
      mmgen/globalvars.py
  4. 6 8
      mmgen/protocol.py
  5. 15 0
      mmgen/util.py
  6. 3 0
      test/unit_tests.py

+ 4 - 12
mmgen/addr.py

@@ -84,12 +84,8 @@ class AddrGeneratorEthereum(AddrGenerator):
 
 	def __init__(self,proto,addr_type):
 
-		try:
-			assert not g.use_internal_keccak_module
-			from sha3 import keccak_256
-		except:
-			from .keccak import keccak_256
-		self.keccak_256 = keccak_256
+		from .util import get_keccak
+		self.keccak_256 = get_keccak()
 
 		from .protocol import hash256
 		self.hash256 = hash256
@@ -144,12 +140,8 @@ class AddrGeneratorMonero(AddrGenerator):
 
 	def __init__(self,proto,addr_type):
 
-		try:
-			assert not g.use_internal_keccak_module
-			from sha3 import keccak_256
-		except:
-			from .keccak import keccak_256
-		self.keccak_256 = keccak_256
+		from .util import get_keccak
+		self.keccak_256 = get_keccak()
 
 		from .protocol import hash256
 		self.hash256 = hash256

+ 10 - 11
mmgen/altcoins/eth/contract.py

@@ -30,20 +30,14 @@ from mmgen.obj import MMGenObject,CoinAddr,TokenAddr,CoinTxID
 from mmgen.util import msg
 from .obj import ETHAmt
 
-try:
-	assert not g.use_internal_keccak_module
-	from sha3 import keccak_256
-except:
-	from mmgen.keccak import keccak_256
-
 def parse_abi(s):
 	return [s[:8]] + [s[8+x*64:8+(x+1)*64] for x in range(len(s[8:])//64)]
 
-def create_method_id(sig):
-	return keccak_256(sig.encode()).hexdigest()[:8]
-
 class TokenBase(MMGenObject): # ERC20
 
+	def create_method_id(self,sig):
+		return self.keccak_256(sig.encode()).hexdigest()[:8]
+
 	def transferdata2sendaddr(self,data): # online
 		return CoinAddr(self.proto,parse_abi(data)[1][-40:])
 
@@ -51,7 +45,7 @@ class TokenBase(MMGenObject): # ERC20
 		return ETHAmt(int(parse_abi(data)[-1],16) * self.base_unit)
 
 	async def do_call(self,method_sig,method_args='',toUnit=False):
-		data = create_method_id(method_sig) + method_args
+		data = self.create_method_id(method_sig) + method_args
 		if g.debug:
 			msg('ETH_CALL {}:  {}'.format(
 				method_sig,
@@ -104,7 +98,7 @@ class TokenBase(MMGenObject): # ERC20
 		from_arg = from_addr.rjust(64,'0') if from_addr else ''
 		to_arg = to_addr.rjust(64,'0')
 		amt_arg = '{:064x}'.format( int(amt / self.base_unit) )
-		return create_method_id(method_sig) + from_arg + to_arg + amt_arg
+		return self.create_method_id(method_sig) + from_arg + to_arg + amt_arg
 
 	def make_tx_in( self,from_addr,to_addr,amt,start_gas,gasPrice,nonce,
 					method_sig='transfer(address,uint256)',from_addr2=None):
@@ -157,6 +151,9 @@ class TokenBase(MMGenObject): # ERC20
 class Token(TokenBase):
 
 	def __init__(self,proto,addr,decimals,rpc=None):
+		if type(self).__name__ == 'Token':
+			from mmgen.util import get_keccak
+			self.keccak_256 = get_keccak()
 		self.proto = proto
 		self.addr = TokenAddr(proto,addr)
 		assert isinstance(decimals,int),f'decimals param must be int instance, not {type(decimals)}'
@@ -167,6 +164,8 @@ class Token(TokenBase):
 class TokenResolve(TokenBase,metaclass=AsyncInit):
 
 	async def __init__(self,proto,rpc,addr):
+		from mmgen.util import get_keccak
+		self.keccak_256 = get_keccak()
 		self.proto = proto
 		self.rpc = rpc
 		self.addr = TokenAddr(proto,addr)

+ 3 - 3
mmgen/globalvars.py

@@ -155,11 +155,11 @@ class GlobalContext(Lockable):
 
 	# global var sets user opt:
 	global_sets_opt = (
-		'minconf','usr_randchars','debug', 'quiet','tx_confs','tx_fee_adj','key_generator' )
+		'use_internal_keccak_module',
+		'minconf','usr_randchars','debug','quiet','tx_confs','tx_fee_adj' )
 
 	# user opt sets global var:
-	opt_sets_global = (
-		'use_internal_keccak_module','subseeds','cached_balances' )
+	opt_sets_global = ( 'subseeds','cached_balances' )
 
 	# 'long' opts - opt sets global var
 	common_opts = (

+ 6 - 8
mmgen/protocol.py

@@ -115,6 +115,10 @@ class CoinProtocol(MMGenObject):
 			if self.tokensym:
 				assert isinstance(self,CoinProtocol.Ethereum), 'CoinProtocol.Base_chk1'
 
+			if self.base_coin in ('ETH','XMR'):
+				from .util import get_keccak
+				self.keccak_256 = get_keccak()
+
 		@property
 		def dcoin(self):
 			return self.coin
@@ -429,8 +433,7 @@ class CoinProtocol(MMGenObject):
 
 		@classmethod
 		def checksummed_addr(cls,addr):
-			from .keccak import keccak_256
-			h = keccak_256(addr.encode()).digest().hex()
+			h = self.keccak_256(addr.encode()).digest().hex()
 			return ''.join(addr[i].upper() if int(h[i],16) > 7 else addr[i] for i in range(len(addr)))
 
 		def pubhash2addr(self,pubkey_hash,p2sh):
@@ -520,13 +523,8 @@ class CoinProtocol(MMGenObject):
 
 			ret = b58dec(addr)
 
-			try:
-				assert not g.use_internal_keccak_module
-				from sha3 import keccak_256
-			except:
-				from .keccak import keccak_256
+			chk = self.keccak_256(ret[:-4]).digest()[:4]
 
-			chk = keccak_256(ret[:-4]).digest()[:4]
 			assert ret[-4:] == chk, f'{ret[-4:].hex()}: incorrect checksum.  Correct value: {chk.hex()}'
 
 			return self.parse_addr_bytes(ret)

+ 15 - 0
mmgen/util.py

@@ -175,6 +175,21 @@ def warn_altcoins(coinsym,trust_level):
 	if not keypress_confirm(m,default_yes=True):
 		sys.exit(0)
 
+def get_keccak():
+
+	from .opts import opt
+	if getattr(opt,'use_internal_keccak_module',False):
+		from .keccak import keccak_256
+		qmsg('Using internal keccak module by user request')
+		return keccak_256
+
+	try:
+		from sha3 import keccak_256
+	except:
+		from .keccak import keccak_256
+
+	return keccak_256
+
 def set_for_type(val,refval,desc,invert_bool=False,src=None):
 
 	if type(refval) == bool:

+ 3 - 0
test/unit_tests.py

@@ -50,6 +50,9 @@ If no test is specified, all available tests are run
 }
 
 sys.argv.insert(1,'--skip-cfg-file')
+
+opts.UserOpts._reset_ok += ('use_internal_keccak_module',)
+
 cmd_args = opts.init(opts_data)
 
 file_pfx = 'nt_' if opt.node_tools else 'ut_'