federation_client.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  1. import nacl.signing
  2. import json
  3. import base64
  4. import requests
  5. import sys
  6. import srvlookup
  7. def encode_base64(input_bytes):
  8. """Encode bytes as a base64 string without any padding."""
  9. input_len = len(input_bytes)
  10. output_len = 4 * ((input_len + 2) // 3) + (input_len + 2) % 3 - 2
  11. output_bytes = base64.b64encode(input_bytes)
  12. output_string = output_bytes[:output_len].decode("ascii")
  13. return output_string
  14. def decode_base64(input_string):
  15. """Decode a base64 string to bytes inferring padding from the length of the
  16. string."""
  17. input_bytes = input_string.encode("ascii")
  18. input_len = len(input_bytes)
  19. padding = b"=" * (3 - ((input_len + 3) % 4))
  20. output_len = 3 * ((input_len + 2) // 4) + (input_len + 2) % 4 - 2
  21. output_bytes = base64.b64decode(input_bytes + padding)
  22. return output_bytes[:output_len]
  23. def encode_canonical_json(value):
  24. return json.dumps(
  25. value,
  26. # Encode code-points outside of ASCII as UTF-8 rather than \u escapes
  27. ensure_ascii=False,
  28. # Remove unecessary white space.
  29. separators=(',',':'),
  30. # Sort the keys of dictionaries.
  31. sort_keys=True,
  32. # Encode the resulting unicode as UTF-8 bytes.
  33. ).encode("UTF-8")
  34. def sign_json(json_object, signing_key, signing_name):
  35. signatures = json_object.pop("signatures", {})
  36. unsigned = json_object.pop("unsigned", None)
  37. signed = signing_key.sign(encode_canonical_json(json_object))
  38. signature_base64 = encode_base64(signed.signature)
  39. key_id = "%s:%s" % (signing_key.alg, signing_key.version)
  40. signatures.setdefault(signing_name, {})[key_id] = signature_base64
  41. json_object["signatures"] = signatures
  42. if unsigned is not None:
  43. json_object["unsigned"] = unsigned
  44. return json_object
  45. NACL_ED25519 = "ed25519"
  46. def decode_signing_key_base64(algorithm, version, key_base64):
  47. """Decode a base64 encoded signing key
  48. Args:
  49. algorithm (str): The algorithm the key is for (currently "ed25519").
  50. version (str): Identifies this key out of the keys for this entity.
  51. key_base64 (str): Base64 encoded bytes of the key.
  52. Returns:
  53. A SigningKey object.
  54. """
  55. if algorithm == NACL_ED25519:
  56. key_bytes = decode_base64(key_base64)
  57. key = nacl.signing.SigningKey(key_bytes)
  58. key.version = version
  59. key.alg = NACL_ED25519
  60. return key
  61. else:
  62. raise ValueError("Unsupported algorithm %s" % (algorithm,))
  63. def read_signing_keys(stream):
  64. """Reads a list of keys from a stream
  65. Args:
  66. stream : A stream to iterate for keys.
  67. Returns:
  68. list of SigningKey objects.
  69. """
  70. keys = []
  71. for line in stream:
  72. algorithm, version, key_base64 = line.split()
  73. keys.append(decode_signing_key_base64(algorithm, version, key_base64))
  74. return keys
  75. def lookup(destination, path):
  76. if ":" in destination:
  77. return "https://%s%s" % (destination, path)
  78. else:
  79. try:
  80. srv = srvlookup.lookup("matrix", "tcp", destination)[0]
  81. return "https://%s:%d%s" % (srv.host, srv.port, path)
  82. except:
  83. return "https://%s:%d%s" % (destination, 8448, path)
  84. def get_json(origin_name, origin_key, destination, path):
  85. request_json = {
  86. "method": "GET",
  87. "uri": path,
  88. "origin": origin_name,
  89. "destination": destination,
  90. }
  91. signed_json = sign_json(request_json, origin_key, origin_name)
  92. authorization_headers = []
  93. for key, sig in signed_json["signatures"][origin_name].items():
  94. authorization_headers.append(bytes(
  95. "X-Matrix origin=%s,key=\"%s\",sig=\"%s\"" % (
  96. origin_name, key, sig,
  97. )
  98. ))
  99. result = requests.get(
  100. lookup(destination, path),
  101. headers={"Authorization": authorization_headers[0]},
  102. verify=False,
  103. )
  104. return result.json()
  105. def main():
  106. origin_name, keyfile, destination, path = sys.argv[1:]
  107. with open(keyfile) as f:
  108. key = read_signing_keys(f)[0]
  109. result = get_json(
  110. origin_name, key, destination, "/_matrix/federation/v1/" + path
  111. )
  112. json.dump(result, sys.stdout)
  113. if __name__ == "__main__":
  114. main()