Make handling of node attributes more flexible.

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.
rewrite
Jan-Philipp Litza 2014-02-21 15:27:19 +01:00
parent 0ab4ce6e2b
commit 6fc1423124
8 changed files with 85 additions and 70 deletions

View File

@ -12,16 +12,11 @@ class alfred:
alias = {} alias = {}
for mac,node in alfred_data.items(): for mac,node in alfred_data.items():
node_alias = {} node_alias = {}
if 'location' in node: for key in node:
try: node_alias[key] = node[key]
node_alias['gps'] = str(node['location']['latitude']) + ' ' + str(node['location']['longitude'])
except:
pass
try: if 'location' in node:
node_alias['firmware'] = node['software']['firmware']['release'] node_alias['geo'] = [node['location']['latitude'], node['location']['longitude']]
except KeyError:
pass
try: try:
node_alias['id'] = node['network']['mac'] node_alias['id'] = node['network']['mac']
@ -30,8 +25,6 @@ class alfred:
if 'hostname' in node: if 'hostname' in node:
node_alias['name'] = node['hostname'] node_alias['name'] = node['hostname']
elif 'name' in node:
node_alias['name'] = node['name']
if len(node_alias): if len(node_alias):
alias[mac] = node_alias alias[mac] = node_alias
return alias return alias

View File

@ -4,12 +4,13 @@ import json
import fileinput import fileinput
import argparse import argparse
import os import os
import datetime
from batman import batman from batman import batman
from alfred import alfred from alfred import alfred
from rrd import rrd from rrd import rrd
from nodedb import NodeDB from nodedb import NodeDB
from d3mapbuilder import D3MapBuilder from json_encoder import CustomJSONEncoder
# Force encoding to UTF-8 # Force encoding to UTF-8
import locale # Ensures that subsequent open()s import locale # Ensures that subsequent open()s
@ -71,11 +72,12 @@ if options['obscure']:
scriptdir = os.path.dirname(os.path.realpath(__file__)) scriptdir = os.path.dirname(os.path.realpath(__file__))
m = D3MapBuilder(db) exported = db.export()
exported['meta'] = {'timestamp': datetime.datetime.utcnow().replace(microsecond=0).isoformat()}
#Write nodes json #Write nodes json
nodes_json = open(options['destination_directory'] + '/nodes.json.new','w') nodes_json = open(options['destination_directory'] + '/nodes.json.new','w')
nodes_json.write(m.build()) json.dump(exported, nodes_json, cls=CustomJSONEncoder)
nodes_json.close() nodes_json.close()
#Move to destination #Move to destination

View File

@ -1,36 +0,0 @@
import json
import datetime
class D3MapBuilder:
def __init__(self, db):
self._db = db
def build(self):
output = dict()
now = datetime.datetime.utcnow().replace(microsecond=0)
nodes = self._db.get_nodes()
output['nodes'] = [{'name': x.name, 'id': x.id,
'macs': ', '.join(x.macs),
'geo': [float(x) for x in x.gps.split(" ")] if x.gps else None,
'firmware': x.firmware,
'flags': x.flags,
'clientcount': x.clientcount
} for x in nodes]
links = self._db.get_links()
output['links'] = [{'source': x.source.id, 'target': x.target.id,
'quality': x.quality,
'type': x.type,
'id': x.id
} for x in links]
output['meta'] = {
'timestamp': now.isoformat()
}
return json.dumps(output)

View File

@ -71,7 +71,7 @@ def import_wikigps(url):
mac = data[0].strip() mac = data[0].strip()
if data[1]: if data[1]:
alias['gps'] = data[1].strip() alias['geo'] = [float(x) for x in data[1].strip().split(' ')]
if data[2]: if data[2]:
alias['name'] = data[2].strip() alias['name'] = data[2].strip()

13
json_encoder.py Normal file
View File

@ -0,0 +1,13 @@
from json import JSONEncoder
class CustomJSONEncoder(JSONEncoder):
"""
JSON encoder that uses an object's __json__() method to convert it to
something JSON-compatible.
"""
def default(self, obj):
try:
return obj.__json__()
except AttributeError:
pass
return super().default(obj)

13
link.py
View File

@ -1,11 +1,20 @@
class Link(): class Link():
def __init__(self): def __init__(self):
self.id = None self.id = None
self.source = None self.source = LinkConnector()
self.target = None self.target = LinkConnector()
self.quality = None self.quality = None
self.type = None self.type = None
def export(self):
return {
'source': self.source.id,
'target': self.target.id,
'quality': self.quality,
'type': self.type,
'id': self.id
}
class LinkConnector(): class LinkConnector():
def __init__(self): def __init__(self):
self.id = None self.id = None

48
node.py
View File

@ -1,4 +1,31 @@
class Node(): from collections import defaultdict
class NoneDict:
"""
A NoneDict acts like None but returns a NoneDict for every item in it.
This is similar to the behaviour of collections.defaultdict in that even
previously inexistent keys can be accessed, but there is nothing stored
permanently.
"""
__repr__ = lambda self: 'NoneDict()'
__bool__ = lambda self: False
__getitem__ = lambda self, k: NoneDict()
__json__ = lambda self: None
def __setitem__(self, key, value):
raise RuntimeError("NoneDict is readonly")
class casualdict(defaultdict):
"""
This special defaultdict returns a NoneDict for inexistent items. Also, its
items can be accessed as attributed as well.
"""
def __init__(self):
super().__init__(NoneDict)
__getattr__ = defaultdict.__getitem__
__setattr__ = defaultdict.__setitem__
class Node(casualdict):
def __init__(self): def __init__(self):
self.name = "" self.name = ""
self.id = "" self.id = ""
@ -9,9 +36,7 @@ class Node():
"gateway": False, "gateway": False,
"client": False "client": False
}) })
self.gps = None super().__init__()
self.firmware = None
self.clientcount = 0
def add_mac(self, mac): def add_mac(self, mac):
mac = mac.lower() mac = mac.lower()
@ -25,7 +50,20 @@ class Node():
def __repr__(self): def __repr__(self):
return self.macs.__repr__() return self.macs.__repr__()
def export(self):
"""
Return a dict that contains all attributes of the Node that are supposed to
be exported to other applications.
"""
return {
"name": self.name,
"id": self.id,
"macs": list(self.macs),
"geo": self.geo,
"firmware": self.software['firmware']['release'],
"flags": self.flags
}
class Interface(): class Interface():
def __init__(self): def __init__(self):
self.vpn = False self.vpn = False

View File

@ -1,4 +1,3 @@
import json
from functools import reduce from functools import reduce
from collections import defaultdict from collections import defaultdict
from node import Node, Interface from node import Node, Interface
@ -18,6 +17,12 @@ class NodeDB:
def get_nodes(self): def get_nodes(self):
return self._nodes 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): def maybe_node_by_fuzzy_mac(self, mac):
mac_a = mac.lower() mac_a = mac.lower()
@ -179,21 +184,12 @@ class NodeDB:
node.add_mac(mac) node.add_mac(mac)
self._nodes.append(node) self._nodes.append(node)
if 'name' in alias: for key in alias:
node.name = alias['name'] node[key] = alias[key]
if 'vpn' in alias and alias['vpn'] and mac and node.interfaces and mac in node.interfaces: if 'vpn' in alias and alias['vpn'] and mac and node.interfaces and mac in node.interfaces:
node.interfaces[mac].vpn = True node.interfaces[mac].vpn = True
if 'gps' in alias:
node.gps = alias['gps']
if 'firmware' in alias:
node.firmware = alias['firmware']
if 'id' in alias:
node.id = alias['id']
# list of macs # list of macs
# if options['gateway']: # if options['gateway']:
# mark_gateways(options['gateway']) # mark_gateways(options['gateway'])