import urllib.parse import httpx from flask_restx import Resource from pydantic import BaseModel, Field import services from controllers.common import helpers from controllers.common.errors import ( FileTooLargeError, RemoteFileUploadError, UnsupportedFileTypeError, ) from controllers.common.schema import register_schema_models from core.file import helpers as file_helpers from core.helper import ssrf_proxy from extensions.ext_database import db from fields.file_fields import FileWithSignedUrl, RemoteFileInfo from libs.login import current_account_with_tenant from services.file_service import FileService from . import console_ns register_schema_models(console_ns, RemoteFileInfo, FileWithSignedUrl) @console_ns.route("/remote-files/") class RemoteFileInfoApi(Resource): @console_ns.response(200, "Remote file info", console_ns.models[RemoteFileInfo.__name__]) def get(self, url): decoded_url = urllib.parse.unquote(url) resp = ssrf_proxy.head(decoded_url) if resp.status_code != httpx.codes.OK: # failed back to get method resp = ssrf_proxy.get(decoded_url, timeout=3) resp.raise_for_status() info = RemoteFileInfo( file_type=resp.headers.get("Content-Type", "application/octet-stream"), file_length=int(resp.headers.get("Content-Length", 0)), ) return info.model_dump(mode="json") class RemoteFileUploadPayload(BaseModel): url: str = Field(..., description="URL to fetch") console_ns.schema_model( RemoteFileUploadPayload.__name__, RemoteFileUploadPayload.model_json_schema(ref_template="#/definitions/{model}"), ) @console_ns.route("/remote-files/upload") class RemoteFileUploadApi(Resource): @console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__]) @console_ns.response(201, "Remote file uploaded", console_ns.models[FileWithSignedUrl.__name__]) def post(self): args = RemoteFileUploadPayload.model_validate(console_ns.payload) url = args.url try: resp = ssrf_proxy.head(url=url) if resp.status_code != httpx.codes.OK: resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True) if resp.status_code != httpx.codes.OK: raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}") except httpx.RequestError as e: raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}") file_info = helpers.guess_file_info_from_response(resp) if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size): raise FileTooLargeError content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content try: user, _ = current_account_with_tenant() upload_file = FileService(db.engine).upload_file( filename=file_info.filename, content=content, mimetype=file_info.mimetype, user=user, source_url=url, ) except services.errors.file.FileTooLargeError as file_too_large_error: raise FileTooLargeError(file_too_large_error.description) except services.errors.file.UnsupportedFileTypeError: raise UnsupportedFileTypeError() payload = FileWithSignedUrl( id=upload_file.id, name=upload_file.name, size=upload_file.size, extension=upload_file.extension, url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id), mime_type=upload_file.mime_type, created_by=upload_file.created_by, created_at=int(upload_file.created_at.timestamp()), ) return payload.model_dump(mode="json"), 201