Add package and file download button

inspector/main.py

@@ -1,7 +1,12 @@
 import os
 import urllib.parse
 
-import gunicorn.http.errors
+try:
+    import gunicorn.http.errors
+    GUNICORN_AVAILABLE = True
+except ImportError:
+    GUNICORN_AVAILABLE = False
+
 import sentry_sdk
 
 from flask import Flask, Response, abort, redirect, render_template, request, url_for
@@ -42,9 +47,23 @@ def traces_sampler(sampling_context):
 app.jinja_env.lstrip_blocks = True
 
 
-@app.errorhandler(gunicorn.http.errors.ParseException)
-def handle_bad_request(e):
-    return abort(400)
+if GUNICORN_AVAILABLE:
+    @app.errorhandler(gunicorn.http.errors.ParseException)
+    def handle_bad_request(e):
+        return abort(400)
+
+
+def _get_file_from_dist(project_name, first, second, rest, distname, filepath):
+    if project_name != canonicalize_name(project_name):
+        return None, None, True
+
+    dist = _get_dist(first, second, rest, distname)
+
+    if not dist:
+        return None, abort(404), False
+
+    contents = dist.contents(filepath)
+    return contents, dist, False
 
 
 @app.route("/")
@@ -162,9 +181,11 @@ def distribution(project_name, version, first, second, rest, distname):
         file_urls = [
             "./" + urllib.parse.quote(filename) for filename in dist.namelist()
         ]
+        download_url = f"/project/{project_name}/{version}/packages/{first}/{second}/{rest}/{distname}/download"
         return render_template(
             "links.html",
             links=file_urls,
+            download_url=download_url,
             h2=f"{project_name}",
             h2_link=f"/project/{project_name}",
             h2_paren=h2_paren,
@@ -184,7 +205,16 @@ def distribution(project_name, version, first, second, rest, distname):
     "/project/<project_name>/<version>/packages/<first>/<second>/<rest>/<distname>/<path:filepath>"  # noqa
 )
 def file(project_name, version, first, second, rest, distname, filepath):
-    if project_name != canonicalize_name(project_name):
+    try:
+        contents, dist, should_redirect = _get_file_from_dist(
+            project_name, first, second, rest, distname, filepath
+        )
+    except FileNotFoundError:
+        return abort(404)
+    except InspectorError:
+        return abort(400)
+
+    if should_redirect:
         return redirect(
             url_for(
                 "file",
@@ -198,6 +228,9 @@ def file(project_name, version, first, second, rest, distname, filepath):
             ),
             301,
         )
+
+    if contents is None:
+        return dist
 
     h2_paren = "View this project on PyPI"
     resp = requests_session().get(f"https://pypi.org/pypi/{project_name}/json")
@@ -210,57 +243,49 @@ def file(project_name, version, first, second, rest, distname, filepath):
     )
     if resp.status_code == 404:
         h3_paren = "❌ Release no longer on PyPI"
+
+    file_extension = filepath.split(".")[-1]
+    report_link = pypi_report_form(project_name, version, filepath, request.url)
+    download_url = f"/project/{project_name}/{version}/packages/{first}/{second}/{rest}/{distname}/{filepath}/download"
+
+    details = [detail.html() for detail in basic_details(dist, filepath)]
+    common_params = {
+        "file_details": details,
+        "mailto_report_link": report_link,
+        "download_url": download_url,
+        "h2": f"{project_name}",
+        "h2_link": f"/project/{project_name}",
+        "h2_paren": h2_paren,
+        "h2_paren_link": f"https://pypi.org/project/{project_name}",
+        "h3": f"{project_name}=={version}",
+        "h3_link": f"/project/{project_name}/{version}",
+        "h3_paren": h3_paren,
+        "h3_paren_link": f"https://pypi.org/project/{project_name}/{version}",
+        "h4": distname,
+        "h4_link": f"/project/{project_name}/{version}/packages/{first}/{second}/{rest}/{distname}/",  # noqa
+        "h5": filepath,
+        "h5_link": f"/project/{project_name}/{version}/packages/{first}/{second}/{rest}/{distname}/{filepath}",  # noqa
+    }
 
-    dist = _get_dist(first, second, rest, distname)
-    if dist:
+    if file_extension in ["pyc", "pyo"]:
+        disassembly = disassemble(contents)
+        decompilation = decompile(contents)
+        return render_template(
+            "disasm.html",
+            disassembly=disassembly,
+            decompilation=decompilation,
+            **common_params,
+        )
+
+    if isinstance(contents, bytes):
         try:
-            contents = dist.contents(filepath)
-        except FileNotFoundError:
-            return abort(404)
-        except InspectorError:
-            return abort(400)
-        file_extension = filepath.split(".")[-1]
-        report_link = pypi_report_form(project_name, version, filepath, request.url)
-
-        details = [detail.html() for detail in basic_details(dist, filepath)]
-        common_params = {
-            "file_details": details,
-            "mailto_report_link": report_link,
-            "h2": f"{project_name}",
-            "h2_link": f"/project/{project_name}",
-            "h2_paren": h2_paren,
-            "h2_paren_link": f"https://pypi.org/project/{project_name}",
-            "h3": f"{project_name}=={version}",
-            "h3_link": f"/project/{project_name}/{version}",
-            "h3_paren": h3_paren,
-            "h3_paren_link": f"https://pypi.org/project/{project_name}/{version}",
-            "h4": distname,
-            "h4_link": f"/project/{project_name}/{version}/packages/{first}/{second}/{rest}/{distname}/",  # noqa
-            "h5": filepath,
-            "h5_link": f"/project/{project_name}/{version}/packages/{first}/{second}/{rest}/{distname}/{filepath}",  # noqa
-        }
-
-        if file_extension in ["pyc", "pyo"]:
-            disassembly = disassemble(contents)
-            decompilation = decompile(contents)
-            return render_template(
-                "disasm.html",
-                disassembly=disassembly,
-                decompilation=decompilation,
-                **common_params,
-            )
-
-        if isinstance(contents, bytes):
-            try:
-                contents = contents.decode()
-            except UnicodeDecodeError:
-                return "Binary files are not supported."
+            contents = contents.decode()
+        except UnicodeDecodeError:
+            return "Binary files are not supported."
 
-        return render_template(
-            "code.html", code=contents, name=file_extension, **common_params
-        )  # noqa
-    else:
-        return "Distribution type not supported"
+    return render_template(
+        "code.html", code=contents, name=file_extension, **common_params
+    )  # noqa
 
 
 @app.route("/_health/")
@@ -271,3 +296,77 @@ def health():
 @app.route("/robots.txt")
 def robots():
     return Response("User-agent: *\nDisallow: /", mimetype="text/plain")
+
+
+@app.route(
+    "/project/<project_name>/<version>/packages/<first>/<second>/<rest>/<distname>/download"
+)
+def download_distribution(project_name, version, first, second, rest, distname):
+    if project_name != canonicalize_name(project_name):
+        return redirect(
+            url_for(
+                "download_distribution",
+                project_name=canonicalize_name(project_name),
+                version=version,
+                first=first,
+                second=second,
+                rest=rest,
+                distname=distname,
+            ),
+            301,
+        )
+
+    try:
+        dist = _get_dist(first, second, rest, distname)
+    except InspectorError:
+        return abort(400)
+
+    if not dist:
+        return abort(404)
+
+    url = f"https://files.pythonhosted.org/packages/{first}/{second}/{rest}/{distname}"
+    return redirect(url, 307)
+
+
+@app.route(
+    "/project/<project_name>/<version>/packages/<first>/<second>/<rest>/<distname>/<path:filepath>/download"
+)
+def download_file(project_name, version, first, second, rest, distname, filepath):
+    try:
+        contents, dist, should_redirect = _get_file_from_dist(
+            project_name, first, second, rest, distname, filepath
+        )
+    except FileNotFoundError:
+        return abort(404)
+    except InspectorError:
+        return abort(400)
+
+    if should_redirect:
+        return redirect(
+            url_for(
+                "download_file",
+                project_name=canonicalize_name(project_name),
+                version=version,
+                first=first,
+                second=second,
+                rest=rest,
+                distname=distname,
+                filepath=filepath,
+            ),
+            301,
+        )
+
+    if contents is None:
+        return dist
+
+    filename = filepath.split("/")[-1]
+
+    return Response(
+        contents,
+        mimetype="application/octet-stream",
+        headers={
+            "Content-Disposition": f"attachment; filename={filename}",
+            "Content-Length": str(len(contents)),
+        },
+    )
+

inspector/static/style.css

@@ -15,6 +15,30 @@
     color: red;
 }
 
+.download-anchor {
+    display: inline;
+    color: #0066cc;
+    text-decoration: none;
+}
+
+.download-anchor:hover {
+    text-decoration: underline;
+}
+
+.download-button {
+    display: inline-block;
+    padding: 10px 20px;
+    background-color: #0066cc;
+    color: white;
+    text-decoration: none;
+    border-radius: 5px;
+    font-weight: bold;
+}
+
+.download-button:hover {
+    background-color: #0052a3;
+}
+
 table, th, td {
     border: 1px solid black;
     border-collapse: collapse;

inspector/templates/code.html

@@ -17,7 +17,11 @@
 {% endblock %}
 
 {% block body %}
-<a href="{{ mailto_report_link }}" class="report-anchor"> <strong>Report Malicious Package</strong> </a>
+<div style="margin-bottom: 15px;">
+    <a href="{{ mailto_report_link }}" class="report-anchor"> <strong>Report Malicious Package</strong> </a>
+    <span style="margin: 0 10px;">|</span>
+    <a href="{{ download_url }}" class="download-anchor" download> <strong>⬇️ Download File</strong> </a>
+</div>
 <pre id="line" class="line-numbers linkable-line-numbers language-{{ name }}">
 {# Indenting the below <code> tag will cause rendering issues! #}
 <code class="language-{{ name }}">{{- code }}</code>

inspector/templates/disasm.html

@@ -17,8 +17,11 @@
 {% endblock %}
 
 {% block body %}
-    <a href="{{ mailto_report_link }}" style="color:red"> <strong>Report Malicious Package</strong> </a>
-    <br>
+    <div style="margin-bottom: 15px;">
+        <a href="{{ mailto_report_link }}" style="color:red"> <strong>Report Malicious Package</strong> </a>
+        <span style="margin: 0 10px;">|</span>
+        <a href="{{ download_url }}" class="download-anchor" download> <strong>⬇️ Download File</strong> </a>
+    </div>
     <br>
     <div>
 

inspector/templates/links.html

@@ -1,6 +1,13 @@
 {% extends 'base.html' %}
 
 {% block body %}
+{% if download_url %}
+<div style="margin-bottom: 20px;">
+  <a href="{{ download_url }}" class="download-button" download>
+    <strong>⬇️ Download Package</strong>
+  </a>
+</div>
+{% endif %}
 <ul>
 {% for link in links %}
   <li><a href="{{ link }}">{{ link|unquote }}</a></li>

requirements/main.txt

@@ -24,6 +24,9 @@ click==8.1.3 \
     --hash=sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e \
     --hash=sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48
     # via flask
+colorama==0.4.6 \
+    --hash=sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6
+    # via click
 flask==3.1.0 \
     --hash=sha256:5f873c5184c897c8d9d1b05df1e3d01b14910ce69607a117bd3277098a5836ac \
     --hash=sha256:d667207822eb83f1c4b50949b1623c8fc8d51f2341d65f72e1a1815397551136