From 6fc1423124def6343a0aeb456c2af269e744b18d Mon Sep 17 00:00:00 2001 From: Jan-Philipp Litza Date: Fri, 21 Feb 2014 15:27:19 +0100 Subject: [PATCH] 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. --- alfred.py | 15 ++++----------- bat2nodes.py | 8 +++++--- d3mapbuilder.py | 36 ------------------------------------ ffhlwiki.py | 2 +- json_encoder.py | 13 +++++++++++++ link.py | 13 +++++++++++-- node.py | 48 +++++++++++++++++++++++++++++++++++++++++++----- nodedb.py | 20 ++++++++------------ 8 files changed, 85 insertions(+), 70 deletions(-) delete mode 100644 d3mapbuilder.py create mode 100644 json_encoder.py diff --git a/alfred.py b/alfred.py index 6d926bb..b8aa1e2 100755 --- a/alfred.py +++ b/alfred.py @@ -12,16 +12,11 @@ class alfred: alias = {} for mac,node in alfred_data.items(): node_alias = {} - if 'location' in node: - try: - node_alias['gps'] = str(node['location']['latitude']) + ' ' + str(node['location']['longitude']) - except: - pass + for key in node: + node_alias[key] = node[key] - try: - node_alias['firmware'] = node['software']['firmware']['release'] - except KeyError: - pass + if 'location' in node: + node_alias['geo'] = [node['location']['latitude'], node['location']['longitude']] try: node_alias['id'] = node['network']['mac'] @@ -30,8 +25,6 @@ class alfred: if 'hostname' in node: node_alias['name'] = node['hostname'] - elif 'name' in node: - node_alias['name'] = node['name'] if len(node_alias): alias[mac] = node_alias return alias diff --git a/bat2nodes.py b/bat2nodes.py index 921b548..e1fde6e 100755 --- a/bat2nodes.py +++ b/bat2nodes.py @@ -4,12 +4,13 @@ import json import fileinput import argparse import os +import datetime from batman import batman from alfred import alfred from rrd import rrd from nodedb import NodeDB -from d3mapbuilder import D3MapBuilder +from json_encoder import CustomJSONEncoder # Force encoding to UTF-8 import locale # Ensures that subsequent open()s @@ -71,11 +72,12 @@ if options['obscure']: 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 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() #Move to destination diff --git a/d3mapbuilder.py b/d3mapbuilder.py deleted file mode 100644 index ff7589f..0000000 --- a/d3mapbuilder.py +++ /dev/null @@ -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) - diff --git a/ffhlwiki.py b/ffhlwiki.py index c1ba01e..588ae72 100755 --- a/ffhlwiki.py +++ b/ffhlwiki.py @@ -71,7 +71,7 @@ def import_wikigps(url): mac = data[0].strip() if data[1]: - alias['gps'] = data[1].strip() + alias['geo'] = [float(x) for x in data[1].strip().split(' ')] if data[2]: alias['name'] = data[2].strip() diff --git a/json_encoder.py b/json_encoder.py new file mode 100644 index 0000000..8d62771 --- /dev/null +++ b/json_encoder.py @@ -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) diff --git a/link.py b/link.py index 896079b..b161608 100644 --- a/link.py +++ b/link.py @@ -1,11 +1,20 @@ class Link(): def __init__(self): self.id = None - self.source = None - self.target = None + self.source = LinkConnector() + self.target = LinkConnector() self.quality = 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(): def __init__(self): self.id = None diff --git a/node.py b/node.py index 0fe35fb..504768a 100644 --- a/node.py +++ b/node.py @@ -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): self.name = "" self.id = "" @@ -9,9 +36,7 @@ class Node(): "gateway": False, "client": False }) - self.gps = None - self.firmware = None - self.clientcount = 0 + super().__init__() def add_mac(self, mac): mac = mac.lower() @@ -25,7 +50,20 @@ class Node(): def __repr__(self): 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(): def __init__(self): self.vpn = False - diff --git a/nodedb.py b/nodedb.py index fa9caed..e5ff30e 100644 --- a/nodedb.py +++ b/nodedb.py @@ -1,4 +1,3 @@ -import json from functools import reduce from collections import defaultdict from node import Node, Interface @@ -18,6 +17,12 @@ class NodeDB: 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() @@ -179,21 +184,12 @@ class NodeDB: node.add_mac(mac) self._nodes.append(node) - if 'name' in alias: - node.name = alias['name'] + 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 - 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 # if options['gateway']: # mark_gateways(options['gateway'])