| 1 | import base64 |
| 2 | from io import BytesIO |
| 3 | import mimetypes |
| 4 | import os |
| 5 | from pathlib import Path |
| 6 | |
| 7 | from flask import Response |
| 8 | from helpers.api import ApiHandler, Input, Output, Request |
| 9 | from helpers import files, runtime |
| 10 | from api import file_info |
| 11 | from urllib.parse import quote |
| 12 | |
| 13 | |
| 14 | |
| 15 | def stream_file_download(file_source, download_name, chunk_size=8192): |
| 16 | """ |
| 17 | Create a streaming response for file downloads that shows progress in browser. |
| 18 | |
| 19 | Args: |
| 20 | file_source: Either a file path (str) or BytesIO object |
| 21 | download_name: Name for the downloaded file |
| 22 | chunk_size: Size of chunks to stream (default 8192 bytes) |
| 23 | |
| 24 | Returns: |
| 25 | Flask Response object with streaming content |
| 26 | """ |
| 27 | # Calculate file size for Content-Length header |
| 28 | if isinstance(file_source, str): |
| 29 | # File path - get size from filesystem |
| 30 | file_size = os.path.getsize(file_source) |
| 31 | elif isinstance(file_source, BytesIO): |
| 32 | # BytesIO object - get size from buffer |
| 33 | current_pos = file_source.tell() |
| 34 | file_source.seek(0, 2) # Seek to end |
| 35 | file_size = file_source.tell() |
| 36 | file_source.seek(current_pos) # Restore original position |
| 37 | else: |
| 38 | raise ValueError(f"Unsupported file source type: {type(file_source)}") |
| 39 | |
| 40 | def generate(): |
| 41 | if isinstance(file_source, str): |
| 42 | # File path - open and stream from disk |
| 43 | with open(file_source, 'rb') as f: |
| 44 | while True: |
| 45 | chunk = f.read(chunk_size) |
| 46 | if not chunk: |
| 47 | break |
| 48 | yield chunk |
| 49 | elif isinstance(file_source, BytesIO): |
| 50 | # BytesIO object - stream from memory |
| 51 | file_source.seek(0) # Ensure we're at the beginning |
| 52 | while True: |
| 53 | chunk = file_source.read(chunk_size) |
| 54 | if not chunk: |
| 55 | break |
| 56 | yield chunk |
| 57 | |
| 58 | # Detect content type based on file extension |
| 59 | content_type, _ = mimetypes.guess_type(download_name) |
| 60 | if not content_type: |
| 61 | content_type = 'application/octet-stream' |
| 62 | |
| 63 | # Create streaming response with proper headers for immediate streaming |
| 64 | response = Response( |
| 65 | generate(), |
| 66 | content_type=content_type, |
| 67 | direct_passthrough=True, # Prevent Flask from buffering the response |
| 68 | headers={ |
| 69 | 'Content-Disposition': make_disposition(download_name), |
| 70 | 'Content-Length': str(file_size), # Critical for browser progress bars |
| 71 | 'Cache-Control': 'no-cache', |
| 72 | 'X-Accel-Buffering': 'no', # Disable nginx buffering |
| 73 | 'Accept-Ranges': 'bytes' # Allow browser to resume downloads |
| 74 | } |
| 75 | ) |
| 76 | |
| 77 | return response |
| 78 | |
| 79 | |
| 80 | def make_disposition(download_name: str) -> str: |
| 81 | # Basic ASCII fallback (strip or replace weird chars) |
| 82 | ascii_fallback = download_name.encode("ascii", "ignore").decode("ascii") or "download" |
| 83 | utf8_name = quote(download_name) # URL-encode UTF-8 bytes |
| 84 | |
| 85 | # RFC 5987: filename* with UTF-8 |
| 86 | return f'attachment; filename="{ascii_fallback}"; filename*=UTF-8\'\'{utf8_name}' |
| 87 | |
| 88 | |
| 89 | def resolve_download_path(path: str) -> str: |
| 90 | """Resolve a requested download path from the File Browser root.""" |
| 91 | base_dir = Path("/") |
| 92 | candidate = Path(path) |
| 93 | |
| 94 | if candidate.is_absolute(): |
| 95 | resolved = candidate.resolve() |
| 96 | else: |
| 97 | resolved = (base_dir / candidate).resolve() |
| 98 | |
| 99 | try: |
| 100 | resolved.relative_to(base_dir) |
| 101 | except ValueError as exc: |
| 102 | raise ValueError("Invalid file path") from exc |
| 103 | |
| 104 | return str(resolved) |
| 105 | |
| 106 | |
| 107 | class DownloadFile(ApiHandler): |
| 108 | |
| 109 | @classmethod |
| 110 | def get_methods(cls): |
| 111 | return ["GET"] |
| 112 | |
| 113 | async def process(self, input: Input, request: Request) -> Output: |
| 114 | file_path = request.args.get("path", input.get("path", "")) |
| 115 | if not file_path: |
| 116 | raise ValueError("No file path provided") |
| 117 | if not file_path.startswith("/"): |
| 118 | file_path = f"/{file_path}" |
| 119 | |
| 120 | try: |
| 121 | file_path = await runtime.call_development_function( |
| 122 | resolve_download_path, file_path |
| 123 | ) |
| 124 | except ValueError as exc: |
| 125 | return Response(str(exc), status=400) |
| 126 | |
| 127 | file = await runtime.call_development_function( |
| 128 | file_info.get_file_info, file_path |
| 129 | ) |
| 130 | |
| 131 | if not file["exists"]: |
| 132 | raise Exception(f"File {file_path} not found") |
| 133 | |
| 134 | if file["is_dir"]: |
| 135 | zip_file = await runtime.call_development_function(files.zip_dir, file["abs_path"]) |
| 136 | directory_name = os.path.basename(file_path.rstrip("/")) or "directory" |
| 137 | download_name = f"{directory_name}.zip" |
| 138 | if runtime.is_development(): |
| 139 | b64 = await runtime.call_development_function(fetch_file, zip_file) |
| 140 | file_data = BytesIO(base64.b64decode(b64)) |
| 141 | return stream_file_download( |
| 142 | file_data, |
| 143 | download_name=download_name |
| 144 | ) |
| 145 | else: |
| 146 | return stream_file_download( |
| 147 | zip_file, |
| 148 | download_name=download_name |
| 149 | ) |
| 150 | elif file["is_file"]: |
| 151 | if runtime.is_development(): |
| 152 | b64 = await runtime.call_development_function(fetch_file, file["abs_path"]) |
| 153 | file_data = BytesIO(base64.b64decode(b64)) |
| 154 | return stream_file_download( |
| 155 | file_data, |
| 156 | download_name=os.path.basename(file_path) |
| 157 | ) |
| 158 | else: |
| 159 | return stream_file_download( |
| 160 | file["abs_path"], |
| 161 | download_name=os.path.basename(file["file_name"]) |
| 162 | ) |
| 163 | raise Exception(f"File {file_path} not found") |
| 164 | |
| 165 | |
| 166 | async def fetch_file(path): |
| 167 | with open(path, "rb") as file: |
| 168 | file_content = file.read() |
| 169 | return base64.b64encode(file_content).decode("utf-8") |