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.
This commit is contained in:
parent
0ab4ce6e2b
commit
6fc1423124
15
alfred.py
15
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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
@ -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()
|
||||
|
|
13
json_encoder.py
Normal file
13
json_encoder.py
Normal 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
13
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
|
||||
|
|
48
node.py
48
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
|
||||
|
||||
|
|
20
nodedb.py
20
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'])
|
||||
|
|
Loading…
Reference in a new issue