Spaces:
Running
Running
| 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 | |