1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
import re
from enum import IntEnum
from typing import List, Union
from hashlib import sha256
from .common import serialize_str, AddressType, write_varint
from .merkle import MerkleTree, element_hash
class WalletType(IntEnum):
WALLET_POLICY_V1 = 1
WALLET_POLICY_V2 = 2
# should not be instantiated directly
class WalletPolicyBase:
def __init__(self, name: str, version: WalletType) -> None:
self.name = name
self.version = version
if (version != WalletType.WALLET_POLICY_V1 and version != WalletType.WALLET_POLICY_V2):
raise ValueError("Invalid wallet policy version")
def serialize(self) -> bytes:
return b"".join([
self.version.value.to_bytes(1, byteorder="big"),
serialize_str(self.name)
])
@property
def id(self) -> bytes:
return sha256(self.serialize()).digest()
class WalletPolicy(WalletPolicyBase):
"""
Represents a wallet stored with a wallet policy.
For version V2, the wallet is serialized as follows:
- 1 byte : wallet version
- 1 byte : length of the wallet name (max 64)
- (var) : wallet name (ASCII string)
- (varint) : length of the descriptor template
- 32-bytes : sha256 hash of the descriptor template
- (varint) : number of keys (not larger than 252)
- 32-bytes : root of the Merkle tree of all the keys information.
The specific format of the keys is deferred to subclasses.
"""
def __init__(self, name: str, descriptor_template: str, keys_info: List[str], version: WalletType = WalletType.WALLET_POLICY_V2):
super().__init__(name, version)
self.descriptor_template = descriptor_template
self.keys_info = keys_info
@property
def n_keys(self) -> int:
return len(self.keys_info)
def serialize(self) -> bytes:
keys_info_hashes = map(
lambda k: element_hash(k.encode()), self.keys_info)
descriptor_template_sha256 = sha256(
self.descriptor_template.encode()).digest()
return b"".join([
super().serialize(),
write_varint(len(self.descriptor_template.encode())),
self.descriptor_template.encode(
) if self.version == WalletType.WALLET_POLICY_V1 else descriptor_template_sha256,
write_varint(len(self.keys_info)),
MerkleTree(keys_info_hashes).root
])
def get_descriptor(self, change: Union[bool, None]) -> str:
"""
Generates a descriptor string based on the wallet's descriptor template and keys.
Args:
change (bool | None): Indicates whether the descriptor is for a change address.
- If None, returns the BIP-389 multipath address for both the receive and change address.
- If True, the descriptor is for a change address.
- If False, the descriptor is for a non-change address.
Returns:
str: The generated descriptor.
"""
desc = self.descriptor_template
for i in reversed(range(self.n_keys)):
key = self.keys_info[i]
desc = desc.replace(f"@{i}", key)
# in V1, /** is part of the key; in V2, it's part of the policy map. This handles either
if change is not None:
desc = desc.replace("/**", f"/{1 if change else 0}/*")
else:
desc = desc.replace("/**", f"/<0;1>/*")
if self.version == WalletType.WALLET_POLICY_V2:
# V2, the /<M;N> syntax is supported. Replace with M if not change, or with N if change
if change is not None:
desc = re.sub(r"/<(\d+);(\d+)>", "/\\2" if change else "/\\1", desc)
return desc
class MultisigWallet(WalletPolicy):
def __init__(self, name: str, address_type: AddressType, threshold: int, keys_info: List[str], sorted: bool = True, version: WalletType = WalletType.WALLET_POLICY_V2) -> None:
n_keys = len(keys_info)
if not (1 <= threshold <= n_keys <= 16):
raise ValueError("Invalid threshold or number of keys")
multisig_op = "sortedmulti" if sorted else "multi"
if (address_type == AddressType.LEGACY):
policy_prefix = f"sh({multisig_op}("
policy_suffix = f"))"
elif address_type == AddressType.WIT:
policy_prefix = f"wsh({multisig_op}("
policy_suffix = f"))"
elif address_type == AddressType.SH_WIT:
policy_prefix = f"sh(wsh({multisig_op}("
policy_suffix = f")))"
else:
raise ValueError(f"Unexpected address type: {address_type}")
key_placeholder_suffix = "/**" if version == WalletType.WALLET_POLICY_V2 else ""
descriptor_template = "".join([
policy_prefix,
str(threshold) + ",",
",".join("@" + str(l) + key_placeholder_suffix for l in range(n_keys)),
policy_suffix
])
super().__init__(name, descriptor_template, keys_info, version)
self.threshold = threshold
|
