From d1d7132a7cc4bccb3886e68d35b1479932647bb2 Mon Sep 17 00:00:00 2001 From: The MMGen Project Date: Tue, 9 Jun 2020 14:17:46 +0000 Subject: [PATCH] PubKey: add privkey attribute, perform pubkey_type matching checks --- mmgen/addr.py | 50 +++++++++++++++++------------ mmgen/obj.py | 11 +++---- test/objtest_py_d/ot_btc_mainnet.py | 5 +-- test/test_py_d/ts_main.py | 4 +-- 4 files changed, 40 insertions(+), 30 deletions(-) diff --git a/mmgen/addr.py b/mmgen/addr.py index 0e5f6ac9..56c242d9 100755 --- a/mmgen/addr.py +++ b/mmgen/addr.py @@ -24,7 +24,7 @@ from hashlib import sha256,sha512 from .common import * from .obj import * from .baseconv import * -from .protocol import init_proto +from .protocol import init_proto,hash160 pnm = g.proj_name @@ -34,30 +34,32 @@ def dmsg_sc(desc,data): class AddrGenerator(MMGenObject): def __new__(cls,proto,addr_type): - if type(addr_type) == str: # allow override w/o check - gen_method = addr_type + + if type(addr_type) == str: + addr_type = MMGenAddrType(proto=proto,id_str=addr_type) elif type(addr_type) == MMGenAddrType: assert addr_type in proto.mmtypes, f'{addr_type}: invalid address type for coin {proto.coin}' - gen_method = addr_type.gen_method else: raise TypeError(f'{type(addr_type)}: incorrect argument type for {cls.__name__}()') - gen_methods = { + + addr_generators = { 'p2pkh': AddrGeneratorP2PKH, 'segwit': AddrGeneratorSegwit, 'bech32': AddrGeneratorBech32, 'ethereum': AddrGeneratorEthereum, 'zcash_z': AddrGeneratorZcashZ, - 'monero': AddrGeneratorMonero} - assert gen_method in gen_methods - me = super(cls,cls).__new__(gen_methods[gen_method]) - me.desc = gen_methods + 'monero': AddrGeneratorMonero, + } + me = super(cls,cls).__new__(addr_generators[addr_type.gen_method]) + me.desc = type(me).__name__ me.proto = proto + me.addr_type = addr_type + me.pubkey_type = addr_type.pubkey_type return me class AddrGeneratorP2PKH(AddrGenerator): def to_addr(self,pubhex): - from .protocol import hash160 - assert type(pubhex) == PubKey + assert pubhex.privkey.pubkey_type == self.pubkey_type return CoinAddr(self.proto,self.proto.pubhash2addr(hash160(pubhex),p2sh=False)) def to_segwit_redeem_script(self,pubhex): @@ -65,6 +67,7 @@ class AddrGeneratorP2PKH(AddrGenerator): class AddrGeneratorSegwit(AddrGenerator): def to_addr(self,pubhex): + assert pubhex.privkey.pubkey_type == self.pubkey_type assert pubhex.compressed,'Uncompressed public keys incompatible with Segwit' return CoinAddr(self.proto,self.proto.pubhex2segwitaddr(pubhex)) @@ -74,8 +77,8 @@ class AddrGeneratorSegwit(AddrGenerator): class AddrGeneratorBech32(AddrGenerator): def to_addr(self,pubhex): + assert pubhex.privkey.pubkey_type == self.pubkey_type assert pubhex.compressed,'Uncompressed public keys incompatible with Segwit' - from .protocol import hash160 return CoinAddr(self.proto,self.proto.pubhash2bech32addr(hash160(pubhex))) def to_segwit_redeem_script(self,pubhex): @@ -96,7 +99,7 @@ class AddrGeneratorEthereum(AddrGenerator): self.hash256 = hash256 def to_addr(self,pubhex): - assert type(pubhex) == PubKey + assert pubhex.privkey.pubkey_type == self.pubkey_type return CoinAddr(self.proto,self.keccak_256(bytes.fromhex(pubhex[2:])).hexdigest()[24:]) def to_wallet_passwd(self,sk_hex): @@ -116,6 +119,7 @@ class AddrGeneratorZcashZ(AddrGenerator): return Sha256(s,preprocess=False).digest() def to_addr(self,pubhex): # pubhex is really privhex + assert pubhex.privkey.pubkey_type == self.pubkey_type key = bytes.fromhex(pubhex) assert len(key) == 32, f'{len(key)}: incorrect privkey length' from nacl.bindings import crypto_scalarmult_base @@ -173,6 +177,7 @@ class AddrGeneratorMonero(AddrGenerator): return a + b def to_addr(self,sk_hex): # sk_hex instead of pubhex + assert sk_hex.privkey.pubkey_type == self.pubkey_type # Source and license for scalarmultbase function: # https://github.com/bigreddmachine/MoneroPy/blob/master/moneropy/crypto/ed25519.py @@ -277,21 +282,26 @@ class KeyGeneratorPython(KeyGenerator): def to_pubhex(self,privhex): assert type(privhex) == PrivKey - return PubKey(self.privnum2pubhex( - int(privhex,16),compressed=privhex.compressed),compressed=privhex.compressed) + return PubKey( + s = self.privnum2pubhex(int(privhex,16),compressed=privhex.compressed), + privkey = privhex ) class KeyGeneratorSecp256k1(KeyGenerator): desc = 'mmgen-secp256k1' def to_pubhex(self,privhex): assert type(privhex) == PrivKey from .secp256k1 import priv2pub - return PubKey(priv2pub(bytes.fromhex(privhex),int(privhex.compressed)).hex(),compressed=privhex.compressed) + return PubKey( + s = priv2pub(bytes.fromhex(privhex),int(privhex.compressed)).hex(), + privkey = privhex ) class KeyGeneratorDummy(KeyGenerator): desc = 'mmgen-dummy' def to_pubhex(self,privhex): assert type(privhex) == PrivKey - return PubKey(privhex,compressed=privhex.compressed) + return PubKey( + s = privhex, + privkey = privhex ) class AddrListEntryBase(MMGenListItem): invalid_attrs = {'proto'} @@ -605,9 +615,9 @@ Removed {{}} duplicate WIF key{{}} from keylist (also in {pnm} key-address file def generate_addrs_from_keys(self): # assume that the first listed mmtype is valid for flat key list - t = self.proto.addr_type(self.proto.mmtypes[0]) - kg = KeyGenerator(self.proto,t.pubkey_type) - ag = AddrGenerator(self.proto,t.gen_method) + at = self.proto.addr_type(self.proto.mmtypes[0]) + kg = KeyGenerator(self.proto,at.pubkey_type) + ag = AddrGenerator(self.proto,at) d = self.data for n,e in enumerate(d,1): qmsg_r('\rGenerating addresses from keylist: {}/{}'.format(n,len(d))) diff --git a/mmgen/obj.py b/mmgen/obj.py index 30bf1a9a..b14ac959 100755 --- a/mmgen/obj.py +++ b/mmgen/obj.py @@ -746,15 +746,14 @@ class WifKey(str,Hilite,InitErrors): return cls.init_fail(e,wif) class PubKey(HexStr,MMGenObject): # TODO: add some real checks - def __new__(cls,s,compressed): + def __new__(cls,s,privkey): try: - assert type(compressed) == bool,"'compressed' must be of type bool" + me = HexStr.__new__(cls,s,case='lower') + me.privkey = privkey + me.compressed = privkey.compressed + return me except Exception as e: return cls.init_fail(e,s) - me = HexStr.__new__(cls,s,case='lower') - if me: - me.compressed = compressed - return me class PrivKey(str,Hilite,InitErrors,MMGenObject): """ diff --git a/test/objtest_py_d/ot_btc_mainnet.py b/test/objtest_py_d/ot_btc_mainnet.py index 6e39fd65..b9f5a9b9 100755 --- a/test/objtest_py_d/ot_btc_mainnet.py +++ b/test/objtest_py_d/ot_btc_mainnet.py @@ -16,6 +16,7 @@ proto = init_proto('btc') tw_pfx = proto.base_coin.lower() + ':' ssm = str(SeedShareCount.max_val) +privkey = PrivKey(proto=proto,s=bytes.fromhex('deadbeef'*8),compressed=True,pubkey_type='std') tests = { 'Int': { @@ -224,8 +225,8 @@ tests = { }, 'PubKey': { 'arg1': 's', - 'bad': ({'arg':1,'compressed':False},{'arg':'F00BAA12','compressed':False},), - 'good': ({'arg':'deadbeef','compressed':True},) # TODO: add real pubkeys + 'bad': ({'s':1,'privkey':False},{'s':'F00BAA12','privkey':False},), + 'good': ({'s':'deadbeef','privkey':privkey},) # TODO: add real pubkeys }, 'PrivKey': { 'arg1': 'proto', diff --git a/test/test_py_d/ts_main.py b/test/test_py_d/ts_main.py index e962be28..5d52f322 100755 --- a/test/test_py_d/ts_main.py +++ b/test/test_py_d/ts_main.py @@ -339,7 +339,7 @@ class TestSuiteMain(TestSuiteBase,TestSuiteShared): from mmgen.addr import AddrGenerator,KeyGenerator rand_coinaddr = AddrGenerator( self.proto, - 'p2pkh' + 'compressed' ).to_addr(KeyGenerator(self.proto,'std').to_pubhex(privkey)) of = joinpath(self.cfgs[non_mmgen_input]['tmpdir'],non_mmgen_fn) write_data_to_file( @@ -376,7 +376,7 @@ class TestSuiteMain(TestSuiteBase,TestSuiteShared): def _make_txcreate_cmdline(self,tx_data): from mmgen.obj import PrivKey privkey = PrivKey(self.proto,os.urandom(32),compressed=True,pubkey_type='std') - t = ('p2pkh','segwit')['S' in self.proto.mmtypes] + t = ('compressed','segwit')['S' in self.proto.mmtypes] from mmgen.addr import AddrGenerator,KeyGenerator rand_coinaddr = AddrGenerator(self.proto,t).to_addr(KeyGenerator(self.proto,'std').to_pubhex(privkey))