abstract NodeDB and D3MapBuilder into classes

bat2nodes.py

@@ -8,8 +8,8 @@ import json
 import fileinput
 import argparse
 
-from node import Node
-from link import Link
+from nodedb import NodeDB
+from d3mapbuilder import D3MapBuilder
 
 parser = argparse.ArgumentParser()
 
@@ -26,123 +26,16 @@ args = parser.parse_args()
 
 options = vars(args)
 
-aliases = dict()
-links = set()
-nodes = []
+db = NodeDB()
 
-def maybe_node_by_mac(nodes, macs):
-  for node in nodes:
-    for mac in macs:
-      if mac in node.macs:
-        return node
-
-  raise
-
-lines = list(fileinput.input(options['batmanjson']))
-
-for line in lines:
-  x = json.loads(line)
-
-  if 'of' in x:
-    try:
-      node = maybe_node_by_mac(nodes, (x['of'], x['secondary']))
-    except:
-      node = Node()
-      node.online = True
-      nodes.append(node)
-
-    node.add_mac(x['of'])
-    node.add_mac(x['secondary'])
-
-for line in lines:
-  x = json.loads(line)
-
-  if 'router' in x:
-    try:
-      node = maybe_node_by_mac(nodes, (x['router'], ))
-    except:
-      node = Node()
-      node.online = True
-      node.add_mac(x['router'])
-      nodes.append(node)
-
-    try:
-      if 'gateway' in x:
-        x['neighbor'] = x['gateway']
-
-      node = maybe_node_by_mac(nodes, (x['neighbor'], ))
-    except:
-      node = Node()
-      node.online = True
-      if x['label'] == 'TT':
-        node.group = 3
-
-      node.add_mac(x['neighbor'])
-      nodes.append(node)
-
-for line in lines:
-  x = json.loads(line)
-
-  if 'router' in x:
-    try:
-      if 'gateway' in x:
-        x['neighbor'] = x['gateway']
-
-      router = maybe_node_by_mac(nodes, (x['router'], ))
-      neighbor = maybe_node_by_mac(nodes, (x['neighbor'], ))
-    except:
-      continue
-
-    a = nodes.index(router)
-    b = nodes.index(neighbor)
-
-    links.add(tuple(sorted((a,b))))
+db.import_batman(list(fileinput.input(options['batmanjson'])))
 
 if options['aliases']:
-  aliases = json.load(open(options['aliases']))
-
-  for mac, alias in aliases.items():
-    try:
-      node = maybe_node_by_mac(nodes, (mac, ))
-    except:
-      continue
-
-    node.name = alias['name']
-    if 'group' in alias:
-      node.group = alias['group']
+  db.import_aliases(json.load(open(options['aliases'])))
 
 if options['gateway']:
-  for gateway in options['gateway']:
-    try:
-      node = maybe_node_by_mac(nodes, (gateway, ))
-    except:
-      continue
+  db.mark_gateways(options['gateway'])
 
-    node.group = 2
+m = D3MapBuilder(db)
 
-def map_link(nodes, pair):
-  distance = 80
-  strength = 0.2
-  if any(filter(lambda x: nodes[x].group == 3, pair)):
-    distance = 10
-    strength = 1
-
-  link = Link()
-  link.pair = pair
-  link.distance = distance
-  link.strength = strength
-  return link
-
-links = [map_link(nodes, x) for x in links]
-
-output = dict()
-
-output['nodes'] = [{'group': x.group, 'name': x.name,
-                    'macs': ', '.join(x.macs)
-                   } for x in nodes if x.online]
-output['links'] = [{'source': x.pair[0], 'target': x.pair[1],
-                    'distance': x.distance,
-                    'strength': x.strength
-                   } for x in links]
-
-print(json.dumps(output))
+print(m.build())

d3mapbuilder.py

@@ -0,0 +1,21 @@
+import json
+
+class D3MapBuilder:
+  db = None
+
+  def __init__(self, db):
+    self.db = db
+
+  def build(self):
+    output = dict()
+
+    output['nodes'] = [{'group': x.group, 'name': x.name,
+                        'macs': ', '.join(x.macs)
+                       } for x in self.db.get_nodes() if x.online]
+    output['links'] = [{'source': x.pair[0], 'target': x.pair[1],
+                        'distance': x.distance,
+                        'strength': x.strength
+                       } for x in self.db.get_links()]
+
+    return json.dumps(output)
+

nodedb.py

@@ -0,0 +1,122 @@
+import json
+from node import Node
+from link import Link
+
+class NodeDB:
+  def __init__(self):
+    self._nodes = []
+    self._links = set()
+
+  # fetch list of links
+  def get_links(self):
+    return [self.map_link(x) for x in self._links]
+
+  # fetch list of nodes
+  def get_nodes(self):
+    return self._nodes
+
+  def maybe_node_by_mac(self, macs):
+    for node in self._nodes:
+      for mac in macs:
+        if mac in node.macs:
+          return node
+
+    raise
+
+  # import_batman(list(fileinput.input(options['batmanjson'])))
+  def import_batman(self, lines):
+
+    for line in lines:
+      x = json.loads(line)
+
+      if 'of' in x:
+        try:
+          node = self.maybe_node_by_mac((x['of'], x['secondary']))
+        except:
+          node = Node()
+          node.online = True
+          self._nodes.append(node)
+
+        node.add_mac(x['of'])
+        node.add_mac(x['secondary'])
+
+    for line in lines:
+      x = json.loads(line)
+
+      if 'router' in x:
+        try:
+          node = self.maybe_node_by_mac((x['router'], ))
+        except:
+          node = Node()
+          node.online = True
+          node.add_mac(x['router'])
+          self._nodes.append(node)
+
+        try:
+          if 'gateway' in x:
+            x['neighbor'] = x['gateway']
+
+          node = self.maybe_node_by_mac((x['neighbor'], ))
+        except:
+          node = Node()
+          node.online = True
+          if x['label'] == 'TT':
+            node.group = 3
+
+          node.add_mac(x['neighbor'])
+          self._nodes.append(node)
+
+    for line in lines:
+      x = json.loads(line)
+
+      if 'router' in x:
+        try:
+          if 'gateway' in x:
+            x['neighbor'] = x['gateway']
+
+          router = self.maybe_node_by_mac((x['router'], ))
+          neighbor = self.maybe_node_by_mac((x['neighbor'], ))
+        except:
+          continue
+
+        a = self._nodes.index(router)
+        b = self._nodes.index(neighbor)
+
+        self._links.add(tuple(sorted((a,b))))
+
+  def import_aliases(self, aliases):
+    for mac, alias in aliases.items():
+      try:
+        node = self.maybe_node_by_mac((mac, ))
+      except:
+        continue
+
+      node.name = alias['name']
+      if 'group' in alias:
+        node.group = alias['group']
+
+  # list of macs
+  # if options['gateway']:
+  #   mark_gateways(options['gateway'])
+  def mark_gateways(self, gateways):
+    for gateway in gateways:
+      try:
+        node = self.maybe_node_by_mac((gateway, ))
+      except:
+        continue
+
+      node.group = 2
+
+  def map_link(self, pair):
+    distance = 80
+    strength = 0.2
+    if any(filter(lambda x: self._nodes[x].group == 3, pair)):
+      distance = 10
+      strength = 1
+
+    link = Link()
+    link.pair = pair
+    link.distance = distance
+    link.strength = strength
+    return link
+