72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
|
import json
|
||
|
|
|
||
|
|
__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 Exporter:
|
||
|
|
def __init__(self, filepath="nodes.json"):
|
||
|
|
self.filepath = filepath
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def generate(nodedb):
|
||
|
|
indexes = {}
|
||
|
|
nodes = []
|
||
|
|
count = 0
|
||
|
|
for node in set(nodedb.values()):
|
||
|
|
nodes.append(node.export())
|
||
|
|
indexes[node.id] = count
|
||
|
|
count += 1
|
||
|
|
|
||
|
|
links = []
|
||
|
|
for node in set(nodedb.values()):
|
||
|
|
if "neighbors" in node:
|
||
|
|
links.extend(
|
||
|
|
{
|
||
|
|
"source": indexes[node.id],
|
||
|
|
"target": indexes[neighbor["neighbor"].id],
|
||
|
|
"quality": neighbor["metric"],
|
||
|
|
"type": "vpn" if neighbor["neighbor"]["vpn"] else None,
|
||
|
|
"id": "-".join((node.id, neighbor["neighbor"].id)),
|
||
|
|
} for neighbor in node["neighbors"]
|
||
|
|
)
|
||
|
|
if "clients" in node:
|
||
|
|
for client in node["clients"]:
|
||
|
|
if not client in indexes:
|
||
|
|
nodes.append({
|
||
|
|
"id": client,
|
||
|
|
})
|
||
|
|
indexes[client] = count
|
||
|
|
count += 1
|
||
|
|
|
||
|
|
links.append({
|
||
|
|
"source": indexes[node.id],
|
||
|
|
"target": indexes[client],
|
||
|
|
"quality": "TT",
|
||
|
|
"type": "client",
|
||
|
|
"id": "-".join((node.id, client)),
|
||
|
|
})
|
||
|
|
|
||
|
|
return {
|
||
|
|
"nodes": nodes,
|
||
|
|
"links": links,
|
||
|
|
}
|
||
|
|
|
||
|
|
def export(self, nodedb):
|
||
|
|
with open(self.filepath, "w") as nodes_json:
|
||
|
|
json.dump(
|
||
|
|
self.generate(nodedb),
|
||
|
|
nodes_json,
|
||
|
|
cls=CustomJSONEncoder
|
||
|
|
)
|