replace segwit wallet class with custom xpub type

lib/base_wizard.py

@@ -26,7 +26,7 @@
 import os
 import bitcoin
 import keystore
-from wallet import Wallet, Imported_Wallet, Standard_Wallet, Segwit_Wallet, Multisig_Wallet, WalletStorage, wallet_types
+from wallet import Wallet, Imported_Wallet, Standard_Wallet, Multisig_Wallet, WalletStorage, wallet_types
 from i18n import _
 from plugins import run_hook
 
@@ -339,9 +339,6 @@ class BaseWizard(object):
         if self.wallet_type == 'standard':
             self.storage.put('seed_type', self.seed_type)
             self.storage.put('keystore', k.dump())
-            if self.seed_type == 'segwit':
-                self.wallet = Segwit_Wallet(self.storage)
-            else:
             self.wallet = Standard_Wallet(self.storage)
             self.run('create_addresses')
         elif self.wallet_type == 'multisig':

lib/bitcoin.py

@@ -41,8 +41,8 @@ TESTNET = False
 ADDRTYPE_P2PKH = 0
 ADDRTYPE_P2SH = 5
 ADDRTYPE_P2WPKH = 6
-XPRV_HEADER = "0488ade4"
-XPUB_HEADER = "0488b21e"
+XPRV_HEADER = 0x0488ade4
+XPUB_HEADER = 0x0488b21e
 HEADERS_URL = "https://headers.electrum.org/blockchain_headers"
 
 def set_testnet():
@@ -53,10 +53,13 @@ def set_testnet():
     ADDRTYPE_P2PKH = 111
     ADDRTYPE_P2SH = 196
     ADDRTYPE_P2WPKH = 3
-    XPRV_HEADER = "04358394"
-    XPUB_HEADER = "043587cf"
+    XPRV_HEADER = 0x04358394
+    XPUB_HEADER = 0x043587cf
     HEADERS_URL = "https://headers.electrum.org/testnet_headers"
 
+
+
+
 ################################## transactions
 
 FEE_STEP = 10000
@@ -736,20 +739,35 @@ def _CKD_pub(cK, c, s):
     return cK_n, c_n
 
 
+def xprv_header(xtype):
+    return ("%08x"%(XPRV_HEADER + xtype)).decode('hex')
+
+def xpub_header(xtype):
+    return ("%08x"%(XPUB_HEADER + xtype)).decode('hex')
+
+def serialize_xpub(xtype, c, k):
+    xprv = xprv_header(xtype) + chr(0)*9 + c + chr(0) + k
+    return EncodeBase58Check(xprv)
+
+def serialize_xpub(xtype, c, cK, depth=0, fingerprint=chr(0)*4, child_number=chr(0)*4):
+    xpub = xpub_header(xtype) + chr(depth) + fingerprint + child_number + c + cK
+    return EncodeBase58Check(xpub)
+
 def deserialize_xkey(xkey, prv):
-    header = XPRV_HEADER if prv else XPUB_HEADER
     xkey = DecodeBase58Check(xkey)
-    assert len(xkey) == 78
+    if len(xkey) != 78:
+        raise BaseException('Invalid length')
     depth = ord(xkey[4])
     fingerprint = xkey[5:9]
     child_number = xkey[9:13]
     c = xkey[13:13+32]
-    if xkey[0:4].encode('hex') == header:
+    header = XPRV_HEADER if prv else XPUB_HEADER
+    xtype = int('0x' + xkey[0:4].encode('hex'), 16) - header
+    if xtype not in ([0, 1] if TESTNET else [0]):
+        raise BaseException('Invalid header')
     n = 33 if prv else 32
     K_or_k = xkey[13+n:]
-    else:
-        raise BaseException('wrong key')
-    return depth, fingerprint, child_number, c, K_or_k
+    return xtype, depth, fingerprint, child_number, c, K_or_k
 
 def deserialize_xpub(xkey):
     return deserialize_xkey(xkey, False)
@@ -757,36 +775,46 @@ def deserialize_xpub(xkey):
 def deserialize_xprv(xkey):
     return deserialize_xkey(xkey, True)
 
+def is_xpub(text):
+    try:
+        deserialize_xpub(text)
+        return True
+    except:
+        return False
+
+def is_xprv(text):
+    try:
+        deserialize_xprv(text)
+        return True
+    except:
+        return False
+
 
 def xpub_from_xprv(xprv):
-    depth, fingerprint, child_number, c, k = deserialize_xprv(xprv)
+    xtype, depth, fingerprint, child_number, c, k = deserialize_xprv(xprv)
     K, cK = get_pubkeys_from_secret(k)
-    xpub = XPUB_HEADER.decode('hex') + chr(depth) + fingerprint + child_number + c + cK
-    return EncodeBase58Check(xpub)
+    return serialize_xpub(xtype, c, cK, depth, fingerprint, child_number)
 
 
-def bip32_root(seed):
+def bip32_root(seed, xtype):
     I = hmac.new("Bitcoin seed", seed, hashlib.sha512).digest()
     master_k = I[0:32]
     master_c = I[32:]
     K, cK = get_pubkeys_from_secret(master_k)
-    xprv = (XPRV_HEADER + "00" + "00000000" + "00000000").decode("hex") + master_c + chr(0) + master_k
-    xpub = (XPUB_HEADER + "00" + "00000000" + "00000000").decode("hex") + master_c + cK
-    return EncodeBase58Check(xprv), EncodeBase58Check(xpub)
-
+    xprv = serialize_xprv(xtype, master_c, master_k)
+    xpub = serialize_xpub(xtype, master_c + cK)
+    return xprv, xpub
 
-def xpub_from_pubkey(cK):
+def xpub_from_pubkey(xtype, cK):
     assert cK[0] in ['\x02','\x03']
-    master_c = chr(0)*32
-    xpub = (XPUB_HEADER + "00" + "00000000" + "00000000").decode("hex") + master_c + cK
-    return EncodeBase58Check(xpub)
+    return serialize_xpub(xtype, chr(0)*32, cK)
 
 
 def bip32_private_derivation(xprv, branch, sequence):
     assert sequence.startswith(branch)
     if branch == sequence:
         return xprv, xpub_from_xprv(xprv)
-    depth, fingerprint, child_number, c, k = deserialize_xprv(xprv)
+    xtype, depth, fingerprint, child_number, c, k = deserialize_xprv(xprv)
     sequence = sequence[len(branch):]
     for n in sequence.split('/'):
         if n == '': continue
@@ -799,13 +827,13 @@ def bip32_private_derivation(xprv, branch, sequence):
     fingerprint = hash_160(parent_cK)[0:4]
     child_number = ("%08X"%i).decode('hex')
     K, cK = get_pubkeys_from_secret(k)
-    xprv = XPRV_HEADER.decode('hex') + chr(depth) + fingerprint + child_number + c + chr(0) + k
-    xpub = XPUB_HEADER.decode('hex') + chr(depth) + fingerprint + child_number + c + cK
+    xprv = xprv_header(xtype) + chr(depth) + fingerprint + child_number + c + chr(0) + k
+    xpub = xpub_header(xtype) + chr(depth) + fingerprint + child_number + c + cK
     return EncodeBase58Check(xprv), EncodeBase58Check(xpub)
 
 
 def bip32_public_derivation(xpub, branch, sequence):
-    depth, fingerprint, child_number, c, cK = deserialize_xpub(xpub)
+    xtype, depth, fingerprint, child_number, c, cK = deserialize_xpub(xpub)
     assert sequence.startswith(branch)
     sequence = sequence[len(branch):]
     for n in sequence.split('/'):
@@ -816,7 +844,7 @@ def bip32_public_derivation(xpub, branch, sequence):
         depth += 1
     fingerprint = hash_160(parent_cK)[0:4]
     child_number = ("%08X"%i).decode('hex')
-    xpub = XPUB_HEADER.decode('hex') + chr(depth) + fingerprint + child_number + c + cK
+    xpub = xpub_header(xtype) + chr(depth) + fingerprint + child_number + c + cK
     return EncodeBase58Check(xpub)
 
 

lib/keystore.py

@@ -241,7 +241,7 @@ class Xpub:
 
     @classmethod
     def get_pubkey_from_xpub(self, xpub, sequence):
-        _, _, _, c, cK = deserialize_xpub(xpub)
+        _, _, _, _, c, cK = deserialize_xpub(xpub)
         for i in sequence:
             cK, c = CKD_pub(cK, c, i)
         return cK.encode('hex')
@@ -298,7 +298,7 @@ class BIP32_KeyStore(Deterministic_KeyStore, Xpub):
 
     def check_password(self, password):
         xprv = pw_decode(self.xprv, password)
-        if deserialize_xprv(xprv)[3] != deserialize_xpub(self.xpub)[3]:
+        if deserialize_xprv(xprv)[4] != deserialize_xpub(self.xpub)[4]:
             raise InvalidPassword()
 
     def update_password(self, old_password, new_password):
@@ -322,14 +322,14 @@ class BIP32_KeyStore(Deterministic_KeyStore, Xpub):
         self.xprv = xprv
         self.xpub = bitcoin.xpub_from_xprv(xprv)
 
-    def add_xprv_from_seed(self, bip32_seed, derivation):
-        xprv, xpub = bip32_root(bip32_seed)
+    def add_xprv_from_seed(self, bip32_seed, xtype, derivation):
+        xprv, xpub = bip32_root(bip32_seed, xtype)
         xprv, xpub = bip32_private_derivation(xprv, "m/", derivation)
         self.add_xprv(xprv)
 
     def get_private_key(self, sequence, password):
         xprv = self.get_master_private_key(password)
-        _, _, _, c, k = deserialize_xprv(xprv)
+        _, _, _, _, c, k = deserialize_xprv(xprv)
         pk = bip32_private_key(sequence, k, c)
         return pk
 
@@ -622,20 +622,6 @@ def is_old_mpk(mpk):
         return False
     return len(mpk) == 128
 
-def is_xpub(text):
-    try:
-        deserialize_xpub(text)
-        return True
-    except:
-        return False
-
-def is_xprv(text):
-    try:
-        deserialize_xprv(text)
-        return True
-    except:
-        return False
-
 def is_address_list(text):
     parts = text.split()
     return bool(parts) and all(bitcoin.is_address(x) for x in parts)
@@ -669,7 +655,8 @@ def from_seed(seed, passphrase):
         keystore.add_seed(seed)
         keystore.passphrase = passphrase
         bip32_seed = Mnemonic.mnemonic_to_seed(seed, passphrase)
-        keystore.add_xprv_from_seed(bip32_seed, "m/")
+        xtype = 0 if t == 'standard' else 1
+        keystore.add_xprv_from_seed(bip32_seed, xtype, "m/")
     return keystore
 
 def from_private_key_list(text):

lib/wallet.py

@@ -1516,6 +1516,7 @@ class Simple_Wallet(Abstract_Wallet):
 
     def load_keystore(self):
         self.keystore = load_keystore(self.storage, 'keystore')
+        self.xpub_type = deserialize_xpub(self.keystore.xpub)[0]
 
     def get_pubkey(self, c, i):
         pubkey_list = self.change_pubkeys if c else self.receiving_pubkeys
@@ -1622,26 +1623,28 @@ class P2SH:
     def pubkeys_to_redeem_script(self, pubkeys):
         raise NotImplementedError()
 
-    def pubkeys_to_address(self, pubkeys):
-        redeem_script = self.pubkeys_to_redeem_script(pubkeys)
+    def pubkeys_to_address(self, pubkey):
+        redeem_script = self.pubkeys_to_redeem_script(pubkey)
         return bitcoin.hash160_to_p2sh(hash_160(redeem_script.decode('hex')))
 
 
-class P2PKH:
+class Standard_Wallet(Simple_Deterministic_Wallet):
+    wallet_type = 'standard'
+
+    def pubkeys_to_redeem_script(self, pubkeys):
+        if self.xpub_type == 1:
+            return transaction.segwit_script(pubkey)
 
     def pubkeys_to_address(self, pubkey):
+        if self.xpub_type == 0:
             return bitcoin.public_key_to_p2pkh(pubkey.decode('hex'))
+        elif self.xpub_type == 1 and TESTNET:
+            redeem_script = self.pubkeys_to_redeem_script(pubkeys)
+            return bitcoin.hash160_to_p2sh(hash_160(redeem_script.decode('hex')))
+        else:
+            raise NotImplementedError()
 
 
-class Standard_Wallet(Simple_Deterministic_Wallet, P2PKH):
-    wallet_type = 'standard'
-
-
-class Segwit_Wallet(Simple_Deterministic_Wallet, P2SH):
-    wallet_type = 'segwit'
-
-    def pubkeys_to_redeem_script(self, pubkey):
-        return transaction.segwit_script(pubkey)
 
 
 
@@ -1736,8 +1739,7 @@ wallet_constructors = {
     'standard': Standard_Wallet,
     'old': Standard_Wallet,
     'xpub': Standard_Wallet,
-    'imported': Imported_Wallet,
-    'segwit': Segwit_Wallet
+    'imported': Imported_Wallet
 }
 
 def register_constructor(wallet_type, constructor):

plugins/trustedcoin/trustedcoin.py

@@ -285,19 +285,16 @@ def get_user_id(storage):
     return long_id, short_id
 
 def make_xpub(xpub, s):
-    _, _, _, c, cK = deserialize_xpub(xpub)
+    version, _, _, _, c, cK = deserialize_xpub(xpub)
     cK2, c2 = bitcoin._CKD_pub(cK, c, s)
-    xpub2 = ("0488B21E" + "00" + "00000000" + "00000000").decode("hex") + c2 + cK2
-    return EncodeBase58Check(xpub2)
-
+    return bitcoin.serialize_xpub(version, c2, cK2)
 
 def make_billing_address(wallet, num):
     long_id, short_id = wallet.get_user_id()
     xpub = make_xpub(billing_xpub, long_id)
-    _, _, _, c, cK = deserialize_xpub(xpub)
+    version, _, _, _, c, cK = deserialize_xpub(xpub)
     cK, c = bitcoin.CKD_pub(cK, c, num)
-    address = public_key_to_bc_address( cK )
-    return address
+    return bitcoin.public_key_to_p2pkh(cK)
 
 
 class TrustedCoinPlugin(BasePlugin):
@@ -524,7 +521,7 @@ class TrustedCoinPlugin(BasePlugin):
         challenge = r.get('challenge')
         message = 'TRUSTEDCOIN CHALLENGE: ' + challenge
         def f(xprv):
-            _, _, _, c, k = deserialize_xprv(xprv)
+            _, _, _, _, c, k = deserialize_xprv(xprv)
             pk = bip32_private_key([0, 0], k, c)
             key = regenerate_key(pk)
             compressed = is_compressed(pk)