Underground-Digital/Workflow-Engine

1import urllib.parse2 3from flask_restful import marshal_with, reqparse4 5from controllers.common import helpers6from controllers.web.wraps import WebApiResource7from core.file import helpers as file_helpers8from core.helper import ssrf_proxy9from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields10from services.file_service import FileService11 12 13class RemoteFileInfoApi(WebApiResource):14    @marshal_with(remote_file_info_fields)15    def get(self, url):16        decoded_url = urllib.parse.unquote(url)17        try:18            response = ssrf_proxy.head(decoded_url)19            return {20                "file_type": response.headers.get("Content-Type", "application/octet-stream"),21                "file_length": int(response.headers.get("Content-Length", -1)),22            }23        except Exception as e:24            return {"error": str(e)}, 40025 26 27class RemoteFileUploadApi(WebApiResource):28    @marshal_with(file_fields_with_signed_url)29    def post(self, app_model, end_user):  # Add app_model and end_user parameters30        parser = reqparse.RequestParser()31        parser.add_argument("url", type=str, required=True, help="URL is required")32        args = parser.parse_args()33 34        url = args["url"]35 36        response = ssrf_proxy.head(url)37        response.raise_for_status()38 39        file_info = helpers.guess_file_info_from_response(response)40 41        if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):42            return {"error": "File size exceeded"}, 40043 44        response = ssrf_proxy.get(url)45        response.raise_for_status()46        content = response.content47 48        try:49            upload_file = FileService.upload_file(50                filename=file_info.filename,51                content=content,52                mimetype=file_info.mimetype,53                user=end_user,  # Use end_user instead of current_user54                source_url=url,55            )56        except Exception as e:57            return {"error": str(e)}, 40058 59        return {60            "id": upload_file.id,61            "name": upload_file.name,62            "size": upload_file.size,63            "extension": upload_file.extension,64            "url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),65            "mime_type": upload_file.mime_type,66            "created_by": upload_file.created_by,67            "created_at": upload_file.created_at,68        }, 20169