Inside Bitcoin

1#!/usr/bin/env python3
2
3import sys
4import json
5import textwrap
6
7max_pubkeys = 0
8
9if len(sys.argv) < 2:
10    print(
11        "This script converts BIP MuSig2 test vectors in a given directory to a C file that can be used in the test framework."
12    )
13    print("Usage: %s <dir>" % sys.argv[0])
14    sys.exit(1)
15
16
17def hexstr_to_intarray(str):
18    return ", ".join([f"0x{b:02X}" for b in bytes.fromhex(str)])
19
20
21def create_init(name):
22    return """
23static const struct musig_%s_vector musig_%s_vector = {
24""" % (
25        name,
26        name,
27    )
28
29
30def init_array(key):
31    return textwrap.indent("{ %s },\n" % hexstr_to_intarray(data[key]), 4 * " ")
32
33
34def init_arrays(key):
35    s = textwrap.indent("{\n", 4 * " ")
36    s += textwrap.indent(
37        ",\n".join(["{ %s }" % hexstr_to_intarray(x) for x in data[key]]), 8 * " "
38    )
39    s += textwrap.indent("\n},\n", 4 * " ")
40    return s
41
42
43def init_indices(array):
44    return " %d, { %s }" % (
45        len(array),
46        ", ".join(map(str, array) if len(array) > 0 else "0"),
47    )
48
49
50def init_is_xonly(case):
51    if len(case["tweak_indices"]) > 0:
52        return ", ".join(map(lambda x: "1" if x else "0", case["is_xonly"]))
53    return "0"
54
55
56def init_optional_expected(case):
57    return hexstr_to_intarray(case["expected"]) if "expected" in case else 0
58
59
60def init_cases(cases, f):
61    s = textwrap.indent("{\n", 4 * " ")
62    for (i, case) in enumerate(cases):
63        s += textwrap.indent("%s\n" % f(case), 8 * " ")
64    s += textwrap.indent("},\n", 4 * " ")
65    return s
66
67
68def finish_init():
69    return "};\n"
70
71
72s = (
73    """/**
74 * Automatically generated by %s.
75 *
76 * The test vectors for the KeySort function are included in this file. They can
77 * be found in src/modules/extrakeys/tests_impl.h. */
78"""
79    % sys.argv[0]
80)
81
82
83s += """
84enum MUSIG_ERROR {
85    MUSIG_PUBKEY,
86    MUSIG_TWEAK,
87    MUSIG_PUBNONCE,
88    MUSIG_AGGNONCE,
89    MUSIG_SECNONCE,
90    MUSIG_SIG,
91    MUSIG_SIG_VERIFY,
92    MUSIG_OTHER
93};
94"""
95
96# key agg vectors
97with open(sys.argv[1] + "/key_agg_vectors.json", "r") as f:
98    data = json.load(f)
99
100    max_key_indices = max(
101        len(test_case["key_indices"]) for test_case in data["valid_test_cases"]
102    )
103    max_tweak_indices = max(
104        len(test_case["tweak_indices"]) for test_case in data["error_test_cases"]
105    )
106    num_pubkeys = len(data["pubkeys"])
107    max_pubkeys = max(num_pubkeys, max_pubkeys)
108    num_tweaks = len(data["tweaks"])
109    num_valid_cases = len(data["valid_test_cases"])
110    num_error_cases = len(data["error_test_cases"])
111
112    # Add structures for valid and error cases
113    s += (
114        """
115struct musig_key_agg_valid_test_case {
116    size_t key_indices_len;
117    size_t key_indices[%d];
118    unsigned char expected[32];
119};
120"""
121        % max_key_indices
122    )
123    s += """
124struct musig_key_agg_error_test_case {
125    size_t key_indices_len;
126    size_t key_indices[%d];
127    size_t tweak_indices_len;
128    size_t tweak_indices[%d];
129    int is_xonly[%d];
130    enum MUSIG_ERROR error;
131};
132""" % (
133        max_key_indices,
134        max_tweak_indices,
135        max_tweak_indices,
136    )
137
138    # Add structure for entire vector
139    s += """
140struct musig_key_agg_vector {
141    unsigned char pubkeys[%d][33];
142    unsigned char tweaks[%d][32];
143    struct musig_key_agg_valid_test_case valid_case[%d];
144    struct musig_key_agg_error_test_case error_case[%d];
145};
146""" % (
147        num_pubkeys,
148        num_tweaks,
149        num_valid_cases,
150        num_error_cases,
151    )
152
153    s += create_init("key_agg")
154    # Add pubkeys and tweaks to the vector
155    s += init_arrays("pubkeys")
156    s += init_arrays("tweaks")
157
158    # Add valid cases to the vector
159    s += init_cases(
160        data["valid_test_cases"],
161        lambda case: "{ %s, { %s }},"
162        % (init_indices(case["key_indices"]), hexstr_to_intarray(case["expected"])),
163    )
164
165    def comment_to_error(case):
166        comment = case["comment"]
167        if "public key" in comment.lower():
168            return "MUSIG_PUBKEY"
169        elif "tweak" in comment.lower():
170            return "MUSIG_TWEAK"
171        else:
172            sys.exit("Unknown error")
173
174    # Add error cases to the vector
175    s += init_cases(
176        data["error_test_cases"],
177        lambda case: "{ %s, %s, { %s }, %s },"
178        % (
179            init_indices(case["key_indices"]),
180            init_indices(case["tweak_indices"]),
181            init_is_xonly(case),
182            comment_to_error(case),
183        ),
184    )
185
186    s += finish_init()
187
188# nonce gen vectors
189with open(sys.argv[1] + "/nonce_gen_vectors.json", "r") as f:
190    data = json.load(f)
191
192    # The MuSig2 implementation only allows messages of length 32
193    data["test_cases"] = list(
194        filter(lambda c: c["msg"] is None or len(c["msg"]) == 64, data["test_cases"])
195    )
196
197    num_tests = len(data["test_cases"])
198
199    s += """
200struct musig_nonce_gen_test_case {
201    unsigned char rand_[32];
202    int has_sk;
203    unsigned char sk[32];
204    unsigned char pk[33];
205    int has_aggpk;
206    unsigned char aggpk[32];
207    int has_msg;
208    unsigned char msg[32];
209    int has_extra_in;
210    unsigned char extra_in[32];
211    unsigned char expected_secnonce[97];
212    unsigned char expected_pubnonce[66];
213};
214"""
215
216    s += (
217        """
218struct musig_nonce_gen_vector {
219    struct musig_nonce_gen_test_case test_case[%d];
220};
221"""
222        % num_tests
223    )
224
225    s += create_init("nonce_gen")
226
227    def init_array_maybe(array):
228        return "%d , { %s }" % (
229            0 if array is None else 1,
230            hexstr_to_intarray(array) if array is not None else 0,
231        )
232
233    s += init_cases(
234        data["test_cases"],
235        lambda case: "{ { %s },  %s, { %s }, %s, %s, %s, { %s }, { %s } },"
236        % (
237            hexstr_to_intarray(case["rand_"]),
238            init_array_maybe(case["sk"]),
239            hexstr_to_intarray(case["pk"]),
240            init_array_maybe(case["aggpk"]),
241            init_array_maybe(case["msg"]),
242            init_array_maybe(case["extra_in"]),
243            hexstr_to_intarray(case["expected_secnonce"]),
244            hexstr_to_intarray(case["expected_pubnonce"]),
245        ),
246    )
247
248    s += finish_init()
249
250# nonce agg vectors