torrent: send jwt in announce url

util/registry/torrent.py

@@ -1,8 +1,15 @@
 import time
 import hashlib
+import urllib
 
 import bencode
 import resumablehashlib
+import jwt
+
+from cryptography.x509 import load_pem_x509_certificate
+from cryptography.hazmat.backends import default_backend
+
+from cachetools import lru_cache
 
 from app import app
 
@@ -10,6 +17,55 @@ from app import app
 ANNOUNCE_URL = app.config.get('TORRENT_ANNOUNCE_URL')
 NAMING_SALT = app.config.get('TORRENT_NAMING_SALT')
 REGISTRY_TITLE = app.config.get('REGISTRY_TITLE')
+JWT_ISSUER = app.config.get('JWT_AUTH_TOKEN_ISSUER')
+
+
+@lru_cache(maxsize=1)
+def _load_certificate_bytes(certificate_file_path):
+  with open(certificate_file_path) as cert_file:
+    return load_pem_x509_certificate(cert_file.read(), default_backend()).public_key()
+
+@lru_cache(maxsize=1)
+def _load_private_key(private_key_file_path):
+  with open(private_key_file_path) as private_key_file:
+    return private_key_file.read()
+
+def _torrent_jwt(info_dict):
+  token_data = {
+    'iss': JWT_ISSUER,
+    'aud': ANNOUNCE_URL,
+    'infohash': _infohash(info_dict),
+  }
+
+  certificate = _load_certificate_bytes(app.config['JWT_AUTH_CERTIFICATE_PATH'])
+  token_headers = {
+    'x5c': [certificate],
+  }
+
+  private_key = _load_private_key(app.config['JWT_AUTH_PRIVATE_KEY_PATH'])
+  return jwt.encode(token_data, private_key, 'RS256', headers=token_headers)
+
+def _infohash(infodict):
+  digest = hashlib.sha1()
+  digest.update(bencode.bencode(infodict))
+  return urllib.quote(digest.digest())
+
+def make_torrent(name, webseed, length, piece_length, pieces):
+  info_dict = {
+    'name': name,
+    'length': length,
+    'piece length': piece_length,
+    'pieces': pieces,
+  }
+
+  return bencode.bencode({
+    'announce': ANNOUNCE_URL + "?jwt=" + _torrent_jwt(info_dict),
+    'url-list': webseed,
+    'encoding': 'UTF-8',
+    'created by': REGISTRY_TITLE,
+    'creation date': int(time.time()),
+    'info': info_dict,
+  })
 
 def public_torrent_filename(blob_uuid):
   return hashlib.sha256(blob_uuid).hexdigest()
@@ -18,22 +74,6 @@ def per_user_torrent_filename(user_uuid, blob_uuid):
   return hashlib.sha256(blob_uuid + user_uuid + NAMING_SALT).hexdigest()
 
 
-def make_torrent(name, webseed, length, piece_length, pieces):
-  return bencode.bencode({
-    'announce': ANNOUNCE_URL,
-    'url-list': webseed,
-    'encoding': 'UTF-8',
-    'created by': REGISTRY_TITLE,
-    'creation date': int(time.time()),
-    'info': {
-      'name': name,
-      'length': length,
-      'piece length': piece_length,
-      'pieces': pieces,
-    },
-  })
-
-
 class PieceHasher(object):
   def __init__(self, piece_size, starting_offset=0, starting_piece_hash_bytes='',
                hash_fragment_to_resume=None):