From bf090c0fab81384d2a95c60cd2be98c9a5f3f8b4 Mon Sep 17 00:00:00 2001
From: TheAMM <the.actual.amm@gmail.com>
Date: Sat, 21 Jul 2018 22:01:19 +0300
Subject: [PATCH] torrents: clean up & optimize magnets further

Removes the specialized template ES magnet creator, since create_magnet()
can use both Torrents and ES objects. Search results will get the
properly escaped magnets, now.
Slightly optimizes the tracker adding and string joins.
RIP base32, wonder how many bad clients will break with sha1.
---
 nyaa/template_utils.py             | 21 ++++-----------------
 nyaa/templates/rss.xml             |  2 +-
 nyaa/templates/search_results.html |  2 +-
 nyaa/torrents.py                   | 19 ++++++++++++++-----
 4 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/nyaa/template_utils.py b/nyaa/template_utils.py
index 3fc11fa..f7e64af 100644
--- a/nyaa/template_utils.py
+++ b/nyaa/template_utils.py
@@ -1,15 +1,13 @@
 import os.path
 import re
-from base64 import b32encode
 from datetime import datetime
 from email.utils import formatdate
-from urllib.parse import urlencode
 
 import flask
 from werkzeug.urls import url_encode
 
 from nyaa.backend import get_category_id_map
-from nyaa.torrents import get_default_trackers
+from nyaa.torrents import create_magnet
 
 app = flask.current_app
 bp = flask.Blueprint('template-utils', __name__)
@@ -20,20 +18,9 @@ _static_cache = {}  # For static_cachebuster
 
 # For processing ES links
 @bp.app_context_processor
-def create_magnet_from_es_info():
-    def _create_magnet_from_es_info(display_name, info_hash, max_trackers=5, trackers=None):
-        if trackers is None:
-            trackers = get_default_trackers()
-
-        magnet_parts = [
-            ('dn', display_name)
-        ]
-        for tracker in trackers[:max_trackers]:
-            magnet_parts.append(('tr', tracker))
-
-        b32_info_hash = b32encode(bytes.fromhex(info_hash)).decode('utf-8')
-        return 'magnet:?xt=urn:btih:' + b32_info_hash + '&' + urlencode(magnet_parts)
-    return dict(create_magnet_from_es_info=_create_magnet_from_es_info)
+def create_magnet_from_es_torrent():
+    # Since ES entries look like ducks, we can use the create_magnet as-is
+    return dict(create_magnet_from_es_torrent=create_magnet)
 
 
 # ######################### TEMPLATE GLOBALS #########################
diff --git a/nyaa/templates/rss.xml b/nyaa/templates/rss.xml
index d9c8732..0b38019 100644
--- a/nyaa/templates/rss.xml
+++ b/nyaa/templates/rss.xml
@@ -12,7 +12,7 @@
 				{% if torrent.has_torrent and not magnet_links %}
 				<link>{{ url_for('torrents.download', torrent_id=torrent.meta.id, _external=True) }}</link>
 				{% else %}
-				<link>{{ create_magnet_from_es_info(torrent.display_name, torrent.info_hash) }}</link>
+				<link>{{ create_magnet_from_es_torrent(torrent) }}</link>
 				{% endif %}
 				<guid isPermaLink="true">{{ url_for('torrents.view', torrent_id=torrent.meta.id, _external=True) }}</guid>
 				<pubDate>{{ torrent.created_time|rfc822_es }}</pubDate>
diff --git a/nyaa/templates/search_results.html b/nyaa/templates/search_results.html
index 7101456..0859e4b 100644
--- a/nyaa/templates/search_results.html
+++ b/nyaa/templates/search_results.html
@@ -86,7 +86,7 @@
 				<td class="text-center" style="white-space: nowrap;">
 					{% if torrent.has_torrent %}<a href="{{ url_for('torrents.download', torrent_id=torrent_id) }}"><i class="fa fa-fw fa-download"></i></a>{% endif %}
 					{% if use_elastic %}
-					<a href="{{ create_magnet_from_es_info(torrent.display_name, torrent.info_hash) }}"><i class="fa fa-fw fa-magnet"></i></a>
+					<a href="{{ create_magnet_from_es_torrent(torrent) }}"><i class="fa fa-fw fa-magnet"></i></a>
 					{% else %}
 					<a href="{{ torrent.magnet_uri }}"><i class="fa fa-fw fa-magnet"></i></a>
 					{% endif %}
diff --git a/nyaa/torrents.py b/nyaa/torrents.py
index b5668e9..8655366 100644
--- a/nyaa/torrents.py
+++ b/nyaa/torrents.py
@@ -1,4 +1,3 @@
-import base64
 import os
 from urllib.parse import quote, urlencode
 
@@ -84,11 +83,21 @@ def create_magnet(torrent, max_trackers=5, trackers=None):
     magnet_parts = [
         ('dn', torrent.display_name)
     ]
-    for tracker in trackers[:max_trackers]:
-        magnet_parts.append(('tr', tracker))
+    magnet_parts.extend(
+        ('tr', tracker_url)
+        for tracker_url in trackers[:max_trackers]
+    )
 
-    b32_info_hash = base64.b32encode(torrent.info_hash).decode('utf-8')
-    return 'magnet:?xt=urn:btih:' + b32_info_hash + '&' + urlencode(magnet_parts, quote_via=quote)
+    # Since we accept both models.Torrents and ES objects,
+    # we need to make sure the info_hash is a hex string
+    info_hash = torrent.info_hash
+    if isinstance(info_hash, (bytes, bytearray)):
+        info_hash = info_hash.hex()
+
+    return ''.join([
+        'magnet:?xt=urn:btih:', info_hash,
+        '&', urlencode(magnet_parts, quote_via=quote)
+    ])
 
 
 def create_default_metadata_base(torrent, trackers=None, webseeds=None):