Underground-Digital/Workflow-Engine

1import urllib.parse2from typing import cast3 4from flask_login import current_user5from flask_restful import Resource, marshal_with, reqparse6 7from controllers.common import helpers8from core.file import helpers as file_helpers9from core.helper import ssrf_proxy10from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields11from models.account import Account12from services.file_service import FileService13 14 15class RemoteFileInfoApi(Resource):16    @marshal_with(remote_file_info_fields)17    def get(self, url):18        decoded_url = urllib.parse.unquote(url)19        try:20            response = ssrf_proxy.head(decoded_url)21            return {22                "file_type": response.headers.get("Content-Type", "application/octet-stream"),23                "file_length": int(response.headers.get("Content-Length", 0)),24            }25        except Exception as e:26            return {"error": str(e)}, 40027 28 29class RemoteFileUploadApi(Resource):30    @marshal_with(file_fields_with_signed_url)31    def post(self):32        parser = reqparse.RequestParser()33        parser.add_argument("url", type=str, required=True, help="URL is required")34        args = parser.parse_args()35 36        url = args["url"]37 38        response = ssrf_proxy.head(url)39        response.raise_for_status()40 41        file_info = helpers.guess_file_info_from_response(response)42 43        if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):44            return {"error": "File size exceeded"}, 40045 46        response = ssrf_proxy.get(url)47        response.raise_for_status()48        content = response.content49 50        try:51            user = cast(Account, current_user)52            upload_file = FileService.upload_file(53                filename=file_info.filename,54                content=content,55                mimetype=file_info.mimetype,56                user=user,57                source_url=url,58            )59        except Exception as e:60            return {"error": str(e)}, 40061 62        return {63            "id": upload_file.id,64            "name": upload_file.name,65            "size": upload_file.size,66            "extension": upload_file.extension,67            "url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),68            "mime_type": upload_file.mime_type,69            "created_by": upload_file.created_by,70            "created_at": upload_file.created_at,71        }, 20172