From 56b3c9833267e888c205cd3f67e95c59e150ff68 Mon Sep 17 00:00:00 2001
From: ThomasV <thomasv@gitorious>
Date: Fri, 26 Jun 2015 14:29:26 +0200
Subject: [PATCH] generic m of n multisig

---
 gui/qt/installwizard.py | 114 +++++++++++++++++++++++++++++++++-------
 gui/qt/main_window.py   |   1 +
 lib/__init__.py         |   2 +-
 lib/account.py          |  36 ++++---------
 lib/transaction.py      |  62 ++++++++--------------
 lib/wallet.py           |  79 +++++++++++-----------------
 plugins/trustedcoin.py  |  14 +++--
 7 files changed, 169 insertions(+), 139 deletions(-)

diff --git a/gui/qt/installwizard.py b/gui/qt/installwizard.py
index e611b3e3..003117f5 100644
--- a/gui/qt/installwizard.py
+++ b/gui/qt/installwizard.py
@@ -1,10 +1,14 @@
+import re
+import sys
+import threading
+
 from PyQt4.QtGui import *
 from PyQt4.QtCore import *
 import PyQt4.QtCore as QtCore
 
 import electrum
 from electrum.i18n import _
-from electrum import Wallet, Wallet_2of2, Wallet_2of3
+from electrum import Wallet
 from electrum import bitcoin
 from electrum import util
 
@@ -13,8 +17,6 @@ from network_dialog import NetworkDialog
 from util import *
 from amountedit import AmountEdit
 
-import sys
-import threading
 from electrum.plugins import always_hook, run_hook
 from electrum.mnemonic import prepare_seed
 
@@ -26,6 +28,42 @@ MSG_ENTER_SEED_OR_MPK = _("Please enter a wallet seed, BIP32 private key, or mas
 MSG_VERIFY_SEED       = _("Your seed is important!") + "\n" + _("To make sure that you have properly saved your seed, please retype it here.")
 
 
+class CosignWidget(QWidget):
+    size = 120
+
+    def __init__(self, m, n):
+        QWidget.__init__(self)
+        self.R = QRect(0, 0, self.size, self.size)
+        self.setGeometry(self.R)
+        self.m = m
+        self.n = n
+
+    def set_n(self, n):
+        self.n = n
+        self.update()
+
+    def set_m(self, m):
+        self.m = m
+        self.update()
+
+    def paintEvent(self, event):
+        import math
+        bgcolor = self.palette().color(QPalette.Background)
+        pen = QPen(bgcolor, 7, QtCore.Qt.SolidLine)
+        qp = QPainter()
+        qp.begin(self)
+        qp.setPen(pen)
+        qp.setRenderHint(QPainter.Antialiasing)
+        qp.setBrush(Qt.gray)
+        for i in range(self.n):
+            alpha = int(16* 360 * i/self.n)
+            alpha2 = int(16* 360 * 1/self.n)
+            qp.setBrush(Qt.green if i<self.m else Qt.gray)
+            qp.drawPie(self.R, alpha, alpha2)
+        qp.end()
+
+
+
 class InstallWizard(QDialog):
 
     def __init__(self, config, network, storage, app):
@@ -36,7 +74,7 @@ class InstallWizard(QDialog):
         self.storage = storage
         self.setMinimumSize(575, 400)
         self.setMaximumSize(575, 400)
-        self.setWindowTitle('Electrum' + '  -  ' + os.path.basename(self.storage.path))
+        self.setWindowTitle('Electrum' + '  -  ' + _('Install Wizard'))
         self.connect(self, QtCore.SIGNAL('accept'), self.accept)
         self.stack = QStackedLayout()
         self.setLayout(self.stack)
@@ -283,6 +321,48 @@ class InstallWizard(QDialog):
         return wallet_type
 
 
+    def multisig_choice(self):
+
+        vbox = QVBoxLayout()
+        self.set_layout(vbox)
+        vbox.addWidget(QLabel(_("Multi Signature Wallet")))
+
+        cw = CosignWidget(2, 2)
+        vbox.addWidget(cw, 1)
+        vbox.addWidget(QLabel(_("Please choose the number of signatures needed to unlock funds in your wallet") + ':'))
+
+        m_edit = QSpinBox()
+        n_edit = QSpinBox()
+        m_edit.setValue(2)
+        n_edit.setValue(2)
+        n_edit.setMinimum(2)
+        n_edit.setMaximum(15)
+        m_edit.setMinimum(1)
+        m_edit.setMaximum(2)
+        n_edit.valueChanged.connect(m_edit.setMaximum)
+
+        n_edit.valueChanged.connect(cw.set_n)
+        m_edit.valueChanged.connect(cw.set_m)
+
+        hbox = QHBoxLayout()
+        hbox.addWidget(QLabel(_('Require')))
+        hbox.addWidget(m_edit)
+        hbox.addWidget(QLabel(_('of')))
+        hbox.addWidget(n_edit)
+        hbox.addWidget(QLabel(_('signatures')))
+        hbox.addStretch(1)
+
+        vbox.addLayout(hbox)
+        vbox.addStretch(1)
+        vbox.addLayout(Buttons(CancelButton(self), OkButton(self, _('Next'))))
+        if not self.exec_():
+            return
+        m = int(m_edit.value())
+        n = int(n_edit.value())
+        wallet_type = '%dof%d'%(m,n)
+        return wallet_type
+
+
     def question(self, msg, yes_label=_('OK'), no_label=_('Cancel'), icon=None):
         vbox = QVBoxLayout()
         self.set_layout(vbox)
@@ -319,7 +399,7 @@ class InstallWizard(QDialog):
 
         if action in ['create', 'restore']:
             if wallet_type == 'multisig':
-                wallet_type = self.choice(_("Multi Signature Wallet"), 'Select wallet type', [('2of2', _("2 of 2")),('2of3',_("2 of 3"))])
+                wallet_type = self.multisig_choice()
                 if not wallet_type:
                     return
             elif wallet_type == 'hardware':
@@ -363,22 +443,14 @@ class InstallWizard(QDialog):
                 wallet.add_seed(seed, password)
                 wallet.create_master_keys(password)
 
-            elif action == 'add_cosigner':
+            elif action == 'add_cosigners':
+                n = int(re.match('(\d+)of(\d+)', wallet.wallet_type).group(2))
                 xpub1 = wallet.master_public_keys.get("x1/")
-                r = self.multi_mpk_dialog(xpub1, 1)
+                r = self.multi_mpk_dialog(xpub1, n - 1)
                 if not r:
                     return
-                xpub2 = r[0]
-                wallet.add_master_public_key("x2/", xpub2)
-
-            elif action == 'add_two_cosigners':
-                xpub1 = wallet.master_public_keys.get("x1/")
-                r = self.multi_mpk_dialog(xpub1, 2)
-                if not r:
-                    return
-                xpub2, xpub3 = r
-                wallet.add_master_public_key("x2/", xpub2)
-                wallet.add_master_public_key("x3/", xpub3)
+                for i, xpub in enumerate(r):
+                    wallet.add_master_public_key("x%d/"%(i+2), xpub)
 
             elif action == 'create_accounts':
                 wallet.create_main_account(password)
@@ -443,8 +515,10 @@ class InstallWizard(QDialog):
                 else:
                     raise BaseException('unknown wallet type')
 
-            elif t in ['2of2', '2of3']:
-                key_list = self.multi_seed_dialog(1 if t == '2of2' else 2)
+            elif re.match('(\d+)of(\d+)', t):
+                n = int(re.match('(\d+)of(\d+)', t).group(2))
+                print t, n
+                key_list = self.multi_seed_dialog(n - 1)
                 if not key_list:
                     return
                 password = self.password_dialog() if any(map(lambda x: Wallet.is_seed(x) or Wallet.is_xprv(x), key_list)) else None
diff --git a/gui/qt/main_window.py b/gui/qt/main_window.py
index 29a8c2fb..3260721c 100644
--- a/gui/qt/main_window.py
+++ b/gui/qt/main_window.py
@@ -276,6 +276,7 @@ class ElectrumWindow(QMainWindow):
         try:
             wallet = Wallet(storage)
         except BaseException as e:
+            traceback.print_exc(file=sys.stdout)
             QMessageBox.warning(None, _('Warning'), str(e), _('OK'))
             return
         action = wallet.get_action()
diff --git a/lib/__init__.py b/lib/__init__.py
index 3f2e88d0..b8119854 100644
--- a/lib/__init__.py
+++ b/lib/__init__.py
@@ -1,7 +1,7 @@
 from version import ELECTRUM_VERSION
 from util import format_satoshis, print_msg, print_json, print_error, set_verbosity
 from wallet import WalletSynchronizer, WalletStorage
-from wallet import Wallet, Wallet_2of2, Wallet_2of3, Imported_Wallet
+from wallet import Wallet, Imported_Wallet
 from network import Network, DEFAULT_SERVERS, DEFAULT_PORTS, pick_random_server
 from interface import Interface
 from simple_config import SimpleConfig, get_config, set_config
diff --git a/lib/account.py b/lib/account.py
index be914e25..a50d3f03 100644
--- a/lib/account.py
+++ b/lib/account.py
@@ -366,15 +366,17 @@ class BIP32_Account(Account):
 
 
 
-class BIP32_Account_2of2(BIP32_Account):
+class Multisig_Account(BIP32_Account):
 
     def __init__(self, v):
-        BIP32_Account.__init__(self, v)
-        self.xpub2 = v['xpub2']
+        self.m = v.get('m', 2)
+        Account.__init__(self, v)
+        self.xpub_list = v['xpubs']
 
     def dump(self):
-        d = BIP32_Account.dump(self)
-        d['xpub2'] = self.xpub2
+        d = Account.dump(self)
+        d['xpubs'] = self.xpub_list
+        d['m'] = self.m
         return d
 
     def get_pubkeys(self, for_change, n):
@@ -385,10 +387,10 @@ class BIP32_Account_2of2(BIP32_Account):
 
     def redeem_script(self, for_change, n):
         pubkeys = self.get_pubkeys(for_change, n)
-        return Transaction.multisig_script(sorted(pubkeys), 2)
+        return Transaction.multisig_script(sorted(pubkeys), self.m)
 
     def pubkeys_to_address(self, pubkeys):
-        redeem_script = Transaction.multisig_script(sorted(pubkeys), 2)
+        redeem_script = Transaction.multisig_script(sorted(pubkeys), self.m)
         address = hash_160_to_bc_address(hash_160(redeem_script.decode('hex')), 5)
         return address
 
@@ -396,25 +398,9 @@ class BIP32_Account_2of2(BIP32_Account):
         return self.pubkeys_to_address(self.get_pubkeys(for_change, n))
 
     def get_master_pubkeys(self):
-        return [self.xpub, self.xpub2]
+        return self.xpub_list
 
     def get_type(self):
-        return _('Multisig 2 of 2')
+        return _('Multisig %d of %d'%(self.m, len(self.xpub_list)))
 
 
-class BIP32_Account_2of3(BIP32_Account_2of2):
-
-    def __init__(self, v):
-        BIP32_Account_2of2.__init__(self, v)
-        self.xpub3 = v['xpub3']
-
-    def dump(self):
-        d = BIP32_Account_2of2.dump(self)
-        d['xpub3'] = self.xpub3
-        return d
-
-    def get_master_pubkeys(self):
-        return [self.xpub, self.xpub2, self.xpub3]
-
-    def get_type(self):
-        return _('Multisig 2 of 3')
diff --git a/lib/transaction.py b/lib/transaction.py
index 7382b776..600b6e98 100644
--- a/lib/transaction.py
+++ b/lib/transaction.py
@@ -368,32 +368,29 @@ def parse_scriptSig(d, bytes):
         d['address'] = address
         return
 
-    # p2sh transaction, 2 of n
+    # p2sh transaction, m of n
     match = [ opcodes.OP_0 ] + [ opcodes.OP_PUSHDATA4 ] * (len(decoded) - 1)
-
     if not match_decoded(decoded, match):
         print_error("cannot find address in input script", bytes.encode('hex'))
         return
-
     x_sig = [x[1].encode('hex') for x in decoded[1:-1]]
-    d['signatures'] = parse_sig(x_sig)
-    d['num_sig'] = 2
-
     dec2 = [ x for x in script_GetOp(decoded[-1][1]) ]
-    match_2of2 = [ opcodes.OP_2, opcodes.OP_PUSHDATA4, opcodes.OP_PUSHDATA4, opcodes.OP_2, opcodes.OP_CHECKMULTISIG ]
-    match_2of3 = [ opcodes.OP_2, opcodes.OP_PUSHDATA4, opcodes.OP_PUSHDATA4, opcodes.OP_PUSHDATA4, opcodes.OP_3, opcodes.OP_CHECKMULTISIG ]
-    if match_decoded(dec2, match_2of2):
-        x_pubkeys = [ dec2[1][1].encode('hex'), dec2[2][1].encode('hex') ]
-    elif match_decoded(dec2, match_2of3):
-        x_pubkeys = [ dec2[1][1].encode('hex'), dec2[2][1].encode('hex'), dec2[3][1].encode('hex') ]
-    else:
+    m = dec2[0][0] - opcodes.OP_1 + 1
+    n = dec2[-2][0] - opcodes.OP_1 + 1
+    op_m = opcodes.OP_1 + m - 1
+    op_n = opcodes.OP_1 + n - 1
+    match_multisig = [ op_m ] + [opcodes.OP_PUSHDATA4]*n + [ op_n, opcodes.OP_CHECKMULTISIG ]
+    if not match_decoded(dec2, match_multisig):
         print_error("cannot find address in input script", bytes.encode('hex'))
         return
-
-    d['x_pubkeys'] = x_pubkeys
+    x_pubkeys = map(lambda x: x[1].encode('hex'), dec2[1:-2])
     pubkeys = [parse_xpub(x)[0] for x in x_pubkeys]     # xpub, addr = parse_xpub()
+    redeemScript = Transaction.multisig_script(pubkeys, m)
+    # write result in d
+    d['num_sig'] = m
+    d['signatures'] = parse_sig(x_sig)
+    d['x_pubkeys'] = x_pubkeys
     d['pubkeys'] = pubkeys
-    redeemScript = Transaction.multisig_script(pubkeys,2)
     d['redeemScript'] = redeemScript
     d['address'] = hash_160_to_bc_address(hash_160(redeemScript.decode('hex')), 5)
 
@@ -535,31 +532,14 @@ class Transaction:
         return self
 
     @classmethod
-    def multisig_script(klass, public_keys, num=None):
+    def multisig_script(klass, public_keys, m):
         n = len(public_keys)
-        if num is None: num = n
-
-        assert num <= n and n in [2,3] , 'Only "2 of 2", and "2 of 3" transactions are supported'
-
-        if num==2:
-            s = '52'
-        elif num == 3:
-            s = '53'
-        else:
-            raise
-
-        for k in public_keys:
-            s += op_push(len(k)/2) + k
-        if n==2:
-            s += '52'
-        elif n==3:
-            s += '53'
-        else:
-            raise
-        s += 'ae'
-
-        return s
-
+        assert n <= 15
+        assert m <= n
+        op_m = format(opcodes.OP_1 + m - 1, 'x')
+        op_n = format(opcodes.OP_1 + n - 1, 'x')
+        keylist = [op_push(len(k)/2) + k for k in public_keys]
+        return op_m + ''.join(keylist) + op_n + 'ae'
 
     @classmethod
     def pay_script(self, output_type, addr):
@@ -617,7 +597,7 @@ class Transaction:
                 script += push_script(x_pubkey)
             else:
                 script = '00' + script          # put op_0 in front of script
-                redeem_script = self.multisig_script(pubkeys,2)
+                redeem_script = self.multisig_script(pubkeys, num_sig)
                 script += push_script(redeem_script)
 
         elif for_sig==i:
diff --git a/lib/wallet.py b/lib/wallet.py
index 1c78bf80..7c3194a7 100644
--- a/lib/wallet.py
+++ b/lib/wallet.py
@@ -240,7 +240,6 @@ class Abstract_Wallet(object):
 
     def load_accounts(self):
         self.accounts = {}
-
         d = self.storage.get('accounts', {})
         for k, v in d.items():
             if self.wallet_type == 'old' and k in [0, '0']:
@@ -248,10 +247,6 @@ class Abstract_Wallet(object):
                 self.accounts['0'] = OldAccount(v)
             elif v.get('imported'):
                 self.accounts[k] = ImportedAccount(v)
-            elif v.get('xpub3'):
-                self.accounts[k] = BIP32_Account_2of3(v)
-            elif v.get('xpub2'):
-                self.accounts[k] = BIP32_Account_2of2(v)
             elif v.get('xpub'):
                 self.accounts[k] = BIP32_Account(v)
             elif v.get('pending'):
@@ -942,7 +937,7 @@ class Abstract_Wallet(object):
 
         if redeemScript:
             txin['redeemScript'] = redeemScript
-            txin['num_sig'] = 2
+            txin['num_sig'] = account.m
         else:
             txin['redeemPubkey'] = account.get_pubkey(*sequence)
             txin['num_sig'] = 1
@@ -1732,55 +1727,45 @@ class NewWallet(BIP32_Wallet, Mnemonic):
         self.add_account('0', account)
 
 
-class Wallet_2of2(BIP32_Wallet, Mnemonic):
-    # Wallet with multisig addresses.
+class Multisig_Wallet(BIP32_Wallet, Mnemonic):
+    # generic m of n
     root_name = "x1/"
     root_derivation = "m/"
-    wallet_type = '2of2'
+
+    def __init__(self, storage):
+        BIP32_Wallet.__init__(self, storage)
+        self.wallet_type = storage.get('wallet_type')
+        m = re.match('(\d+)of(\d+)', self.wallet_type)
+        self.m = int(m.group(1))
+        self.n = int(m.group(2))
+
+    def load_accounts(self):
+        self.accounts = {}
+        d = self.storage.get('accounts', {})
+        v = d.get('0')
+        if v:
+            if v.get('xpub3'):
+                v['xpubs'] = [v['xpub'], v['xpub2'], v['xpub3']]
+            elif v.get('xpub2'):
+                v['xpubs'] = [v['xpub'], v['xpub2']]
+            self.accounts = {'0': Multisig_Account(v)}
 
     def create_main_account(self, password):
-        xpub1 = self.master_public_keys.get("x1/")
-        xpub2 = self.master_public_keys.get("x2/")
-        account = BIP32_Account_2of2({'xpub':xpub1, 'xpub2':xpub2})
+        account = Multisig_Account({'xpubs': self.master_public_keys.values(), 'm': self.m})
         self.add_account('0', account)
 
     def get_master_public_keys(self):
         return self.master_public_keys
 
     def get_action(self):
-        xpub1 = self.master_public_keys.get("x1/")
-        xpub2 = self.master_public_keys.get("x2/")
-        if xpub1 is None:
-            return 'create_seed'
-        if xpub2 is None:
-            return 'add_cosigner'
+        for i in range(self.n):
+            if self.master_public_keys.get("x%d/"%(i+1)) is None:
+                return 'create_seed' if i == 0 else 'add_cosigners'
         if not self.accounts:
             return 'create_accounts'
 
 
 
-class Wallet_2of3(Wallet_2of2):
-    # multisig 2 of 3
-    wallet_type = '2of3'
-
-    def create_main_account(self, password):
-        xpub1 = self.master_public_keys.get("x1/")
-        xpub2 = self.master_public_keys.get("x2/")
-        xpub3 = self.master_public_keys.get("x3/")
-        account = BIP32_Account_2of3({'xpub':xpub1, 'xpub2':xpub2, 'xpub3':xpub3})
-        self.add_account('0', account)
-
-    def get_action(self):
-        xpub1 = self.master_public_keys.get("x1/")
-        xpub2 = self.master_public_keys.get("x2/")
-        xpub3 = self.master_public_keys.get("x3/")
-        if xpub1 is None:
-            return 'create_seed'
-        if xpub2 is None or xpub3 is None:
-            return 'add_two_cosigners'
-        if not self.accounts:
-            return 'create_accounts'
-
 
 class OldWallet(Deterministic_Wallet):
     wallet_type = 'old'
@@ -1859,8 +1844,8 @@ wallet_types = [
     ('standard', 'xpub',     ("BIP32 Import"),             BIP32_Simple_Wallet),
     ('standard', 'standard', ("Standard wallet"),          NewWallet),
     ('standard', 'imported', ("Imported wallet"),          Imported_Wallet),
-    ('multisig', '2of2',     ("Multisig wallet (2 of 2)"), Wallet_2of2),
-    ('multisig', '2of3',     ("Multisig wallet (2 of 3)"), Wallet_2of3)
+    ('multisig', '2of2',     ("Multisig wallet (2 of 2)"), Multisig_Wallet),
+    ('multisig', '2of3',     ("Multisig wallet (2 of 3)"), Multisig_Wallet)
 ]
 
 # former WalletFactory
@@ -1898,7 +1883,10 @@ class Wallet(object):
                     WalletClass = c
                     break
             else:
-                raise BaseException('unknown wallet type', wallet_type)
+                if re.match('(\d+)of(\d+)', wallet_type):
+                    WalletClass = Multisig_Wallet
+                else:
+                    raise BaseException('unknown wallet type', wallet_type)
         else:
             if seed_version == OLD_SEED_VERSION:
                 WalletClass = OldWallet
@@ -2012,10 +2000,7 @@ class Wallet(object):
 
     @classmethod
     def from_multisig(klass, key_list, password, storage):
-        if len(key_list) == 2:
-            self = Wallet_2of2(storage)
-        elif len(key_list) == 3:
-            self = Wallet_2of3(storage)
+        self = Multisig_Wallet(storage)
         key_list = sorted(key_list, key = lambda x: klass.is_xpub(x))
         for i, text in enumerate(key_list):
             assert klass.is_seed(text) or klass.is_xprv(text) or klass.is_xpub(text)
diff --git a/plugins/trustedcoin.py b/plugins/trustedcoin.py
index 004ded7f..7d51f1e6 100644
--- a/plugins/trustedcoin.py
+++ b/plugins/trustedcoin.py
@@ -34,7 +34,7 @@ from electrum import bitcoin
 from electrum.bitcoin import *
 from electrum.mnemonic import Mnemonic
 from electrum import version
-from electrum.wallet import Wallet_2of3
+from electrum.wallet import Multisig_Wallet, BIP32_Wallet
 from electrum.i18n import _
 from electrum.plugins import BasePlugin, run_hook, hook
 
@@ -170,9 +170,13 @@ class TrustedCoinCosignerClient(object):
 server = TrustedCoinCosignerClient(user_agent="Electrum/" + version.ELECTRUM_VERSION)
 
 
-class Wallet_2fa(Wallet_2of3):
+class Wallet_2fa(Multisig_Wallet):
 
-    wallet_type = '2fa'
+    def __init__(self, storage):
+        BIP32_Wallet.__init__(self, storage)
+        self.wallet_type = '2fa'
+        self.m = 2
+        self.n = 3
 
     def get_action(self):
         xpub1 = self.master_public_keys.get("x1/")
@@ -191,13 +195,13 @@ class Wallet_2fa(Wallet_2of3):
         return Mnemonic('english').make_seed(num_bits=256, prefix=SEED_PREFIX)
 
     def estimated_fee(self, tx):
-        fee = Wallet_2of3.estimated_fee(self, tx)
+        fee = Multisig_Wallet.estimated_fee(self, tx)
         x = run_hook('extra_fee', tx)
         if x: fee += x
         return fee
 
     def get_tx_fee(self, tx):
-        fee = Wallet_2of3.get_tx_fee(self, tx)
+        fee = Multisig_Wallet.get_tx_fee(self, tx)
         x = run_hook('extra_fee', tx)
         if x: fee += x
         return fee