From e83068cd73abb1567d9e6e641049acd0a99ce8c0 Mon Sep 17 00:00:00 2001
From: Mohanson <mohanson@outlook.com>
Date: Thu, 16 May 2024 18:21:02 +0800
Subject: [PATCH] Replace p2as with 2-of-2 p2ms

---
 README.md          |  2 +-
 btc/__init__.py    |  1 +
 btc/taproot.py     | 22 ++++++++++++++++++
 example/taproot.py | 58 +++++++++++++++++++---------------------------
 4 files changed, 48 insertions(+), 35 deletions(-)
 create mode 100644 btc/taproot.py

diff --git a/README.md b/README.md
index bcaa17c..5eb9b60 100644
--- a/README.md
+++ b/README.md
@@ -32,7 +32,7 @@ $ python example/addr.py --net mainnet --prikey 1
 
 **example/taproot.py**
 
-This example demonstrates how to create a P2TR script with two script spending paths: p2pk and p2as.
+This example demonstrates how to create a P2TR script with two script spending paths: p2pk and p2ms(2-of-2 multisig).
 
 ```sh
 $ python example/taproot.py
diff --git a/btc/__init__.py b/btc/__init__.py
index e1a68f7..74d72f5 100644
--- a/btc/__init__.py
+++ b/btc/__init__.py
@@ -9,4 +9,5 @@
 from . import rpc
 from . import schnorr
 from . import secp256k1
+from . import taproot
 from . import wallet
diff --git a/btc/taproot.py b/btc/taproot.py
new file mode 100644
index 0000000..2a06cf9
--- /dev/null
+++ b/btc/taproot.py
@@ -0,0 +1,22 @@
+import btc
+import typing
+
+
+class Leaf:
+    def __init__(self, script: bytearray):
+        data = bytearray()
+        data.append(0xc0)
+        data.extend(btc.core.compact_size_encode(len(script)))
+        data.extend(script)
+        self.hash = btc.core.hashtag('TapLeaf', data)
+        self.script = script
+
+
+class Node:
+    def __init__(self, l: typing.Self | Leaf, r: typing.Self | Leaf):
+        if l.hash < r.hash:
+            self.hash = btc.core.hashtag('TapBranch', l.hash + r.hash)
+        else:
+            self.hash = btc.core.hashtag('TapBranch', r.hash + l.hash)
+        self.l = l
+        self.r = r
diff --git a/example/taproot.py b/example/taproot.py
index c15c402..5abe89f 100644
--- a/example/taproot.py
+++ b/example/taproot.py
@@ -1,39 +1,22 @@
 import btc
-import typing
-Self = typing.Self
 
-# This example shows how to create a P2TR script with two unlock conditions: p2pk and p2as.
-
-
-class TapLeaf:
-    def __init__(self, script: bytearray):
-        data = bytearray()
-        data.append(0xc0)
-        data.extend(btc.core.compact_size_encode(len(script)))
-        data.extend(script)
-        self.hash = btc.core.hashtag('TapLeaf', data)
-        self.script = script
-
-
-class TapNode:
-    def __init__(self, l: Self | TapLeaf, r: Self | TapLeaf):
-        if l.hash < r.hash:
-            self.hash = btc.core.hashtag('TapBranch', l.hash + r.hash)
-        else:
-            self.hash = btc.core.hashtag('TapBranch', r.hash + l.hash)
-        self.l = l
-        self.r = r
+# This example shows how to create a P2TR script with two unlock conditions: p2pk and p2ms.
 
 
 # Here created two scripts, one of which is a p2pk script, which requires that it can only be unlocked by private key 2,
-# and the other is an p2as(always success) script, which means that anyone can spend the utxo.
-mast = TapNode(
-    TapLeaf(btc.core.script([
+# and the other is an 2-of-2 multisig script.
+mast = btc.taproot.Node(
+    btc.taproot.Leaf(btc.core.script([
         btc.opcode.op_pushdata(btc.core.PriKey(2).pubkey().sec()[1:]),
         btc.opcode.op_checksig,
     ])),
-    TapLeaf(btc.core.script([
-        btc.opcode.op_1
+    btc.taproot.Leaf(btc.core.script([
+        btc.opcode.op_pushdata(btc.core.PriKey(3).pubkey().sec()[1:]),
+        btc.opcode.op_checksig,
+        btc.opcode.op_pushdata(btc.core.PriKey(4).pubkey().sec()[1:]),
+        btc.opcode.op_checksigadd,
+        btc.opcode.op_n(2),
+        btc.opcode.op_equal,
     ]))
 )
 
@@ -64,7 +47,7 @@ def txin(self, op: btc.core.OutPoint):
         ])
 
 
-class Tp2trp2as:
+class Tp2trp2ms:
     def __init__(self, pubkey: btc.core.PubKey):
         self.pubkey = pubkey
         self.addr = btc.core.address_p2tr(pubkey, mast.hash)
@@ -77,10 +60,17 @@ def __init__(self, pubkey: btc.core.PubKey):
             self.prefix = 0xc0
 
     def sign(self, tx: btc.core.Transaction):
-        return tx
+        for i, e in enumerate(tx.vin):
+            m = btc.secp256k1.Fr(int.from_bytes(tx.digest_segwit_v1(i, btc.core.sighash_all, mast.r.script)))
+            r, s = btc.schnorr.sign(btc.secp256k1.Fr(4), m)
+            e.witness[0] = bytearray(r.x.x.to_bytes(32) + s.x.to_bytes(32)) + bytearray([btc.core.sighash_all])
+            r, s = btc.schnorr.sign(btc.secp256k1.Fr(3), m)
+            e.witness[1] = bytearray(r.x.x.to_bytes(32) + s.x.to_bytes(32)) + bytearray([btc.core.sighash_all])
 
     def txin(self, op: btc.core.OutPoint):
         return btc.core.TxIn(op, bytearray(), 0xffffffff, [
+            bytearray(65),
+            bytearray(65),
             mast.r.script,
             bytearray([self.prefix]) + self.pubkey.sec()[1:] + mast.l.hash,
         ])
@@ -115,8 +105,8 @@ def txin(self, op: btc.core.OutPoint):
 # Spending by script path: always success(op_1).
 mate.transfer(user_p2tr.script, 1 * btc.denomination.bitcoin)
 assert user_p2tr.balance() == btc.denomination.bitcoin
-user_p2as = btc.wallet.Wallet(Tp2trp2as(user_p2tr.signer.pubkey))
-print('main: spending by script path p2as')
-user_p2as.transfer_all(mate.script)
+user_p2ms = btc.wallet.Wallet(Tp2trp2ms(user_p2tr.signer.pubkey))
+print('main: spending by script path p2ms')
+user_p2ms.transfer_all(mate.script)
 assert user_p2tr.balance() == 0
-print('main: spending by script path p2as done')
+print('main: spending by script path p2ms done')