6fc1423124
This commit makes Nodes special dicts that return None-like objects for inexistent keys, making it a dynamic attribute store. Also, it removes the D3MapBuilder and moves its logic to the Node and Link classes' newly introduced export() method. Only they need to be changed to populate the final nodes.json with more attributes.
382 lines
9 KiB
Python
382 lines
9 KiB
Python
from functools import reduce
|
|
from collections import defaultdict
|
|
from node import Node, Interface
|
|
from link import Link, LinkConnector
|
|
|
|
class NodeDB:
|
|
def __init__(self):
|
|
self._nodes = []
|
|
self._links = []
|
|
|
|
# fetch list of links
|
|
def get_links(self):
|
|
self.update_vpn_links()
|
|
return self.reduce_links()
|
|
|
|
# fetch list of nodes
|
|
def get_nodes(self):
|
|
return self._nodes
|
|
|
|
def export(self):
|
|
return {
|
|
'nodes': [node.export() for node in self.get_nodes()],
|
|
'links': [link.export() for link in self.get_links()],
|
|
}
|
|
|
|
def maybe_node_by_fuzzy_mac(self, mac):
|
|
mac_a = mac.lower()
|
|
|
|
for node in self._nodes:
|
|
for mac_b in node.macs:
|
|
if is_derived_mac(mac_a, mac_b):
|
|
return node
|
|
|
|
raise KeyError
|
|
|
|
def maybe_node_by_mac(self, macs):
|
|
for node in self._nodes:
|
|
for mac in macs:
|
|
if mac.lower() in node.macs:
|
|
return node
|
|
|
|
raise KeyError
|
|
|
|
def maybe_node_by_id(self, mac):
|
|
for node in self._nodes:
|
|
if mac.lower() == node.id:
|
|
return node
|
|
|
|
raise KeyError
|
|
|
|
def parse_vis_data(self,vis_data):
|
|
for x in vis_data:
|
|
|
|
if 'of' in x:
|
|
try:
|
|
node = self.maybe_node_by_mac((x['of'], x['secondary']))
|
|
except:
|
|
node = Node()
|
|
node.flags['online'] = True
|
|
if 'legacy' in x:
|
|
node.flags['legacy'] = True
|
|
self._nodes.append(node)
|
|
|
|
node.add_mac(x['of'])
|
|
node.add_mac(x['secondary'])
|
|
|
|
for x in vis_data:
|
|
|
|
if 'router' in x:
|
|
try:
|
|
node = self.maybe_node_by_mac((x['router'], ))
|
|
except:
|
|
node = Node()
|
|
node.flags['online'] = True
|
|
if 'legacy' in x:
|
|
node.flags['legacy'] = True
|
|
node.add_mac(x['router'])
|
|
self._nodes.append(node)
|
|
|
|
# If it's a TT link and the MAC is very similar
|
|
# consider this MAC as one of the routers
|
|
# MACs
|
|
if 'gateway' in x and x['label'] == "TT":
|
|
if is_similar(x['router'], x['gateway']):
|
|
node.add_mac(x['gateway'])
|
|
|
|
# skip processing as regular link
|
|
continue
|
|
|
|
try:
|
|
if 'neighbor' in x:
|
|
try:
|
|
node = self.maybe_node_by_mac((x['neighbor']))
|
|
except:
|
|
continue
|
|
|
|
if 'gateway' in x:
|
|
x['neighbor'] = x['gateway']
|
|
|
|
node = self.maybe_node_by_mac((x['neighbor'], ))
|
|
except:
|
|
node = Node()
|
|
node.flags['online'] = True
|
|
if x['label'] == 'TT':
|
|
node.flags['client'] = True
|
|
|
|
node.add_mac(x['neighbor'])
|
|
self._nodes.append(node)
|
|
|
|
for x in vis_data:
|
|
|
|
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
|
|
|
|
# filter TT links merged in previous step
|
|
if router == neighbor:
|
|
continue
|
|
|
|
link = Link()
|
|
link.source = LinkConnector()
|
|
link.source.interface = x['router']
|
|
link.source.id = self._nodes.index(router)
|
|
link.target = LinkConnector()
|
|
link.target.interface = x['neighbor']
|
|
link.target.id = self._nodes.index(neighbor)
|
|
link.quality = x['label']
|
|
link.id = "-".join(sorted((link.source.interface, link.target.interface)))
|
|
|
|
if x['label'] == "TT":
|
|
link.type = "client"
|
|
|
|
self._links.append(link)
|
|
|
|
for x in vis_data:
|
|
|
|
if 'primary' in x:
|
|
try:
|
|
node = self.maybe_node_by_mac((x['primary'], ))
|
|
except:
|
|
continue
|
|
|
|
node.id = x['primary']
|
|
|
|
def reduce_links(self):
|
|
tmp_links = defaultdict(list)
|
|
|
|
for link in self._links:
|
|
tmp_links[link.id].append(link)
|
|
|
|
links = []
|
|
|
|
def reduce_link(a, b):
|
|
a.id = b.id
|
|
a.source = b.source
|
|
a.target = b.target
|
|
a.type = b.type
|
|
a.quality = ", ".join([x for x in (a.quality, b.quality) if x])
|
|
|
|
return a
|
|
|
|
for k, v in tmp_links.items():
|
|
new_link = reduce(reduce_link, v, Link())
|
|
links.append(new_link)
|
|
|
|
return links
|
|
|
|
def import_aliases(self, aliases):
|
|
for mac, alias in aliases.items():
|
|
try:
|
|
node = self.maybe_node_by_mac([mac])
|
|
except:
|
|
try:
|
|
node = self.maybe_node_by_fuzzy_mac(mac)
|
|
except:
|
|
# create an offline node
|
|
node = Node()
|
|
node.add_mac(mac)
|
|
self._nodes.append(node)
|
|
|
|
for key in alias:
|
|
node[key] = alias[key]
|
|
|
|
if 'vpn' in alias and alias['vpn'] and mac and node.interfaces and mac in node.interfaces:
|
|
node.interfaces[mac].vpn = True
|
|
|
|
# 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:
|
|
print("WARNING: did not find gateway '",gateway,"' in node list")
|
|
continue
|
|
|
|
node.flags['gateway'] = True
|
|
|
|
def update_vpn_links(self):
|
|
changes = 1
|
|
while changes > 0:
|
|
changes = 0
|
|
for link in self._links:
|
|
if link.type == "client":
|
|
continue
|
|
|
|
source_interface = self._nodes[link.source.id].interfaces[link.source.interface]
|
|
target_interface = self._nodes[link.target.id].interfaces[link.target.interface]
|
|
if source_interface.vpn or target_interface.vpn:
|
|
source_interface.vpn = True
|
|
target_interface.vpn = True
|
|
if link.type != "vpn":
|
|
changes += 1
|
|
|
|
link.type = "vpn"
|
|
|
|
def count_clients(self):
|
|
for link in self._links:
|
|
try:
|
|
a = self.maybe_node_by_id(link.source.interface)
|
|
b = self.maybe_node_by_id(link.target.interface)
|
|
|
|
if a.flags['client']:
|
|
client = a
|
|
node = b
|
|
elif b.flags['client']:
|
|
client = b
|
|
node = a
|
|
else:
|
|
continue
|
|
|
|
node.clientcount += 1
|
|
except:
|
|
pass
|
|
|
|
def obscure_clients(self):
|
|
|
|
globalIdCounter = 0
|
|
nodeCounters = {}
|
|
clientIds = {}
|
|
|
|
for node in self._nodes:
|
|
if node.flags['client']:
|
|
node.macs = set()
|
|
clientIds[node.id] = None
|
|
|
|
for link in self._links:
|
|
ids = link.source.interface
|
|
idt = link.target.interface
|
|
|
|
try:
|
|
node_source = self.maybe_node_by_fuzzy_mac(ids)
|
|
node_target = self.maybe_node_by_id(idt)
|
|
|
|
if not node_source.flags['client'] and not node_target.flags['client']:
|
|
# if none of the nodes associated with this link are clients,
|
|
# we do not want to obscure
|
|
continue
|
|
|
|
if ids in clientIds and idt in clientIds:
|
|
# This is for corner cases, when a client
|
|
# is linked to another client.
|
|
clientIds[ids] = str(globalIdCounter)
|
|
ids = str(globalIdCounter)
|
|
globalIdCounter += 1
|
|
|
|
clientIds[idt] = str(globalIdCounter)
|
|
idt = str(globalIdCounter)
|
|
globalIdCounter += 1
|
|
|
|
elif ids in clientIds:
|
|
newId = generateId(idt)
|
|
clientIds[ids] = newId
|
|
ids = newId
|
|
|
|
link.source.interface = ids;
|
|
node_source.id = ids;
|
|
|
|
elif idt in clientIds:
|
|
newId = generateId(ids,nodeCounters)
|
|
clientIds[idt] = newId
|
|
idt = newId
|
|
|
|
link.target.interface = idt;
|
|
node_target.id = idt;
|
|
|
|
link.id = ids + "-" + idt
|
|
|
|
except KeyError:
|
|
pass
|
|
|
|
# extends node id by incremented node counter
|
|
def generateId(nodeId,nodeCounters):
|
|
if nodeId in nodeCounters:
|
|
n = nodeCounters[nodeId]
|
|
nodeCounters[nodeId] = n + 1
|
|
else:
|
|
nodeCounters[nodeId] = 1
|
|
n = 0
|
|
|
|
return nodeId + "_" + str(n)
|
|
|
|
# compares two MACs and decides whether they are
|
|
# similar and could be from the same node
|
|
def is_similar(a, b):
|
|
if a == b:
|
|
return True
|
|
|
|
try:
|
|
mac_a = list(int(i, 16) for i in a.split(":"))
|
|
mac_b = list(int(i, 16) for i in b.split(":"))
|
|
except ValueError:
|
|
return False
|
|
|
|
# first byte must only differ in bit 2
|
|
if mac_a[0] | 2 == mac_b[0] | 2:
|
|
# count different bytes
|
|
c = [x for x in zip(mac_a[1:], mac_b[1:]) if x[0] != x[1]]
|
|
else:
|
|
return False
|
|
|
|
# no more than two additional bytes must differ
|
|
if len(c) <= 2:
|
|
delta = 0
|
|
|
|
if len(c) > 0:
|
|
delta = sum(abs(i[0] -i[1]) for i in c)
|
|
|
|
# These addresses look pretty similar!
|
|
return delta < 8
|
|
|
|
def is_derived_mac(a, b):
|
|
if a == b:
|
|
return True
|
|
|
|
try:
|
|
mac_a = list(int(i, 16) for i in a.split(":"))
|
|
mac_b = list(int(i, 16) for i in b.split(":"))
|
|
except ValueError:
|
|
return False
|
|
|
|
if mac_a[4] != mac_b[4] or mac_a[2] != mac_b[2] or mac_a[1] != mac_b[1]:
|
|
return False
|
|
|
|
x = list(mac_a)
|
|
x[5] += 1
|
|
x[5] %= 255
|
|
if mac_b == x:
|
|
return True
|
|
|
|
x[0] |= 2
|
|
if mac_b == x:
|
|
return True
|
|
|
|
x[3] += 1
|
|
x[3] %= 255
|
|
if mac_b == x:
|
|
return True
|
|
|
|
x = list(mac_a)
|
|
x[0] |= 2
|
|
x[5] += 2
|
|
x[5] %= 255
|
|
if mac_b == x:
|
|
return True
|
|
|
|
x = list(mac_a)
|
|
x[0] |= 2
|
|
x[3] += 1
|
|
x[3] %= 255
|
|
if mac_b == x:
|
|
return True
|
|
|
|
return False
|