synapse/graph/graph.py

import sqlite3
import pydot
import cgi
import json
import datetime
import argparse
import urllib2


def make_name(pdu_id, origin):
    return "%s@%s" % (pdu_id, origin)


def make_graph(pdus, room, filename_prefix):
    pdu_map = {}
    node_map = {}

    origins = set()
    colors = set(("red", "green", "blue", "yellow", "purple"))

    for pdu in pdus:
        origins.add(pdu.get("origin"))

    color_map = {color: color for color in colors if color in origins}
    colors -= set(color_map.values())

    color_map[None] = "black"

    for o in origins:
        if o in color_map:
            continue
        try:
            c = colors.pop()
            color_map[o] = c
        except:
            print "Run out of colours!"
            color_map[o] = "black"

    graph = pydot.Dot(graph_name="Test")

    for pdu in pdus:
        name = make_name(pdu.get("pdu_id"), pdu.get("origin"))
        pdu_map[name] = pdu

        t = datetime.datetime.fromtimestamp(
            float(pdu["ts"]) / 1000
        ).strftime('%Y-%m-%d %H:%M:%S,%f')

        label = (
            "<"
            "<b>%(name)s </b><br/>"
            "Type: <b>%(type)s </b><br/>"
            "State key: <b>%(state_key)s </b><br/>"
            "Content: <b>%(content)s </b><br/>"
            "Time: <b>%(time)s </b><br/>"
            "Depth: <b>%(depth)s </b><br/>"
            ">"
        ) % {
            "name": name,
            "type": pdu.get("pdu_type"),
            "state_key": pdu.get("state_key"),
            "content": cgi.escape(json.dumps(pdu.get("content")), quote=True),
            "time": t,
            "depth": pdu.get("depth"),
        }

        node = pydot.Node(
            name=name,
            label=label,
            color=color_map[pdu.get("origin")]
        )
        node_map[name] = node
        graph.add_node(node)

    for pdu in pdus:
        start_name = make_name(pdu.get("pdu_id"), pdu.get("origin"))
        for i, o in pdu.get("prev_pdus", []):
            end_name = make_name(i, o)

            if end_name not in node_map:
                print "%s not in nodes" % end_name
                continue

            edge = pydot.Edge(node_map[start_name], node_map[end_name])
            graph.add_edge(edge)

        # Add prev_state edges, if they exist
        if pdu.get("prev_state_id") and pdu.get("prev_state_origin"):
            prev_state_name = make_name(
                pdu.get("prev_state_id"), pdu.get("prev_state_origin")
            )

            if prev_state_name in node_map:
                state_edge = pydot.Edge(
                    node_map[start_name], node_map[prev_state_name],
                    style='dotted'
                )
                graph.add_edge(state_edge)

    graph.write('%s.dot' % filename_prefix, format='raw', prog='dot')
    graph.write_png("%s.png" % filename_prefix, prog='dot')
    graph.write_svg("%s.svg" % filename_prefix, prog='dot')


def get_pdus(host, room):
    transaction = json.loads(
        urllib2.urlopen(
            "http://%s/context/%s/" % (host, room)
        ).read()
    )

    return transaction["pdus"]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate a PDU graph for a given room by talking "
                    "to the given homeserver to get the list of PDUs. \n"
                    "Requires pydot."
    )
    parser.add_argument(
        "-p", "--prefix", dest="prefix",
        help="String to prefix output files with"
    )
    parser.add_argument('host')
    parser.add_argument('room')

    args = parser.parse_args()

    host = args.host
    room = args.room
    prefix = args.prefix if args.prefix else "%s_graph" % (room)

    pdus = get_pdus(host, room)

    make_graph(pdus, room, prefix)