From bdcc6c531b4be3f03bddd95aa51b4c7aac0ea5e8 Mon Sep 17 00:00:00 2001
From: Timothy Rogers <github@hack13.me>
Date: Mon, 26 May 2025 14:38:13 -0400
Subject: [PATCH] Addressing feature request for Content-Disposition

---
 app.py                      | 48 ++++++++++++++++++++++++++++++++++++-
 storage.py                  | 17 ++++++++++++-
 templates/asset_detail.html |  7 +++---
 templates/edit_asset.html   |  3 +--
 4 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/app.py b/app.py
index 2178f57..d223557 100644
--- a/app.py
+++ b/app.py
@@ -1,6 +1,7 @@
 import os
 import uuid
-from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
+import mimetypes
+from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, send_from_directory, send_file, Response, stream_with_context
 from werkzeug.utils import secure_filename
 from config import Config
 from flask_migrate import Migrate
@@ -263,5 +264,50 @@ def delete_asset_file(id):
         flash('Failed to delete file: ' + str(e), 'error')
         return redirect(url_for('asset_detail', id=asset_id))
 
+@app.route('/download/<int:file_id>')
+def download_file(file_id):
+    """Download a file with its original filename"""
+    try:
+        asset_file = AssetFile.query.get_or_404(file_id)
+        filename = asset_file.filename
+        download_name = asset_file.original_filename or filename
+
+        # Guess the mime type
+        mime_type, _ = mimetypes.guess_type(download_name)
+        if mime_type is None:
+            mime_type = 'application/octet-stream'
+
+        app.logger.debug(f"Starting download of {filename} as {download_name} with type {mime_type}")
+
+        try:
+            file_stream = app.storage.get_file_stream(filename)
+            
+            def generate():
+                try:
+                    while True:
+                        chunk = file_stream.read(8192)  # Read in 8KB chunks
+                        if not chunk:
+                            break
+                        yield chunk
+                finally:
+                    file_stream.close()
+
+            response = Response(
+                stream_with_context(generate()),
+                mimetype=mime_type
+            )
+            response.headers['Content-Disposition'] = f'attachment; filename="{download_name}"'
+            return response
+
+        except Exception as e:
+            app.logger.error(f"Error streaming file {filename}: {str(e)}", exc_info=True)
+            flash('Error downloading file. Please try again.', 'error')
+            return redirect(url_for('asset_detail', id=asset_file.asset_id))
+
+    except Exception as e:
+        app.logger.error(f"Error in download_file: {str(e)}", exc_info=True)
+        flash('File not found or error occurred.', 'error')
+        return redirect(url_for('index'))
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5432, debug=True)
diff --git a/storage.py b/storage.py
index 1792d39..6886564 100644
--- a/storage.py
+++ b/storage.py
@@ -172,4 +172,19 @@ class StorageBackend:
         full_path = self._get_full_path(filename)
         if self.protocol == 's3':
             return self.fs.exists(f"{self.bucket}/{full_path}")
-        return self.fs.exists(full_path) 
\ No newline at end of file
+        return self.fs.exists(full_path)
+
+    def get_file_stream(self, filename: str):
+        """Get a file stream from storage"""
+        try:
+            if self.protocol == 's3':
+                s3_path = f"{self.bucket}/{self._get_full_path(filename)}"
+                self.logger.debug(f"Opening S3 file stream: {s3_path}")
+                return self.fs.open(s3_path, 'rb')
+            else:
+                full_path = self._get_full_path(filename)
+                self.logger.debug(f"Opening local file stream: {full_path}")
+                return open(full_path, 'rb')
+        except Exception as e:
+            self.logger.error(f"Failed to get file stream for {filename}: {str(e)}", exc_info=True)
+            raise 
\ No newline at end of file
diff --git a/templates/asset_detail.html b/templates/asset_detail.html
index 6d32fe7..95d1011 100644
--- a/templates/asset_detail.html
+++ b/templates/asset_detail.html
@@ -64,15 +64,14 @@
             </div>
             {% endif %}
 
-            <div class="attached-files content-box">
-                <h2>Attached Files</h2>
+            <div class="files-section">
+                <h2>Files</h2>
                 {% if asset.files %}
                 <ul class="files-list">
                     {% for file in asset.files %}
                     <li class="file-item">
                         <a
-                            href="{{ file.file_url }}"
-                            target="_blank"
+                            href="{{ url_for('download_file', file_id=file.id) }}"
                             class="file-link"
                         >
                             <i class="fas fa-file"></i>
diff --git a/templates/edit_asset.html b/templates/edit_asset.html
index db9937f..4364416 100644
--- a/templates/edit_asset.html
+++ b/templates/edit_asset.html
@@ -109,8 +109,7 @@
             {% for file in asset.files %}
             <li class="file-item">
                 <a
-                    href="{{ file.file_url }}"
-                    target="_blank"
+                    href="{{ url_for('download_file', file_id=file.id) }}"
                     class="file-link"
                 >
                     <i class="fas fa-file"></i>