1  #! /usr/bin/env python3
  2  
  3  import matplotlib.pyplot as plt
  4  import networkx as nx
  5  import random, math
  6  import json
  7  import sys, os
  8  import string
  9  import typer
 10  from enum import Enum
 11  
 12  
 13  # Enums & Consts
 14  
 15  # To add a new node type, add appropriate entries to the nodeType and nodeTypeSwitch
 16  class nodeType(Enum):
 17      DESKTOP = "desktop" # waku desktop config
 18      MOBILE  = "mobile"  # waku mobile config
 19  
 20  nodeTypeSwitch = {
 21          nodeType.DESKTOP : "rpc-admin = true\nkeep-alive = true\n",
 22          nodeType.MOBILE  : "rpc-admin = true\nkeep-alive = true\n"
 23      }
 24  
 25  
 26  # To add a new network type, add appropriate entries to the networkType and networkTypeSwitch
 27  # the networkTypeSwitch is placed before generate_network(): fwd declaration mismatch with typer/python :/
 28  class networkType(Enum):
 29      CONFIGMODEL         = "configmodel"
 30      SCALEFREE           = "scalefree"           # power law
 31      NEWMANWATTSSTROGATZ = "newmanwattsstrogatz" # mesh, smallworld
 32      BARBELL             = "barbell"             # partition
 33      BALANCEDTREE        = "balancedtree"        # committees?
 34      STAR                = "star"                # spof
 35  
 36  
 37  NW_DATA_FNAME = "network_data.json"
 38  NODE_PREFIX = "waku"
 39  SUBNET_PREFIX = "subnetwork"
 40  
 41  
 42  ### I/O related fns ##############################################################
 43  
 44  # Dump to a json file
 45  def write_json(dirname, json_dump):
 46      fname = os.path.join(dirname, NW_DATA_FNAME)
 47      with open(fname, "w") as f:
 48          json.dump(json_dump, f, indent=2)
 49  
 50  
 51  def write_toml(dirname, node_name, toml):
 52      fname = os.path.join(dirname, f"{node_name}.toml")
 53      with open(fname, "w") as f:
 54          f.write(toml)
 55  
 56  
 57  # Draw the network and output the image to a file; does not account for subnets yet
 58  def draw(dirname, H):
 59      nx.draw(H, pos=nx.kamada_kawai_layout(H), with_labels=True)
 60      fname = os.path.join(dirname, NW_DATA_FNAME)
 61      plt.savefig(f"{os.path.splitext(fname)[0]}.png", format="png")
 62      plt.show()
 63  
 64  
 65  # Has trouble with non-integer/non-hashable keys 
 66  def read_json(fname):
 67      with open(fname) as f:
 68          jdata = json.load(f)
 69      return nx.node_link_graph(jdata)
 70  
 71  
 72  # check if the required dir can be created
 73  def exists_or_nonempty(dirname):
 74      if not os.path.exists(dirname):
 75          return False
 76      elif not os.path.isfile(dirname) and os.listdir(dirname):
 77          print(f"{dirname}: exists and not empty")
 78          return True
 79      elif os.path.isfile(dirname):
 80          print(f"{dirname}: exists but not a directory")
 81          return True
 82      else:
 83          return False
 84  
 85  
 86  ### topics related fns #############################################################
 87  
 88  # Generate a random string of upper case chars
 89  def generate_random_string(n):
 90      return "".join(random.choice(string.ascii_uppercase) for _ in range(n))
 91  
 92  
 93  # Generate the topics - topic followed by random UC chars - Eg, topic_XY"
 94  def generate_topics(num_topics):
 95      topic_len = int(math.log(num_topics)/math.log(26)) + 1  # base is 26 - upper case letters
 96      topics = {i: f"topic_{generate_random_string(topic_len)}" for i in range(num_topics)}
 97      return topics
 98  
 99  
100  # Get a random sub-list of topics
101  def get_random_sublist(topics):
102      n = len(topics)
103      lo = random.randint(0, n - 1)
104      hi = random.randint(lo + 1, n)
105      sublist = []
106      for i in range(lo, hi):
107          sublist.append(topics[i])
108      return sublist
109  
110  
111  ### network processing related fns #################################################
112  
113  # Network Types
114  def generate_config_model(n):
115      #degrees = nx.random_powerlaw_tree_sequence(n, tries=10000)
116      degrees = [random.randint(1, n) for i in range(n)]
117      if (sum(degrees)) % 2 != 0:         # adjust the degree to be even
118          degrees[-1] += 1
119      return nx.configuration_model(degrees) # generate the graph
120  
121  
122  def generate_scalefree_graph(n):
123      return  nx.scale_free_graph(n)
124  
125  
126  # n must be larger than k=D=3
127  def generate_newmanwattsstrogatz_graph(n):
128      return nx.newman_watts_strogatz_graph(n, 3, 0.5)
129  
130  
131  def generate_barbell_graph(n):
132      return nx.barbell_graph(int(n/2), 1)
133  
134  
135  def generate_balanced_tree(n, fanout=3):
136      height = int(math.log(n)/math.log(fanout))
137      return nx.balanced_tree(fanout, height)
138  
139  
140  def generate_star_graph(n):
141      return nx.star_graph(n)
142  
143  
144  networkTypeSwitch = {
145          networkType.CONFIGMODEL : generate_config_model,
146          networkType.SCALEFREE   : generate_scalefree_graph,
147          networkType.NEWMANWATTSSTROGATZ : generate_newmanwattsstrogatz_graph,
148          networkType.BARBELL     : generate_barbell_graph,
149          networkType.BALANCEDTREE: generate_balanced_tree,
150          networkType.STAR        : generate_star_graph
151      }
152  
153  
154  # Generate the network from nw type
155  def generate_network(n, network_type):
156      return postprocess_network(networkTypeSwitch.get(network_type)(n))
157  
158  
159  # Label the generated network with prefix
160  def postprocess_network(G):
161      G = nx.Graph(G)                             # prune out parallel/multi edges
162      G.remove_edges_from(nx.selfloop_edges(G))   # remove the self-loops
163      mapping = {i: f"{NODE_PREFIX}_{i}" for i in range(len(G))}
164      return nx.relabel_nodes(G, mapping)         # label the nodes
165  
166  
167  def generate_subnets(G, num_subnets):
168      n = len(G.nodes)
169      if num_subnets == n:   # if num_subnets == size of the network
170          return {f"{NODE_PREFIX}_{i}": f"{SUBNET_PREFIX}_{i}" for i in range(n)}
171  
172      lst = list(range(n))
173      random.shuffle(lst)
174      offsets = sorted(random.sample(range(0, n), num_subnets - 1))
175      offsets.append(n-1)
176  
177      start = 0
178      subnets = {}
179      subnet_id =  0
180      for end in offsets:
181          for i in range(start, end+1):
182              subnets[f"{NODE_PREFIX}_{lst[i]}"] = f"{SUBNET_PREFIX}_{subnet_id}"
183          start = end
184          subnet_id += 1
185      return subnets 
186  
187  
188  ### file format related fns ###########################################################
189  #Generate per node toml configs
190  def generate_toml(topics, node_type=nodeType.DESKTOP):
191      topic_str = " ".join(get_random_sublist(topics))   # space separated topics
192      return f"{nodeTypeSwitch.get(node_type)}topics = \"{topic_str}\"\n"
193  
194  
195  # Generates network-wide json and per-node toml and writes them 
196  def generate_and_write_files(dirname, num_topics, num_subnets, G):
197      topics = generate_topics(num_topics)
198      subnets = generate_subnets(G, num_subnets)
199      json_dump = {}
200      for node in G.nodes:
201          write_toml(dirname, node, generate_toml(topics))        # per node toml
202          json_dump[node] = {}
203          json_dump[node]["static-nodes"] = []
204          for edge in G.edges(node):
205              json_dump[node]["static-nodes"].append(edge[1])
206          json_dump[node][SUBNET_PREFIX] = subnets[node]
207      write_json(dirname, json_dump)                              # network wide json
208  
209  
210  ### the main ##########################################################################
211  def main(
212          dirname: str = "WakuNetwork", num_nodes: int = 4, num_topics: int = 1, 
213          network_type: networkType = networkType.NEWMANWATTSSTROGATZ.value, 
214          node_type: nodeType = nodeType.DESKTOP.value,
215          num_subnets: int = -1,
216          num_partitions: int = 1):
217  
218      # sanity checks
219      if num_partitions > 1:
220          raise ValueError(f"--num-partitions {num_partitions}, Sorry, we do not yet support partitions")
221      if num_subnets > num_nodes:
222          raise ValueError(f"num_subnets must be <= num_nodes: num_subnets={num_subnets}, num_nodes={num_nodes}")
223      if num_subnets == -1:
224          num_subnets = num_nodes
225  
226      # Generate the network
227      G = generate_network(num_nodes, network_type)
228  
229      # Refuse to overwrite non-empty dirs
230      if exists_or_nonempty(dirname) :
231          sys.exit(1)
232      os.makedirs(dirname, exist_ok=True)
233  
234      # Generate file format specific data structs and write the files; optionally, draw the network
235      generate_and_write_files(dirname, num_topics, num_subnets, G)
236      #draw(dirname, G)
237  
238  
239  if __name__ == "__main__":
240      typer.run(main)