graph.py 4.3 KB

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