139 lines
5.3 KiB
Python
139 lines
5.3 KiB
Python
"""
|
||
Download Blueprint
|
||
Handles file download endpoints.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from flask import Blueprint, send_file, current_app
|
||
|
||
from ..shared import _format_response, log_performance, logger
|
||
from ..auth import require_api_key
|
||
|
||
|
||
def _mark_task_downloaded(task_id):
|
||
"""Mark task as downloaded and schedule deletion in database."""
|
||
import sqlite3
|
||
from pathlib import Path
|
||
|
||
# Use independent database connection (not from flask.g which may be closed)
|
||
from ..db import get_db_path as get_config_db_path
|
||
db_path = get_config_db_path(current_app)
|
||
|
||
# Get cleanup age for successful tasks from config (in seconds)
|
||
successful_task_cleanup_age = current_app.config.get('SUCCESSFUL_TASK_CLEANUP_AGE', 3600)
|
||
|
||
try:
|
||
conn = sqlite3.connect(str(db_path), check_same_thread=False)
|
||
conn.execute("PRAGMA foreign_keys=ON")
|
||
conn.execute("PRAGMA busy_timeout=3000")
|
||
|
||
# Update downloaded timestamp and set deletion time based on config
|
||
conn.execute("""
|
||
UPDATE tasks
|
||
SET downloaded_at = datetime('now', '+8 hours'),
|
||
delete_after_at = datetime('now', '+8 hours', '+' || ? || ' seconds')
|
||
WHERE task_id = ?
|
||
""", (successful_task_cleanup_age, task_id))
|
||
|
||
conn.commit()
|
||
logger.info(f"Task {task_id} marked as downloaded, scheduled for deletion in {successful_task_cleanup_age} seconds")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to mark task {task_id} as downloaded: {str(e)}", exc_info=True)
|
||
finally:
|
||
if 'conn' in locals():
|
||
conn.close()
|
||
|
||
# Create blueprint
|
||
download_bp = Blueprint('download', __name__, url_prefix='/download')
|
||
|
||
|
||
@download_bp.route('/<path:filename>')
|
||
@log_performance
|
||
def download_file(filename):
|
||
"""Download a processed file."""
|
||
from flask import request
|
||
|
||
logger.info(f"Download request for file: {filename} from IP {request.remote_addr}")
|
||
|
||
# Check API key from header or query parameter
|
||
provided_key = request.headers.get('X-API-Key') or request.args.get('api_key')
|
||
if not provided_key:
|
||
logger.warning(f"API key missing from request: POST /download/{filename}")
|
||
return _format_response(401, "API key required")
|
||
|
||
# Validate API key
|
||
from ..auth import verify_api_key, get_db
|
||
db = get_db()
|
||
row = db.execute(
|
||
"SELECT key_hash, salt, revoked FROM api_keys WHERE key_id = ?",
|
||
(provided_key,)
|
||
).fetchone()
|
||
|
||
if not row:
|
||
logger.warning(f"Unknown API key used: POST /download/{filename}")
|
||
return _format_response(401, "Invalid API key")
|
||
|
||
if row['revoked']:
|
||
logger.warning(f"Revoked API key used: POST /download/{filename}")
|
||
return _format_response(401, "API key revoked")
|
||
|
||
if not verify_api_key(provided_key, row['key_hash'], row['salt']):
|
||
logger.warning(f"Invalid API key hash: POST /download/{filename}")
|
||
return _format_response(401, "Invalid API key")
|
||
|
||
try:
|
||
# 支持两种路径格式：
|
||
# 1. 绝对路径（以 / 开头，如 /full/path/to/file）
|
||
# 2. 相对路径（task_id/filename）
|
||
if filename.startswith('/'):
|
||
# 绝对路径 - 直接使用
|
||
file_path = Path(filename)
|
||
task_id = None # Can't determine task_id from absolute path
|
||
else:
|
||
# 相对路径 - 相对于 OUTPUT_FOLDER
|
||
output_folder = Path(current_app.config.get('OUTPUT_FOLDER') or '')
|
||
if not output_folder:
|
||
logger.error("OUTPUT_FOLDER not configured")
|
||
return _format_response(500, "服务器配置错误")
|
||
|
||
# 解析 task_id/filename 格式
|
||
parts = filename.split('/', 1)
|
||
if len(parts) != 2:
|
||
logger.warning(f"Invalid relative path format: {filename}")
|
||
return _format_response(400, "无效的文件路径")
|
||
|
||
task_id, filename_part = parts
|
||
file_path = output_folder / task_id / filename_part
|
||
|
||
# Security check - ensure file is within output folder
|
||
file_path = file_path.resolve()
|
||
output_folder = Path(current_app.config.get('OUTPUT_FOLDER') or '').resolve()
|
||
if output_folder and not str(file_path).startswith(str(output_folder)):
|
||
logger.warning(f"Security violation: Attempted to access file outside output folder: {filename}")
|
||
return _format_response(403, "访问被拒绝")
|
||
|
||
if not file_path.exists():
|
||
logger.warning(f"File not found: {filename}")
|
||
return _format_response(404, "文件未找到")
|
||
|
||
if not file_path.is_file():
|
||
logger.warning(f"Path is not a file: {filename}")
|
||
return _format_response(400, "不是文件")
|
||
|
||
file_size = file_path.stat().st_size
|
||
logger.info(f"Serving file: {filename} ({file_size} bytes)")
|
||
|
||
# Mark download immediately before sending file
|
||
if task_id:
|
||
try:
|
||
_mark_task_downloaded(task_id)
|
||
except Exception as e:
|
||
logger.error(f"Failed to mark download for task {task_id}: {str(e)}")
|
||
|
||
response = send_file(file_path)
|
||
return response
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error serving file {filename}: {str(e)}", exc_info=True)
|
||
return _format_response(500, "内部服务器错误") |