Update utils.py

utils.py

@@ -1,5 +1,8 @@
 import os
 import re
+import ipaddress
+import urllib.parse
+from shlex import quote as shqt
 import gradio as gr
 from constants import (
     DIFFUSERS_FORMAT_LORAS,
@@ -20,7 +23,6 @@ from diffusers.pipelines.pipeline_loading_utils import variant_compatible_siblin
 from stablepy.diffusers_vanilla.utils import checkpoint_model_type
 from pathlib import PosixPath
 from unidecode import unidecode
-import urllib.parse
 import copy
 import requests
 from requests.adapters import HTTPAdapter
@@ -39,8 +41,42 @@ MODEL_ARCH = {
 }
 
 
+def validate_url(url: str) -> str:
+    """Validate URL protocol and block SSRF (localhost, private & cloud metadata IPs)."""
+    url = url.strip()
+    if not url:
+        raise ValueError("URL cannot be empty.")
+
+    parsed = urllib.parse.urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        raise ValueError(f"Invalid protocol '{parsed.scheme}'. Only HTTP/HTTPS are allowed.")
+
+    hostname = (parsed.hostname or "").lower()
+    if not hostname:
+        raise ValueError("Invalid URL: missing hostname.")
+
+    # SSRF Protection: Block loopback and local hosts
+    if hostname in ("localhost", "0.0.0.0", "127.0.0.1", "::1"):
+        raise ValueError("Access to local/loopback address is blocked.")
+
+    # SSRF Protection: Block private and link-local / cloud metadata IPs (e.g. 169.254.169.254)
+    try:
+        ip = ipaddress.ip_address(hostname)
+        if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved:
+            raise ValueError(f"Access to internal IP ({hostname}) is blocked.")
+    except ValueError:
+        pass  # Standard domain name
+
+    # Normalize Civitai domain alias
+    if hostname == "civitai.red" or hostname.endswith(".civitai.red"):
+        url = url.replace("civitai.red", "civitai.com")
+
+    return url
+
+
 def read_safetensors_header_from_url(url: str):
     """Read safetensors header from a remote Hugging Face file."""
+    url = validate_url(url)
     meta = get_hf_file_metadata(url)
 
     # Step 1: first 8 bytes → header length
@@ -248,7 +284,7 @@ def civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
     curl_command = (
         f'curl -L -sI --connect-timeout 5 --max-time 5 '
         f'-H "Content-Type: application/json" '
-        f'-H "Authorization: Bearer {civitai_api_key}" "{url}"'
+        f'-H {shqt(f"Authorization: Bearer {civitai_api_key}")} {shqt(url)}'
     )
 
     headers = os.popen(curl_command).read()
@@ -277,7 +313,7 @@ def civ_redirect_down(url, dir_, civitai_api_key, romanize, alternative_name):
 
     wget_command = (
         f'wget -c -nv '
-        f'-O "{os.path.join(dir_, filename_base)}" "{redirect_url}"'
+        f'-O {shqt(os.path.join(dir_, filename_base))} {shqt(redirect_url)}'
     )
     r_code = os.system(wget_command)  # noqa
 
@@ -303,7 +339,7 @@ def civ_api_down(url, dir_, civitai_api_key, civ_filename):
     if not civ_filename:
         wget_command = (
             f'wget -c -nv '
-            f'-P "{dir_}" "{url_dl}"'
+            f'-P {shqt(dir_)} {shqt(url_dl)}'
         )
         os.system(wget_command)
     
@@ -313,12 +349,13 @@ def civ_api_down(url, dir_, civitai_api_key, civ_filename):
         if not os.path.exists(output_path):
             wget_command = (
                 f'wget -c -nv '
-                f'-O "{output_path}" "{url_dl}"'
+                f'-O {shqt(output_path)} {shqt(url_dl)}'
             )
             os.system(wget_command)
     
     return output_path
 
+
 def drive_down(url, dir_):
     import gdown
 
@@ -361,13 +398,13 @@ def hf_down(url, dir_, hf_token, romanize):
     if hf_token:
         os.system(
             f'wget -c -nv '
-            f'--header="Authorization: Bearer {hf_token}" '
-            f'-O "{os.path.join(dir_, filename)}" "{url}"'
+            f'--header={shqt(f"Authorization: Bearer {hf_token}")} '
+            f'-O {shqt(output_path)} {shqt(url)}'
         )
     else:
         os.system(
             f'wget -c -nv '
-            f'-O "{os.path.join(dir_, filename)}" "{url}"'
+            f'-O {shqt(output_path)} {shqt(url)}'
         )
 
     return output_path
@@ -375,13 +412,26 @@ def hf_down(url, dir_, hf_token, romanize):
 
 def download_things(directory, url, hf_token="", civitai_api_key="", romanize=False):
     url = url.strip()
+    if not url:
+        return None
+
+    # SSRF & protocol validation
+    try:
+        url = validate_url(url)
+    except Exception as e:
+        msg = f"Download blocked: {e}"
+        print(f"\033[91m{msg}\033[0m")
+        gr.Warning(msg)
+        return None
+
+    hostname = (urllib.parse.urlparse(url).hostname or "").lower()
     downloaded_file_path = None
 
-    if "drive.google.com" in url:
+    if hostname == "drive.google.com" or hostname.endswith(".drive.google.com"):
         downloaded_file_path = drive_down(url, directory)
-    elif "huggingface.co" in url:
+    elif hostname in ("huggingface.co", "hf.co") or hostname.endswith((".huggingface.co", ".hf.co")):
         downloaded_file_path = hf_down(url, directory, hf_token, romanize)
-    elif "civitai." in url:
+    elif hostname == "civitai.com" or hostname.endswith(".civitai.com"):
         url = url.replace("civitai.red", "civitai.com")
         if not civitai_api_key:
             msg = "You need an API key to download Civitai models."
@@ -407,7 +457,7 @@ def download_things(directory, url, hf_token="", civitai_api_key="", romanize=Fa
     else:
         os.system(
             f'wget -c -nv '
-            f'-P "{directory}" "{url}"'
+            f'-P {shqt(directory)} {shqt(url)}'
         )
 
     return downloaded_file_path
@@ -726,4 +776,4 @@ def html_template_message(msg):
 
 def escape_html(text):
     """Escapes HTML special characters in the input text."""
-    return text.replace("<", "&lt;").replace(">", "&gt;").replace("\n", "<br>")
+    return text.replace("<", "&lt;").replace(">", "&gt;").replace("\n", "<br>")