diff --git a/api/controllers/files/image_preview.py b/api/controllers/files/image_preview.py index 4b2d61e7c3..6b3ac93cdf 100644 --- a/api/controllers/files/image_preview.py +++ b/api/controllers/files/image_preview.py @@ -1,5 +1,5 @@ from flask import Response, request -from flask_restful import Resource +from flask_restful import Resource, reqparse from werkzeug.exceptions import NotFound import services @@ -41,24 +41,39 @@ class FilePreviewApi(Resource): def get(self, file_id): file_id = str(file_id) - timestamp = request.args.get("timestamp") - nonce = request.args.get("nonce") - sign = request.args.get("sign") + parser = reqparse.RequestParser() + parser.add_argument("timestamp", type=str, required=True, location="args") + parser.add_argument("nonce", type=str, required=True, location="args") + parser.add_argument("sign", type=str, required=True, location="args") + parser.add_argument("as_attachment", type=bool, required=False, default=False, location="args") - if not timestamp or not nonce or not sign: + args = parser.parse_args() + + if not args["timestamp"] or not args["nonce"] or not args["sign"]: return {"content": "Invalid request."}, 400 try: - generator, mimetype = FileService.get_signed_file_preview( + generator, upload_file = FileService.get_file_generator_by_file_id( file_id=file_id, - timestamp=timestamp, - nonce=nonce, - sign=sign, + timestamp=args["timestamp"], + nonce=args["nonce"], + sign=args["sign"], ) except services.errors.file.UnsupportedFileTypeError: raise UnsupportedFileTypeError() - return Response(generator, mimetype=mimetype) + response = Response( + generator, + mimetype=upload_file.mime_type, + direct_passthrough=True, + headers={}, + ) + if upload_file.size > 0: + response.headers["Content-Length"] = str(upload_file.size) + if args["as_attachment"]: + response.headers["Content-Disposition"] = f"attachment; filename={upload_file.name}" + + return response class WorkspaceWebappLogoApi(Resource): diff --git a/api/services/file_service.py b/api/services/file_service.py index 22ea923f6b..6193a39669 100644 --- a/api/services/file_service.py +++ b/api/services/file_service.py @@ -1,7 +1,6 @@ import datetime import hashlib import uuid -from collections.abc import Generator from typing import Literal, Union from flask_login import current_user @@ -132,7 +131,7 @@ class FileService: return upload_file @staticmethod - def get_file_preview(file_id: str) -> str: + def get_file_preview(file_id: str): upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first() if not upload_file: @@ -171,7 +170,7 @@ class FileService: return generator, upload_file.mime_type @staticmethod - def get_signed_file_preview(file_id: str, timestamp: str, nonce: str, sign: str): + def get_file_generator_by_file_id(file_id: str, timestamp: str, nonce: str, sign: str): result = file_helpers.verify_file_signature(upload_file_id=file_id, timestamp=timestamp, nonce=nonce, sign=sign) if not result: raise NotFound("File not found or signature is invalid") @@ -183,10 +182,10 @@ class FileService: generator = storage.load(upload_file.key, stream=True) - return generator, upload_file.mime_type + return generator, upload_file @staticmethod - def get_public_image_preview(file_id: str) -> tuple[Generator, str]: + def get_public_image_preview(file_id: str): upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first() if not upload_file: