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'])