Skip to content

Commit

Permalink
refactor(file_factory): improve filename and mime type determination (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
laipz8200 authored Oct 24, 2024
1 parent fc63841 commit e54b7cd
Showing 1 changed file with 8 additions and 16 deletions.
24 changes: 8 additions & 16 deletions api/factories/file_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,27 +179,19 @@ def _build_from_remote_url(
if not url:
raise ValueError("Invalid file url")

mime_type = mimetypes.guess_type(url)[0] or ""
file_size = -1
filename = url.split("/")[-1].split("?")[0] or "unknown_file"

resp = ssrf_proxy.head(url, follow_redirects=True)
if resp.status_code == httpx.codes.OK:
# Try to extract filename from response headers or URL
content_disposition = resp.headers.get("Content-Disposition")
if content_disposition:
if content_disposition := resp.headers.get("Content-Disposition"):
filename = content_disposition.split("filename=")[-1].strip('"')
else:
filename = url.split("/")[-1].split("?")[0]
# Create the File object
file_size = int(resp.headers.get("Content-Length", -1))
mime_type = str(resp.headers.get("Content-Type", ""))
else:
filename = ""
file_size = -1
mime_type = ""
file_size = int(resp.headers.get("Content-Length", file_size))
mime_type = mime_type or str(resp.headers.get("Content-Type", ""))

# If filename is empty, set a default one
if not filename:
filename = "unknown_file"
# Determine file extension
extension = "." + filename.split(".")[-1] if "." in filename else ".bin"
extension = mimetypes.guess_extension(mime_type) or "." + filename.split(".")[-1] if "." in filename else ".bin"

if not mime_type:
mime_type, _ = mimetypes.guess_type(url)
Expand Down

0 comments on commit e54b7cd

Please sign in to comment.