ffmap-backend/ffmap/outputs/d3json.py

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()):
            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,
            "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
            )