Files
dify/api/controllers/files/storage_download.py

57 lines
1.8 KiB
Python

from urllib.parse import quote, unquote
from flask import Response, request
from flask_restx import Resource
from pydantic import BaseModel, Field
from werkzeug.exceptions import Forbidden, NotFound
from controllers.files import files_ns
from extensions.ext_storage import storage
from extensions.storage.file_presign_storage import FilePresignStorage
DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}"
class StorageDownloadQuery(BaseModel):
timestamp: str = Field(..., description="Unix timestamp used in the signature")
nonce: str = Field(..., description="Random string for signature")
sign: str = Field(..., description="HMAC signature")
files_ns.schema_model(
StorageDownloadQuery.__name__,
StorageDownloadQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0),
)
@files_ns.route("/storage/<path:filename>/download")
class StorageFileDownloadApi(Resource):
def get(self, filename: str):
filename = unquote(filename)
args = StorageDownloadQuery.model_validate(request.args.to_dict(flat=True))
if not FilePresignStorage.verify_signature(
filename=filename,
timestamp=args.timestamp,
nonce=args.nonce,
sign=args.sign,
):
raise Forbidden("Invalid or expired download link")
try:
generator = storage.load_stream(filename)
except FileNotFoundError:
raise NotFound("File not found")
encoded_filename = quote(filename.split("/")[-1])
return Response(
generator,
mimetype="application/octet-stream",
direct_passthrough=True,
headers={
"Content-Disposition": f"attachment; filename*=UTF-8''{encoded_filename}",
},
)