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,10 @@ 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)
 
 
 @app.route("/")
@@ -162,9 +168,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,
@@ -221,11 +229,13 @@ def file(project_name, version, first, second, rest, distname, filepath):
             return abort(400)
         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,
@@ -271,3 +281,99 @@ 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):
+    """Download entire distribution package."""
+    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)
+
+    # Get the raw distribution file from PyPI
+    url = f"https://files.pythonhosted.org/packages/{first}/{second}/{rest}/{distname}"
+    try:
+        resp = requests_session().get(url, stream=True)
+        resp.raise_for_status()
+    except Exception:
+        return abort(404)
+
+    # Return the file with proper headers to trigger download
+    return Response(
+        resp.content,
+        mimetype="application/octet-stream",
+        headers={
+            "Content-Disposition": f"attachment; filename={distname}",
+            "Content-Length": str(len(resp.content)),
+        },
+    )
+
+
+@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):
+    """Download individual file from a distribution."""
+    if project_name != canonicalize_name(project_name):
+        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,
+        )
+
+    try:
+        dist = _get_dist(first, second, rest, distname)
+    except InspectorError:
+        return abort(400)
+
+    if not dist:
+        return abort(404)
+
+    try:
+        contents = dist.contents(filepath)
+    except FileNotFoundError:
+        return abort(404)
+    except InspectorError:
+        return abort(400)
+
+    # Extract just the filename from the path
+    filename = filepath.split("/")[-1]
+
+    # Return the file with proper headers to trigger download
+    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