1
+ import os
2
+ import random
3
+ from importlib .resources import files
1
4
from pathlib import Path
2
5
3
6
import click
7
+ import inquirer
8
+ import yaml
9
+
10
+ from .constants import DEFAULT_TAG , SUPPORTED_TAGS
4
11
5
12
6
13
@click .group (name = "graph" , hidden = True )
@@ -17,3 +24,188 @@ def import_json(infile: Path, outfile: Path, cb: str, ln_image: str):
17
24
Returns XML file as string with or without --outfile option.
18
25
"""
19
26
raise Exception ("Not Implemented" )
27
+
28
+
29
+ def custom_graph (
30
+ num_nodes : int ,
31
+ num_connections : int ,
32
+ version : str ,
33
+ datadir : Path ,
34
+ fork_observer : bool ,
35
+ fork_obs_query_interval : int ,
36
+ ):
37
+ datadir .mkdir (parents = False , exist_ok = False )
38
+ # Generate network.yaml
39
+ nodes = []
40
+ connections = set ()
41
+
42
+ for i in range (num_nodes ):
43
+ node = {"name" : f"tank-{ i :04d} " , "connect" : [], "image" : {"tag" : version }}
44
+
45
+ # Add round-robin connection
46
+ next_node = (i + 1 ) % num_nodes
47
+ node ["connect" ].append (f"tank-{ next_node :04d} " )
48
+ connections .add ((i , next_node ))
49
+
50
+ # Add random connections
51
+ available_nodes = list (range (num_nodes ))
52
+ available_nodes .remove (i )
53
+ if next_node in available_nodes :
54
+ available_nodes .remove (next_node )
55
+
56
+ for _ in range (min (num_connections - 1 , len (available_nodes ))):
57
+ random_node = random .choice (available_nodes )
58
+ # Avoid circular loops of A -> B -> A
59
+ if (random_node , i ) not in connections :
60
+ node ["connect" ].append (f"tank-{ random_node :04d} " )
61
+ connections .add ((i , random_node ))
62
+ available_nodes .remove (random_node )
63
+
64
+ nodes .append (node )
65
+
66
+ network_yaml_data = {"nodes" : nodes }
67
+ network_yaml_data ["fork_observer" ] = {
68
+ "enabled" : fork_observer ,
69
+ "configQueryInterval" : fork_obs_query_interval ,
70
+ }
71
+
72
+ with open (os .path .join (datadir , "network.yaml" ), "w" ) as f :
73
+ yaml .dump (network_yaml_data , f , default_flow_style = False )
74
+
75
+ # Generate node-defaults.yaml
76
+ default_yaml_path = files ("resources.networks" ).joinpath ("node-defaults.yaml" )
77
+ with open (str (default_yaml_path )) as f :
78
+ defaults_yaml_content = f .read ()
79
+
80
+ with open (os .path .join (datadir , "node-defaults.yaml" ), "w" ) as f :
81
+ f .write (defaults_yaml_content )
82
+
83
+ click .echo (
84
+ f"Project '{ datadir } ' has been created with 'network.yaml' and 'node-defaults.yaml'."
85
+ )
86
+
87
+
88
+ def inquirer_create_network (project_path : Path ):
89
+ # Custom network configuration
90
+ questions = [
91
+ inquirer .Text (
92
+ "network_name" ,
93
+ message = click .style ("Enter your network name" , fg = "blue" , bold = True ),
94
+ validate = lambda _ , x : len (x ) > 0 ,
95
+ ),
96
+ inquirer .List (
97
+ "nodes" ,
98
+ message = click .style ("How many nodes would you like?" , fg = "blue" , bold = True ),
99
+ choices = ["8" , "12" , "20" , "50" , "other" ],
100
+ default = "12" ,
101
+ ),
102
+ inquirer .List (
103
+ "connections" ,
104
+ message = click .style (
105
+ "How many connections would you like each node to have?" ,
106
+ fg = "blue" ,
107
+ bold = True ,
108
+ ),
109
+ choices = ["0" , "1" , "2" , "8" , "12" , "other" ],
110
+ default = "8" ,
111
+ ),
112
+ inquirer .List (
113
+ "version" ,
114
+ message = click .style (
115
+ "Which version would you like nodes to run by default?" , fg = "blue" , bold = True
116
+ ),
117
+ choices = SUPPORTED_TAGS ,
118
+ default = DEFAULT_TAG ,
119
+ ),
120
+ ]
121
+
122
+ net_answers = inquirer .prompt (questions )
123
+ if net_answers is None :
124
+ click .secho ("Setup cancelled by user." , fg = "yellow" )
125
+ return False
126
+
127
+ if net_answers ["nodes" ] == "other" :
128
+ custom_nodes = inquirer .prompt (
129
+ [
130
+ inquirer .Text (
131
+ "nodes" ,
132
+ message = click .style ("Enter the number of nodes" , fg = "blue" , bold = True ),
133
+ validate = lambda _ , x : int (x ) > 0 ,
134
+ )
135
+ ]
136
+ )
137
+ if custom_nodes is None :
138
+ click .secho ("Setup cancelled by user." , fg = "yellow" )
139
+ return False
140
+ net_answers ["nodes" ] = custom_nodes ["nodes" ]
141
+
142
+ if net_answers ["connections" ] == "other" :
143
+ custom_connections = inquirer .prompt (
144
+ [
145
+ inquirer .Text (
146
+ "connections" ,
147
+ message = click .style ("Enter the number of connections" , fg = "blue" , bold = True ),
148
+ validate = lambda _ , x : int (x ) >= 0 ,
149
+ )
150
+ ]
151
+ )
152
+ if custom_connections is None :
153
+ click .secho ("Setup cancelled by user." , fg = "yellow" )
154
+ return False
155
+ net_answers ["connections" ] = custom_connections ["connections" ]
156
+ fork_observer = click .prompt (
157
+ click .style (
158
+ "\n Would you like to enable fork-observer on the network?" , fg = "blue" , bold = True
159
+ ),
160
+ type = bool ,
161
+ default = True ,
162
+ )
163
+ fork_observer_query_interval = 20
164
+ if fork_observer :
165
+ fork_observer_query_interval = click .prompt (
166
+ click .style (
167
+ "\n How often would you like fork-observer to query node status (seconds)?" ,
168
+ fg = "blue" ,
169
+ bold = True ,
170
+ ),
171
+ type = int ,
172
+ default = 20 ,
173
+ )
174
+ custom_network_path = project_path / "networks" / net_answers ["network_name" ]
175
+ click .secho ("\n Generating custom network..." , fg = "yellow" , bold = True )
176
+ custom_graph (
177
+ int (net_answers ["nodes" ]),
178
+ int (net_answers ["connections" ]),
179
+ net_answers ["version" ],
180
+ custom_network_path ,
181
+ fork_observer ,
182
+ fork_observer_query_interval ,
183
+ )
184
+ return custom_network_path
185
+
186
+
187
+ @click .command ()
188
+ def create ():
189
+ """Create a new warnet network"""
190
+ try :
191
+ project_path = Path (os .getcwd ())
192
+ # Check if the project has a networks directory
193
+ if not (project_path / "networks" ).exists ():
194
+ click .secho (
195
+ "The current directory does not have a 'networks' directory. Please run 'warnet init' or 'warnet create' first." ,
196
+ fg = "red" ,
197
+ bold = True ,
198
+ )
199
+ return False
200
+ custom_network_path = inquirer_create_network (project_path )
201
+ click .secho ("\n New network created successfully!" , fg = "green" , bold = True )
202
+ click .echo ("\n Run the following command to deploy this network:" )
203
+ click .echo (f"warnet deploy { custom_network_path } " )
204
+ except Exception as e :
205
+ click .echo (f"{ e } \n \n " )
206
+ click .secho (f"An error occurred while creating a new network:\n \n { e } \n \n " , fg = "red" )
207
+ click .secho (
208
+ "Please report the above context to https://github.com/bitcoin-dev-project/warnet/issues" ,
209
+ fg = "yellow" ,
210
+ )
211
+ return False
0 commit comments