update quickstart

Author: willcl-ark

Diff for: src/warnet/main.py

@@ -1,4 +1,5 @@
 import os
+import random
 import subprocess
 import sys
 from importlib.resources import files
@@ -16,6 +17,7 @@
 from .image import image
 from .network import copy_network_defaults, copy_scenario_defaults
 from .status import status as status_command
+from .util import SUPPORTED_TAGS
 
 QUICK_START_PATH = files("resources.scripts").joinpath("quick_start.sh")
 
@@ -57,19 +59,36 @@ def quickstart():
             return False
 
         create_project = click.confirm("Do you want to create a new project?", default=True)
+        if not create_project:
+            click.echo("Setup completed successfully!")
+            return True
+
+        default_path = os.path.abspath(os.getcwd())
+        project_path = click.prompt(
+            "Enter the project directory path",
+            default=default_path,
+            type=click.Path(file_okay=False, dir_okay=True, resolve_path=True),
+        )
 
-        if create_project:
-            default_path = os.path.abspath(os.getcwd())
-            project_path = click.prompt(
-                "Enter the project directory path",
-                default=default_path,
-                type=click.Path(file_okay=False, dir_okay=True, resolve_path=True),
-            )
-
+        custom_network = click.confirm("Do you want to create a custom network?", default=True)
+        if not custom_network:
             create_warnet_project(Path(project_path))
+            click.echo("Setup completed successfully!")
+            return True
+
+        network_name = click.prompt(
+            "Enter the network name",
+            type=str,
+        )
+        nodes = click.prompt("How many nodes would you like?", type=int)
+        connections = click.prompt("How many connects would you like each node to have?", type=int)
+        version = click.prompt(
+            "Which version would you like nodes to have by default?",
+            type=click.Choice(SUPPORTED_TAGS, case_sensitive=False),
+        )
 
-        click.echo("Setup completed successfully!")
-        return True
+        create_warnet_project(Path(project_path))
+        custom_graph(nodes, connections, version, Path(project_path) / "networks" / network_name)
 
     except Exception as e:
         print(f"An error occurred while running the quick start script:\n\n{e}\n\n")
@@ -91,6 +110,8 @@ def create_warnet_project(directory: Path, check_empty: bool = False):
         richprint(f"[green]Created warnet project structure in {directory}[/green]")
     except Exception as e:
         richprint(f"[red]Error creating project: {e}[/red]")
+        raise e
+
 
 @cli.command()
 @click.argument(
@@ -103,6 +124,7 @@ def create(directory: Path):
         return
     create_warnet_project(directory)
 
+
 @cli.command()
 def init():
     """Initialize a warnet project in the current directory"""
@@ -149,3 +171,71 @@ def auth(kube_config: str) -> None:
 
 if __name__ == "__main__":
     cli()
+
+
+def custom_graph(num_nodes: int, num_connections: int, version: str, datadir: Path):
+    datadir.mkdir(parents=False, exist_ok=False)
+    # Generate network.yaml
+    nodes = []
+
+    for i in range(num_nodes):
+        node = {"name": f"tank-{i:04d}", "connect": []}
+
+        # Add round-robin connection
+        next_node = (i + 1) % num_nodes
+        node["connect"].append(f"tank-{next_node:04d}")
+
+        # Add random connections
+        available_nodes = list(range(num_nodes))
+        available_nodes.remove(i)
+        if next_node in available_nodes:
+            available_nodes.remove(next_node)
+
+        for _ in range(min(num_connections - 1, len(available_nodes))):
+            random_node = random.choice(available_nodes)
+            node["connect"].append(f"tank-{random_node:04d}")
+            available_nodes.remove(random_node)
+
+        nodes.append(node)
+
+    # Add image tag to the first node
+    nodes[0]["image"] = {"tag": "v0.20.0"}
+
+    network_yaml_data = {"nodes": nodes}
+
+    with open(os.path.join(datadir, "network.yaml"), "w") as f:
+        yaml.dump(network_yaml_data, f, default_flow_style=False)
+
+    # Generate defaults.yaml
+    defaults_yaml_content = """
+chain: regtest
+
+collectLogs: true
+metricsExport: true
+
+resources: {}
+  # We usually recommend not to specify default resources and to leave this as a conscious
+  # choice for the user. This also increases chances charts run on environments with little
+  # resources, such as Minikube. If you do want to specify resources, uncomment the following
+  # lines, adjust them as necessary, and remove the curly braces after 'resources:'.
+  # limits:
+  #   cpu: 100m
+  #   memory: 128Mi
+  # requests:
+  #   cpu: 100m
+  #   memory: 128Mi
+
+image:
+  repository: bitcoindevproject/bitcoin
+  pullPolicy: IfNotPresent
+  # Overrides the image tag whose default is the chart appVersion.
+  tag: "27.0"
+
+config: |
+  dns=1
+"""
+
+    with open(os.path.join(datadir, "defaults.yaml"), "w") as f:
+        f.write(defaults_yaml_content.strip())
+
+    click.echo(f"Project '{datadir}' has been created with 'network.yaml' and 'defaults.yaml'.")