ffmap-backend/ffmap/outputs/d3json.py

91 lines
2.9 KiB
Python

import json
from datetime import datetime
__all__ = ["Exporter"]
class CustomJSONEncoder(json.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)
class Output:
def __init__(self, filepath="nodes.json"):
self.filepath = filepath
@staticmethod
def generate(nodedb):
indexes = {}
nodes = []
count = 0
for node in set(nodedb.values()):
node_export = node.export()
node_export["flags"] = {
"gateway": "vpn" in node and node["vpn"],
"client": False,
"online": True
}
nodes.append(node_export)
indexes[node.id] = count
count += 1
links = {}
for node in set(nodedb.values()):
for neighbor in node.get("neighbors", []):
key = (neighbor["neighbor"].id, node.id)
rkey = tuple(reversed(key))
if rkey in links:
links[rkey]["quality"] += ","+neighbor["metric"]
else:
links[key] = {
"source": indexes[node.id],
"target": indexes[neighbor["neighbor"].id],
"quality": neighbor["metric"],
"type": "vpn" if neighbor["neighbor"]["vpn"] or node["vpn"] else None,
"id": "-".join((node.id, neighbor["neighbor"].id)),
}
for client in node.get("clients", []):
if not client in indexes:
nodes.append({
"id": client,
"flags": {
"client": True,
"online": True,
"gateway": False
}
})
indexes[client] = count
count += 1
links[(node.id, client)] = {
"source": indexes[node.id],
"target": indexes[client],
"quality": "TT",
"type": "client",
"id": "-".join((node.id, client)),
}
return {
"nodes": nodes,
"links": list(links.values()),
"meta": {
"timestamp": datetime.utcnow()
.replace(microsecond=0)
.isoformat()
}
}
def output(self, nodedb):
with open(self.filepath, "w") as nodes_json:
json.dump(
self.generate(nodedb),
nodes_json,
cls=CustomJSONEncoder
)