1
0

graph.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. import argparse
  2. import cgi
  3. import datetime
  4. import json
  5. import pydot
  6. import urllib2
  7. # Copyright 2014-2016 OpenMarket Ltd
  8. #
  9. # Licensed under the Apache License, Version 2.0 (the "License");
  10. # you may not use this file except in compliance with the License.
  11. # You may obtain a copy of the License at
  12. #
  13. # http://www.apache.org/licenses/LICENSE-2.0
  14. #
  15. # Unless required by applicable law or agreed to in writing, software
  16. # distributed under the License is distributed on an "AS IS" BASIS,
  17. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  18. # See the License for the specific language governing permissions and
  19. # limitations under the License.
  20. def make_name(pdu_id, origin):
  21. return "%s@%s" % (pdu_id, origin)
  22. def make_graph(pdus, room, filename_prefix):
  23. pdu_map = {}
  24. node_map = {}
  25. origins = set()
  26. colors = {"red", "green", "blue", "yellow", "purple"}
  27. for pdu in pdus:
  28. origins.add(pdu.get("origin"))
  29. color_map = {color: color for color in colors if color in origins}
  30. colors -= set(color_map.values())
  31. color_map[None] = "black"
  32. for o in origins:
  33. if o in color_map:
  34. continue
  35. try:
  36. c = colors.pop()
  37. color_map[o] = c
  38. except Exception:
  39. print("Run out of colours!")
  40. color_map[o] = "black"
  41. graph = pydot.Dot(graph_name="Test")
  42. for pdu in pdus:
  43. name = make_name(pdu.get("pdu_id"), pdu.get("origin"))
  44. pdu_map[name] = pdu
  45. t = datetime.datetime.fromtimestamp(float(pdu["ts"]) / 1000).strftime(
  46. "%Y-%m-%d %H:%M:%S,%f"
  47. )
  48. label = (
  49. "<"
  50. "<b>%(name)s </b><br/>"
  51. "Type: <b>%(type)s </b><br/>"
  52. "State key: <b>%(state_key)s </b><br/>"
  53. "Content: <b>%(content)s </b><br/>"
  54. "Time: <b>%(time)s </b><br/>"
  55. "Depth: <b>%(depth)s </b><br/>"
  56. ">"
  57. ) % {
  58. "name": name,
  59. "type": pdu.get("pdu_type"),
  60. "state_key": pdu.get("state_key"),
  61. "content": cgi.escape(json.dumps(pdu.get("content")), quote=True),
  62. "time": t,
  63. "depth": pdu.get("depth"),
  64. }
  65. node = pydot.Node(name=name, label=label, color=color_map[pdu.get("origin")])
  66. node_map[name] = node
  67. graph.add_node(node)
  68. for pdu in pdus:
  69. start_name = make_name(pdu.get("pdu_id"), pdu.get("origin"))
  70. for i, o in pdu.get("prev_pdus", []):
  71. end_name = make_name(i, o)
  72. if end_name not in node_map:
  73. print("%s not in nodes" % end_name)
  74. continue
  75. edge = pydot.Edge(node_map[start_name], node_map[end_name])
  76. graph.add_edge(edge)
  77. # Add prev_state edges, if they exist
  78. if pdu.get("prev_state_id") and pdu.get("prev_state_origin"):
  79. prev_state_name = make_name(
  80. pdu.get("prev_state_id"), pdu.get("prev_state_origin")
  81. )
  82. if prev_state_name in node_map:
  83. state_edge = pydot.Edge(
  84. node_map[start_name], node_map[prev_state_name], style="dotted"
  85. )
  86. graph.add_edge(state_edge)
  87. graph.write("%s.dot" % filename_prefix, format="raw", prog="dot")
  88. # graph.write_png("%s.png" % filename_prefix, prog='dot')
  89. graph.write_svg("%s.svg" % filename_prefix, prog="dot")
  90. def get_pdus(host, room):
  91. transaction = json.loads(
  92. urllib2.urlopen(
  93. "http://%s/_matrix/federation/v1/context/%s/" % (host, room)
  94. ).read()
  95. )
  96. return transaction["pdus"]
  97. if __name__ == "__main__":
  98. parser = argparse.ArgumentParser(
  99. description="Generate a PDU graph for a given room by talking "
  100. "to the given homeserver to get the list of PDUs. \n"
  101. "Requires pydot."
  102. )
  103. parser.add_argument(
  104. "-p", "--prefix", dest="prefix", help="String to prefix output files with"
  105. )
  106. parser.add_argument("host")
  107. parser.add_argument("room")
  108. args = parser.parse_args()
  109. host = args.host
  110. room = args.room
  111. prefix = args.prefix if args.prefix else "%s_graph" % (room)
  112. pdus = get_pdus(host, room)
  113. make_graph(pdus, room, prefix)