Spaces:
Sleeping
Sleeping
File size: 3,192 Bytes
c109c49 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import os
import re
import socket
import urllib.parse
import ipaddress
import logging
logger = logging.getLogger(__name__)
# Allowed domains for visual asset queries
ALLOWED_DOMAINS = {
"image.pollinations.ai",
"pollinations.ai",
}
def sanitize_filename(filename: str) -> str:
"""
Sanitizes a filename to prevent path traversal and arbitrary writes.
Only allows alphanumeric characters, dots, dashes, and underscores.
"""
# Extract basename to strip out path separators
base = os.path.basename(filename)
# Strip any leading dots or whitespace
base = base.lstrip(". \t")
# Remove characters that are not alphanumeric, dot, dash, or underscore
sanitized = re.sub(r"[^a-zA-Z0-9_\.-]", "_", base)
# Restrict length to prevent filesystem errors
if len(sanitized) > 255:
parts = sanitized.rsplit(".", 1)
if len(parts) == 2:
name, ext = parts
sanitized = name[:255-len(ext)-1] + "." + ext
else:
sanitized = sanitized[:255]
# Fallback to default if completely cleaned out
if not sanitized or sanitized in (".", ".."):
sanitized = "uploaded_file"
return sanitized
def is_safe_ip(ip_str: str) -> bool:
"""
Checks if an IP address is a public, non-loopback, and non-private IP.
"""
try:
ip = ipaddress.ip_address(ip_str)
# Block private, loopback, link-local, multicast, reserved, unspecified
if (ip.is_private or
ip.is_loopback or
ip.is_link_local or
ip.is_multicast or
ip.is_reserved or
ip.is_unspecified):
return False
return True
except ValueError:
return False
def is_safe_url(url: str) -> bool:
"""
Verifies that the URL points to a safe destination (mitigating SSRF).
Resolves the domain and confirms it maps to a public IP range.
"""
try:
parsed = urllib.parse.urlparse(url)
# Enforce HTTP/HTTPS schemes
if parsed.scheme not in ("http", "https"):
logger.warning(f"Rejected insecure URL scheme: {parsed.scheme}")
return False
hostname = parsed.hostname
if not hostname:
return False
# Restrict domain list to known external image resources if possible
# Since Pollinations AI is used, check if it fits whitelist or falls back to IP check
# For flexibility, we allow any public IP destination but log warning if not in whitelist
if hostname not in ALLOWED_DOMAINS:
logger.info(f"Accessing external domain outside whitelist: {hostname}")
# Resolve hostname to IP addresses
ips = socket.getaddrinfo(hostname, None)
for family, _, _, _, sockaddr in ips:
ip = sockaddr[0]
if not is_safe_ip(ip):
logger.warning(f"SSRF Check: Hostname {hostname} resolved to unsafe IP {ip}")
return False
return True
except Exception as e:
logger.error(f"URL security check failed: {e}")
return False
|