1+ import os
2+ import random
3+ from importlib .resources import files
14from pathlib import Path
25
36import click
7+ import inquirer
8+ import yaml
9+
10+ from .constants import DEFAULT_TAG , SUPPORTED_TAGS
411
512
613@click .group (name = "graph" , hidden = True )
@@ -17,3 +24,188 @@ def import_json(infile: Path, outfile: Path, cb: str, ln_image: str):
1724 Returns XML file as string with or without --outfile option.
1825 """
1926 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