graph: add generator

delete random AS graphs of size `n` using either default tag (latest) or
random tags from SUPPORTED_TAGS

Author: willcl-ark

src/warnet/cli/graph.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+from typing import List
+
+import click
+from rich import print
+
+from warnet.cli.rpc import rpc_call
+from warnet.utils import DEFAULT_TAG
+
+
+@click.group(name="graph")
+def graph():
+    """Graph commands"""
+
+
+@graph.command()
+@click.argument("params", nargs=-1, type=str)
+@click.option("--outfile", type=Path)
+@click.option("--version", type=str, default=DEFAULT_TAG)
+@click.option("--bitcoin_conf", type=Path)
+@click.option("--random", is_flag=True)
+def create(params: List[str], outfile: Path, version: str, bitcoin_conf: Path, random: bool = False):
+    """
+    Create a graph file of type random AS graph with [params]
+    """
+    try:
+        result = rpc_call(
+            "graph_generate",
+            {
+                "params": params,
+                "outfile": str(outfile) if outfile else "",
+                "version": version,
+                "bitcoin_conf": str(bitcoin_conf),
+                "random": random
+            },
+        )
+        print(result)
+    except Exception as e:
+        print(f"Error generating graph: {e}")

src/warnet/cli/main.py

@@ -2,6 +2,7 @@
 from rich import print as richprint
 
 from warnet.cli.debug import debug
+from warnet.cli.graph import graph
 from warnet.cli.network import network
 from warnet.cli.rpc import rpc_call
 from warnet.cli.scenarios import scenarios
@@ -12,6 +13,7 @@ def cli():
 
 
 cli.add_command(debug)
+cli.add_command(graph)
 cli.add_command(scenarios)
 cli.add_command(network)
 

src/warnet/server.py

@@ -6,15 +6,20 @@
 import sys
 import threading
 from datetime import datetime
+from io import BytesIO
 from logging.handlers import RotatingFileHandler
 from logging import StreamHandler
-from typing import List, Dict
+from pathlib import Path
+from typing import List, Dict, Optional
 from flask import Flask, request
 from flask_jsonrpc.app import JSONRPC
 
+import networkx as nx
+
 import scenarios
 from warnet.warnet import Warnet
 from warnet.utils import (
+    create_graph_with_probability,
     gen_config_dir,
 )
 
@@ -81,6 +86,8 @@ def setup_rpc(self):
         self.jsonrpc.register(self.network_down)
         self.jsonrpc.register(self.network_info)
         self.jsonrpc.register(self.network_status)
+        # Graph
+        self.jsonrpc.register(self.graph_generate)
         # Debug
         self.jsonrpc.register(self.generate_deployment)
         # Server
@@ -271,6 +278,20 @@ def thread_start(wn):
         t.start()
         return f"Starting warnet network named '{network}' with the following parameters:\n{wn}"
 
+    def graph_generate(self, params: List[str], outfile: str, version: str, bitcoin_conf: Optional[str] = None, random: bool = False) -> str:
+        graph_func = nx.generators.random_internet_as_graph
+
+        graph = create_graph_with_probability(graph_func, params, version, bitcoin_conf, random)
+
+        if outfile:
+            file_path = Path(outfile)
+            nx.write_graphml(graph, file_path)
+            return f"Generated graph written to file: {outfile}"
+        bio = BytesIO()
+        nx.write_graphml(graph, bio)
+        xml_data = bio.getvalue()
+        return f"Generated graph:\n\n{xml_data.decode('utf-8')}"
+
     def network_down(self, network: str = "warnet") -> str:
         """
         Stop all containers in <network>.

src/warnet/utils.py

@@ -11,19 +11,25 @@
 import time
 from io import BytesIO
 from pathlib import Path
+from typing import Dict, List, Optional
+
+import networkx as nx
 
 from test_framework.p2p import MESSAGEMAP
 from test_framework.messages import ser_uint256
 
 
 logger = logging.getLogger("utils")
 
+
 SUPPORTED_TAGS = [
     "25.1",
     "24.2",
     "23.2",
     "22.2",
 ]
+DEFAULT_TAG = SUPPORTED_TAGS[0]
+WEIGHTED_TAGS = [tag for index, tag in enumerate(reversed(SUPPORTED_TAGS)) for _ in range(index + 1)]
 
 
 def exponential_backoff(max_retries=5, base_delay=1, max_delay=32):
@@ -187,7 +193,7 @@ def parse_bitcoin_conf(file_content):
     return result
 
 
-def dump_bitcoin_conf(conf_dict):
+def dump_bitcoin_conf(conf_dict, for_graph=False):
     """
     Converts a dictionary representation of bitcoin.conf content back to INI-style string.
 
@@ -213,6 +219,9 @@ def dump_bitcoin_conf(conf_dict):
         for sub_key, sub_value in values:
             result.append(f"{sub_key}={sub_value}")
 
+    if for_graph:
+        return ",".join(result)
+
     # Terminate file with newline
     return "\n".join(result) + "\n"
 
@@ -408,3 +417,94 @@ def default_bitcoin_conf_args() -> str:
             conf_args.append(f"-{key}={value}")
 
     return " ".join(conf_args)
+
+
+def create_graph_with_probability(graph_func, params: List, version: str, bitcoin_conf: Optional[str], random_version: bool):
+    kwargs = {}
+    for param in params:
+        try:
+            key, value = param.split("=")
+            kwargs[key] = value
+        except ValueError:
+            msg = f"Invalid parameter format: {param}"
+            logger.error(msg)
+            return msg
+
+    # Attempt to convert numerical values from string to their respective numerical types
+    for key in kwargs:
+        try:
+            kwargs[key] = int(kwargs[key])
+        except ValueError:
+            try:
+                kwargs[key] = float(kwargs[key])
+            except ValueError:
+                pass
+
+    logger.debug(f"Parsed params: {kwargs}")
+
+    try:
+        graph = graph_func(**kwargs)
+    except TypeError as e:
+        msg = f"Failed to create graph: {e}"
+        logger.error(msg)
+        return msg
+
+    # calculate degree
+    degree_dict = dict(graph.degree(graph.nodes()))
+    nx.set_node_attributes(graph, degree_dict, 'degree')
+
+    # add a default layout
+    pos = nx.spring_layout(graph)
+    for node in graph.nodes():
+        graph.nodes[node]['x'] = float(pos[node][0])
+        graph.nodes[node]['y'] = float(pos[node][1])
+
+    # parse and process conf file
+    conf_contents = ""
+    if bitcoin_conf is not None:
+        conf = Path(bitcoin_conf)
+        if conf.is_file():
+            with open(conf, 'r') as f:
+                # parse INI style conf then dump using for_graph
+                conf_dict = parse_bitcoin_conf(f.read())
+                conf_contents = dump_bitcoin_conf(conf_dict, for_graph=True)
+
+    # populate our custom fields
+    for node in graph.nodes():
+        if random_version:
+            graph.nodes[node]['version'] = random.choice(WEIGHTED_TAGS)
+        else:
+            graph.nodes[node]['version'] = version
+        graph.nodes[node]['bitcoin_config'] = conf_contents
+        graph.nodes[node]['tc_netem'] = ""
+
+    # remove type and customer fields from edges as we don't need 'em!
+    for edge in graph.edges():
+        del graph.edges[edge]["customer"]
+        del graph.edges[edge]["type"]
+
+    convert_unsupported_attributes(graph)
+    return graph
+
+
+def convert_unsupported_attributes(graph):
+    # Sometimes networkx complains about invalid types when writing the graph
+    # (it just generated itself!). Try to convert them here just in case.
+    for _, node_data in graph.nodes(data=True):
+        for key, value in node_data.items():
+            if isinstance(value, set):
+                node_data[key] = list(value)
+            elif isinstance(value, (int, float, str)):
+                continue
+            else:
+                node_data[key] = str(value)
+
+    for _, _, edge_data in graph.edges(data=True):
+        for key, value in edge_data.items():
+            if isinstance(value, set):
+                edge_data[key] = list(value)
+            elif isinstance(value, (int, float, str)):
+                continue
+            else:
+                edge_data[key] = str(value)
+