From 51e456a4e5046a6a14509d7898bc6c87d3f9a44e Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Mon, 30 Dec 2024 11:31:39 +0000
Subject: [PATCH] improve load_cryptodome() monkey-patch function

---
 alt-requirements.txt         |  2 +-
 mmgen/util2.py               | 41 ++++++++++++++++++++----------------
 pyproject.toml               |  1 +
 test/gentest.py              |  7 +++---
 test/modtest_d/ut_dep.py     |  3 ++-
 test/modtest_d/ut_testdep.py |  4 ++--
 6 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/alt-requirements.txt b/alt-requirements.txt
index 1a668a67..acdfd202 100644
--- a/alt-requirements.txt
+++ b/alt-requirements.txt
@@ -1 +1 @@
-pycryptodomex
+pycryptodome
diff --git a/mmgen/util2.py b/mmgen/util2.py
index 7e3da4e5..f374d202 100755
--- a/mmgen/util2.py
+++ b/mmgen/util2.py
@@ -30,18 +30,30 @@ def die_pause(ev=0, s=''):
 	input('Press ENTER to exit')
 	sys.exit(ev)
 
-# monkey-patch function for monero-python: permits its use with pycryptodome (e.g. MSYS2)
-# instead of the expected pycryptodomex
-def load_cryptodomex():
-	try:
-		import Cryptodome # cryptodomex
-	except ImportError:
+def cffi_override_fixup():
+	from cffi import FFI
+	class FFI_override:
+		def cdef(self, csource, override=False, packed=False, pack=None):
+			self._cdef(csource, override=True, packed=packed, pack=pack)
+	FFI.cdef = FFI_override.cdef
+
+# monkey-patch function: makes modules pycryptodome and pycryptodomex available to packages that
+# expect them (monero-python, eth-keys), regardless of which one is installed on system
+def load_cryptodome(called=[]):
+	if not called:
+		cffi_override_fixup()
 		try:
-			import Crypto # cryptodome
+			import Crypto # Crypto == pycryptodome
 		except ImportError:
-			die(2, 'Unable to import either the ‘pycryptodomex’ or ‘pycryptodome’ package')
+			try:
+				import Cryptodome # Crypto == pycryptodome
+			except ImportError:
+				die(2, 'Unable to import the ‘pycryptodome’ or ‘pycryptodomex’ package')
+			else:
+				sys.modules['Crypto'] = Cryptodome # Crypto == pycryptodome
 		else:
-			sys.modules['Cryptodome'] = Crypto
+			sys.modules['Cryptodome'] = Crypto # Cryptodome == pycryptodomex
+		called.append(True)
 
 # called with no arguments by pyethereum.utils:
 def get_keccak(cfg=None, cached_ret=[]):
@@ -51,15 +63,8 @@ def get_keccak(cfg=None, cached_ret=[]):
 			cfg._util.qmsg('Using internal keccak module by user request')
 			from .contrib.keccak import keccak_256
 		else:
-			try:
-				from Cryptodome.Hash import keccak
-			except ImportError as e:
-				try:
-					from Crypto.Hash import keccak
-				except ImportError as e2:
-					msg(f'{e2} and {e}')
-					die('MMGenImportError',
-						'Please install the ‘pycryptodome’ or ‘pycryptodomex’ package on your system')
+			load_cryptodome()
+			from Crypto.Hash import keccak
 			def keccak_256(data):
 				return keccak.new(data=data, digest_bytes=32)
 		cached_ret.append(keccak_256)
diff --git a/pyproject.toml b/pyproject.toml
index 9c2251f1..f0874546 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -91,4 +91,5 @@ ignored-classes = [ # ignored for no-member, otherwise checked
 	"SwapMgrBase",
 	"Opts",
 	"Help",
+	"FFI_override",
 ]
diff --git a/test/gentest.py b/test/gentest.py
index 0cb9c852..36076682 100755
--- a/test/gentest.py
+++ b/test/gentest.py
@@ -565,7 +565,6 @@ from mmgen.key import PrivKey
 from mmgen.addr import MMGenAddrType
 from mmgen.addrgen import KeyGenerator, AddrGenerator
 from mmgen.keygen import get_backends
-from mmgen.util2 import load_cryptodomex
 from test.include.common import getrand, get_ethkey, set_globals
 
 gtr = namedtuple('gen_tool_result', ['wif', 'addr', 'viewkey'])
@@ -582,9 +581,9 @@ vmsg = cfg._util.vmsg
 
 proto = cfg._proto
 
-if proto.coin == 'XMR':
-	# This must be done at top level, not in monero tool __init__
-	load_cryptodomex()
+if proto.coin in ('XMR', 'ETH', 'ETC'):
+	from mmgen.util2 import load_cryptodome
+	load_cryptodome()
 
 if __name__ == '__main__':
 	from mmgen.main import launch
diff --git a/test/modtest_d/ut_dep.py b/test/modtest_d/ut_dep.py
index 04cea944..88383a08 100755
--- a/test/modtest_d/ut_dep.py
+++ b/test/modtest_d/ut_dep.py
@@ -32,7 +32,8 @@ class unit_tests:
 	def keccak(self, name, ut): # used by ETH, XMR
 		from mmgen.util2 import get_keccak
 		try:
-			get_keccak()
+			keccak_256 = get_keccak()
+			keccak_256(b'abc')
 		except Exception as e:
 			rmsg(str(e))
 			return False
diff --git a/test/modtest_d/ut_testdep.py b/test/modtest_d/ut_testdep.py
index 35b4b058..46accab1 100755
--- a/test/modtest_d/ut_testdep.py
+++ b/test/modtest_d/ut_testdep.py
@@ -47,8 +47,8 @@ class unit_tests:
 		return True
 
 	def monero_python(self, name, ut):
-		from mmgen.util2 import load_cryptodomex
-		load_cryptodomex()
+		from mmgen.util2 import load_cryptodome
+		load_cryptodome()
 		from monero.seed import Seed
 		Seed('deadbeef' * 8).public_address()
 		return True