diff --git a/bat2nodes.py b/bat2nodes.py index 482915d..d07faed 100644 --- a/bat2nodes.py +++ b/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()) diff --git a/d3mapbuilder.py b/d3mapbuilder.py new file mode 100644 index 0000000..4fd10db --- /dev/null +++ b/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) + diff --git a/nodedb.py b/nodedb.py new file mode 100644 index 0000000..7defde2 --- /dev/null +++ b/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 +