123 lines
2.7 KiB
Python
123 lines
2.7 KiB
Python
|
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
|
||
|
|
|